The Role of Pretrained Representations for the OOD Generalization of RL Agents
Authors: Frederik Träuble
Andrea Dittadi
Manuel Wuthrich
Felix Widmaier
Peter Vincent Gehler
Ole Winther
Francesco Locatello
Olivier Bachem
Bernhard Schölkopf
Stefan Bauer
Venue: ICLR 2022 Poster (2022)
Keywords: representations
out-of-distribution
generalization
deep learning
reinforcement learning

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.

Published as a conference paper at ICLR 2022
THE ROLE OF PRETRAINED REPRESENTATIONS FOR THE OOD GENERALIZATION OF RL AGENTS
Frederik Träuble,' Andrea Dittadi;, Manuel Wüthrich,' Felix Widmaier,' Peter Gehler, Ole Winther, Francesco Locatello, Olivier Bachem," Bernhard Schölkopf, Stefan Bauer 1.5
I Maa Planck Institute for Intelligent Systems, Tübingen, Germany, Technical University of Denmark, AAmazon Lablets, *Google Brain, 'KTH Stockholm
ABSTRACT
Building sample-efficient agents that generalize fit-of-distribution (OOD) in real- world settings remains a fundamental unsolved problem on the path towards achiev- ing 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- representa- tions affect the OOD generalization of downstream agents. We observe that many agents are surprisingly robust realistic distribution shifts, including the challeng- ing 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 thus be used to select pretrained representations that will lead to agents that generalize.
INTRODUCTION
Robust out-of-distributio (OOD) generalization is one of the key open challenges in machine
learning. This is particularly relevant for the deployment of ML models the real world, where
need systems that generalize beyond the i.i.d. (independent and identically distributed) data setting
(Schölkopf al. 2021 Djolonga al. 2020 Koh al. 2021 Barbu et al. 2019 Azulay & Weiss)
2019 Roy et al. 2018 Gulrajani & Lopez- Paz, 20200 Hendrycks & Dietterich 2019] Michaelis
al. 2019 Funk 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
statisties (Zhang et al. 2018 Phister al. 2019 Cobbe et al. 2019 Ahmed et al., 2021 Ke et al.
2021). 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 (Marr
1982, Gordon & Irwin 1996. Lake et al. 20177 Anand et al. 2019 Spelke 1990). Therefore,
particularly promising path is to base predictions and decisions on simil ar low-dimensional represen-
tations of our world (Bengio et al., 2013
Kaiser et al. 2019 Finn et al. 2016 Barreto et al. 2017
Dittadi et al. 202 Stooke al. 2021 Vinyals et al. 2019). 9. The learned representation should
facilitate efficient downstream learning (Hslami et al. 2018 Anand et al. [2099] Stooke et al.
202
Van Steenkiste et al. (2019) and exhibit better generalization (Zhang et al. 20200 Srinivas et al. 2020).
Learning such a representation from scratch for every downstream task and every new variation would be inefficient. If we learned to juggle three balls, we should be able generalize to oranges or apples without learning again from scratch. We could even do it with cherimoyas, a fruit that we might have never seen before. We can effectively reuse our generic representation of the world.
* "gqual contribution. The order was chosen random and can he swapped. Correspondence to: frederi tr..uruulllee. Traeuble@tuuein e@t cerrobinnnen.p.p. mpg. de and adi i.t.tt.. t

---

Published as a conference paper at ICLR 2022
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 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 be sound. We therefore base our study on the robot platform introduced by Wüthhric 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. (2021c) conveniently introduced dataset of simulated and real-world images of this setup with ground ground-truth labels, which can be used to pretrain and evaluate representations. To train downstream agents, we adapted the simulated reinforcement learning benchmark Causal from Ahmed et al. (202T) that was developed for this platform. Building upon these works, we design our experimental study as follows (see Fig. I First, we pretrain representations from static simulated images of the setup and evaluate a collection of representation metrics. Following prior work (Watter et al.) 2015 Van Hoof et al. 2016; Ghadirzadeh et al., 2017 Nair et al., 2018: Ha & Schmidhuber 2018 Eslami et al. 2018), we focus on autoencoder- based representations. Then, we train downstream agents from this
fixed representation on 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.
The goal of this work is 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 their performance under diverse range of distribution shifts
1 and systematically investigate
We extensively analyze the relationship between the generalization of our RL agents and a substantial set of representation metrics.
Notably, we find that specific representation metric that measures the generalization of simple downstream proxy task reliably predicts the generalization of downstream RI. 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 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 domain randomization.
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.
framework for optimizing a latent variable model (x( = Jg Pa (x z z)p z with parameters 8, typically with a fixed prior p(z) = (z; O, I), using amortized stochastic variational inference. A variational distri- bution q (z x) with parameters é approximates the intractable posterior pz x). The approximate posterior and generative model, typically called encoder and decoder and parameterized by neural networks, are jointly optimized by maximizing lower bound to the log likelihood (the ELBO):
Variational autoencoders. VAEs (Kingma & Welling 2014 Rezende et al. 2014)
logpe(()) 2 gx, [log (x z)] DKL (z x))p((z)) CELBO (x)
structured latent space (Higgins) et al. 2017a; Burgess et al. (0118). While VAEs are typically trained without supervision, we also
In B- VAEs, the KL term is modulated by factor / to enforce
employ form of weak supervision (Locatello et al. 2020) that encourages disentanglement.
Reinforcement learning. A Reinforcement Learning (RL.) problem is typically modeled as a Par- tially Observable Markov Decision Process (POMDP) defined tuple (S, A, T, R, Ω, O, H)
1 'Training 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 2 Additional results and videos are provided httpt: h/tp:///sseeeee. googl www/ vi ww/ewwwwoodooooo
2

---

Published as a conference paper at ICLR 2022
1. Pretrain Representations
2. Train Downstream Policy
decoder
freeze trained encoder
MLP Policy
robot
Evaluate Trained Policies
(n) in-distriburion
(b) OODI
0OD2-sim
(d) OOD2-real
Figure 1: Overview of our experimental setup for investigating diitr--distiitttion generalization in down- stream tasks. (1) We train 240 3-VAEs on the robotic dataset from DDttadi al. (202Tc). (2) We then train downstream policies to solve object reaching or pushing, using multiple random RL seeds per VAE. The input to a policy consists of the output of a pretrained encoder and additional task-related observable variables. Crucially, the policy is only trained on subset of the cube colors from the pretraining dataset. (3) Finally, we evaluate these policies their respective tasks four different scenarios: in-distribution, i.e. with cube colors used policy training; (b) OODI, i.c. with cube colours previously seen by the encoder but OOD the policy; (c) 00m2-sim, having cube colours also OOD to the encoder; (d) sim-to-ree no-ato--all zero-shot on the real-world setup.
with states E S, actions E and observations o E Ω determined by the state and action of the en- vironment Oooss a). T88++1 dz) is the transition probability distribution function, [[(s,, (z) is the reward function, 7 is the discount factor, po(s) the initial state distribution at the beginning each episode, and H is the time horizon per episode. The objective in RL is to learn policy [0, 1]. typically parameterized by a neural network, that maximizes the total discounted expected re- ward ((T) = E Σm 7 R(s,, at]] There a broad range of mrdel-free learning algorithms to find n by policy gradient optimization or by learning value functions while trading off exploration and ex- ploitation (Haarnoja et al. (0188) 8b, Schulman et al. 2017)] /utton al.) 1999; (Schulman et al.) 2015b a Silver et 2014] Fujimoto 2 al. 2018). Here, we optimize the objective above with Soft Actor Critic (SAC), an off off-policy method that simultancously maximizes the expected reward and the entropy (T((ss)), and is widely used in control tasks due to its sample efficiency (Haarnoja et al. 2018 8b).
x A
A robotic setup to evaluate ant-of-distribution generalization. Our study is based real robot platform where a robotic finger with three joints manipulates a cube a bowl-shaped stage Wüthrich labeled dataset introduced by Dittadi al. (2021c)
et al. (020). We pretrain representations 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. Some of these factors are correlated (Dittadi et al. 2021ch, which may be problematic for representation learners. especially in the context of disentanglement (Träuble al.) 2021 Chen ct al. 202T). After training the representations, we train downstream agents and evaluate their generalization on an adapted version of the simulated CausalWorld benchmark (Ahmed al., 2021 that was developed for the same setup. Finally, we test sim-to-real generalization on the real robot.
Our experimental setup, illustrated in Fig. M allows to systematically investigate broad range of nnt-of-diitribution scenarios in 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 subset of the cube colors, and evaluate it on the same colors (this is the typical scenario in RL), (2) on the held-out cube colors that are still known to the encoder, or (3) OOD w.r.t. the encoder's training distribution, e.g. novel colors and shapes
the real world.
We closely follow the framework for measuring OOD generalization proposed by Dittadi et al. (2021c). In this framework, representation is initially learned downstream model is trained on a subset D, C to predict the ground -tuth factors from the learned representation. Generalization then evaluated by testing the downstream model set D:
training set D, and a simple

---

Published as a conference paper at ICLR 2022
that differs distributionally from D1, e.g. containing images corresponding to held-out values of chosen factor of variation (FoV), Dittadi et al. (2021cc) consider two flavors of OOD generalization when D2 C D, ie. the OOD test set is subset of the
depending on the choice of D2: First, 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 as OOD2, 'D and D2 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. (202c)), we use a simple 2-layer Multi-Layer Perceptron (MLP) for downstream factor prediction, we train one MLP for each FoV, and report the negative error. This simple and cheap generalization metric could serve convenient proxy for the generalization of more expensive downstream tasks. We refer to these generalization scores as GS-OOD1, GS GS-OOD2-sim, and GS-OOD2-real depending on the scenario.
The focus of Dittadi et al. (222c)) was to scale VAE-based 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 study research questions: To what extent can agents generalize under distribution shift? Do they generalize in different ways depending on the type of 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.) (2021c)?
3 STUDY DESIGN
Robotic setup. Our setup is based on TriFinger (Wüthrich al. 2020) and consists of robotic finger with three joints that can be controlled to manipulate an object (e.g. a cube) in bowl-shaped stage. The agent receives camera observation consistent with the images in Dittadi et al. (2021c) and outputs three-dinn diree-dieniional action. During training, which always happens in simulation, the agent only observes a cube of four possible colors, randomly sampled at every episode (see Fig. 1 step 2).
Distribution shifts. After training, we evaluate these agents in 7 environments: (1) the training environment, which the typical setting in RI, (2) the OODI setting with cube colors that are OOD for the agent but still in-distribution for the encoder, (3) the more challenging OOD2-sim setting where the colors are also OOD for the encoder, (4-6) the OOD2 settings where the object colors are as in the 3 previous settings but the cube is replaced by sphere (a previously unseen shape). (7) the 00D2-real setting, where we evaluate zero-shot sim-to-real transfer on the real robotic platform.
Tasks. We begin our study with the object reaching downstream control task, where the agent has to reach an object placed at an arbitrary random position in the arena. This is significantly more challeng- ing than directly predicting the ground-truth factors, the agent has to learn to reach the cube by act- ing on the joints, with a scalar reward as the only learning signal. Consequently, the compute required to learn this task is about 1, 1,000 times greater than in the simple factor prediction case. We additionally include in our study a pushing task which consists of pushing an object to goal position that is sam- pled at each episode. Learning this task takes one order of magnitude more compute than object reach- ing, likely due to the complex rigid-body dynamics and object interactions. To the best of our knowl- edge, this is the most challenging manipulation task that is currently feasible on our setup. [Ahmed et al.) (2021) report solving similar pushing task, but require the full ground-truth state to be observable.
Training the RL agents. The inputs at time are the camera observation of and a vector of observable variables I2 containing the joint angles and velocities, as well as the target object position in pushing. We then feed the camera observation o, into an encoder e that was pretrained on the dataset in Dittadi et al. c0. The result is concatenated with Hr. yielding state vector 8; [are, e(ox)]. We then use SAC to train the policy with s, as input. The policy, value, and Q networks are implemented as MLPs with 2 hidden layers of size 256. When training the policies, we keep the encoder frozen.
Model sweep. To shed light on the research questions outlined in the previous sections, we perform a large-scale study in which we train 240 representation models and 1 11,520 downstream policies, as described below. See Appendix [foo further implementation details.
We train 120 3-VAEs (Higgins al. 2017a) and 120 Ada-GVAEs (Locatello et al. 2020) with a subset of the hyperparameter configurations and neural architecture from Dittadi et al. (202cc). Specifically, we consider 3 E 21, 2),,
1), B annealing over (0, 50000) steps, with and without
4

