RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability
Authors: Chuning Zhu
Max Simchowitz
Siri Gadipudi
Abhishek Gupta
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Model-Based Reinforcement Learning
Deep Reinforcement Learning

Abstract: Visual model-based RL methods typically encode image observations into low-dimensional representations in a manner that does not eliminate redundant information. This leaves them susceptible to spurious variations -- changes in task-irrelevant components such as background distractors or lighting conditions. In this paper, we propose a visual model-based RL method that learns a latent representation resilient to such spurious variations. Our training objective encourages the representation to be maximally predictive of dynamics and reward, while constraining the information flow from the observation to the latent representation. We demonstrate that this objective significantly bolsters the resilience of visual model-based RL methods to visual distractors, allowing them to operate in dynamic environments. We then show that while the learned encoder is able to operate in dynamic environments, it is not invariant under significant distribution shift. To address this, we propose a simple reward-free alignment procedure that enables test time adaptation of the encoder. This allows for quick adaptation to widely differing environments without having to relearn the dynamics and policy. Our effort is a step towards making model-based RL a practical and useful tool for dynamic, diverse domains and we show its effectiveness in simulation tasks with significant spurious variations.

RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability
Chuning Zhu University of Washington Seattle, WA 98105 zchuning@es. hattiingges.wwahiintton... edu
Max Simchowitz
Massachusetts Institute of Technology
Boston, MA 02139
msimchow@mit edu
Siri Gadipudi University Washington Seattle, WA 98105 sg06@uv.edu
Abhishek Gupta University of Washington Seattle, WA 98105 abhgupta@cs washington, edu
Abstract
Visual model-based RL methods typically encode image observations into low- dimensional representations in manner that does not eliminate redundant in- formation. This leaves them susceptible to spurious variations changes task-irrelevant components such as background distractors or lighting conditions. In this paper, we propose visual model-based RL method that learns latent representation resilient such spurious variations. Our training objective encour- ages the representation be maximally predictive of dynamics and reward, while constraining the information flow from the observation to the latent representation. We demonstrate that this objective significantly bolsters the resilience of visual model- based RL methods to visual distractors, allowing them to operate in dynamic environments. We then show that while the learned encoder is resilient to spirious variations, it not invariant under significant distribution shift. To address this, we propose simple reward-fres alignment procedure that enables test time adaptation of the encoder. This allows for quick adaptation widely differing environments without having to relearn the dynamics and policy. Our effort a step towards mak- ing model-based RL practical and useful tool for dynamic, diverse domains. We show its effectiveness in simulation benchmarks with significant spurious variations well real-world egocentric navigation task with noisy TVs in the background. Videos and code: https: ///ccunninnggg. gi thub. io/repo- httrr/epoobbitee/../ website/.
1 Introduction
Consider the difference between training a single robot arm against a plain background with rein- forcement learning (RL), and learning operate the same amidst of plentiful dynamic distractors uncontrollable elements such as changing lighting and disturbances in the scene. The latter must contend with spurious variations differences in environments which are irrelevant for the task but potentially confusing for a vision-based RI. agent resilience to which indispensable for truly versatile embodied agents deployed in real world settings.
Standard end-to-end techniques for visual RL struggle in the presence of spurious variations [64, 48), in part because they fail to discard fask-irrelevant elements. To improve generalization [38, 59], self-supervised representation learning methods [23, 39, 55, 54, 17, 31] pre-train visual encoders that compress visual observations. These methods aim for lossless compression of how image observations evolve in time (e.g. by minimizing reconstruction error). Unaware the demands of downstream
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

tasks, these methods also cannot determine which elements of an environment can be discarded. As such, they often struggle in dynamic and diverse scenes [64, 48, 17] ones where significant portions of the observations are both unpredictable and irrelevant despite being remarkably successful in static domains.
This paper proposes Resilient Model-Based RL by Regularizing Posteior Predictability (RePo) - an algorithm for learning lossy la- tent representations resilient to spurious variations. A representation satisfactory if it (a) predicts its own dynamics and (b) accurately predicts the reward. To satisfy these criteria, RePo jointly learns (i) visual encoder mapping high-dimg nigh-dimensional observations to inter- mediate image "encodings" (ii) latent encoder which compresses histories of intermediate image encodings into compressed latent rep- resentations (iii) dynamics model in the latent representation space, and (iv) a reward predictor to most accurately predict current and future rewards. What distinguishes us from past work [63, 12, 17] a new desideratum of predictability; that, conditioned on past latents and actions, future latent dynamics should look determin- istic as possible. This is because an agent should try to maximize its control over task-relevant parts of the state, whilst neglecting aspects of the environment that it cannot influence [20, 60]. RePo optimizes novel loss which encourages predictability, thereby discarding a broad range of spurious variations in aspects of the environment which out of the agents control (e.g. changes in background, lighting, or visual traffic in the background). At the same time, by penalizing reward prediction error, we capture the task-relevant aspects of the dynamics necessary for learning performant policies.
Figure 1: Reinforcement lear- ing in environments with spurious variations including dynamic elements like humans, changes in lighting and training across a range of visual appearances.
RePo implements deceptively simple modification to recurrent state-space models for model-based RL [17, 59, 46J. We maximize mutual information (MI) between the current representation and all future rewards, while minimizing the mutual information between the representation and observation. Instead of minimizing image reconstruction error, optimize variational lower bound on the MI-objective which tractably enforces that the learned observation encoder, latent dynamics and reward predictors are highly informative of reward, while ensuring latents are predictable as possible (in the sense described above). We demonstrate that the representations, and the policies built thereupon, learned through RePo succeed in environments with significant amounts of dynamic and uncontrollable distractors, as well as across domains with significant amounts of variability and complexity. Through ablations, also validate the necessity of our careful algorithm design and optimization decisions.
While these learned representations enable more effective reinforcement learning in dynamic, complex environments, the visual encoders (point (i) above) mapping from observations into intermediate encodings suffer from distribution shift in new environments with novel visual features (c.g. new background not seen at train time.) We propose a simple test-time adaptation scheme which uses (mostly) unlabeled test-time data to adapt the visual encoders only, whilst keeping all other aspects of the RePo model fixed. Because RePo ensures resilience of the compressed latent representation at training time, modifying only the test-time visual encoders to match training time representations allows representations to recover optimal performance with only minor amounts of adaptation.
Concretely, the key contributions of this work are: (1) We propose simple representation learning algorithm RePo for learning representations that are informative of rewards, while being as predictable possible. This allows model-based RL to scale to dynamic, cluttered environments, avoiding reconstruction. (2) We show that while the learned encoders may be susceptible to distribution shift, they are amenable simple test-time adaptation scheme that can allow for quick adaptation in new environments. (3) We demonstrate the efficacy of RePo on a number of simulation and real-world domains with dynamic and diverse environments.
Related Work
Our work is related to a number of techniques for visual model-based reinforcement learning, but differs in crucial elements that allow it to scale to dynamic environments with spurious variations.

---

Model-Based RL. Though model- based RL began with low-dimensional, compact state spaces [26, 37, 27, 57], advances in visual model-based reinforcement learning [17, 19, 18, 44, 42, 21] learn latent representations and dynamies models from high dimensional visual feedback (typically via recurrent state-space models). Perhaps most relevant to RePo DREAMER [17]. Section 4 explains the salient differences between DREAMER and RePo; notably, we eschew reconstruction loss in pursuit of resilience to spurious variations. A closely related work is TD-MPC [22]. which learns a task-oriented latent representation by predicting the value function. However, its representation may not discard irrelevant information and necessarily contains information about the policy.
Representation Learning for Control. There is a plethora of techniques for pretraining visual representations using unsupervised learning objectives [38, 34, 30, 32, 41, 49, 47, 13]. While these can be effective on certain domains, they do not take downstream tasks into account. Task-relevant representation learning for RL. uses the reward function to guide representation learning, typically in pursuit of value-equivalence (e.g. via bisimulation) [63, 8, 62, 12, 50, 22]. However, these approaches do little to explicitly counteract spurious variations. Our work aligns with a line of work that disentangles task-relevant and task-irrelevant components of the MDP. [7, 6] obtain provable guarantees for representation learning with exogeneous distractors parts of the state space whose dynamics is independent of the agent's s actions. [56] introduces more granular decomposition of the MDP across the task relevance and controllability axes. Our work, in contrast, does not impose a specific form on the spurious variations.
Domain Adaptation. Unsupervised domain adaptation adapts representations across visually dif- ferent source and target domains [66, 58, 45, 11, 25]. These techniques predominantly adapt visual encoders by minimizing a distribution measure across source and training distributions, such as MMI) [5, 33, 53, 24], KL divergence [67, 35] or
divergence [11, In [61], distribution matching was extended to sequential decision making. While domain adaptation settings typically assume that the source and target share underlying marginal or joint distribution in latent space, this assumption does not hold in online RL because the data is being collected incrementally through exploration, and hence the marginals may not match. Hence, our test-time adaptation technique, as outlined in Section 4.1, introduces a novel support matching objective that enforces the test distribution to be in support of the train distribution, without trying to make the distributions identical.
3 Preliminaries
MDPs. A (discounted) MDP .M- = (S, P. Po. consists of a state-space S, action space A. discount factor, y E (0, 1). transition and P(..)) : S A + (SS), initial state distribu- tion Po E ASS), and reward function r(.-)) S x A + [J (assumed deterministic for sim- plicity). A policy π : :S + ((A)
mapping from states to distributions over actions. We let EA denote expectations under So ~ Po. ag ~ Tst),, and 8t41 P(st, at); the value is (x) = EM EEE >"rss,, at) 8p s]. and VM = Po [VM(())] The goal is to learn a policy that maximizes the sum of expected returns Ey [E ( a:) | 80 s], as in most RL problems, but we do so based on a belief state as explained below.
Visual RL and Representations. For our purposes, we take states 8t to be visual observations 8, = O2 E O; for simplicity, we avoid explicitly describing POMDP formulation this can be subsumed either by introducing a belief-state [68], or by assuming that images (or sequences thereof, e.g. to estimate velocities) are sufficient to determine rewards and transitions [36]. The states o, may be digh-dimensional, so we learn encoders h O + X' to an encoding space X. We compress these encodings x, further into latent states 2t, described at length in our method in Section
By spurious variation, we informally mean the presence of features of the states s, which are irrelevant our task, but which do vary across trajectories. These can take the form of explicit distractors- either static objects (e.g. background wall-paper) or dynamic processes (e.g. video coming from a television) that do not affect the part of the state space involved in our task [7, 6]. Spurious variation can also encompass processes which are not easy to disentangle with the state: for example, lighting conditions will affect all observations, and hence will affect the appearance of transition dynamics.
Spurious variation.
Consider the following canonical example: MDP with state space S x S2, where for (s(1), (22) S x Sz, the reward r(s, a) is function F(s(1), a) only of the projection onto S². Moreover, suppose that (())))()))))))))))))))))))) E I s,a,, where st ~ P(s, a), is distribution P(s(1), a) again only depending on (1)), Then, the states (22) can be viewed spuriously various. For example, if
3

---

