The Role of Pretrained Representations for the OOD Generalization of RL Agents

Building sample-efficient agents that generalize out-of-distribution ( OOD ) in realworld 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 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 simto-real case . In addition , we find that the generalization performance of a simple downstream proxy task reliably predicts the generalization performance of our reinforcement learning control tasks under a wide range of practically relevant OOD settings . Such proxy tasks can thus be used to select pretrained representations that will lead to agents that generalize out-of-distribution . 1 INTRODUCTION . Robust out-of-distribution ( OOD ) generalization is one of the key open challenges in machine learning . This is particularly relevant for the deployment of ML models to the real world , where we need systems that generalize beyond the i.i.d . ( independent and identically distributed ) data setting ( Azulay & Weiss , 2019 ; Barbu et al. , 2019 ; Djolonga et al. , 2020 ; Funk et al. , 2021 ; Gulrajani & Lopez-Paz , 2020 ; Hendrycks & Dietterich , 2019 ; Koh et al. , 2021 ; Michaelis et al. , 2019 ; Roy et al. , 2018 ; Schölkopf et al. , 2021 ) . One instance of such models are agents that learn by interacting with a training environment and we would like them to generalize to other environments with different statistics ( Ahmed et al. , 2021 ; Cobbe et al. , 2019 ; Ke et al. , 2021 ; Pfister et al. , 2019 ; Zhang et al. , 2018 ) . Consider the example of a robot with the task of moving a cube to a target position : Such an agent can easily fail as soon as some aspects of the environment differ from the training setup , e.g . the shape , color , and other object properties , or when transferring from simulation to real world . Humans do not suffer from these pitfalls when transferring learned skills beyond a narrow training domain , presumably because they represent visual sensory data in a concise and useful manner ( Anand et al. , 2019 ; Gordon & Irwin , 1996 ; Lake et al. , 2017 ; Marr , 1982 ; Spelke , 1990 ) . Therefore , a particularly promising path is to base decisions and predictions on similar low-dimensional representations of our world ( Barreto et al. , 2017 ; Bengio et al. , 2013 ; Finn et al. , 2016 ; Kaiser et al. , 2019 ) . The learned representation should facilitate efficient downstream learning ( Anand et al. , 2019 ; Eslami et al. , 2018 ; Stooke et al. , 2020 ) and exhibit better generalization ( Srinivas et al. , 2020 ; Zhang et al. , 2020 ) . Learning such a representation from scratch for every downstream task and every new variation would be highly inefficient . If we learned to juggle three balls , we should be able to generalize to oranges or apples without learning again from scratch . We could even do it with cherimoyas , a type of fruit that we might have never seen before . We can effectively reuse our generic representation of the world . We thus consider deep learning agents trained from pretrained representations and ask the following questions : To what extent do they generalize under distribution shifts similar to those mentioned above ? Do they generalize in different ways or to different degrees depending on the type of distribution shift , including sim-to-real ? Can we predict the OOD generalization of downstream agents from properties of the pretrained representations ? To answer the questions above , we need our experimental setting to be realistic , diverse , and challenging , but also controlled enough for the conclusions to be sound . We therefore base our study on the robot platform introduced by Wüthrich et al . ( 2020 ) . The scene comprises a robot finger with three joints that can be controlled to manipulate a cube in a bowl-shaped stage . Dittadi et al . ( 2021 ) conveniently introduced a dataset of simulated and real-world images of this setup with ground-truth labels , which can be used to pretrain and evaluate representations . To train downstream agents , we adapted the simulated reinforcement learning benchmark CausalWorld from Ahmed et al . ( 2021 ) that was also developed for this robotic setup . Building upon these works , we can design our experimental study as follows : First , we pretrain representations from static simulated images of the setup and evaluate a collection of representation metrics . Following prior work ( Eslami et al. , 2018 ; Ghadirzadeh et al. , 2017 ; Ha & Schmidhuber , 2018 ; Nair et al. , 2018 ; Van Hoof et al. , 2016 ; Watter et al. , 2015 ) , we focus on autoencoder-based representations . Then , we train downstream agents from this fixed representation on a set of environments . Finally , we investigate the zero-shot generalization of these agents to new environments that are out of the training distribution , including the real robot . Our setup is summarized in Fig . 1 . The goal of this work is to provide the first systematic and extensive account of the OOD generalization of downstream RL agents in a robotic setup , and how this is affected by characteristics of the upstream pretrained representations . We summarize our contributions as follows : • We train 240 representations and 11,520 downstream policies,1 and systematically investigate the performance of these policies under a diverse range of distribution shifts.2 • We extensively analyze the relationship between the generalization of our RL agents and a substantial set of representation metrics . • Notably , we find that a specific representation metric that measures the generalization of a simple downstream proxy task reliably predicts the generalization of downstream RL agents under the broad spectrum of OOD settings considered here . This metric can thus be used to select pretrained representations that will lead to more robust downstream policies . • In the most challenging of our OOD scenarios , we deploy a subset of the trained policies to the corresponding real-world robotic platform , and observe surprising zero-shot sim-to-real generalization without any fine-tuning or domain randomization . 2 BACKGROUND . In this section , we provide relevant background on the methods for representation learning and reinforcement learning , and on the robotic setup to evaluate out-of-distribution generalization . Variational autoencoders . VAEs ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) are a framework for optimizing a latent variable model pθ ( x ) = ∫ z pθ ( x | z ) p ( z ) dz with parameters θ , typically with a fixed prior p ( z ) = N ( z ; 0 , I ) , using amortized stochastic variational inference . A variational distribution qφ ( z |x ) with parameters φ approximates the intractable posterior pθ ( z |x ) . The approximate posterior and generative model , typically called encoder and decoder and parameterized by neural networks , are jointly optimized by maximizing a lower bound to the log likelihood ( the ELBO ) : log pθ ( x ) ≥ Eqφ ( z |x ) [ log pθ ( x | z ) ] −DKL ( qφ ( z |x ) ‖p ( z ) ) = L ELBO θ , φ ( x ) . ( 1 ) In β-VAEs , the KL term is modulated by a factor β to enforce a more structured latent space ( Burgess et al. , 2018 ; Higgins et al. , 2017a ) . While VAEs are typically trained without supervision , we also employ a form of weak supervision proposed by Locatello et al . ( 2020 ) that encourages disentanglement ( Higgins et al. , 2018 ) . Reinforcement learning . A Reinforcement Learning ( RL ) problem is typically modeled as a Partially Observable Markov Decision Process ( POMDP ) defined as a tuple ( S , A , T , R , Ω , O , γ , ρ0 , H ) with states s ∈ S , actions a ∈ A and observations o ∈ Ω determined by the state and action of the environment O ( o|s , a ) . T ( st+1|st , at ) is the transition probability distribution function , R ( st , at ) is the reward function , γ is the discount factor , ρ0 ( s ) is the initial state distribution at the beginning of each episode , and H is the time horizon per episode . The objective in RL is to learn a policy 1Training the representations required approximately 0.62 GPU years on NVIDIA Tesla V100 . Training and evaluating the downstream policies required about 86.8 CPU years on Intel Platinum 8175M processors . 2Additional results and videos are provided at https : //sites.google.com/view/ood-rl . π : S ×A→ [ 0 , 1 ] , typically parameterized by a neural network , that maximizes the total discounted expected reward J ( π ) = E [ ∑H t=0 γ tR ( st , at ) ] . There is a broad range of proposed model-free learning algorithms to find π∗ by policy gradient optimization or learning value functions while trading off exploration and exploitation ( Fujimoto et al. , 2018 ; Haarnoja et al. , 2018b ; Schulman et al. , 2015a ; b ; 2017 ; Silver et al. , 2014 ; Sutton et al. , 1999 ) . Here , we optimize the objective above with Soft Actor Critic ( SAC ) , a widely used off-policy method in control tasks due to its sample efficiency ( Haarnoja et al. , 2018b ) . SAC aims to improve sample-inefficiency and convergence in RL by simultaneously maximizing the expected reward and the entropy H ( π ( ·|st ) ) . A robotic setup to evaluate out-of-distribution generalization . Our study is based on a real robot platform where a robotic finger with three joints manipulates a cube in a bowl-shaped stage ( Wüthrich et al. , 2020 ) . We pretrain representations on a labeled dataset introduced by Dittadi et al . ( 2021 ) which consists of simulated and real-world images of this setup . This dataset has 7 underlying factors of variation ( FoV ) : angles of the three joints , and position ( x and y ) , orientation , and color of the cube . After training the representations , we train downstream agents and evaluate their generalization on an adapted version of the simulated CausalWorld benchmark ( Ahmed et al. , 2021 ) that was developed for the same setup . Finally , we can test sim-to-real generalization on the real robot . Our experimental setup , illustrated in Fig . 1 , allows us to systematically investigate a broad range of out-of-distribution scenarios in a controlled way . We pretrain our representations from this simulated dataset that covers 8 distinct cube colors . We then train an agent from this fixed representation on a subset of the cube colors , and evaluate it ( 1 ) on the same colors ( this is the typical scenario in RL ) , ( 2 ) on the held-out cube colors known to the representation , or ( 3 ) OOD w.r.t . the representation training distribution , e.g . on novel colors and shapes or on the real world . We closely follow the framework for measuring OOD generalization proposed by Dittadi et al . ( 2021 ) . In this framework , a representation is initially learned on a training set D , and a simple downstream model is trained on a subset D1 ⊂ D to predict the ground-truth factors from the learned representation . Generalization is then evaluated by testing the downstream model on a set D2 that differs distributionally from D1 , e.g . containing images corresponding to held-out values of a chosen factor of variation ( FoV ) . Dittadi et al . ( 2021 ) consider two flavors of OOD generalization depending on the choice of D2 : First , the case when D2 ⊂ D , i.e . the OOD test set is a subset of the dataset for representation learning . This is denoted by OOD1 and corresponds to the scenario ( 2 ) from the previous paragraph . In the other scenario , referred to as OOD2 , D and D2 are disjoint and distributionally different . This even stronger OOD shift corresponds to case ( 3 ) above . The generalization score for D2 is then measured by the ( normalized ) mean absolute prediction error across all FoVs except for the one that is OOD . Following Dittadi et al . ( 2021 ) , we use a simple 2-layer Multi-Layer Perceptron ( MLP ) for downstream factor prediction . This simple and cheap generalization metric could serve as a convenient proxy for the generalization of more expensive downstream tasks . From now on , we will report the negative error for interpretability , and decompose this metric in addition on a per-factor level . We refer to these generalization scores as GS-OOD1 , GS-OOD2-sim , and GS-OOD2-real accordingly . The focus of Dittadi et al . ( 2021 ) was to scale VAE-based learning approaches to more realistic scenarios and study the generalization of these simple downstream tasks , with a particular emphasis on disentanglement . Building upon their contributions , we can leverage the broader potential of this robotic setup with many more OOD2 scenarios to study our research questions : To what extent can agents generalize under distribution shift ? Do they generalize in different ways depending on the type of distribution shift ( including sim-to-real ) ? Can we predict the OOD generalization of downstream agents from properties of the pretrained representations such as the GS metrics from Dittadi et al . ( 2021 ) ?

Summary: The paper discusses a study of different learned autoencoder based representations in the context of generalisation (in and out of distribution) in reinforcement learning. It focuses on two types of VAE models and a robotic block reaching/pushing setting. Comparisons are performed to separately test OOD data for the representation and learned policy as well as generalisation to a real world setting after training in simulation. The analysis shows correlations between various properties and metrics and final downstream RL performance