---

Published as a conference paper at ICLR 2022
imput noise, and 10 random seeds per configuration. The latent space size is fixed to 10 following
prior work (Kim & Mnih] 2018: Chen et al. 2018 Locatello et al. 2020 Träuble et al. )0
For object reaching, we train 20 downstream policies (varying random seed) for each of the 240
VAEs. The resulting 4,800 policies are trained for 400k steps (approximately 2,400 episodes).
Since pushing takes substantially longer to train, we limit the number of policies trained on this task: We choose subset of 96 VAEs corresponding to only 4 seeds, and then use 10 seeds per representation. The resulting 960 policies are trained for 3M steps (about 9,000 episodes). Finally. for both tasks we also investigate the role of regularization on the policy. More specifically, we repeat the two training sweeps from above (5,760 policies), with the difference that now the policies are trained with 1.1 regularization on the first layer.
Limitations of our study, Although we aim to provide a sound and extensive empirical study, such studies are inevitably computationally demanding. Thus, we found it necessary to make certain design choices. For each of these choices, we attempted to follow common practice, in order to maintain our study relevant, general, and useful as possible. One such decision is that of focusing on autoencoder- based representations. To answer our questions on the effect of upstream representations on the generalization of downstream policies, these representations are obtained is not directly relevant to answer our research question. Following Dittadi al. (202Tc), we chose to focus on 3-VAE and Ada-GVAF models, as they were shown to provide a broad set of representations, including fully disentangled ones. Although we conjecture that other classes of representation learning algorithms should generally reveal similar trends those found in our study, this is undoubtedly an interesting extension. As for the RL algorithm used in this work, SAC is known to be a particularly cample-efficiese model- free RL method that is a popular choice in robotics (Haarnoja et al.] 2018a| Kiran et al.] [202]] Singh et al.) (2119). Extensive results on pushing from ground-truti features on the same setup in Ahmed et al. (2021) indicate that methods like TD3 (Fujimoto et al., (018)) or PPO (Schulman et al., 2017) perform very similarly to SAC under the same reward structure and observation space. Thus, we expect the results of our study to hold beyond SAC. Another interesting direction is the study of additional regularization schemes on the policy network, an aspect that is often overlooked in RL. We expect the potential insights from extending the study along these axes to not justify the additional compute costs and corresponding carbon footprint. However, with improving efficiency and decreasing costs, we believe these could become worthwhile validation experiments in the future.
need diverse range of representations. How
RESULTS
We discuss our results in three parts: In Section 4.1
we present the training results of our large-scale sweep, and how policy regularization and different properties of the pretrained representations affect in-distribution reward. Section 4.2 gives an extensive account of which metries of the pretrained representations predict OOD generalization of the agents in simulated environments. Finally, in Section 4.3 we perform a similar evaluation on the real robot, a zero-shot sim-to-real scenario.
4.1 RESULTS IN THE TRAINING ENVIRONMENT
Fig. shows the training curves of all policies for object reaching and pushing in terms of the lusk-specific success metric. Here we use success metrics for interpretability, as their range is always [0, 1]. In object reaching, the success metric indicates progress from the initial end effector position to the optimal distance from the center of the cube. It is 0 if the final distance is not smaller than the initial distance, and if the end effector is touching the center
face the cube. pushing, the success metric is defined as the volumetric overlap of the cube with the goal cube, and the task can be visually considered solved with a score around 80%.
From the training curves we can conclude that both tasks can be consistently solved from pixels using pretrained representations. In particular, all policies on object reaching attain almost perfect scores. Unsurprisingly, the more complex pushing task requires significantly more training, and the variance across policies is larger. Nonetheless, almost all policies learn to solve the task satisfactorily.
To investigate the effect of representations the training reward, we now compute its Spearman rank correlations with various supervised and unsupervised metrics of the representations (Fig 2 bottom). By training reward, we mean the average reward of a fully trained policy over 200 episodes
5

---