Representation learning
Pollcy learning with prior
marK Toes
Prior Ngltee.1Ry, (R.)))
MLE
Anwoard Madea (eloaalBes)
Dy, S
4
Absteriar Rr,kees
Inference with posterior
PC
Eirader Me (Pees)
Latents
Data
P
ALL
Objertisss
Op-t
9111
Figure 2: RePo learns a latent representation resilient to spurious variations by predicting the dynamics and the reward while constraining the information flow from images.
g(1) is a Lagrangian state and s(2) is static background, then it clear that transitions of Lagrangian state and reward do not depend on s(2), Our template also encompasses dynamic distractors; e-g. television show in the background has its own dynamics, and these also do not affect reward or physical dynamics. Even varying lighting conditions can be encompassed in this framework: the shadows in a scene or brightness of the environment should not affect reward or physics, even though these visual features themselves evolve dynamically in response to actions and changes in state. That is, there are examples of spurious variation where (1) (e.g. Lagrangian state) affect (2)) (e.g. certain visual features), but not the other way round. In all cases, "spurious" implies that states (s) (2) J120.
and their possible variations due to different environments, have no bearing on optimal actions.
4 RePo: Parsimonious Representation Learning without Reconstruction
We propose a simple technique for learning task-relevant representations that encourages parsimony by removing all information that is neither pertinent the reward nor the dynamics. Such representations discard information about spurious variations, while retaining the information actually needed for decision making.
To describe our method formally, we introduce some notation (which is also shown in Fig 2). Let O be the space of image observations, X' the space of encoded observations, where f O + X represents the encoding function from images observations to encoded observations, and Z the space of latent representations. Note that r++1 is simply the instantancous encoding of the image o++1 as It-1 = (0(++1), but the latent representation 2t-1 at time step t + 1 is an aggregation of the current encoding I++1 and previous latent 2 and action a₃ Let 3 post denote the space of "posteriors" on latent dynamics: 2 of the form p( ( E 22, a, I+++), where 2,, 2 E 2, a, E A, Lt+1 E X, and where and 20 Pn has some initial distribution pn- In words, the latent posterior use past latent state and action, in addition to current encoding to determine current latent. Control policies and learned dynamics models act this latent representation 2,+1, and not simply the image encoding x++1 so as to incorporate historical information.
Let Douf denote the distribution over experienced actions, observations and rewards from the environment ((11T, 01 T Dbut). For p E Posst let Ep,s denote expectation of (aI.T, 01:7, 71:r) Dbuts Jt h(o;) and the latents 21+1 p(- 24, ae. 2++1) drawn from the latent posterior, with the initial latent žy ~ Pu- Our starting proposal is to optimize the latent posterior p and image encoder h such that information between the latent representation and future reward is maximized, while bottlenecking [1] the information between the latent and the observation:
max Lp.h (o.n(21:17 (Z1:1; F1. d1:T) s.t. Ip.A (21:7 01:2 (1:T) < e.
(4.1)
n.h
Above, (21T) (1:r) denotes mutual information between latents and rewards conditioned actions under the Ep, distribution, and distribution Lp.h (21:3) 01.7 (1:r) measures information

---

between latents and observations under Ep. as well. Thus, (4.1) aims to preserve large mutual information with rewards whilst minimizing information stored from observations.
Optimizing mutual information is intractable in general, so we propose two variational relaxations of both objects (proven in Appendix B)
(auaZrITTT (LTT) 2 Ep,h
- log (rr 2i] Er=o DKI(p( zer ue T441) ge(. ze, at))]
(4.2)
(2.0( 01r (1:r) < Ep,
(4.3)
where 4r and lx are variational families representing beliefs over rewards r; and latent representations Et+1, respectively. We refer to zt+1 p(- 2t, (t, (t+)) as the latent posterior, because it conditions on the latest encoded observation .++1 = h((041). We call the variational approximation % 2e, 4z)
the latent prior because it does not use the current observation o++11 (or it's encoding xx+1) to determine 2-+-1- Note that the right hand side of Eq. (4.3) depends on h through 1t41 = h() and thus gradients of this expression incorporate gradients through h.
The magic of Eq. (4.3). The upper bound (4.3) reveals a striking feature which is the core of our method: that, in order to reduce extraneous information in the latents 2, about observations Ot. it is enough to match the latent posterior 2t+1 ~ p(- 2,, at, ft+11) to our latent prior %x(- 2t, az) that does not condition on current T++1- Elements that are spurious variations can be captured by p( 2t, at, It-1), ) but not by 9z(- Et, at), since dz is not informed by the latest observation encoding TT-1, and spurious variations are not predictable. To match the latent posterior and the latent prior, the latent representation must omit these spurious variations. For example, in an environment with a TV in the background, removing the TV images reduces next-step stochasticity of the environment. Thus, (4.3) encourages representations to omit television images.
The relaxed bottleneck. The above discussion may make it seem as if we suffer in the presence of task-relevant stochasticity. However, by replacing the terms in Eq. (4.1) with their relaxations in Eqs. (4.2) and (4.3), we only omit the stochasticity that is not useful for reward-predictte these substitutions, and move to a penalty-formulatin like dual-gradiee descent [2]. The resulting objective we optimize to learn the latent posterior p. latent prior lz. reward predictor dr and observation encoder jointly is:
We make amenable to constrained optimization methods
max min Ep.. 94
E log gr (re. =+] + (En,A
E/ DKI.(p(- | zz, Or, Te+1) |l (x(- Ze; an))] -o). (4.4)
Implementation details. We parameterize p and q using recurrent state-space model (RSSM)
[17]. The RSSM consists of an encoder ha(a, oc), a latent dynamics model go [=(++11 1 22, az)
corresponding to the prior, a representation model pe pe(2(+1 24, ar, It+1) corresponding to the posterior, and reward predictor (r: zt). We optimize (4.4) using dual gradient descent. In addition, we use the KL balancing technique introduced in Dreamer V2 [19] to balance the learning of the prior and the posterior. Concretely, we compute the KL divergence in Eq. (4.4) as DKL (p q) = dDkL((P) || 9) (1 DKL (p I| [g]), where |-] denotes the stop gradient operator and a E [o, 1] is the balancing parameter. With the removal of reconstruction, the KL balancing parameters becomes especially important as shown by our ablation in Sec. 5.
Policy learning As is common in the literature on model-based reinforcement learning [19, 17, 18), our training procedure alternates between (1) Representation Learning: learning a representation by solving the optimization problem outlined in Eq. (4.4) to infer a latent posterior p( | 2,, at T2+1).
a latent prior qu (2141 | 24, az), an encoder x; h(o,)) and reward predictor qe (r 22). and (2) Policy Learning: using the inferred representation, dynamics model and reward predictor to learn a policy T (a, 2 for control. With the latent representation and dynamics model, we perform actor-critic policy learning [16, 10] by rolling out trajectories in the latent space. The critic Vz(2) is trained to predict the discounted cumulative reward given a latent state, and the actor (a 2) is trained to take the action that maximizes the critic's 's prediction. While policy learning is carried out entirely using the latent prior as the dynamics model, during policy execution (referred to as inference in Fig. 2), we infer the posterior distribution p(=++1 2t, (x+1) over latent representations from the current observation, and use this to condition the policy acting in the world. We refer readers to Appendix C for further details.
5

---

Training Distribution
Reweighting
(Keratn)
I ('eraned
P(arraia)
Support
Match reweighted distribution == support constraint
Testing Distribution
constraint
p(
Test-time adaptation
Figure 3: Depiction of test time adaptation scheme for latent alignment via support constraints. During exploration. the marginal distributions may not match perfectly, so we match the supports of the latent features instead, using reweighted distribution constraint.
Comparison to DREAMER, DEEPMDP, and BISIMULATION. DREAMER [17] was first derived to optimize pixel-reconstrrucion, leading to high-fidelity dynamics but susceptibility to spurious varia- tions. Naively removing pixel reconstruction from dreamer, however, leads poor performance [17]. Our objective can be interpreted as modifying DREAMER so as to maintain sufficiently accurate dynamics, but without the fragility of pixel-reconstruction. DEEPMDP [12] the latents 2 to exactly the image encodings Ie = h(oo). It learns dynamics P X x A ( such that the distribution Æt+1 P(h(o)). ae) is close to Z(+1 h(o--1), 0t+1 ~ P" (oz, at). where P denotes ground-truth transition dynamics; this enforces consistency of dynamics under encoding. The above distributions are viewed as conditional on past observation and action, and as a result, highly non-parsimonious representations such as the identity are valid under this objective. BISIMULAT ATINOUATION [63] learns an optimal representation the sense that perfect bisimulation metric does not dis- card any relevant information about an MDP. However, there is no guarantee that it will disregard irrlevant information. Indeed, the identity mapping induces a trivial bisimulation metric. Hence, BISIMULATION compress only by reducing the dimensionality of the latent space. In contrast, further compress the encodings z; into latents 2N so as to enforce the latent prior gz ae. =t) is close to the latest observation-depeedenn posterior distribution p(- Ra, az It+1). As mentioned in Eq. (4.3), this ensures information compression and invalidates degenerate representations such as the identity mapping.
4.1 Transferring Invariant Latent Representations via Test-Time Adaptation
While resilient spurious variations seen during training, our learned latents 2; and hence the policies which depend on them may not generalize to new environment which exhibit systematic distribution shift, e.g. lighting changes or background changes. The main source of degradation is that encoder h O + X may observe images that it has not seen at train time; thus the latent, which depend observations through x, h(oo), may behave erratically, even when system dynamics remain unchanged.
Relying the resilience of posteriors p latents 2; introduced by RePo, propose test-time adaption strategy to only adjust the encoder to the new environment, whilst leaving fixed. A natural approach is to apply unsupervised domain adaptation methods [66, 58) to adapt the visual encoder h to Deest These domain adaptation techniques typically operate in supervised learning settings, and impose distributional constraints between source and target domains [61, 25), where the distributions of training and test data are stationary and assumed to be the same in some feature space. A distribution matching constraint would be:
min D(Parain (Pestt) s.t. Prest heest Drest, Purain h Darain- htea()))
(4.5)
In Eq. (4.5), we consider matching the distributions over encodings x of observations o. Specifically, we assume D Dtrain and Drest denote training and test-buffer distributions over observations o, Perain hurain o Divain denotes the distribution of x (orain()) where Dursin is encoded by the train-time

---

Realistic Maniskill
Distracted DeepMind Control
Lazy TurtleBot
Figure 4: Depiction of the environments being used for evaluation. (Left):: the Distracted DeepMind Control suite [64], (Top Right): Maniskil12 [15]] environments with realistic backgrounds from Matterport [3]. (Bottom Right): Turtleoot environment with two TVs playing random videos in the background.
encoder htrain: and Puet heest Dtest denotes encodings under test-time encoder hest (+) over which we optimize. Here, D).-)) denotes an f-divergence, such as the x² -i-vieergence.
Support Constraint. (4.5) fails capture that the encoded distributions at train and test time differ at the start of our adaption phase: suboptimal encoder performance at the start of the adaptation phase causes the policy to visit sub- sub-optimal regions of state space not impossible to match the distribution as in standard unsupervised domain adaptation. We therefore propose to replace (4.5) with support constraint, enforcing that the distribution of htest Drest is We consider the following idealized objective:
train time. Thus, it may be
contained in the support of herain o D Dtrain-
min hitee (-)
D(T Perain Ptest) s.t. EgwPeran [T(x)]
(4.6)
Here, by T- Ptrain. we mean the we-weighted density of Ptrain = htrain o Derain by a function ( and 7(-) > 0 ensures this reweighted distribution is also valid probability distribution. The reweighting operation T Perain seems intractable at first, but we show that if take D)-.)) (- to be the x" divergence, then Eq. (4.6) admits the following tractable Lagrangian formulation (we refer readers to [65] and Appendix B for a thorough derivation)
The constraints Ep,,ai [r(x)]
min max (T [(x)] - -Erca [(()) -(,)) + X(EEPeuin, - 1), (4.7) +() 2 0, (1) 16)A
where above, E RR X IR and the objective depends on heest through the definition Ptest /teest о Drest. This objective is now tractable saddle point optimization, which can be solved with standard stochastic optimization techniques. The optimization alternates between optimizing reweighting and the visual encoder Atstt and the dual variables J, A. Throughout adaptation, we freeze all other parts of the recurrent state space model and only optimize the encoder. We provide more intuition for the support constraint in Appendix E.
Calibration. We note that naively reweighting by 7(-) can cause degenerate encodings that collapse into one point. To prevent this, we regularize the support constraint by also ensuring that some set of paired "calibration" states across training and testing domains share the same encoding. We collect paired trajectories in the training and testing domains using actions generated by an exploration policy, and minimize the & loss between the training and testing encoding of each pair of observations. We defer the details of the complete optimization to Appendix C.
5 Experimental Evaluation
We conduct empirical experiments to answer the following research questions: (1) Does RePo enable learning in dynamic, distracted environments with spurious variations? (2) Do representations learned by RePo quickly adapt to new environments with test time adaptation? (3) Does RePo help learning in static, but diverse and cluttered environments?
7

---

1.00
le3
Walker Stand
1e2
Walker Walk
162
Walker Run
Eplodd Reiumm
0.75
0.50
0.25
2
0.00
1o2
Cheetah Run
1c5
102
Hopper Stand
165
102
Cartpole Swingup le5
Epicade Retum
7.5-
5.0
2.5
0.0
Environment Steps
105
Environment Steps
105
Environment Steps
le5
RoPa (Ours)
TIA
TD-MPC
DBC
Denaised MDP
DeepMDP
RePo Convergence
Figure 5: Results on distracted DeepMind control environments. These environments have spurious variations. and RePo is able to successfully learn in all of them, both faster and achieving higher asymptotic returns than prior representation learning methods.
Evaluation domains We evaluate our method primarily in three different settings. (1) Distracted DeepMind Control Suite [64, 63] is variant of DeepMind Control Suite where the static back- ground is replaced with natural videos (Fig. 4). For adaptation experiments, we train agents on static undistracted backgrounds and adapt them distracted variants. (2) Realistic Maniskill is benchmark we constructed based on the Maniskil12 benchmark [15], but with realistic backgrounds from [3] simulate learning in diverse range of human homes. We solve three tasks LiftCube, PushCube, and TurnFaucet in variety of background settings. (3) Lazy TurtleBot is a real-worke robotic setup where a TurtleBot has to reach some goal location from egocentric observations in furnished room. However, there are two TVs playing random videos to distract the "lazy" robot. We provide more details about evaluation domains in Appendix D.
Baselines We compare our method with a number of techniques that explicitly learn representations and use them for learning control policies. (1) Dreamer [17] is state-of-the-at visual model-based RL method that learns a latent representation by reconstructing images. (2) TIA [9] renders Dreamer more robust to visual distractors by using a separate dynamics model to capture the iask-irrelevant components in the environment. (3) Denoised MDP [56] further learns a factorized latent dynamics model that disentangles controllability and reward relevance.
(4) TD-MPC [22] trains a latent dynamics model to predict the value function and uses a hybrid planning method to extract a policy. (5) DDP [12] is model-free method that learns representation by predicting dynamics and reward, and then performs actor-critic policy learning on the learned representation.
(6) Deep Bisimulation for Control DBC [63] is model-free algorithm which encodes images into a latent space that preserves the bisimulation metric.
We also compare with a number of techniques for test-time adaptation of these representations. (I) calibrated distribution matching, a variant of the method proposed in Section 4. using a distribu- tion matching constraint rather than support matching one, (2) uncalibrated support matching, a variant of the method proposed in Section 4.1, using a support matching constraint but without using paired examples, (3) uncalibrated distribution matching, variant of the method proposed in Section 1, using a distribution matching constraint, but without using paired examples, (4) invariance through latent alignment ILA [61], ], technique for test-time adaptation of representations with distribution matching and enforcing consistency in latent dynamics, (5) calibration, a bascline that only matches the encodings of paired examples.
Does RePo learn behaviors in environments with spurious variations? We evaluate our method's ability to ignore spurious variations on a suite of simulated benchmark environments with dynamic visual backgrounds (Fig. 4); these are challenging because uncontrollable elements of the environment visually dominate a significant portion of the scene. Fig. 5 shows our method outperforms the baselines across six Distracted Deep Control environments, both in terms of learning speed and asymptotic performance. This implies that our method successfully learns latent representations resilient to spurious variations. Dreamer [17] attempts to reconstruct the dynamic visual distractors which is challenging in these domains. TIA [9] and Denoised MDP [56] occasional success when
8

---

le3
Walker Stand
62
W Walker Walk
1e2
Cheetah Run
1.0
0.8
Epitade
0.6
Environment
164
Environment Steps
1c4
Calibrate Support (Ours)
Support
Source
Calibrated Distribuzion
bustiibution
Calbration
Figure 6: Results on adaptation from static environments dynamic environments in Deepmind control. RePo with calibrated support constraints outperforms ablations and previous techniques for domain adaptation.
Push Cubse
Lift Cube
Turn Faucet
Epippde Retuur
20
5.0
0.0
Environment Stops
Environment Steps
Environment Stops
le5
RePo (Ours)
TIA
TD-MPC
DBC
Denoised
DeepMDP
Figure 7: Results of training agents varying static environments in Maniskill [15]. RePo is able to learn quickly and efficiently than alternatives even in static domains.
they dissociate the task-relevant and irrelevant components, but they suffer from high variance and optimization failures. TD-MPC [22] is affected by spurious variations as its representations are not minimal. The model-free baselines DeepMDP [12] and DBC [63] exhibit lower sample efficiency on the more complex domains despite performing well on simpler ones.
To further validate RePo's ability to handle spurious varia- tions in the real world, we evaluate its performance on Lazy TurtleBot, where a mobile robot has navigate around a furnished room to reach the goal from egocentric observa- tions (Fig. 4). To introduce spurious variations, we place two TVs playing random Youtube videos along the critical paths to the goal. As shown in Table. I, RePo is able reach the goal with nontrivial success within 15K environ- ment steps, whereas Dreamer fails to reach the goal. We provide details about the setup in Appendix. D.
Table 1: Results on Lazy TurtbleBot at 15K environment steps. RePo achieves nontrivial success whereas Dreamer fails to reach the goal.
Success
Return
RePo (Ours)
62.5%
-24.3
Dreamer [17]
0.0%
-61.7
Do representations learned by RePo transfer under distribution shift? We evaluate the effective- ness of the test-time adaptation method described in Section 4.1 on three DeepMind Control domains: Walker Stand, Walker Walk, and Cheetah Run. We train the representation environments with static backgrounds, and adapt the representation to domains with natural video distractors (as shown Fig. 4). For methods that use calibration between the source and target environments, we collect 10 trajectories of paired observations. Results are shown in Fig. 6. RePo shows the ability to adapt quickly across all three domains, nearly recovering the full training performance within 50k steps. Performance degrades if we replace the support constraint with a distribution matching objective, as it is infcasible to match distributions with the test-time distribution having insufficient exploration. We also observe that by removing the calibration examples, both support constraint and distribution perform worse as the distributions tend to collapse. We found the addition of dynamics consistency in ILA to be ineffective. Nor is calibaration alone sufficient for adaptation.
Does RePo learn across diverse environments with varying visual features? While the previous two sections studied learning and adaptation in dynamic environments with uncontrollable elements, we also evaluate RePo on it's ability to learn in diverse range of environments, each with a realistic and cluttered static background. Being able to learn more effectively in these domains suggests that RePo focuses it's representation capacity on the important elements of the task across environments, rather than trying to reconstruct the entire background for every environment.

---

Dreamer
Reconstruction
3.-
Observation
-2.
-50
-3.
RaPo Probe
Figure 10: Top two principle components of RePo and Dreamer's latent representa- tions across different backgrounds. RePo's latent representation more compact than Dreamer's s, which enables data sharing.
Figure 9: Probing representations learned by RePo shows that the background is largely ignored, while [17] tries recon- struct the cost of the agent prediction.
We test on three robotic manipulation tasks LiftCube, PushCube, and TurnFaucet with realistic backgrounds depicted in Fig. 4. As shown in Fig. 7, our method achieves saturating performance across all three tasks. Dreamer [17] spends its representation capacity memorizing backgrounds and unable to reach optimal task performance. TIA [9] suffers from high variance and occasionally fails to dissociate task-relevant from irreeerell features. Denoised MDP [56], TD-MPC [22], and DBC [63] learn to ignore the background two of the tasks but generally lag behind RePo in terms of sample efficiency. DeepMDP [12] fails to learn meaningful behavior in any task.
Visualizing representations learned by RePo To decipher our representation learning objective, we probe the learned rep- resentations by post-hoc training separate image decoder to
Walker Walk
Episode Return
reconstruct image observations from the latents. We visualize the results in Fig. 9 and compare them with Dreamer recon- structions [17]. Our representation contains little information about background but is capable of reconstructing the agent, implying that contains only task-relevant information.
4-
2.
2
In addition to probing, states of RePo and Dreamer by visualizing their top two prin- cipal components. We collect the same trajectory across all backgrounds Maniskill and visualize the final recurrent latent state inferred by RePo and Dreamer respectively. As shown in Fig. 10, RePo produces more compact latent representations than Dreamer, meaning the latent states encode less information about background variations. This enables RePo to share data across different backgrounds, explaining its superior sample efficiency compared to baselines.
qualitatively compare the latent
Environment Steps
1e5
RePo (Ours)
KL Balancing
No Info Bottleneck
Figure Ablating objectives show- ing the importance of information bot- tleneck and KL balancing described in Section 4.
Ablation experiments We conduct ablation experiments to determine the effect of hyperparameters in Fig. 8. As can see, the performance is erucially dependent the information bottleneck e, as well as KL balancing. We refer readers to Appendix for more thorough discussion.
Discussion
This work presents RePo, a technique for learning parsimonious representations that are resilient to spurious variations. Our representation is effective on learning in dynamic, distracted environments. And while the representation is subject to degradation under distribution shift, it can be quickly adapted to new domains by a semi-supervised test-time adaptation procedure. A limitation of our method that the learned dynamics model is no longer task-agnostic, as only captures task- relevant information. This can be potentially addressed by simultancously predicting multiple reward objectives. Our framework opens up several interesting directions for future research, such as: can multi-tase tulk variant of RePo allow for representations applicable to some distribution of tasks? Can we apply our algorithm in a continual learning setup? We believe our method holds promise in these more general settings, especially for real robots deployed into dynamic, human centrie environments.
10

---

Acknowledgments and Disclosure of Funding
We would like to thank Marius Memmel, Max Balsells, and many other members of the WHIRD Lab at University of Washington for valuable feedback and discussions.
References
[1] A. A. Alemi, I. Fischer, V. Dillon, and K. Murphy. Deep variational information bottleneck. In International Conference Learning Representations, 2017.
[2] S. Boyd and L. Vandenberghe. Convex Optimization. Cambridge University Press, 2004.
[3] A. Chang, A. Dai, T. Funkhouser, M. Halber, M. Niessner, M. Savva, S. Song, A. Zeng. and Zhang. Matterport3D: Learning from RGB-D data in indoor environments. International Conference 3D Vision (3DV), 2017.
[4] B. Dai, N. He, Y. Pan, B. Boots, and L. Song. Learning from Conditional Distributions via Dual Embeddings.
Singh and Zhu, editors, Proceedings of the 20th International Conference on Artificial Intelligence and Statisties, volume 54 of Proceedings of Machine Learning Research, pages 1458-1467 PMLR, 20-22 Apr 2017.
[5] Dorii and A. Ghodsi. Adapting component analysis. In M. Zaki, A. Siebes, X. Yu, B. Goethals, G. I. Webb, and X. Wu, editors, 12th IEEE International Conference on Data Mining, ICDM 2012, Brussels, Belgium, December 10- 13, 2012, pages 846-851, IEEE Computer Society, 2012.
[6] Y. Efroni, D. J. Foster, D. Misra, A. Krishnamurthy, and J. Langford. Sample-efficient reinforce- ment learning in the presence of exogenous information. P.-L. Loh and M. Raginsky, editors. Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 5062-5127. PMLR, 02-05 Jul 2022.
[7] Y. Efroni, D. Misra, A. Krishnamurthy, A. Agarwal, and J. Langford. Provably filtering exogenous distractors using multistep inverse dynamics. In International Conference on Learning Representations, 2022.
[8] N. Ferns and D. Precup. Bisimulation metrics are optimal value functions. In N. Zhang and J. Tian, editors, Proceedings of the Thirtieth Conference on Uncertainty in Artificial Intelligence, UAI 2014, Quebec City, Quebec, Canada, July 23-27, 2014, pages 210-219, AUAI Press, 2014.
[9] X. Fu, G. Yang, P. Agrawal, and T. Jaakkola. Learning task informed abstractions. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 3480-3491. PMLR, 18-24 Jul 2021.
[10] S. Fujimoto, H. van Hoof, and D. Meger. Addressing function approximation error in actor-critic methods. In J. G. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässsan Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 1582-1591. PMLR, 2018
[11] Y. Ganin, E. Ustinova, H. Ajakan, P. Germain, H. Larochelle, F. Laviolette, M. Marchand, and V. S. Lempitsky. Domain-ade aaversiadversarialll training of neural networks. CoRR, 0bbs//505.078/ 2015.
[12] C. Gelada, S. Kumar, J. Buckman, O. Nachum, and M. G. Bellemare. DeepMDP: Learning con- tinuous latent space models for representation learning. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages
PMLR, 09-15 Jun 2019.
[13] D. Ghosh, A. Gupta, and S. Levine. Learning actionable representations with goal conditioned policies. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
[14] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Bengio. Generative adversarial nets. In Advances neural information processing systems, pages 2672-2680, 2014.

---

[15] J. Gu, F. Xiang, X. Li, Z. Ling, X. Liu. T. Mu. Y. Tang, S. Tao, X. Wei, Y. Yao, X. Yuan, P. Xie, Z. Huang, R. Chen, and H. Su. Maniskil12: A unified benchmark for generalizable manipulation skills. In International Conference
Learning Representations, 2023.
[16] T. Haarmoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off- maximum entropy deep reinforcement learning with a stochastic actor. arXiv preprint arXiv: 1100.001290 2018.
[17] D. Hafner, T. P. Lillicrap, J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. In Sth International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[18] D. Hafner, T. P. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, Calisornia, USA, volume 97 of Proceedings of Machine Learning Research, pages 2555-2565. PMLR, 2019.
[19] D. Hafner, T. P. Lillicrap, M. Norouzi, and Ba. Mastering atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenRevi 2021.
[20] K. Haninger, R. Vicente-Garcia, and J. Krüger. Towards learning controllable representations of physical systems. CoRR, abs/201 1.09906, 2020.
[21] N. Hansen, Y. Lin, H. Su, X. Wang, V. Kumar, and A. Rajeswaran. Modem: Accelerating visual model-based reinforcement learning with demonstrations. CoRR, abs/2212.05699 2022.
[22] N. Hansen, X. Wang, and H. Su. Temporal difference learning for model predictive control. In International Conference on Machine Learning. PMLR, 2022.
[23] K. He, X. Chen, S. Xie, Li, P. Dollár, and R. B. Girshick. Masked autoencoders are scalable vision learners. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 15979-15988. IEEE, 2022.
[24] Hoffman. E. Tzeng, T. Darrell, and K. Saenko. Simultancous deep transfer across domains and tasks. In G. Csurka, editor, Domain Adapiation in Computer Vision Applications, Advances in Computer Vision and Pattern Recognition, pages 173-187, Springer, 2017.
[25] J. Hoffman, E. Tzeng, T. Park. J. Zhu, P. Isola, K. Saenko, A. A. Efros, and T. Darrell. Cycada: Cycle-consistent adversarial domain adaptation. CoRR,
2017.
[26] M. Janner, J. Fu, M. Zhang, and S. Levine. When to trust your model: Model-based policy optimization. In H. M. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Aché-Buu E. B. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 12498-12509 2019.
[27] M. Janner, I. Mordatch, and S. Levine. Gamma-models Generative temporal difference learning for infinite-horizon prediction. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[28) W. Kay, J. Carreira, K. Simonyan, B. Zhang, C. Hillier, S. Vijayanarasimhan, F. Viola, T. Green, T. Back, A. Natsev, M. Suleyman, and A. Zisserman. The kinetics human action video dataset. ArXiv, abs/1 705. 705.06950, 2017.
[29] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2015.
[30] M. Laskin, A. Srinivas, and P. Abbeel. CURL: contrastive unsupervised representations for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 1 119 of Proceedings of Machine Learning Research, pages 5639-56500 PMLR, 2020.
12

---

[31] A. X. Lee, A. Nagabandi, P. Abbeel, and S. Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[32] T. Lesort, N. D. Rodríguez, J. Goudou, and D. Filliat. State representation learning for control: An overview. CoRR, abs/1802.0418)) 2018.
[33] M. Long, J. Wang, G. Ding, J. Sun, and P. S. Yu. Transfer joint matching for unsupervised domain adaptation. In 2014 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2014, Columbus, OHI, USA, June 23-28, 2014, pages 1410-1417. IEEE Computer Society, 2014.
[34] A. Majumdar, K. Yadav, S. Arnaud, Y. J. Ma, C. Chen, S. Silwal, A. Jain, V. Berges, P. Abbeel, J. Malik, D. Batra, Y. Lin, O. Maksymets, A. Rajeswaran, and Meier. Where are we in the search for artificial visual cortex for embodied intelligence? CoRR, 11w/2303.182400 2023.
[35] Z. Meng. J. Li, Y. Gong, and B. Juang. Adversarial teacher-student learning for unsupervised domain adaptation. In 2018 IEEE International Conference on Acousties, Speech and Signal Processing, ICASSP 2018, Calgary, AB, Canada, April 15-20, 2018, pages 5949-5953. IEEE, 2018.
[36] V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. A. Ried- miller. Playing atari with deep reinforcement learning. CoRR, abs/1312.5602, 2013.
[37] A. Nagabandi, G. Kahn, R. S. Fearing, and S. Levine. Neural network dynamics for model- based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE International Conference on Robotics and Automation, ICRA 2018, Brisbane, Australia, May 2 21-25, 2018, pages 7559-7566. IEEE, 2018
[38] S. Nair, A. Rajeswaran, V. Kumar, C. Finn, and A. Gupta. R3M: A universal visual represen- tation for robot manipulation. In K. Liu, D. Kulic, and J. Ichnowski, editors, Conference on Robot Learning, CoRL 2022, 14-18 December 2022, Auckland, New Zealand, volume 205 of Proceedings of Machine Learning Research, pages 992-909.100 PMLR, 2022.
[39] M. Oquab, T. Darcet, T. Moutakanni, H. Vo, M. Szafraniec, V. Khalidov, P. Fernandez, D. Haziza, F. Massa, A. El-Nouby, M. Assran, N. Ballas, W. Galuba, R. Howes, P. Huang. S. Li, I. Misra, M. G. Rabbat, V. Sharma, G. Synnaeve, H. Xu, H. Jégou, Mairal, P. Labatut, A. Joulin, and P. Bojanowski. Dinov2: Learning robust visual features without supervision. CoRR, abs/2304.07193, 2023.
[40] X. B. Peng, A. Kanazawa, S. Toyer, P. Abbeel, and S. Levine. Variational discriminator bottleneck: Improving imitation learning, inverse RI., and GANs by constraining information flow. In International Conference on Learning Representations, 2019.
[41] I. Radosavovic, T. Xiao, S. James, P. Abbecl, J. Malik, and T. Darrell. Real-world robot learning with masked visual pre-training. In K. Liu, D. Kulic, and J. Ichnowski, editors, Conference on Robot Learning, CoRL 2022, 14-18 December 2022, Auckland, New Zealand, volume 205 of Proceedings of Machine Learning Research, pages 416-426. PMLR, 2022.
[42] R. Rafailov, T. Yu, A. Rajeswaran, and C. Finn. Offline reinforcement learning from images with latent space models. In A. Jadbabaie, J. Lygeros, G. J. Pappas, A. Parrilo, B. Recht, C. J. Tomlin, and M. N. Zeilinger, editors, Proceedings of the 3rd Annual Conference on Learning for Dynamics and Control, L4DC 2021, 7-8 June 2021, Virtual Event, Switzerland, volume 144 of Proceedings of Machine Learning Research, pages 1554-166.. PMLR, 2021.
[43] K. Ramakrishnan, A. Gokaslan, E. Wijmans, O. Maksymets, A. Clegg, J. M. Turner, E. Un- dersander, W. Galuba, A. Westbury, A. X. Chang, M. Savva, Y. Zhao, and D. Batra. Habitat- matterport 3d dataset (HM3d): 1000 large-scale 3d environments for embodied AL. In Thirty-ffith Conference Neural Information Processing Systems Datasets and Benchmarks Track, 2021.
13

---

[44] O. Rybkin, C. Zhu, A. Nagabandi, K. Daniliidis, I. Mordatch, and S. Levine. Model-based reinforcement learning via latent-space collocation. M. Meila and T. Zhang, editors. Proceed- ings of the 38th International Conference on Machine Learning, ICML. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 9190-9201. PMLR, 2021.
[45] S. Sankaranarayanan, Y. Balaji, A. Jain, S. Lim, and R. Chellappa. Learning from synthetic data: Addressing domain shift for semantic segmentation. In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pages 3752-3761. Computer Vision Foundation IEER Computer Society, 2018.
[46] R. Sekar, O. Rybkin, K. Danillidis, P. Abbeel, D. Hafner, and D. Pathak. Planning to explore via self-supervised world models. In Proceedings of the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 8583 8592. PMLR, 2020.
[47] Sermanet, C. Lynch, Y. Chebotar, J. Hsu, E. Jang, S. Schaal, and S. Levine. Time-contrastive networks: Self-supervised learning from video. In 2018 IEEE International Conference on Robotics and Automation, ICRA 2018, Brishane, Australia, May 2 21- 2018, pages 1134-1141. IEEE, 2018.
[48] A. Stone, O. Ramirez, K. Konolige, and R. Jonschkowski. The distracting control suite A challenging benchmark for reinforcement learning from pixels. CoRR. abs/2101.00722) 2021.
[49] A. Stooke, K. Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference Machine Learning, ICML 2021, 18-24 July 2021. Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 9870 9879. PMLR, 2021.
[50] J. Subramanian, A. Sinha, R. Seraj, and A. Mahajan. Approximate information state for approximate planning and reinforcement learning in partially observed systems. J. Mach. Learn. Res., 23:12:1-12:88 2022.
[51] Y. Tassa, Y. Doron, A. Muldal, T. Erez, Y. Li, D. de Las Casas, D. Budden, A. Abdolmaleki, Merel, A. Lefrancq, T. Lillicrap, and M. Riedmiller. Deepmind control suite, 2018.
[52] E. Tzeng, J. Hoffman, K. Saenko, and T. Darrell. Adversarial discriminative domain adaptation (workshop extended abstract). In 5th International Conference on Learning Representations, ICLR 2017, Toulon, Frunce, April 24-26, 2017, Workshop Track Proceedings. OpenReview.net, 2017.
[53] E. Tzeng, Hoffman, N. Zhang. K. Saenko, and T. Darrell. Deep domain confusion: Maximiz- ing for domain invariance. CoRR, abs/1412.3474,, 2014.
[54] A. van den Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. CORR, abs/1 807880077777 007433 2018.
[55] A. van den Oord, O. Vinyals, and K. Kavukcuoglu. Neural discrete representation learning. In
I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus. S. V. N. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30: Anmual Conference on Neural Information Processing Systems 2017, December 4.9, 2017, Long Beach, CA, USA, pages 6306 6315,315,550 2017.
[56] T. Wang, S. Du, A. Torralba, P. Isola, A. Zhang, and Y. Tian. Denoised mdps: Learning world models better than the world itself. In International Conference on Machine Learning. PMLR, 2022.
[57] G. Williams, A. Aldrich, and E. A. Theodorou. Model predictive path integral control: From theory to parallel computation. Journal of Guidance, Control, and Dynamics, 40(2):344-357, 2017.
[58] G. Wilson and D. Cook. A survey unsupervised deep domain adaptation. ACM Trans. Intell. Syst. Technol., 116):51:1-51:46, 2020.
14

---

[59] P. Wu, A. Escontrela, D. Hafner, P. Abbeel, and K. Goldberg. Daydreamer: World models for physical robot learning. In K. Liu, D. Kulic, and Ichnowski, editors, Conference on Robot Learning, CORI. 2022, 14-18 December 2022, Auckland, New Zealand, volume 205 of Proceedings of Machine Learning Research, pages 2226 22240220. PMIR, 2022.
[60] M. Yang, D. Schuurmans,
P. Abbeel, and O. Nachum. Dichotomy of control: Separating what you control from what you cannot. CoRR, abs/2210.134355 2022.
[61] T. Yoneda, G. Yang, M. R. Walter, and B. Stadie. Invariance through latent alignment, 2021.
[62] A. Zhang, C. Lyle, S. Sodhani, A. Filos, M. Kwiatkowska, Pineau, Y. Gal, and Precup. Invariant causal prediction for block mdps. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 11214 11224. PMLR, 2020.
[63] A. Zhang, R. T. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant repre- sentations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
[64] A. Zhang, Y. Wu, and J. Pineau. Natural environment benchmarks for reinforcement learning. CoRR, 1tss/8811066022 2018.
[65] R. Zhang, B. Dai, L. Li, and D. Schuurmans. Gendice: Generalized offline estimation of stationary values. In 8th International Conference Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[66] Y. Zhang. A survey of unsupervised domain adaptation for visual recognition. CoRR, abs/2112.06 10./2112.06745, 2021.
[67] Zhuang. X. Cheng, P. Luo, Pan, and Q. He. Supervised representation learning with double encoding-layer autoencoder for transfer learning. ACM Trans. Intell. Syst. Technol., 9()::16:111:17, 2018.
[68] L. M. Zintgraf, K. Shiarlis, M. Igl. S. Schulze, Y. Gal, K. Hofmann, and S. Whiteson. Varibad: A very good method for bayes-adaptive deep RL via meta-learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
15

---

Supplementary Materials for
"RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability"
A Algorithm Pseudocode
Algorithm I Resilient Model-Based RL by Regularizing Posterior Predictability (RePo)
1: Initialize dataset Dout with S random seed episodes.
2: Initialize neural network parameters 0, Φ, v randomly and set dual variable 3 = Bu-
3: while not converged do
4: for update step c= == 1... do
5:
// Dynamics Learning
6:
Draw B data sequences {(or, 4t, rr)))== ~ Dbuf-
7:
Encode images F1.7 (o(0))))
8:
Infer prior and posterior distributions (0(2; | 22-1, at at-1), po 04-1,00-1,30 Sample latent states 2, ~ p=(2; 2r-1, 0c-1, and infer reward distributions gr (rx zz).
9:
10:
Compute the Lagrangian objective in (4.4) Update 0 and using dual gradient descent. // Behavior Learning
12:
13:
Imagine trajectories {(z,,,,) " 1++ from each zt- Predict rewards r, qa (rr 27), values v, (w 2r), actions a, Talar 2T).
15:
Update á and v using actor-critic learning.
// Data Collection
17:
for time step t = 1 N do
1...
Compute 2 po(2t; 2t-1, a:-1, ho():)) from history and current observation. Compute ~ Ta(a 2 from policy.
19:
20:
Execute a, environment and collect or+1, 0c+1-
21:
Add experience to dataset Duuf + Duuf U {or, at, Ti Ftft=1
1.N
Algorithm 2 Supervsue Adaptation of Visual Encoder Using Support Constraint
1: Fix training encoder ha, model ps. policy n. and replay buffer Dtrain-
2: Initialize test-time replay buffer Dusst with S random seed episodes.
3: Initialize test-time encoder hgg by setting 0' 0.
4: Initialize neural networks Tp and fw randomly and set dual variable Ag-
5: Collect N calibration trajectories with paired observations Deal <- {(o oft at rt) using exploration policy.
6: while not converged do
7: for update step 1... C do
8:
// Adaptation
9:
Draw B observations from training and test buffers 1.B D Drrains putt D Drest-
10:
Encode images using respective encoders 1::B (o gles ha (O Compute the Lagrangian objective in (4.7)
11:
Draw pairs of observations from the calibration buffer (0, uaai d (ealt 1 ~ Deal- ,, train = ho (o eai_triin) cal test hg (ogal Encode images using encoders 11:B
12:
13:
14:
Compute MSE loss between calibration encodings and add to the Lagrangian objective. Update 8', p, w and À using dual gradient descent.
15:
16:
// Data Collection
17:
for time step / 1...N N do
18:
Compute 2, (= 2t-1, 0:-1, lag (o.)) from history and current observation. Compute 4; ~ T (4: =x)) from policy.
19:
Execute a, environment and collect oc+1, ac+1.
21:
Add experience to dataset Drest + Dreet L (or, az, rz J.
16

---

B Derivations
B.1 RePo objective
Recall that Duuf denote the distribution over experienced actions, observations and rewards from the environment ((01:T, 01:T; F1:T) but). For p 3 post (our class of posteriors), let pp, denote expectation of (01:7, 01:T, (1:r)) ~ Duuf= It h(oz) and the latents 2++1 p(- 2t, Qg, 2441) drawn from the latent posterior, with the initial latent 2g Po-
We derive variational lower bound for the following information bottleneck objective from eq. (4.1).
max I.. (1:T) s.t. 1p,((Z1:;:1177 O1:T (I:T) <E p.h
where the mutual information is with respect to the distribution described by the expectation Ep, above. The goal is to learn a latent representation that maximally predicts the dynamics and reward while sharing minimal information with the observations. Notice that, as defined, is only latent dynamics model. Thus, for the derivation we will introduce p as corresponding the distribution of all variables under Ep which in particular takes into account the distribution from Duuf-
Definition B.1. We let P denote the joint distribution of (21:T, 01:7; 21:T; 01:T) under Ep.h-
Lower bound reward prediction. We first derive variational lower bound for the objective Ip,((2i2 Fi:r (1:r). For our derivation, we make a simplifying assumption:
Assumption B.1. Assume that there is function (r 24) such that p(r= 21:0, 01:t) = pr zx).
The simplifying assumption enforces that 2, is sufficient for the reward given past latents, actions, and the current action ae. We remark that Assumption B.1 is not strictly necessary, and the following lower bound we derive is still valid if it fails; it may just be very loose. The derivation begins as follows:
(ZITiiii a1r)
Epph Mogg((rin 1irr, (1r)) + H Hr,t(rir (1:))
+
Ep.h Dlog p(ri: 21:r,11:))
T
= Epph
logp,(r 22, at))
(Assumption B.1)
=
Eph
Σ log (r: 2())
+
(n. 2x) ll gr(r, 2))
Ep.h
Σ
log (r. 22))
where above Hp.h denotes an entropy term under distribution Ep and denotes equality up to constant that does not depend on our choice of parametrization p. which only dictates latents 2,- In the fourth line, add and subtract log q. (re(ez) in each term of the summation. The last step the nonnegativity of KL. divergence.
Upper bound dynamic compression. We proceed to derive upper bound for the constraint (2,77 01.7 (1.)). We make a simplifying assumption analogous to Assumption B. again, this assumption can also discarded, but illustrates that our lower bound is sharper when this assumption holds.
Assumption B.2. We assume that there a function Pz(z++1 Ze, at) such that ((ze+1 21:e; (1:t) = (2(411 2z, a.).
17

---

With Assumption B.2, we then invoke variational upper bound which replaces Pa P2 with variational approximation 9 (2 a.):
Ze, (2)
(21:T: 01. (1:T) Ourr, (1.r) - log p(21:1 01:r)]
p. [log p(21:1
=
-
p
log p( Ze, 0e, (t+1)
log ((2):) (1:r)
(definition of latent dynamics)
p
log pggp( 24, (4,
la 2es az)
(Assumption B.2)
T
Epp
log p(2441 2z, ae Ft+1)
log d (zt11 (1)
(RL(Pe(2441 (2.aa) ae(2241 2.,a)))
T
T log Ou (2(21 (2++1
< Egp
Σ
log p(z+11 2t, at, 1141)
24, as)
1=1
where in the last two lines we add and subtract 4s (Zt41 24, ae) and apply the nonnegativity of KL. divergence. In the third line,
the fact that transition dynamics model is Markov, with
p(2+11 01:t, 01:t; 21:+) p(z+11 Mo++1), (t, =x) = p(=++1 Pe41,0(;21)
(B.1)
where under Ep,h. 3y = h(oz) holds by definition. Notice that the same identity would be true instead tried to lower bound the mutual information (p.A(= 01.T (1:T), showing that objective is agnostic to whether we consider compressing information with respect to 21.7 or 01:7. This may be seem suprrising, but is simple consequence of choosing a latent dynamics model which depends on 01 only through T1.2.
Combining the derivations. at the objective in eq. (4.4)
Combining these results and writing in Lagrangian form, arrive
min Ep.t P.
max
ET= log (re =1] + (Ep.A Eno DKL(p(- ze, 04 ((41)) |l aa(- ze: (e)) -e)
B.2 Tractable objective for support constraint
We derive a tractable variant for the support matching objective in eq. (4.6): Pteat) s.t. ErwPeran [T()] 1.
min he (-)
D(T Pirain
where Ptrain denotes the distribution of image encodings training time and Pust denotes the distribution of image encodings at test time. is a reweighting function and the constraint EgrPen [r(x)] ensures that 7- Perain is a valid distribution. Let Do any f-divergence, have
min 7())20,Pears))
DØ(T Purain || Prest)
Purain (x)) T(x)
min +(-)20,kemms)
(+) (2)
(definition of f-divergence)
Prese (z)
Perain (2r) T T(x) [(x) - +()))))(())]
min 7())20
( Presa
max f(-)
(by convex conjugacy)
Prese (x)
Ptrnin (xx)
min 7()20.hka)))
max
[ Prest f())
f(x) ~0'((()))
Prese (x)
(by interchangeability principle [4])
min 7()20.
max PP (T(x) f()
[(x)) E.p.o s ( (f()))
We follow [4] and use the x2 divergence where 6(x) = (x - 1)² and dz()) = y + x. This gives
min
maxEp,.. 0. f(-)
max Ep.... (r(x) f(x)] - Erren f(x) [(1)3
+(-)
18

---

Incorporating the constraint as a Lagrange multiplier, we arrive at the objective in eq. (4.7).
min max Pr [T(x) f(x)] - Epi.s [(x)+ + 1 [ + A(Ep [r(2) - 1).
+(-) 2 0.Rewes)) (.) f(-JA
C Implementation Details
Model architecture We base our implementation of RePo on Dreamer [17]. We fix the image size to 64 x 64 and parameterize the image encoder using a 4-layer CNN with f32, 64, 128, 256} channels, kernel size 4, stride 2, and ReLU activation. This results in an embedding size of 1024. The recurrent state space model is parametrized by a GRU operating on deterministic beliefs. Given the current belief, state (sampled from either prior or posterior), and action, the GRU recurrently predicts the next belief. We parameterize the prior 9a (Rt+1 2t, at) as Gaussian whose mean and standard deviation are predicted by passing the next belief through 2-layer MLP. Similarly, we parametrize the posterior p(z++1 2t, Qt, T++1) as a Gaussian whose mean and standard deviation are predicted by passing the next belief along with the next image embedding through 2-layer MLP. We set the belief size to 200 and the state size to 30. As in Dreamer, we use both deterministic belief and stochastic state inputs to the prediction heads, including the reward model, policy, and value function. The reward model and value function are 4-layer MLPs. The policy is 4- MLP outputting a squashed Gaussian distribution. We use 200 hidden units per layer and ELU activation for all MIPss
Table 2: Hyperparameters for evaluation tasks
30
r
Walker
le-5
3
5
Cheetah
le-5
5
Hopper
le-4
3
Cartpole
le-4
4
Maniskill
le-4
4
TurtleBot
le-5
3
5
Optimization We train the RL agent in an online setting, performing 100 training steps for every 500 environment steps (except for TurtleBot which trains every 100 environment steps). In each training step. we sample 50 trajectories of length 50 from the replay buffer and optimize the RePo objective using the prior and posterior inferred from these trajectories. We then perform behavior learning by using the posterior states as initial latent states and rolling out the policy for 15 steps in the dynamics model. To balance bias and variance, we train the value function predict the generalized value estimate [17] with À = 0.95 and discount factor y = 0.99. We optimize all components with the Adam [29] optimizer. The image encoder, recurrent state space model, and reward model share the same learning rate of 3e-4. The policy and value function use a learning rate of 8e-5. Specific to our method, we initialize the Lagrange multiplier as 3 3g and set its learning rate to le-4. We cast the KL balancing parameter to the ratio r between the number of prior training steps to posterior training steps, where r translates to r_ and KL balancing ratio r on our evaluation tasks and report the best hyperparameters in Table 2.
r+1 We tune the initial Lagrange multiplier ,9. target KL e,
Baselines We implement Dreamer [17] and TIA [9] on top of the RePo codebase and tune their hyperparameters on the evaluation tasks. For DBC and DeepMDP, we the official implementation along with their reported hyperparameters. We use determinisitic transition model for DBC as we find it to perform better than using stochastic transition model.
Adaptation We parametrize the reweighting function 7 as 2-layer MLP with 256 hidden units in each layer and ReLU activation. The dual function f analogous to the discriminator in GAN architecture [14]. To prevent vanishing gradient, we parameterize using variational discriminator bottleneck [40] with hidden layers of size 256 and bottleneck layer of size 64. Prior to training, we collect 10 calibration trajectories with corresponding observations from both the source and the target domains, using an expert policy as an approximation for an exploration policy. We initialize the test- test-time encoder with the weights of the training encoder and adapt it online, performing 100
19

---

adaptation steps for every 500 environment steps taken. In each adaptation step, we sample 2500 observations from the offline source buffer and online target buffer respectively, and optimize the tractable support constraint objective (4.7). In addition, we sample 2500) paired observations from the calibration data and minimize the f2 loss between their encodings. All components are optimized using Adam [29] optimizer. We use learning rate of 3e-4 for heest 5e-5 for le-4 for f, and 5e-3 for the Lagrange multiplier À initialized to le-4.
Adaptation baselines We reuse the same model architecture for all variants of our proposed method. For baselines involving a distribution matching objective, we optimize the standard GAN [14] objective to minimize the Janson- Shamnon divergence between training and test-time encoding distributions. To enforce dynamics consistency for the ILA [61] baseline, we forward the image encodings through the RSSM to get the latent states and compute dynamics violation as the KI. divergence between the prior and the posterior. We minimize dynamics violation by backpropagating the gradient through the model to the encodings.
D Environment Details
Distracted DeepMind Control To evaluate the ability of RePo to learn in dynamic environments, we use the distracted DeepMind Control Suite proposed in [63]. Specifically, we replace the static background in standard DeepMind Control environments [51] with grayscale videos from the Kinectics-400 Dataset [28]. We use time limit of 1000 and an action repeat of 2 for all environments. We evaluate all methods with 4 random seeds.
Realistic Maniskill We evaluate the ability of RePo to learn across diverse scenarios on 3 manipu- lation tasks adapted from the ManiSkill 2 benchmark [15]:
Push Cube: the goal is to push a cube reach a position on the floor.
Lift Cube: the goal is lift cube above certain height.
Turn Faucet: the goal is to turn the faucet to a certain angle.
To simulate real-world scenarios, we replace the default background with realistic scenes from the Habitat Matterport dataset [43]. We curate 90 different scenes and randomly load a new scene the beginning of each episode. We use time limit of 100 and action repeat of 1. We evaluate all methods with random seeds.
To evaluate RePo's 's resilience against spurious variations in the real world, we furnish a room to mimic a typical household setup and train a TurtleBot reach certain goal location from egocentric observations (Fig. 4). We introduce spurious variations by placing two TVs playing random YouTubbe videos along the critical paths to the goal. The robot features a discrete action space with 4 actions: move forward, rotate left, rotate right, and no-op. To minimize algorithmic change, we convert the action space to a continuous 2D action space where the discrete actions correspond to (1, 0), (0, 1), 0), (0, -1) respectively, and continuous actions are mapped back to the closest discrete actions in L2 distance. The reward the negative distance to goal, where the robot's 's state is estimated using overhead camera mounted on the ceiling. We apply time limit of 100 and an action repeat of 1. At the start of training, we prefill the replay buffer with 5K offline transitions collected by sampling and reaching random goals using 10K steps of pure offline training. We then train each method online for 10K environment steps, using the state-based controller for automatic resets. Due to time constraints, evaluate each method for 2 seeds.
Lazy TurtleBot
based controller. This is followed by
E
Additional Experiments
E.1
Results on standard DMC environments
In Fig. 11, we provide additional comparisons to Dreamer on standard DeepMind Control tasks. We note that RePo is able to match Dreamer in asymptotic performance despite not being trained with a reconstruction objective. Our method slightly lags behind Dreamer in terms of sample efficiency, which is due to reward signals being inherently more sparse than pixel reconstruction signals.
20

---

163
Walker Stand
162
Walker Walk
Walker Run
0.75
0.25
Cheetah Run
1:2
Hopper Stand
105
Cartpola Swingup le5
7.5
7.5
5.0
5.0
2.5
2.5
0.91
Envronment Stops
Environmen: Steps
Envirorment Staps
RePo (Ours)
Dreamer
Figure 11: Comparison with Dreamer standard DeepMind control environments. RePo matches Dreamer in asymptotic performance despite not being trained with a reconstruction objective.
E.2 Results Distracted Walker Walk with scoreboard
We evaluate RePo variant of Distracted Walker Walk which consists of scoreboard displaying the cumulative reward of the current episode. The goal is to investigate if RePo's seward-peedictive nature causes its latent representation to collapse onto predicting the visualization in Fig. 12, which shows that the learned latent reconstructs the joint positions of the agent and ignores the score. This is because RePo predicts only the current reward but also the dynamics and future rewards. The empirical performance with the scoreboard is 872.86 = 43.71, while without the scoreboard is 868. 72 53.93.
We present the probing
Observation
363.96
383.80
383.00
308.06
RePo Probe
Figure 12:
Visualization of RePo's representation on distracted DMC with scoreboard. Since RePo predicts not only the current reward but also the dynamics and future rewards, its latent representation does not collapse to predicting the score.
E.3 Additional ablations
We perform additional ablation experiments distracted DMC Walker Walk analyze the effect of hyperparameters. All other hyperparameters are fixed to those reported in Table. 2. Fig. 13 illustrates the effect of (3 learning rate. When the learning rate is too high, the constraint is enforced immediately after training begins, which can lead to ineffective exploration. When the learning rate low, dual gradient descent converges slowly. We find the sweet spot to be a learning rate that relaxes the information bottleneck at the beginning of training and gradually tightens the bottleneck discard redundant information in the representation. Fig. 14 show the result of ablating the KL constraint value c. With too small KL target, the representation becomes degenerate and fails to capture task-relevant information, as indicated by the increase in reward loss. large KL target, on the other hand, can result in the dynamics model being inaccurate, thereby affecting behavior learning. Finally, Fig. 15 explores the effect of KL balancing. accordance with [19], we find that the dynamics model is generally harder to train than the representation model. Hence, larger KI. balancing parameter tends to result better performance. However, we train the prior too aggressively, then the posterior becomes poorly regularized and captures iorkiirelevant information.
21

---

KI Divergence
Reward Loss
Retur
9.19
9.15
2-
Envirommnnt Steps
Environement Staps
Eniiionmmnt
Figure 13: Ablating 3 learning rate, A larger 3 learning rate corresponds to faster convergence of the constraint, which can lead to ineffective exploration.
KL Divergence
Reward Loss
Retum
6-
0.2
Envitoment Steps
Environment Steps
Enviiooment
Figure 14: Ablating information bottleneck e. A tighter information bottleneck generally induces more parsimonious representations and more accurate dynamics models. But too tight a bottleneck can thwart reward learning.
Divergence
1--1
Reward Loss
Retur
9.30
1.0
0.5
9.20
0
Eneionment Steps
Environment Steps
Environment Steps 165
ratio=1
ralioes
rationy
Figure 15: Ablating KL balancing ratio. The dynamics model (prior) is generally harder to train than the representation model (posterior), which suggests training the prior more frequently than the posterior, i.e. using higher KL balancing ratio. However, training the prior too aggressively leads poorly regularized posterior.
E.4 An illustrative example of support constraint
We provide a didactic example to illustrate the intuition behind using support constraint for alignment. Consider two ground truth latent distributions shown in Fig. 16. The training dis- tribution Perain is uniform distribution span- ning [0,
Training Latent Distribution
Testing Latent Distribution
[0, 6). and the test-time distribu- tion Ptest is a uniform distribution covering exactly half the support of Ptrriin This is to simulate insufficient online exploration at test time. We construct nonlinear emission func- tion f X O f(x) 0. to generate observations. For illustration purposes, we use the same emission function for both training and testing domains. Given access to perfect training encoder, our goal is to learn a test-time encoder that recovers the ground truth encoding function (inverse of emission function) and in turn the test-time latent distribution. We compare the support constraint objective (detailed in eq. (4.7)) with distribution matching objective minimizing the Sinson-Shannor divergence.
2-
Figure 16: Ground truth latent distributions training and testing domains.
First, notice in Fig. 17 that naively optimizing either objective fails recover the ground truth test-time latent distribution. The distribution matching objective "stretches" the test-time distribution to match the training distribution, whereas the support constraint only forces the former to be within the support of the latter without any additional stipulation. A common method in unsupervised domain adaptation is enforce dynamics consistency between adjacent states, which in our toy example can be interpreted preserving the pairwise distance between latent states. When we add in
22

---

Vanilla
Dynamics
Calibration
Training Testing
Cal brattor
Distribution
Matching
Support
Constraint
Comparison between distribution matching and support constraint. When there is insufficient exploration, the distribution matching object conflicts with dynamics consistency and calibration objectives, whereas support constraint conforms to these objectives.
Figure 17:
this objective, a crucial difference between support constraint and distribution matching transpires: while the support constraint conforms to the dynammic consistency objective through the reweighting function and recovers the correct shape of the distribution, the distribution matching objective is inherently conflicting with dynamics consistency. Intuitively, the distribution matching objective attempts to stretch the distribution, whereas the dynamies consistency pulls together. This leads to suboptimal results shown in the middle column of the first row in Fig. 17. Still, neither method can recover the ground truth distribution by only using dynamics consistency. the relative positions of latents but not the global position. To address this, we propose to match the ground truth latent state for small number calibration examples. These calibration examples effectively serve an anchor for the distribution. When the calibration samples have good coverage, we see the support matching objective successfully recovers the ground truth encoding function and latent distribution. Distribution matching, the other hand, yields suboptimal solution it still conflicts with the calibration objective.
it only preserves
We further investigate the case where the cali- bration examples do not have good coverage. In Fig. 18, we see that when the calibration exam- ples only cover a small region the test-time distribution, the result of optimizing the support constraint accurate only around the calibration examples. Under this circumstance, the addition of dynamics consistency proves to be quite ef- fective. Since (1) the distribution is anchored around the calibration samples and (2) pairwise distance is preserved, the only valid solution is recovering the ground truth latent distribution. In our DMC benchmark, we do observe degra- dation when we replace the exploration policy with random policy for collecting calibration samples. Yet we find optimizing for dynamics consistency through RSSM be of little avail. We hypothesize this is due to the recurrent archi- tecture interfering with the optimization, and leave it for future work to investigate potential solutions.
Calibration Dynamics
Calibration
Training
Testing
Leamedd
Calibration
Support constraint with skewed calibration examples. When the calibration examples are skewed. the result of optimizing the support constraint is only accurate around the calibration examples. However, adding dynamics consistency effectively propagates the calibration signal other states.
Figure 18:
23


Reviews:
[{"review_id": "l2Ad8aKxhU", "reviewer": "Reviewer_zSkA", "summary": "This work proposes a modification to Dreamer [1], wherein the authors discard the reconstructions objective, and replace it with an objective that maximizes the information contained in the latent representation about the reward, while minimizing the information about the original image. This pushes the model to learn representations that are agnostic to 'spurious variations' - changes in the input that do not affect the reward. \n\nThe authors compare the method to Dreamer [1], TIA [2] and bisimulation [3], and show that RePo deals with distracting background on DMC very well compared to other methods.\n\nAuthors also show that the method can generalize from static training data to test-time data with distractors if an additional test-time adaptation stage is introduced.\n\n[1] Hafner D, Lillicrap T, Ba J and Norouzi M 2019 Dream to Control: Learning Behaviors by Latent Imagination _arXiv_\n[2]  Fu X, Yang G, Agrawal P and Jaakkola T 2021 Learning Task Informed Abstractions\n[3] Zhang A, McAllister R, Calandra R, Gal Y and Levine S 2021 Learning Invariant Representations for Reinforcement Learning without Reconstruction", "questions": "- Can the authors clarify what would happen if the image itself contains information about the score.\n- As a baseline, would be useful to see how Dreamer performs with distribution shfit.\n\n**typos**\n- 22 - learning to operator - > learning to operate\n- 120 - to a an -> to an\n- 135 - effect -> affect\n- 366 - could be used should be removed", "limitations": "The limitations are adequately addressed in the submission.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written\n- The proposed method is simple\n- The method performs better than baselines on DMC with distractions.\n- The method is able to work with distractions introduced at test-time if additional test-time adaptation stage is used", "weaknesses": "- This method, as authors have mentioned, is no longer task-agnostic.\n- Possibly, if there's some information about the reward in the image (like a score), the model will focus on that instead", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments. We address your questions below.\n\n> Can the authors clarify what would happen if the image itself contains information about the score.\n\nSince RePo not only predicts the current reward, but also the dynamics and future rewards, adding the score information would not cause it to only focus on the score. To verify this, we construct a variant of Distracted Walker Walk consisting of a scoreboard displaying the current episodic return. We train a RePo agent on this environment and probe its latent states to reconstruct the observations. As shown in Figure 3, RePo successfully reconstructs the joint positions of the agent and does not only focus on the score.\n\n> As a baseline, would be useful to see how Dreamer performs with distribution shfit.\n\nWe provide the baseline comparison in Table 1. We see that Dreamer experiences a significant performance drop when the background becomes out of distribution.\n"}, {"review_id": "HBhMkAmJ6y", "reviewer": "Reviewer_CzT8", "summary": "The paper has two major contributions. The first one is the introduction of a model-based RL method that learns latent representations of visual RL environments that are resilient to spurious variations. The second contribution is a test-time alignment method that makes the latent representations invariant distribution shifts, thereby enabling the quick adaptation to diverse environments.", "questions": "Should I assume that the method name RePo comes from Resilient Posterior? If so, please clarify that the first time you use it, e.g., line 40.\n \nAre there any constraints on the environment distribution shift? On the experiments you tested on natural video distractors. What if it was something extreme like white noise?\n", "limitations": "I do not foresee any negative societal impact. Limitations were not discussed. The authors should add a small paragraph with the limitations.\n", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper proposes a novel method to address a very important problem in reinforcement learning. The RePo method is simple, elegant, and well-motivated and the discussion around eq. 4.3 makes it easy to intuitively understand it.  The presentation is something to compliment the authors for. The paper is very well written and easy to follow, despite the mathematical complexity of the underlying concepts. The relevant work section is concise, and the authors clearly relate their work with previous papers.\n", "weaknesses": "The only minor weakness that I can note is that the authors did not discuss any subtle limitations of their method, which I am sure it does have. Adding such a discussion would make the paper more credible.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and for finding our paper “simple, elegant, and well-motivated.” We address your comments below.\n\n> The only minor weakness that I can note is that the authors did not discuss any subtle limitations of their method, which I am sure it does have.\n\nWe note in Section 5 that due to its reward-predicting nature, RePo’s latent representation is no longer task-agnostic, which makes it challenging to generalize to new reward functions. This could be potentially addressed by training the model to predict multiple reward objectives at once, where each objective exposes some partial view of the environment dynamics. We also note that paired calibration data used for alignment can be expensive to collect, and we hope to explore more unsupervised approaches to alignment in the future.\n\n> Are there any constraints on the environment distribution shift? On the experiments you tested on natural video distractors. What if it was something extreme like white noise?\n\nWe do not put an explicit constraint on the environment distribution shift. In the case where the background is replaced by white noise, we expect both our method and baselines to perform well. This is because if the noise is random Gaussian and independent at each time step, then it contains no dynamical information and thus will not contribute as a spurious correlation. In Fig. 5 of the supplementary pdf, we compare RePo on Distracted Walker Walk with video distractor and white noise, and we do not see a difference in terms of performance. \n\n> Should I assume that the method name RePo comes from Resilient Posterior? If so, please clarify that the first time you use it, e.g., line 40.\n\nThe name RePo comes from **Re**gularizing **Po**sterior Predictability. We will add a clarification in the revision.\n"}, {"review_id": "xVswlrnrbj", "reviewer": "Reviewer_ZoCZ", "summary": "This paper presents an approach that aims to enhance visual model-based RL by learning state representations that filter out irrelevant information while retaining task-specific information. The proposed method achieves this by maximizing the mutual information (MI) between the current representation and future rewards, while simultaneously minimizing the MI between the representation and observations. In order to facilitate the transfer of the learned representation to new environments with distribution shifts, this paper introduces a support constraint regularization technique. This technique ensures that a specific set of \"calibration\" states, paired across training and testing domains, share the same encoding. The empirical results obtained demonstrate the effectiveness of the proposed method.", "questions": "Most important questions are presented in above. \nFurthermore, what are the theoretical advantages of RePo over bisimulation-based methods? It appears that the proposed method necessitates test-time calibration to adapt to changes in distractors, whereas bisimulation-based methods can transfer their representations directly without any additional modifications. Notably, certain bisimulation approaches [1, 2] have demonstrated impressive performance not only in standard settings but also in distractor settings.\n\n[1]: Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, Mark Rowland: MICo: Improved representations via sampling-based state similarity for Markov decision processes. NeurIPS 2021: 30113-30126\n\n[2]: Hongyu Zang, Xin Li, Mingzhong Wang: SimSR: Simple Distance-Based State Representations for Deep Reinforcement Learning. AAAI 2022: 8997-9005", "limitations": "The method necessitates test-time adaptation instead of direct transfer. This requirement for additional calibration time may increase computational complexity during the evaluation stage.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper addresses an important open problem in RL and is very well motivated.\n2. The logic flow of the paper is good, and the presentation is clear and coherent making it a pleasant read.\n3. Test-time adaptation is a reasonable and interesting topic in this field.\n4. Empirical studies show the performance gain of the proposed method.", "weaknesses": "1. The MDP setting in this paper bears strong resemblance to the exogenous distractors setting discussed in [1, 2, 3], as well as the denoised MDP setting introduced in [4]. It is important to include a discussion of these works in the related work section. Additionally, it is worth noting that the denoised MDP approach presented in [4] also offers a model-based method for learning state representations, which should be considered as a baseline method in the experiments conducted.\n2. Utilizing the recurrent state-space model (RSSM) to parameterize both p and q is a reasonable choice. However, understanding Assumption A.1 is challenging. While $z_t$ can represent past latents and past actions, it raises questions about whether $z_t$ can effectively incorporate the information from $a_t$, given that the RSSM does not explicitly model their corresponding relationship.\n\nTypos: \n1. Equation 4.1 is a bit strange. Seems the equation in Appendix A.1 is more reasonable.\n2. Line 192 \"parametrize\" -> \"parameterize\"\n\n[1]: Yonathan Efroni, Dylan J. Foster, Dipendra Misra, Akshay Krishnamurthy, John Langford: Sample-Efficient Reinforcement Learning in the Presence of Exogenous Information. COLT 2022: 5062-5127\n\n[2]: Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alekh Agarwal, John Langford: Provable RL with Exogenous Distractors via Multistep Inverse Dynamics. CoRR abs/2110.08847 (2021)\n\n[3]: George Trimponias, Thomas G. Dietterich: Reinforcement Learning with Exogenous States and Rewards. CoRR abs/2303.12957 (2023)\n\n[4]: Tongzhou Wang, Simon S. Du, Antonio Torralba, Phillip Isola, Amy Zhang, Yuandong Tian: Denoised MDPs: Learning World Models Better Than the World Itself. ICML 2022: 22591-22612", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments and finding our paper “well-motivated” and a “pleasant read.” We address your questions below.\n\n> The MDP setting in this paper bears strong resemblance to the exogenous distractors setting discussed in [1, 2, 3], as well as the denoised MDP setting introduced in [4].\n- The work on exogenous distractors is very interesting, and we will be sure to include an extended discussion in the revision. While these works provide interesting theoretical insights, they are not deployable at the scale of our experimental environments. These works further assume that the distractor state of the MDP is entirely decoupled from the effect of control actions. This means that the exogenous MDP algorithms will learn the entire controllable part of the MDP even if it contains reward-irrelevant information.\n- The Denoised MDP refines the exogenous MDP perspective by learning a factorized latent state which disentangles the non-controllable, reward-irrelevant controllable, and reward-relevant controllable parts of the MDP. However, unlike RePo, they do not provide a mechanism to enforce minimality of all the factors. Indeed, we show that our method outperforms Denoised MDP in Figs. 1 and 2 of the supplementary pdf. \n\n> While z_t can represent past latents and past actions, it raises questions about whether z_t can effectively incorporate the information from a_t, given that the RSSM does not explicitly model their corresponding relationship.\n\nOur assumption is that the reward can be predicted from states and is independent of current action. We build this assumption into our model, which allows us to predict the dynamics, reward, value function, and policy all from the same latent state. It is straightforward to generalize our model to allow for action-dependent rewards by conditioning the reward predictor on the action. However, in our applications of interest, rewards are state dependent so we build in this information.\n\n> Furthermore, what are the theoretical advantages of RePo over bisimulation-based methods? It appears that the proposed method necessitates test-time calibration to adapt to changes in distractors, whereas bisimulation-based methods can transfer their representations directly without any additional modifications.\n- Bisimulation (Bisim) is optimal in the sense that a perfect bisimulation metric does not discard any relevant information about an MDP. However, there is no guarantee that Bisim will disregard irrelevant information. Indeed, the identity mapping induces a trivial Bisim metric without data compression. Thus, Bisim methods compress only by reducing the dimensionality of the latent variable; RePo, on the other hand, has an explicit term in its loss that encourages discarding irrelevant information by regularizing the posterior predictability of rewards/dynamics [1][2][3].\n- We note that the latent representation of our method is invariant, but the visual encoder does not generalize out of distribution, hence the need for test-time adaptation. As indicated in the invariant risk minimization [4] line of work, any approach that uses a CNN encoder to extract visual features -- Bisim included -- also suffers from the same problem and benefits from test-time adaptation.\n\n[1] Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic Markov Decision Processes. AAAI 2020.\n\n[2] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. ICLR 2021.\n\n[3] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, Doina Precup. Invariant Causal Prediction for Block MDPs. ICML 2020.\n\n[4] Elan Rosenfeld, Pradeep Ravikumar, Andrej Risteski. The Risks of Invariant Risk Minimization. ICLR 2021.\n"}, {"review_id": "h7YE6pTgRg", "reviewer": "Reviewer_z3JG", "summary": "This paper presents a model-based reinforcement learning (RL) method which learns a task-relevant latent-space dynamic model. It learns a latent representation containing task-relevant features by 1) maximising the mutual information between latent states and rewards whilst 2) minimising the mutual information between the latent states and observations. This contrasts approaches which reconstruct the image observations and thus learn task-irrelevant features. They highlight that their method fails to adapt to environment changes which results in a distribution shift to the observations, e.g. changes in lighting. To this end, they present a method for adapting to distribution shifts at test time. They test their method in modified versions of the DeepMind Control Suite and Maniskill with (distracted) video backgrounds.", "questions": "1. Are the authors aware of the literature mentioned above? Do you think you need to compare to these methods?\n2. $q$ is usually used for approximate posteriors and $p$ is used for priors. Is there a reason you have swapped these around?\n", "limitations": "The authors mention that their method is not task agnostic and I think this is sufficient. ", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper addresses an important problem which is needed for deploying model-based RL algorithms in real-world robotics scenarios.\nI am not familiar with test-time adaptation for visual model-based RL, however, I do think this is an imprtant problem and the approach seems sensible.\n\nThe experiments seem to be reproducible from the implementation details in the appendix. However, I have not ran the code to check.\n", "weaknesses": "The paper's biggest weakness is its positioning against prior work.\nThere exist visual model-based RL methods that learn task-relevant latent representations which are not cited in this paper.\nFor example:\n```\n- Denoised MDPs\n    @inproceedings{tongzhouw2022denoisedmdps,\n    title={Denoised MDPs: Learning World Models Better Than The World Itself},\n    author={Wang, Tongzhou and Du, Simon S. and Torralba, Antonio and Isola, Phillip and Zhang, Amy and Tian, Yuandong},\n    booktitle={International Conference on Machine Learning},\n    organization={PMLR},\n    year={2022}\n    }\n- TD-MPC\n    @inproceedings{Hansen2022tdmpc,\n        title={Temporal Difference Learning for Model Predictive Control},\n        author={Nicklas Hansen and Xiaolong Wang and Hao Su},\n        booktitle={ICML},\n        year={2022}\n    }\n- ALM\n    @inproceedings{ghugare2023simplifying,\n      title={Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective},\n      booktitle = {Advances in Neural Information Processing Systems},\n      author={Raj Ghugare and Homanga Bharadhwaj and Benjamin Eysenbach and Sergey Levine and Ruslan Salakhutdinov},\n      year={2022},\n    }\n```\nGiven the similarity of the work, I think it is important to not only discuss them in the related work but to also compare to them in the experiments section.\nEspecially as it is not obvious to me that this is a better approach to learning a task-relevant latent space.\nI am interested to see how this method compares but I do not think this can be published without such comparisons.\n\nThe training curves should show the converged values with dashes lines. In many figures it is not possible to see where algorithms converge. For example, dreamer in Walker Stand/Walker Walk/Walker Run in Figure 5.\n\nThere are a lot of minor corrections and typos.\n\nMinor corrections and typos:\n- How many seeds were used for each experiment?\n- Do the shaded regions on figures show variance or std or something else?\n- Line 22 - operator should be operate\n- Line 22 - the sentence doesn't read properly\n  - Perhaps you could replace \"and learning to operator the same arm amidst of plentiful dynamic distractors\" with \"and learning to operate the same arm amidst dynamic distractors\".\n- Eq 4.1 - $I_{p,h}(z_{1:T};x_{1:T} \\mid a_{1:T})$ should be $I_{p,h}(z_{1:T}; o_{1:T} \\mid a_{1:T})$\n- Figures 4/7/8/9 captions should have full stops.\n- Figures 1/2/3 are never referenced.\n- Figure 1 is not very informative so I am unsure why it is included?\n- Line 113 - brackets in expectation\n- Line 123/132 - What does the $h$ subscript in $s_{h}$ represent? You introduce $s_{t}$ where the subscript is time $t$ but then it looks like you're indexing at the encoder?\n- Line 84-87 - This sentence doesn't read properly.\n- Line 221 - $\\bar{P} : \\mathcal{X} \\times \\mathcal{A} \\rightarrow \\Delta(\\mathcal{X})$ should have $\\times$ not $\\rightarrow$?\n- Line 324 - \"Does representations\" should be \"Do representations\"\n- Line 360 - Incorrect use of semicolon\n- Line 366 - \"Can we apply our algorithm could be used in a continual learning setup?\" doesn't make sense.\n- Eq. 4.1/4.5/4.6 are missing full stops.\n- Be consistent with Eq. vs Equation, vs Eqn. vs Eqs.\n  - Line 197 - Eqn should be Eq. for consistency\n  - Line 197 - Eqn should be Eq. for consistency\n  - Line 177 - (4.3) should be Eq. 4.3?\n  - Line 167 - (4.1) should be Eq. 4.1?\n  - Line 261 - (4.5) should be Eq. 4.5?\n  - Line 265 - (4.5) should be Eq. 4.5?\n- Similarly, with Fig vs Fig. vs Figure.\n- Parsimonious/eschew seem unnecessarily complicated words. I had to look them up... I recommend simpler words to make the paper more accessible. I asked around and 60% of people would need to look up both words.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review and constructive feedback. We address your comments below.\n\n> Given the similarity of the work, I think it is important to not only discuss them in the related work but to also compare to them in the experiments section.\n\nThank you for suggesting these baselines. We provide additional results in Figs. 1 and 2 of the supplementary pdf showing that our method indeed outperforms Denoised MDP and TD-MPC. We hypothesize that our advantage over denoised MDP is (a) we do not get penalized for state reconstruction, (b) we do not spend representation capacity on learning a dynamics model in the task-irrelevant components of the state space, and (c) we do not need to tune multiple loss-weights as hyperparameters. TD-MPC, on the other hand, predicts the value function (but not state reconstruction). However, it has no explicit mechanism to encourage minimality of the representation. By contrast, RePo introduces a simple objective that encourages the learning of a minimal task-relevant latent representation. In addition, we propose a reward-free alignment method to handle test-time distribution shift.\n\nWe omit the comparison with ALM since the authors did not evaluate their method on visual domains.\n\n> The training curves should show the converged values with dashes lines. In many figures it is not possible to see where algorithms converge.\n\nWe provide the convergence returns of RePo for all Distracted DMC environments in Fig. 1. All agents converge after 1M environment steps.\n\n> How many seeds were used for each experiment?\n\nWe run 4 seeds for each method on each environment.\n\n> Do the shaded regions on figures show variance or std or something else?\n\nThe shaded regions denote standard error.\n\n> Figure 1 is not very informative so I am unsure why it is included?\n\nWe include Fig. 1 to illustrate the motivation of our method. We will improve it in the revision by making the illustration more realistic and intuitive.\n\n> Line 123/132 - What does the h subscript in s_h represent? You introduce s_t where the subscript is time t but then it looks like you're indexing at the encoder?\n\nWe made a typo on lines 123 and 132. There is no subscript on $s$. We will fix it in the revision.\n\n> q is usually used for approximate posteriors and p is used for priors. Is there a reason you have swapped these around?\n\nWe follow Dreamer [1] convention and use q to denote the prior and p for the posterior.\n\nThank you again for your thorough review! We will be sure to fix the typos in the revision.\n\n[1] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, Mohammad Norouzi. Dream to Control: Learning Behaviors by Latent Imagination. ICLR 2020.\n"}]

Author rebuttal:
We thank the reviewers for their careful reading and constructive feedback. To address the questions raised in the reviews, we conduct the following additional experiments and report the results in the supplementary pdf:
- We introduce two baselines suggested by the reviewers: Denoised MDP [1] and TD-MPC [2]. Denoised MDP learns a factorized latent dynamics model that disentangles controllability and reward relevance. However, they do not enforce minimality since they still reconstruct the state and do not optimize an information bottleneck objective. TD-MPC trains a latent dynamics model to predict the value function and uses a hybrid planning method to extract a policy. While its latent representation is task-oriented, it is not minimal and necessarily contains information about the policy. As shown in Figs. 1 and 2, RePo generally outperforms both baselines across Distracted DMC and Realistic Maniskill environments. 
- We evaluate RePo on a variant of Distracted Walker Walk which consists of a scoreboard displaying the cumulative reward of the current episode. The purpose is to test if RePo’s reward-predictive nature causes its latent representation to collapse onto predicting the score. We present the probing visualization in Fig. 3, which shows that the learned latent reconstructs the joint positions of the agent and ignores the score. This is because RePo predicts not only the current reward but also the dynamics and future rewards.
- We qualitatively compare the latent states of Dreamer and RePo by visualizing their top two principal components. Specifically, we collect the same trajectory across 90 different backgrounds and visualize the final recurrent latent state inferred by Dreamer and RePo respectively. As shown in Fig. 4. RePo produces more compact latent representations than Dreamer, meaning the latent states encode less information about background variations. This enables RePo to share data across environments, which explains its superior sample efficiency compared to other baselines.
- In Fig.5, we evaluate RePo on a variant of Distracted Walker Walk with white noise as the background and observe no performance difference compared to the video distractor setting. 
- As a baseline comparison, we train Dreamer agents on the original DMC environments and evaluate them on the distracted variants. From Table 1, we can see that Dreamer experiences a significant performance drop when the background becomes out of distribution.

We provide additional clarifications, explanations and discussion in the per-reviewer responses.

[1] Tongzhou Wang, Simon S. Du, Antonio Torralba, Phillip Isola, Amy Zhang, Yuandong Tian. Denoised MDPs: Learning World Models Better Than the World Itself. ICML 2022.

[2] Nicklas Hansen, Xiaolong Wang, Hao Su. Temporal Difference Learning for Model Predictive Control. ICML 2022.


Meta-review: This paper proposes a new objective to train visual model-based RL algorithms with the aim of avoiding modelling spurious variations that occur in the image. This is a much-needed work to improve deep MBRL algorithms. The authors also propose a test-time adaptation of the encoder for handling distribution shifts. The results prove that the proposed approach is better than Dreamer and other methods.

Both reviewers with a score of 5 asked for more comparisons. The authors made some comparisons in the rebuttal pdf and also explained why some comparisons might not be necessary. I am satisfied with the answers.

I think this is an interesting research direction that is worth a spotlight.
Decision: Accept (spotlight)