[Paper Review] The Role of Pretrained Representations for the OOD Generalization of RL Agents

Andrea Dittadi, Frederik Träuble+8 more|arXiv (Cornell University)|Jul 12, 2021

Reinforcement Learning in Robotics63 references4 citations

TL;DR

This paper investigates how pretrained variational autoencoder (VAE)-based representations improve out-of-distribution (OOD) generalization in reinforcement learning (RL) agents. Through extensive experiments on 240 representations and 10,000 RL policies in simulation, it finds that proxy tasks on these representations reliably predict downstream OOD generalization, enabling efficient selection of representations that yield robust agents—even under challenging sim-to-real shifts.

ABSTRACT

Building sample-efficient agents that generalize out-of-distribution (OOD) in real-world settings remains a fundamental unsolved problem on the path towards achieving higher-level cognition. One particularly promising approach is to begin with low-dimensional, pretrained representations of our world, which should facilitate efficient downstream learning and generalization. By training 240 representations and over 10,000 reinforcement learning (RL) policies on a simulated robotic setup, we evaluate to what extent different properties of pretrained VAE-based representations affect the OOD generalization of downstream agents. We observe that many agents are surprisingly robust to realistic distribution shifts, including the challenging sim-to-real case. In addition, we find that the generalization performance of a simple downstream proxy task reliably predicts the generalization performance of our RL agents under a wide range of OOD settings. Such proxy tasks can thus be used to select pretrained representations that will lead to agents that generalize.

Motivation & Objective

• To investigate how pretrained representations enhance out-of-distribution (OOD) generalization in RL agents.
• To evaluate the impact of various VAE-based representation properties on downstream RL policy performance under distribution shifts.
• To determine whether a simple proxy task can reliably predict the OOD generalization capability of RL agents.
• To identify representation characteristics that lead to robust, sample-efficient RL agents in real-world-like settings.

Proposed method

• Training 240 distinct VAE-based representations on a simulated robotic environment to capture diverse world representations.
• Training over 10,000 RL policies on top of each representation to assess downstream generalization under OOD settings.
• Using a proxy task—evaluating representation quality on a simple, non-RL downstream task—to predict OOD generalization performance of RL agents.
• Evaluating generalization across multiple OOD settings, including sim-to-real transfer, distribution shifts in observation space, and action space variations.
• Analyzing the correlation between proxy task performance and actual RL agent performance across diverse representation types.

Experimental results

Research questions

• RQ1 To what extent do pretrained VAE-based representations improve OOD generalization in RL agents?
• RQ2 Can a simple proxy task accurately predict the OOD generalization performance of RL agents trained on different representations?
• RQ3 Which properties of pretrained representations most strongly correlate with robust OOD generalization in downstream RL policies?
• RQ4 How does representation quality affect sample efficiency and robustness under sim-to-real distribution shifts?

Key findings

• Many RL agents trained on pretrained representations exhibit strong robustness to realistic distribution shifts, including sim-to-real transfer, despite no explicit domain adaptation.
• The performance of a simple downstream proxy task on the representations strongly predicts the OOD generalization performance of the corresponding RL agents across diverse OOD settings.
• Representations that perform well on the proxy task consistently yield RL agents with superior generalization, enabling efficient selection of high-performing representations.
• Even without fine-tuning, pretrained representations significantly enhance sample efficiency and OOD generalization, suggesting their value as inductive biases for real-world RL.