Published as a conference paper at ICLR 2022
object reaching
pushing
Average success
0.8
Average success
0.8
0.6
0.4
0.4
0.2
0.2
200000
400000
30
Time step
Time step
1e6
reward in cleato.
roward diamb.
12
(1) cube color
(23-4) pains
GBT
MLP
MLP (1
MLP MIP
MID (0)
MLP (B)
MLP
MPP (T)
ELBO
CST
M.P
(1
MI 22
ML 09
MUP e0
MLP (0)
MLP
ELBO
recons 000
oubc
dube ortentation
Figure 2: Top: Average training success, aggregated over all policies from the sweep (median, quartiles. 5th/95th percentiles). Bottom: Rank correlations between representation metries and in in-distribution reward (evaluated when the policies are fully trained), in the case without regularization. Correlations are color-coded red (positive) or blue (negative) when statistically significant (p<0.05), otherwise they are gray.
in the training environment (see Appendix A On object reaching, the final reward correlates with the ELBO and the reconstruction loss. A simple supervised metric evaluate a representation is how well a small downstream model can predict the ground-trui truth factors of variation. Following Dittadi al. (2021c), we use the MLP10000 and GBT10000 metrics (simply MLP and GBT in the following), where MLPs and Gradient Boosted Trees (GBTs) are trained to predict the FoVs from 10,000 samples. The training reward correlates with these metrics as well, especially with the MLP accuracy. This is not entirely surprising: if MLP predict the Fo from the representations, our policies using the same architecture could in principle retrieve the FoVs relevant for the task. Interestingly, the correlation with the overall MLP metric mostly stems from the cube pose Fo i.e. those that are not included in the ground- trrttthtthhhho state Tt. These results suggest that these metrics can be used to select good representations for downstream RL. On the more challenging task of pushing, the correlations are milder but most of them are still statistically significant.
Summary. Both tasks can be consistently solved from pixels using pretrained representations. Unsu- pervised (ELBO, reconstruction loss) and supervised (ground-truth factor prediction) in-distribution metries of the representations are correlated with reward in the training environment.
4.2 OUT-OF-DISTRIUUTII GENERALIZATION IN SIMULATION
In- and out-of outto--distribution rewards. After training, the in- distribution reward correlates with OODI performance on both tasks (especially with regularization), but not with O0D2 performance (see Fig. 3). Moreover, rewards in OODI and OOD2 environments are moderately correlated across tasks and regularization settings.
object reaching pushing
distrb.
12
00D1
43
30
18
31
00D2-sim
12
31
distnb.
71
17
35
12
Unsupervised metrics and informativeness. In Fig. 4 (left) we as- sess the relation between OOD reward and in-distribution metrics (ELBO, reconstruction loss, MLP, and GBT). Both ELBO and recon- struction loss exhibit a correlation with OOD1 reward, but not with O0D2 reward. These unsupervised metrics can thus be useful for selecting representations that will lead to more robust downstream RL tasks, as long as the encoder is in-distribution. While the GBT score not correlated with reward under distribution shift, we observe a significant correlation between OOD1 reward and the MLP score, which measures downstream factor prediction accuracy of an MLP with the same architecture as the one parameterizing the policies. As in Section 4. we further investigate the source of this correlation, and find it in the pose parameters of the cube. Correlations in the O0D2 setting are much weaker, thus conclude that these metrics do not appear helpful for model selection in this case. Our results on pushing confirm these conclusions although correlations are generally weaker, presumably due to the more complicated nature of this task. An extensive discussion is provided in Appendix B.2
O0D1
71
35
00D2-sim
17 28
12 28
in disrrb..
00D1
00D2 sim
in debiitt
00D1
0002.sm
00D0
Figure 3: Correlations between training (in distrib.) and OOD rewards (p<0.05).

---

Published as a conference paper at ICLR 2022
In-dtstributin metrics
Generalization Scores
-17
-10
10
13
14
14
10
-2
84 57
RS1
M&F
EURO
Reco 1o9
05-0001
G5-000
(1)
12
030000
63000
60-000 (2)
0540001
03
00-000004
08/0002- (1)
08 08:0000-00 @0
4000040 (0)
02 05.000044 (9)
08.0000 (00
06 00000 00 (1)
0010/00 ())
00
00
06 00-0000040
68 0000
Figure 4: Rank correlations of representation propertics with OOD1 and OOD2 reward on object reaching
without regularization. Numbering when splitting metrics by FoV: (1) cube color; (2-4) joint angles; (5-7) cube position and rotation. Correlations are color-coded as described in Fig.
Correlations with generalization scores. Here we analyze the link between generalization in RL and the generalization scores (GS) discussed in Section [ which measure the generalization of downstream FoV predictors out of distribution, as opposed to the MLP and GBT metrics considered above. For both OOD scenarios, the distribution shifts underlying these GS scores are the same as the ones in the RL tasks in simulation. We summarize our findings in Fig. 4 (right) on the object reaching task. Reward in the OODI setting is significantly correlated with the GS-OOD1 metric of the pretrained representation. We observe an even stronger correlation between the reward in the simulated OOD2 setting and the corresponding GS-OOD2-sim and GS-OOD2-real scores. On per-factor level, we that the source of the observed correlations primarily stems from the generalization scores w.r.t. the pose parameters of the cube. The OOD generalization metrics can therefore be used as proxies for the corresponding form of generalization in downstream RL tasks. This has practical implications for the training of RL downstream policies which are generally known to be brittle to distribution shifts, labeled images. This allows for selecting representations that yield more robust downstream policies.
can measure representation' 's generalization score from few
Disentangled representations. Disentanglement has been shown to be helpful for downstream performance and OOD1 generalization even with MI.s (Dittadi et al. 2021c). However, in object reaching, we only observe a weak correlation with some disentanglement metrics (Fig. (5). In agreement with (Dittadi et al. 2021c). disentanglement does not correlate with OOD2 generalization. The same study observed that disentanglement correlates with the informativeness of a representation. To understand if these weak correlations originate from this common confounder, we investigate whether they persist after adjusting for MLP FoV prediction accuracy. Given two representations with similar MLP accuracy, does the more disentangled one exhibit better OOD1 generalization? To measure this we predict success from the MI.P accuracy using kNN (k=5) (Locatello et al.) 2019) and compute the residual reward by subtracting the amount of reward explained by the MLP metric. Fig. 5) shows that this resolves the remaining correlations with disentanglement. Thus, for the RL downstream tasks considered here, disentanglement per se does not seem to be useful for OOD generalization. We present similar results on pushing in Appendix B.2
Policy regularization and observation noise. It might seem unsurprising that disentanglement is not useful for generalization in RL, as MLP policies do not have any explicit inductive bias to exploit it. Thus, we attempt to introduce such inductive bias by repeating all experiments with L1 regularization on the first layer of the policy. Although regularization improves OODI and OOD2 generalization in general (see box plots in Fig. (). we observe no clear link with disentanglement. Furthermore, in accordance with Dittadi et al. (2021c), we find that observation noise when training representations is beneficial for O0D2 generalization. See Appendix B.2 for a detailed discussion.
Stronger OOD shifts: evaluating on novel shape. On object reaching, we also test generalization w.r.t. novel shape by replacing the cube with a sphere. This corresponds strong OOpe-type shift, since shape was never varied when training the representations. Surprisingly, the policies appear to be robust to the novel shape. In fact, when the sphere has the same colors that the cube had during policy training, all policies get closer than 5 cm to the sphere on average, with a mean success metric of 95%. On sphere colors from the OODI split, more than 98.5% move the finger closer than this threshold, and the strongest distribution shift (OOD2-sim colors, and cube replaced by sphere) almost 70% surpass that threshold with an average success metric above 80%.
Summary. (1) In- and out-of-distrine distribution rewards correlated, as long as the representation remains in its training distribution (OODI). (2) Similarly, in-distribution representation metrics (both unsupervised and supervised) predict OOD1 reward, but are not reliable when the representation is
7

---

Published as a conference paper at ICLR 2022
in distribution
O0D1
00D2
reward - 0001
0.95
0.985
disentangeemen eeeetnngeeeenttt
0.94
O0D1
0.85
0.945
10
-28
0.80
00
-13
0.940
0.90
0.75
edi 60n
0.935
0.70
SMP
motulant
medium
0.88
0.930
0.99
0.85
L1
L1 reg.
L1
Figure 5: Box plots: fractional success object reaching split according to low (blue), medium-higg (orange), and almost perfect (green) disentanglement. LI regularization the first layer the MLP policy has a positive effect on OODI and OOD2 generalization with minimal sacrifice terms of training reward (see scale). Correlation matrix (left): although we observe mild correlation between some disentanglement metrics and 00D1 (but not OOD2) generalization, this does not hold when adjusting for representation informativeness. Correlations are color-coded descrihed Fig. We use disentanglement metrics from Fastwood & Williams 22088; Chen al. (018): Kumar et al. (2018)) Ridgeway & Mozer (018))
Average 0OD2-real cube)
Count
8
preen
reward 00D2-real
33
35
13
32
-22
0.8
(red
0.8
noword detrit.
reward 00D1
rewerd OCD2
GS-0002-sim
08-0002-000
MLP
Count
blu
recons,
n0 nolse
noiee
Fractional success OOD2-real
Figure 6: Zero-shot sim-to-real on object reaching on over 2,000 episodes. Left: Rank-corelations on
real platform with red cube (color-coded as described in Fig. 2 Middle: Training encoders with additive noise improves sim-to tim-to-real generalization. Right: Histogram of fractional success in the more challenging 00D2-real- 2oDD-real-(grren..ble bluc) scenario from 50 policies across 4 different goal positions.
OOD (OOD2). (3) Disentanglement does not correlate with generalization in experiments, while (4) imput noise when training representations is beneficial for OOD2 generalization. (5) Most notably, the GS metrics, which measure generalization under distribution shifts, are significantly correlated with RL performance under similar distribution shifts. We thus recommend using these convenient proxy metrics for selecting representations that will yield robust downstream policies.
4.3 DEPLOYING POLICIES TO THE REAL WORLD
We now evaluate large subset of the agents on the real robot without fine-tuning, quantify their zero-shot sim-to-real generalization, and find metrics that correlate with real-world performance.
Reaching, We choose 960 policies trained in simulation, based on 96 representations and 10 random seeds, and evaluate them on two (randomly chosen, but far apart) goal positions using a red cube. While red cube was in the training distribution, we consider this to be O0D2 because real- world images represent a strong distribution shift for the encoder (Dittadi et al.) 2021c Djolonga et al. 2020). Although sim-to-real in robotics is considered to be very challenging without domain randomization or fine-tuning (Tobin et al. 2017: Finn et al. 2017 Rusu et al.) 2017). many of our policies obtain a high fractional success without resorting these methods. In addition, in Fig. G (left) we observe sig- nificant correlations between zero-shot real-world performance and some the previously discussed metrics. First, there is positive correlation with the 0OD2-sim reward: Policies that generalize to un- seen cube colors in simulation also generalize to the real world. Second, representations with high GS- 00D2-sin sim and (especially) GS-OOD2-reas scores are promising candidates for sim-to-real transfer. Third, if no labels are available, the weaker correlation with the reconstruction loss on the simulated images can exploited for representation selection. Finally, as observed by Dittadi et al. (202lc)
8

---

Published as a conference paper at ICLR 2022
for simple downstream tasks, input noise while learning representations is beneficial for sim-to-
real generalization (Fig. 6 middle).
Based on these findings, we select 50 policies with high GS-OOD2-real score, and evaluate them on the real world with a green and a blue cube, which is an even stronger OOD2 distri- bution shift. In Fig. 6 (right), where metrics are averaged over cube positions per policy, we observe that most policies can still solve the task: approximately 80% of them position the finger less than 5 cm from the cube. Lastly, we repeat the evaluations on the green sphere that previously performed in simulation, and ob- serve that many policies successfully reach this completely novel object. See Appendix B.3 and the project website for additional results and videos of deployed policies.
time
Figure We select pushing policies with high GS- OOD2-real score. When deployed on the real robot without fine-tuning, they succeed in pushing the cube to specified goal position (transparent blue cube).
Pushing. We now test whether our real-world findings on object reaching also hold for pushing. We again select policies with high GS-OOD2-rewt score and encoders trained with input noise. We record episodes on diverse goal positions and cube colors to support our finding that pushing policies in simulation can generalize to the real robot. In Fig. 7 show three representative episodes with successful task completions and refer to the project site for video recordings and further episodes.
Summary. Policies trained in simulation can solve the task on the real robot without domain ran- domization or fine-tuning. Reconstruction loss, encoder robustness, and OOD2 reward in simulation are all good predictors of real-world performance. For real-world applications, we recommend using GS-OOD AS-OOOD-siim or GS-OOD2-real for model selection, and training the encoder with noise.
5 OTHER RELATED WORK
A key unsolved challenge in RL is the brittleness of agents to distribution shifts in the environment, even if the underlying structure largely unchanged (Cobbe al. 2019 Ahmed al. 202T). This is related to studies on representation learning and generalization in downstream tasks (Gondal et al.,
2019 Steenbrugge et al. 2018 Dittadi et al.] [021b] Esmaeili et al., 2019 Chaabouni et al., (020), as well as domain generalization (see Wang et al. (2021) )) for an overview). More specifically for RL, Higgins et al. (2017b) focus on domain adaptation and zero-shot transfer DeepMind Lab and MuloCo environments, and claim disentanglement improves robustness. To obtain better transfer capabilities, Asadi et al. (222)) argue for discretizing the state space in continuous control domains by clustering states where the optimal policy is similar. Kulkarni et al. (2015) propose geometric object representations by means of keypoints or image-space coordinates and Wulfmeier et al.) (2021) inves- tigate the effect of different representations Transfer becomes especially challenging from the simulation to the real world, a phenomenon often referred the sim-to-real gap. This is particularly crucial in RL, real-world training is expensive, requires sample-effie -enplc-ffficient methods, and is sometimes unfeasible if the reward structure requires ac- curate ground truth labels (Dullc-Arnodd et al.) 2019: Kormushev al.] 2013). This issue typically tackled with large-scale domain randomization in simulation (Akkaya et al. 2019 James et al. 2019).
the learning and exploration of different robotics tasks.
6 CONCLUSION
Robust -ittt--ddistibution (OOD) generalization is still one of the key open challenges in machine learning. We attempted to answer central questions on the generalization of reinforcement learning agents in a robotics context, and how this is affected by pretrained representations. We presented a large-scale empirical study in which we trained over 10,000 downstream agents given pretrained representations, and extensively tested them under variety of distribution shifts, including sim-to- real. We observed agents that generalize OOD, and found that some properties of the pretrained representations can be useful to predict which agents will generalize better. We believe this work brings us one step closer to understanding the generalization abilities of learning systems, and we hope that it encourages many further important studies in this direction.
9

---

Published as a conference paper at ICLR 2022
ETHICS STATEMENT
Our study is based on synthetic and real data of a robotic setup where a robotic finger interacts with cube. Our study does therefore involve any human subjects leading to discrimination, bias or fairness concerns, or privacy and security issues. Representation learning and generalization is important across many disciplines and upplications and could have harmful consequences without humans in the loop in safety-relevant settings. Having sound understanding of the robustness of given MI system based on such pretrained representations distributions shifts crucial to avoid harmful consequences in potential future high-stake applications to society, such as human-robot interaction (c.g. robotic surgery), autonomous driving, healthcare applications or other fairness- related settings. Here, we investigate narrow aspect of this, that is, learning arguably harmless manipulation skills like reaching or pushing an object with simple robotic finger. Our conclusions for OOD generalization based on this setting and thus cannot be directly transferred to any given application setting of concern.
REPRODUCIBILITY STATEMENT
To make sure our experiments are fully reproducible, we provided full account of all required setup and implementation details in Section and Appendix
ACKNOWLEDGEMENTS
We would like thank Anirudh Goyal, Georg Martius, Nasim Rahaman, Vaibhav Agrawal, Max Horn, and the Causality group at the MPI for useful discussions and feedback. We thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting FT. Part of the experiments were generously supported with compute credits by Amazon Web Services.
REFERENCES
Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Manuel Wüthrich, Yoshua Bengio, Bernhard Schölkopf, and Stefan Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. International Conference for Learning Representations, 2021.
Ilge Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, et al. Solving rubik's cube with robot hand. arXiv preprint arXiv:1910.07111 2019.
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. arXiv preprint arXiv:1906.088 08226, 2019.
Kavosh Asadi, David Abel, and Michael Littman. Learning state abstractions for transfer in continuous control. arXiv preprint arXiv:2002.055/8 2020.
Aharon Azulay and Yair Weiss. Why do deep convolutional networks generalize so poorly to small image transformations? Journal of Machine Learning Research, 1114)144-22. 2019.
Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Dan Gutfreund, Josh Tenenbaum, and Boris Katz. Objectnet: A large-scale bias-controlled dataset for pushing the limits of object recognition models. In Advances in Neural Information Processing Systems, pp. 9448-9444 9458, 2019.
André Barreto, Will Dabney. Rémi Munos. Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver. Successor features for transfer reinforcement learning. Advances neural information processing systems, pp. 4065-4065, 2017.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence, 3588):1798-1888, 2013.
10

---

Published as a conference paper at ICLR 2022
Christopher P Burgess, Irina Higgins, Arka Pal. Loic Matthey, Nick Watters, Guillaume Des- jardins, and Alexander Lerchner. Understanding disentangling in beta- VAE. arXiv preprint (rXi::1804003599 2018.
Rahma Chaabouni, Eugene Kharitonov, Diane Bouchacourt, Emmanuel Dupoux, and Marco Baroni. Compositionality and generalization in emergent languages. ACL 2020-8th annual meeting of the Association for Computational Linguistics, 2020.
Tian Qi Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disentanglement in variational autoencoders. In Advances Neural Information Processing Systems, 2018.
Yukun Chen, Andrea Dittadi, Frederik Träuble, Stefan Bauer, and Bernhard Schölkopf. Boxhead: A dataset for learning hierarchical representations. In SVRHM 2021 Workshop @ NewrlPS, 2021.
Karl Cobbe, Oleg Klimov, Chris Hesse, Tachoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In International Conference Machine Learning, pp. 1282-1289. PMLR, 2019.
Andrea Dittadi, Frederik K Drachmann, and Thomas Bolander. Planning from pixels in atari with learned symbolic representations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp.
2021a.
Andrea Dittadi, Samuele Papa, Michele De Vita, Bernhard Schölkopf, Ole Winther, and Francesco Locatello. Generalization and robustness implications in object-centric learning, arXiv preprint arXiv:2107.00637, 2021b.
Andrea Dittadi, Frederik Träuble, Francesco Locatello, Manuel Wüthrich, Vaibhav Agrawal, Ole Winther, Stefan Bauer, and Bernhard Schölkopf. On the Transfer of Disentangled Representations in Realistic Settings. In International Conference on Learning Representations, 2021c
Josip Djolonga, Jessica Yung. Michael Tschannen, Rob Romijnders, Lucas Beyer, Alexander Kolesnikov, Joan Puigcerver, Matthias Minderer, Alexander D' Amour, Dan Moldovan, et al. On robustness and transferability of convolutional neural networks. arXiv preprint arXiv:2007.08858, 2020.
Gabriel Dulac-Arnold Daniel Mankowitz, and Todd Hester. Challenges of real-world reinforcement learning. arXiv preprint
2019.
Cian Eastwood and Christopher KI Williams. A framework for the quantitative evaluation of disentangled representations. International Conference on Learning Representations, 2018.
SM Ali Eslami, Danilo Jimenez Rezende, Frederic Besse, Fabio Viola, Ari S Morcos, Marta Gar- nelo, Avraham Ruderman, Andrei A Rusu, Ivo Danihelka, Karol Gregor, et al. Neural scene representation and rendering. Science, 36066394):1204-1220,, 2018.
Babak Esmaeili, Hao Wu, Sarthak Jain, Alican Bozkurt, N Siddharth, Brooks Paige, Dana H Brooks, Jennifer Dy, and Jan Willem Meent. Structured disentangled representations. In The 22nd International Conference Artificial Intelligence and Statistics, pp.
2019.
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pp. 512-519. IEEE, 2016.
Chelsea Finn, Tianhe Yu, Justin Fu, Pieter Abbeel, and Sergey Levine. Generalizing skills with semi-supervised reinforcement learning. In International Conference on Learning Representations, 2017.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International Conference Machine Learning, pp. 1587-1596. PMLR, 2018.
Niklas Funk, Charles Schaff, Rishabh Madan, Takuma Yoneda, Julen Urain De Jesus, Joe Watson, Ethan K. Gordon, Felix Widmaier, Stefan Bauer, Siddhartha S. Srinivasa, Tapomayukh Bhattachar- jee, Matthew R. Walter, and Jan Peters. Benchmarking structured policies and policy optimization for worldworld dexterous object manipulation. IEEE Robotics and Automation Letters Special Issue: Robotic Grasping and Manipulation Challenges and Progress, 2021.
11

---

Published as a conference paper at ICLR 2022
Ali Ghadirzadeh, Atsuto Maki, Danica Kragic, and Märten Björkman. Deep predictive policy training using reinforcement learning. In 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 2351-2358. IEEE, 2017.
Muhammad Waleed Gondal, Manuel Wüthrich, Djordje Miladinović, Francesco Locatello, Martin Breidt, Valentin Volchkov, Joel Akpo, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. On the transfer of inductive bias from simulation to the real world: a new Misentanglement dataset. In Advances in Neural Information Processing Systems, 2019.
Robert D Gordon and David E Irwin. What's an object file? evidence from priming studies. Perception & Psychophysics, (888))1260-1277)) 1996.
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.01434. 2020.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution, arXiv preprint arXiv:1809. 01999, 2018.
Tuomas Haarnoja, Sehoon Ha, Aurick Zhou, Jie Tan, George Tucker, and Sergey Levine. Learning to walk via deep reinforcement learning. arXiv preprint arXiv: 2018a.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning. pp. 1861-1870. PMLR, 2018b.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corrup- tions and perturbations. Proceedings of the International Conference on Learning Representations, 2019.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. VAE: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017a.
Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In International Conference on Machine Learning, 2017b.
Ashley Hill, Antonin Raffin, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, Rene Traore, Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, and Yuhuai Wu. Stable baselines. https://githnn com/ 2018
Ell-a/stable http///aaaabbbbllleeeeeeeeebeeee
Stephen James, Paul Wohlhart, Mrinal Kalakrishnan, Dmitry Kalashnikov, Alex Irpan, Julian Ibarz, Sergey Levine, Raia Hadsell, and Konstantinos Bousmalis. Sim-to-real via sim-to-sim: Data- efficient robotic grasping via randomized-to tead-to-ccaaa canonical adaptation networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 12627-12637, 2019.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint arXiv: /900::9900..00 0e0..000300077 2019.
Nan Rosemary Ke, Aniket Rajiv Didolkar, Sarthak Mittal, Anirudh Goyal, Guillaume Lajoie, Stefan Bauer, Danilo Jimenez Rezende, Michael Curtis Mozer, Yoshua Bengio, and Christopher Pal. Systematic evaluation of causal discovery in visual model based reinforcement learning, 2021.
net. URL htttp:://pppereeaeeee hwt.:/eee///outitaeeeee Eorum? pomumiddgppaa id=gp50zb1= -C-
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference on Machine Learning, 2018
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. In International Conference Learning Representations, 2014.
12

---

Published as a conference paper at ICLR 2022
B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. IEEE Transactions Intelligent Transportation Systems, 2021.
Pang Wei Koh, Shiori Sagawa, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, et al. Wilds: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning, pp. 5637-5664. PMLR, 2021.
Petar Kormushev, Sylvain Calinon, and Darwin G Caldwell. Reinforcement learning in robotics: Applications and real-world challenges. Robotics, 2(3): 122-148, 2013.
Tejas D Kulkarni, William F Whitney, Pushmeet Kohli, and Josh Tenenbaum. Deep convolutional inverse graphies network. In Advances in Neural Information Processing Systems, 2015.
Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational inference of disen- tangled latent concepts from unlabeled observations. In International Conference on Learning Representations, 2018.
Brenden M Lake, Tomer D Ullman, Joshua Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and Brain Sciences, 40, 2017.
Francesco Locatello, Gabriele Abbati, Thomas Rainforth, Stefan Bauer, Bernhard Schölkopf, and Olivier Bachem. On fairnsss of disentangled representations. In Advances in Neural Information Processing Systems, pp. 14611-14624, 2019.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentanglement without compromises. arXiv preprint urXiv:2002.008880 2020.
David Marr. Vision: A Computational Investigation into the Human Representation and Processing of Visual Information. Henry Holt and Co., Inc., 1982.
Claudio Michaelis, Benjamin Mitzkus, Robert Geirhos, Evgenia Rusak, Oliver Bringmann, Alexan- der Ecker, Matthias Bethge, and Wieland Brendel. Benchmarking robustness in object detection: Autonomous driving when winter is coming. arXiv preprint 1907.07484, 2019.
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. In Advances in Neural Information Processing Systems, 2018.
Niklas Pfister, Stefan Bauer, and Jonas Peters. Learning stable and predictive structures in kinetic systems. Proceedings of the National Academy of Sciences, (6651):25405-2544 2019.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. arXiv preprint urXiv: 14XX::4401.4082 2014.
Karl Ridgeway and Michael C Mozer. Learning deep disentangled embeddings with the f-statistic loss. In Advances in Neural Information Processing Systems, 2018.
Prasun Roy, Subhankar Ghosh, Saumik Bhattacharya, and Umapada Pal. Effects of degradations deep neural network architectures. arXiv preprint 1807 10108, 2018.
Andrei A Rusu, Matej Večerík, Thomas Rothörl, Nicolas Heess, Razvan Pascanu, and Raia Hadsell. Sim-to-real robot learning from pixels with progressive nets. In Conference on Robot Learning, pp. 262-270, 2017.
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 109(5):612-63) 2021.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889-1897. PMLR, 2015a.
13

---

Published as a conference paper at ICLR 2022
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. arXiv preprint 1rXiv:150600244: 02438, 2015b.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint (oXiv:1707.06477 2017.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, pp. 387-395, PMLR, 2014.
Avi Singh, Larry Yang, Kristian Hartikainen, Chelsea Finn, and Sergey Levine. End-to-end robotic reinforcement learning without reward engineering. arXiv preprint 10X1v:1904.07854, 2019.
Elizabeth S Spelke. Principles of object perception. Cognitive science, 14(1):29-56, 1990.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. arXiv preprint 2Xiv::000.041366, 2020.
Xander Steenbrugge, Sam Leroux, Tim Verbelen, and Bart Dhoedt. Improving generalization for abstract reasoning tasks using disentangled feature representations. In Workshop on Relational Representation Learning at NeurIPS, 2018.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learming, pp. 9870-9879. P, 2021.
Richard S Sutton, David A McAllester, Satinder P Singh, Yishay Mansour, et al. Policy gradient methods for reinforcement learning with function approximation. NIPs, volume 99, pp. 1057- 1063. Citeseer, 1999.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ international conference on intelligent robots and systems (IROS), pp. 23-30. IEEE, 2017.
Frederik Träuble, Elliot Creager, Niki Kilbertus, Francesco Locatello, Andrea Dittadi, Anirudh Goyal, Bernhard Schölkopf, and Stefan Bauer. On disentangled representations learned from correlated data. International Conference on Machine Learning (ICML.) 2021.
Herke Van Hoof, Nutan Chen, Maximilian Karl, Patrick van der Smagt, and Jan Peters. Stable rein- forcement learning with autoencoders for tactile and visual data. In 2016 IEEE/RSJ international conference on intelligent robots and systems (IROS), pp. 3928-3934. IEEE, 2016.
Sjoerd Van Steenkiste, Francesco Locatello, Jürgen Schmidhuber, and Olivier Bachem. Are disen- tangled representations helpful for abstract visual reasoning? Advances Neural Information Processing Systems, 32, 2019.
Oriol Vinyals, Igor Babuschkin, Wojciech M Czameeki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782)3350-355) 2019.
Jindong Wang, Cuiling Lan, Chang Liu, Yidong Ouyang, Wenjun Zeng, and Tao Qin. Generalizing to unseen domains: A survey on domain generalization. arXiv preprint arXiv:2 2021.
Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. arXiv preprint S55V::1506.036655 2015.
Markus Wulfmeier, Arunkumar Byravan, Tim Hertweck, Irina Higgins, Ankush Gupta, Tejas Kulka- mi, Malcolm Reynolds, Denis Teplyashin, Roland Hafner, Thomas Lampe, et al. Representation matters: Improving perception and exploration for robotics. In 202/ IEEE International Conference Robotics and Automation (ICRA), pp. 6512-6519. IEEF, 2021.
14

---

Published as a conference paper at ICLR 2022
Manuel Wüthrich, Felix Widmaier, Felix Grimminger, Joel Akpo, Shruti Joshi, Vaibhav Agrawal, Bilal Hammoud, Majid Khadiv, Miroslav Bogdanovic, Vincent Berenz, et al. Trifinger: An open-source robot for learning dexterity. arXiv preprint arXhv:2008.033500 2020.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742. 2020.
Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. study on overfitting in deep reinforcement learning. arXiv preprint arXiv:/804.0668 2018.
15

---

Published as a conference paper at ICLR 2022
A IMPLEMENTATION DETAILS
Task definitions and reward structure. We derive both tasks, object reaching and pushing, from the World environments introduced by Ahmed et al. (202T). pretrain representations on dataset introduced by Dittadi et al. (2021c). and allow only one finger to move in our RL experiments. We introduce the object reaching environment that involves an unmovable cube. We used reward structures similar to those in Ahmed et al. (202I):
object reaching: Ft -550 [d ex) d(gc-1,e--1).
pushing: rz -50 [d(or: et) -((o--1,(--)) 250 [d(or, d(o1-119t-1) Pt
where t. denotes the time step, Pr [0, 1] is the fractional overlap with the goal cube time t, E RR is the end-effector position, E Rd the cube position, 9t E R the goal position, and d(-.-) denotes the Euclidean distance. The cube in object reaching is fixed, i.e. O2 - gr for all The time limit is 2 seconds in object reaching and 4 seconds in pushing.
Success metries. Besides the accumulated reward along episodes, that is determined by the reward normalized success definitions for better inter-
function, we also report two reward-indepente pretability: In object reaching, the success metric indicates progress from the initial end effector position to optimal distance from the center of the cube. is if the final distance greater than equal to the initial distance, and 1 if the end effector is touching the center of a face of the cube. In pushing, the success metric is defined as the volumetric overlap of the cube with the goal cube, and the task be visually considered solved with a score around 80%. We observed that accumulated reward and success are very strongly correlated, thus allowing us to use one or the other for measuring performance.
Training and evaluation details. During training, the goal position is randomly sampled every episode. Similarly, the object color is sampled from the specified training colors from Dj that correspond to the OOD1-B split from Dittadi al. (202lc).
For each policy evaluation (in-distribution and mulated reward and final success over 200 episodes with randomly sampled cube positions and the respective object color in both tasks.
variants), we average the accu-
SAC implementation. Our implementation of SAC builds upon the stable basebaaeeeeeeee inseeiimn inne pack- age (Hill et al. 2018). We use the same SAC hyperparameters used for pushing Ahmed al. (2021). When using L1 regularization, all MLPs, scaled by a coefficient a. We gradually increase regularization by linearly annealing a from 0 to 5 10-7 over 200,000 time steps in object reaching, and from 0 to 6 10-8 over 3,000,000 time steps in pushing.
add to the loss function the L1 norm of the first layers of
B ADDITIONAL RESULTS
B. TRAINING ENVIRONMENT
Fig. 2 in Section 4.1 shows correlations of unsupervised and supervised metrics with in-distribution reward for object reaching and pushing, only in the case without regularization. In Fig. & we also show these results in the with regularization, well when adjusting for MLP informativeness.
B.2 OUT-OF-DISTRIBUTIOO GENERALIZATION IN SIMULATION
In Section 4.2
we discussed rank-correlations OOD rewards with unsupervised, informativeness and generalization scores object reaching without regularization. In Fig, 9 we also show these results for the case with regularization and on pushing. as well as when adjusting for MLP informa- tiveness. Without regularization. as we observed on object reaching. confirming our conclusions on this more complex task. When using regularization, rank corrlationns very similar across both tasks. Interestingly, the correlation
observe pushing very similar correlations along all metrics
16

---

Published as a conference paper at ICLR 2022
object nuadiigg regariaicaaiinn
sbartb.
20
24
13
16
40
36
34
Ond reward in diarta
pushiing no www.aannntooooooaaaatatoooo
reward in cletto.
-1
-10
-10
regularizatio
-35
(ad) Mull roward in civeris
-18
pushing L1 regrrrizztion
neward in ciarro
31
31
22
28
16
12
-12
16
MEPM
EL9O
09-000
009900 20
1000111
090000 99
030000 00
020000
6500
(1)
GP 08-0000044
06 00000000-00 11
Sammoo P1
08-00000
GS 05:0000-1111 35
08-0000-411 30
02-00004111
G5-0003-11
05.0022-00 (1)
63.002--00
20
08.0002000
650002-000
00
030002200
(0
05-0002-053
(1)
00
196
090
08 050000
Figure 8: Rank correlations between metrics and in-distribution reward, with and without adjusting for informa-
tiveness. Correlations are color-coded as described in Fig.
object neuching. requlariostion
roward 0001
25
43
10
16
MiD wward Di0D1
40
21
18
28
18
21
37
38
44
25
2
-17
12
10
11
-3
-14
12
10
.25
reward
aee MLP) reward Oc02
10
10
10
41
45
25
46
46
pushing
DO rogultrizal on
reward 0001
15
ML.") saward O0D1
-10
13
13
10
10
-18
36
-18
125
12
-18
-20
-14
-27
-5
-15
-10
-6
15
-22
roward Moi Ten)
2000
object renching
regukattztion
iwwnrd OCTD1
11
1000 roward Me We
-10
-13
-7
2
-19
-15
10
-11
-5
-10
4
-15
-6
+15
-10
-14
40
-15
roward 2000
pustrry regufarical
Pnwwrd 00D1
12
Heannd OCD1
-15
-111
19
13
4
2
17
9
-11
16
11
18
reward 0c103
ORI
MLP
ML
EU80
68-0001
06-0000 m
08,0000
88 220110
0900001 11
06-0001
05-2001
05-20011 (7)
050002255
05-0000-mm11
05-0302-5i
68-0322
08-0302aan15
08:002/ann
09-0333-M111
56-0000--a
05 m0000
68 60000
68 0000 (7
MEG
00000
Figure 9; Rank correlations between metries and OOD reward, with and without adjusting for informativeness.
Correlations are color-coded deseribed in Fig.
between GS-OOD2 scores and OOD2 generalization of the policy is even stronger when using LI reg- ularization. In contrast to our observations without regularization, find that the correlation between GS-OOD1 and OOD1 generalization of the policy vanishes when adjusting for the MLP metric.
B.2.1 DISENTANGLED REPRESENTATIONS
As discussed in Section 4.2 for object reaching without regularization, we observe in Fig. D weak correlation between some disentanglement metrics and OOD1 reward, which however vanishes when adjusting for MLP informativeness. In agreement with Dittadi et al. (202cc), we observe no significant correlation between disentanglement and O0D2 generalization, for both tasks, with and without regularization. From Fig. I0 we see that in some cases, especially without regularization, very high DCI score seems to lead better performance average. However, this behavior is
17

---

Published as a conference paper at ICLR 2022
distribution
O0D1
OOD2
0.955
0.94
0.90
0.950
0.85
0.92
0.945
0.80
0.940
0.90
0.75
0.935
low
0.70
medium
0.88
0.930
0.99
reg.
reg.
L1 reg.
rog
reg.
reg.
in distribution
00D1
0OD2
0.85
0.7
0.7
0.6
0.6
0.75
0.5
0.70
0.3
0.3
0.65
medium
0,2
0.2
0.99
0.60
0.1
reg.
no reg.
L1 reg.
no reg.
reg.
no reg.
Figure 10: Fractional success on object reaching (top) and pushing (bottom). split according to low (blue). medium-higg high (orange), and almost perfect (green) disentanglement. Results for object reaching are also reported Fig. in Section 4.2
significant (within error bars), as opposed to the results shown simpler downstream tasks by Dittadi et al.) (202cc). Furthermore, this trend is likely due to representation informativeness, since the correlations with disentanglement disappear when adjusting for the MLP score, as discussed above.
B.2.2 REGULARIZATION
As seen in Fig. 10 regularization generally has a positive effect on OODI and OOD2 generalization, which is particularly prominent in the OODI setting. On the other hand, it leads to lower training rewards both object reaching and in pushing. In the latter, the performance drop is particularly significant, while in object reaching is negligible.
B.2.3 SAMPLE EFFICIENCY
In addition to the analysis reported in the main paper, we investigate how representation prop- erties affect sample efficiency. Specifically. we store checkpoints of our policies t (20k, 50k, 100k, 400k) for object reaching and / E 1200k,5004 3M3 for pushing. We then evaluate policies these time step on the same three environments as before: (1) on the cube colors from training; (2) the OODI cube colors; and (3) on the 0OD2-sim cube colors. Results are
summarized in Fig. II for object reaching and Fig. 12 for pushing.
On object reaching (Fig. Unsupervised metrics (ELBO and reconstruction loss) display a correlation with the training reward, do the supervised informativeness metrics (GBT and MLP). This is strongest early timesteps. meaning these scores could important for sample efficiency. Similarly, we observe a correlation with the disentanglement scores DCI, MIG and SAP. With the help of the additional evaluation of
we observe very similar trends with and without regularization:
18

---

Published as a conference paper at ICLR 2022
MOOTO Tup
05-0001
05-2000
CD
1rS 14, PereeNn D002
144 55, FINN Eibicn
48
FENAND D002
P46 55 Piree 29029
15
11
42 10-99-0 -2
Figure 11: Sample efficiency analysis for object reaching. Rank correlations of rewards with relevant metrics along multiple time steps. Correlations are color-coded as described in Fig. 2
CCI
ML
14 -133
14
-15-12
15 25
02
-3
10
<7 -13
13
12
116 00-45--1
-12
11
21
Figure 12: Sample efficiency analysis for pushing. Rank correlations of rewards with relevant metrics along multiple time steps. Correlations are color-coded as described in Fig. 2
rewards adjusted for MLP informativeness, we can attribute this correlation again to this common confounder. Crucially, we see that the generalization scores (GS) are correlated with generalization of the corresponding policies under OODI and O0D2 shifts for all recorded time steps, confirming the results in the main text.
On pushing (Fig. 12p. many correlations at early checkpoints are significantly reduced, especially with regularization. This behavior might be due to the more complicated nature of the task, which involves learning to reach the cube first, and then push it to the goal. Correlations are primarily seen towards the end of training, with similar spurious correlations with disentanglement as elaborated above. Importantly, correlations between generalization scores (GS) and policy generalization under the same distribution shifts remain strong and statistically significant, corroborating the analysis in the main text.
B.2.4 GENERALIZATION TO A NOVEL SHAPF
As mentioned in Section 4.2 on the object reaching task, we also test generalization w.r.t. novel object shape by replacing the cube with an unmovable sphere. This corresponds to a strong OOD2- type shift, since shape was never varied when training the representations. We then evaluate a subset of 960 trained policies as before, with the same color splits. Surprisingly. the policies appear to handle the novel shape as we see from the histograms in Fig. 13 terms of success and final distance. In fact, when the sphere has the same colors that the cube had during policy training, all policies get closer than 5 cm to the sphere on average, with a mean success metric of about 95%. On sphere colors from the OODI split, more than 98.5% move the finger closer than this threshold, and on the strongest distribution shift (OOD2-sim colors and cube replaced by sphere) almost 70% surpass that threshold with an average success metric above 80%.
19

---

Published as a conference paper at ICLR 2022
Training colors
00D1 colors
0002 colors
300
Count
200
200
100
0.80
0.85
0.90
0.95
0.6
0.7
0.8
0.9
0.5
0.6
0.7
0.8
09
Fractional
Fractional
Fractional success
Count
200
100
100
100
50
0.02
0.03
0.04
0.02
0.04
0.06
0.08
0.02
0.04
0.06
0.08
0.10
0.12
Distance
Distance
Distance
Figure 13: Testing policies for object reaching under same in-distribution, OOD1, and OOD2 evaluation protocols regarding object color in simulation, but replacing the cube with a sphere, which was never used in training.
time
Transferring policies for object reaching the real robot setup without any fine- fine-tuning green sphere (unseen shape and color). Correlations are color-coded described Fig.
Figure 14:
B.3 DEPLOYING POLICIES TO THE REAL WORLD
In Fig. 14 we show three representative episodes of testing a reaching policy on the real robot for the strong OOD shift with novel sphere object shape instead of the cube from training. We present the respective videos in the project page. There we also present videos of additional real-world episodes on pushing and reaching cubes of different colors.
20


Reviews:
[{"review_id": "y1RDKNO1eAp", "reviewer": "Reviewer_AwPM", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written and intuitive. Further studies of generalisation in robot learning are overall highly relevant and this study contributes to better understanding. That being said, it seems to be very close to prior existing work Dittadi et al 2021 (citation in paper) and some insights are close to common knowledge. \n\nThe evaluation of generalisation also remains limited. After discussing the 7 different factors of variation in the domain, the full investigation purely focuses on color which limits the meaningfulness of the evaluation. \n\nOn the positive side, it’s great to see the focus on investigating proxy metrics which could be evaluated as a replacement to full reinforcement learning and this part of the paper is highly relevant.\n\nIt’s interesting that a sparsity inducing regularisation on the first policy layer has no effect on the usefulness of disentangled representations. Maybe it would be useful here to evaluate with sparsity regularisation across the whole policy.\n\nMinor:\nThe initial section only cites papers starting 2019 regarding the problem of generalisation. This question has a much longer history in ML research which should be taken into account.\nTable 1 on page 2 includes acronyms which are explained many pages later. Please make sure that an acronym is explained before using it on its own.\nFigure 1 likely presents representative view points and not agent inputs. Additionally providing these would be more helpful in understanding the sim to real gap.\nWhy is there no OOD1 in real?\nPlease add an argument for the specific choice of the 2 types of VAEs. What is the benefit of adding Ada-GVAE in addition to beta-VAE?\nWhy not evaluate with beta=0 to include the regular autoencoder case?\nThe dimensionality of the latent space could have a big impact. I’d recommend at least a minor ablation with the best known parameters and sweeping over the size of the latent space.\nWhy are the correlations ‘milder’ on the more complex pushing task? It does not seem directly intuitive and it would be useful to add a hypothesis.\nSome related work on representation learning for reinforcement learning is missing [1,2,3]\n\n[1] Kulkarni, Tejas D. et al. “Unsupervised Learning of Object Keypoints for Perception and Control.” ArXiv abs/1906.11883 (2019): n. pag.\n[2] Wulfmeier, Markus, et al. \"Representation matters: Improving perception and exploration for robotics.\" 2021 IEEE International Conference on Robotics and Automation (ICRA). IEEE, 2021.\n[3] Higgins, Irina et al. “Towards a Definition of Disentangled Representations.” ArXiv abs/1812.02230 (2018): n. pag.\n", "summary_of_the_paper": "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\n\n", "summary_of_the_review": "Overall a relevant investigation into the effect of representations on generalisation in RL. While some insights are trivial and the only investigated factor of variation for generalisation is color, there are aspects in the paper worth distributing and discussing (in particular the evaluation of proxy metrics for downstream performance).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "OBk_B4w4ln2", "reviewer": "Reviewer_BEPU", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths \n- OoD generalization of RL agents is a very active area of research. The paper provides a systematic and thorough empirical study on the impact of pretrained representations.\n- The correlation analysis between different proxy metrics and the OoD generalization performance is extensive and sound.\n\nWeakness\n- While the key claim of the paper lies in OoD generalization, a large part of experiments are essentially not OoD. Among the four scenarios (IID, OoD1, OoD2, and real-world), only the OoD2 presents a clear domain shift to the pre-trained encoder. \n- It seems not very clear to which degree the *agents generalize OoD*. A more direct comparison between IID, OoD1 and OoD2 would perhaps make claim 1 about OoD generalization more convincing.  \n- The core of claim 2, i.e., the GS metric can be used for representation selection, looks over-stretched to me since:\n  - it requires labeled samples from the OoD domain, which is often not realistic\n  - it requires ground truth values of the underlying factors of variation, which is often not realistic neither\n- Finally, the experiment design overall looks quite incremental to [A Dittadi et al. ICLR'21].", "summary_of_the_paper": "The paper presents an extensive empirical evaluation of the VAE-based pretrained representations for OoD generalization of RL agents in a robotic setup.\n\nFrom the reported experiment results, it draws two conclusions:\n1) agents build on top of the pretrained representations can generalize to some challenging OoD scenarios; \n2) the prediction error of the ground truth values of the underlying factors of variations can be used as a proxy metric for selecting representation for OoD agents.", "summary_of_the_review": "While the empirical evaluation is extensive and rigorous, the two key claims made in the paper are debatable. I, therefore, recommend reject. I'd be happy to change my rating if the author can better support the two claims during rebuttal. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "KDVmCTucCLh", "reviewer": "Reviewer_vE69", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Primary confusion/clarity concern: I had trouble understanding the rank-correlation results.  I think the clarity issues begin at the following quote in Section 2: “As factor prediction mode...as GS-OOD1, GS-OOD2-sim, and GS-OOD2-real accordingly.”  After reading these 4 sentences many times, I still think they are very difficult to understand (primarily on my first and second readthrough, but even now that I think I mostly understand the approach, I still think these sentences are confusing).  Similarly, sentences like “A simple supervised metric to evaluate a representation is how well a small downstream model can predict the ground-truth factors of variation” and “we use the MLP10000 and GBT10000 metrics (simply MLP and GBT in the following), where MLPs and GBTs are trained to predict the FoVs from 10,000 samples” and “rank correlations between representation metrics and training reward” need to be explained better.  (If space is a factor, then pointing to an appendix for more detailed/intuitive explanations is a good strategy.)\n\nMinor problem: “Correlations with p<0.05 are color-coded.”  How are they color-coded?  This is probably obvious to a reader without any of the primary confusion above (I think it’s red, but even then, which shade of red?), but making it explicit would be better.\n\nMinor grammar edit: “As factor prediction model, we will likewise...” missing article adjective?\n\nStrengths:\n- Outstanding limitations discussion; this frank discussion of limitations strengthens the paper.\n- Thorough overall results with some fascinating insights.\n- The sim-to-real experiments were very nice.\n\nQuestions:\n- MLP stands for multi-layer perceptron, correct?  This should be defined.  GLB should be defined.\n- What are the differences between MLP vs MLP(1-7)?\n- Can you give any intuition about why MLP(1) might have a negative correlation in Figure 2?  Is this just noise?\n\n**Update**:  After reading the authors' responses and other reviews, most of my concerns have been addressed (assuming the two changes I suggest are made).  I am updating my score to an 8; I feel that this is a strong paper. Regarding other reviewers' concerns:\n\nI understand reviewer BEPU's concerns and agree with some of them; but do not think they are deal-breakers.  The paper has its limitations, but I think it is extremely strong despite these.  Robotics papers do tend to have more limitations than papers based exclusively on simulated/game environments, since real-world robotics experiments are much more time-consuming to perform (but they are also often much more convincing, since they are much closer to real-world problems of interests, and not games or simplistic simulations).\n\nI respectfully disagree with many of Reviewer YvcV's concerns, many of which amount to \"the authors should have done [some other nice thing]\".  I think we can always ask authors to do more, and think the many of the reviewer's suggestions would be nice to do, but I think the contribution is sufficient in this case, and that rejecting because the authors did not do even more is not the correct response.  This paper's results are impressive and extensive despite the limitations that Reviewer YvcV points out.  I think part of the disconnect may be non-robotics RL researchers forgetting how time-consuming and worthwhile real-world robotics experiments are (compared to simulated or game experiments).", "summary_of_the_paper": "The authors study the problem of generalizing out of distribution for RL.  They conduct an extensive empirical study, using simulation and real robotics, and observed properties that predict how well an agent will generalize.\n", "summary_of_the_review": "This empirical paper has some clarity issues, but is exceptionally strong otherwise.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Ck2HfaHgJf", "reviewer": "Reviewer_YvcV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n* The paper studies the effect of pretraining visual representation on the downstream out-of-distribution task performance.\n* The paper conducted a relatively thorough analysis on this issue. More specifically, the paper studies the effect of VAE pretraining on the TriFinger environment. The studied factors include the measuring metric, network regularization, input noise, different OOD distributions.\n* While most of the conclusions are intuitive to understand, it's good to see that the paper has many experiments to support them.\n\nWeaknesses:\n\n* While the paper did a thorough analysis on VAE and in the TriFinger environment, I found these two are limiting. First, there are many different pretraining techniques such as contrastive learning and using auxilliary loss. This paper only studies VAE. Therefore, the paper title is a bit over-claimed. Second, TriFinger is a challenging control task but not necessarily a difficult task in terms of perception. As we can see in the example environment figures (Figure 1), the image input to the policy is relatively simple and has a clean background. Therefore, it is possible that the perception network can easily learn to only focus on the object in the image and ignore other things during training. And such a network can be more robust than expected and can transfer to the real world better. Therefore, the conclusions drawn in this paper only apply to VAE pre-training and TriFinger environments. It remains unclear whether the conclusions still hold in visually more complex environments such as ProcGen environments.\n\n* Even just in the TriFinger environments, a more systematic analysis should include the testings on the TriFinger environments with different backgrounds. As we can see in Figure 1, the bowl environments are very clean and look similar in simulation and in the real world. Since the pretraining is mainly about training the perception module, OOD tests should also change the bowl texture, color during the testing.\n", "summary_of_the_paper": "This paper is an analysis paper that uses well-known existing pretraining techniques to study the effect of pretrained representations on out-of-distribution generalization performance. The paper mainly uses the VAEs to train the visual representation network. While I believe this is an important topic, the paper only analyzed one pretraining technique and in a visually simple environment.", "summary_of_the_review": "I believe the paper is analyzing an important issue in reinforcement learning. However, more experiments on TriFinger with different visual backgrounds and other visually complex environments such as Procgen, and experiments with different pretraining techniques are needed to make the analysis more thorough and convincing.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: It can be prohibitively expensive to train a reinforcement learner from scratch &mdash; particularly in cases where experience is expensive to obtain, such as with a physical robot. So, we might hope to speed up RL in a couple of ways: first, by pre-training a representation that makes subsequent RL need less data; and second by running our RL on a cheaper proxy environment such as a simulator. For pre-training, we hope to be able to take advantage of available pre-collected data, and we hope to be able to use supervised learning or reconstruction tasks since they can be cheaper than RL. For either pre-training or a proxy environment, we have to deal with distribution shifts: the properties of the environment may change between pre-training and RL, and between RL and testing the learned policy.

The paper presents an empirical study of how different pre-trained representations and different distribution shifts affect RL performance. It evaluates a number of representations trained by different VAEs (differing in aspects such as loss and hyperparameter settings) under various scenarios of distribution shift. It also asks whether we can predict the performance of the learned policies from properties of the representations, before going to the expense of training and evaluating our reinforcement learner.

The paper concludes that it is possible to significantly reduce RL data requirements using pre-trained representations, even in the presence of significant distribution shifts &mdash; including demonstrating zero-shot sim2real transfer. And, the paper concludes that inexpensive measurements of OOD performance on supervised tasks can at least partially predict success in generalization.

The reviewers praised the extensive experimental evaluation, including a large number of experiments on a physical robot, as well as the investigation of less-expensive ways to predict generalization.

Some reviewers were concerned that the choice of environments was limiting &mdash; e.g., that the distributional distance between in-distribution and out-of-distribution tests was limited, or that the results might not generalize to other related robotic environments. However, in the end there was support for the conclusion that the experiments cover a sufficiently general and interesting question.
Decision: Accept (Poster)