SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Comparison to DREAMER, DEEPMDP, and BISIMULATION. DREAMER [17] was first derived to optimize pixel-reconstruction, leading to high-fidelity dynamics but susceptibility to spurious variations. Naively removing pixel reconstruction from dreamer, however, leads to 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] sets the latents  $z_{t}$  to exactly the image encodings  $x_{t} = h(o_{t})$ . It learns a dynamics  $\bar{P}: \mathcal{X} \times \mathcal{A} \to \triangle(\mathcal{X})$  such that the distribution  $\bar{x}_{t+1} \sim \bar{P}(h(o_{t}), a_{t})$  is close to  $x_{t+1} \sim h(o_{t+1})$ ,  $o_{t+1} \sim P^{\star}(o_{t}, a_{t})$ , where  $P^{\star}$  denotes a 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. BISIMULATION [63] learns an optimal representation in the sense that a perfect bisimulation metric does not discard any relevant information about an MDP. However, there is no guarantee that it will disregard irrelevant information. Indeed, the identity mapping induces a trivial bisimulation metric. Hence, BISIMULATION compress only by reducing the dimensionality of the latent space. In contrast, we further compress the encodings  $x_{t}$  into latents  $z_{t}$  so as to enforce the latent prior  $q_{z}(\cdot | a_{t}, z_{t})$  is close to the latest observation-dependent posterior distribution  $p(\cdot | z_{t}, a_{t}, x_{t+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 to spurious variations seen during training, our learned latents  $z_{t}$  - 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: \mathcal{O} \to \mathcal{X}$  may observe images that it has not seen at train time; thus the latent, which depend on observations through  $x_{t} = h(o_{t})$ , may behave erratically, even when system dynamics remain unchanged.
Relying on the resilience of our posteriors  $p$  over latents  $z_{t}$  introduced by RePo, we propose a test-time adaption strategy to only adjust the encoder  $h$  to the new environment, whilst leaving  $p$  fixed. A natural approach is to apply unsupervised domain adaptation methods [66, 58] to adapt the visual encoder  $h$  to  $h_{\mathrm{test}}$ . 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  _ {h _ {\text {t e s t}} (\cdot)} \mathrm {D} \left(\mathcal {P} _ {\text {t r a i n}} \| \mathcal {P} _ {\text {t e s t}}\right) \text {s . t .} \mathcal {P} _ {\text {t e s t}} = h _ {\text {t e s t}} \circ \mathcal {D} _ {\text {t e s t}}, \mathcal {P} _ {\text {t r a i n}} = h \circ \mathcal {D} _ {\text {t r a i n}}. \tag {4.5}
$$
In Eq. (4.5), we consider matching the distributions over encodings  $x$  of observations  $o$ . Specifically, we assume  $\mathcal{D}_{\mathrm{train}}$  and  $\mathcal{D}_{\mathrm{test}}$  denote training and test-buffer distributions over observations  $o$ ,  $\mathcal{P}_{\mathrm{train}} = h_{\mathrm{train}} \circ \mathcal{D}_{\mathrm{train}}$  denotes the distribution of  $x = h_{\mathrm{train}}(o)$  where  $o \sim \mathcal{D}_{\mathrm{train}}$  is encoded by the train-time
Distracted DeepMind Control
Figure 4: Depiction of the environments being used for evaluation. (Left): the Distracted DeepMind Control suite [64], (Top Right): Maniskill2 [15] environments with realistic backgrounds from Matterport [3]. (Bottom Right): TurtleBot environment with two TVs playing random videos in the background.
Realistic Maniskill
Lazy TurtleBot
encoder  $h_{\mathrm{train}}$ , and  $\mathcal{P}_{\mathrm{test}} = h_{\mathrm{test}} \circ \mathcal{D}_{\mathrm{test}}$  denotes encodings under a test-time encoder  $h_{\mathrm{test}}(\cdot)$  over which we optimize. Here,  $\mathrm{D}(\cdot, \cdot)$  denotes an  $f$ -divergence, such as the  $\chi^2$ -divergence.
Support Constraint. (4.5) fails to 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-optimal regions of state space not seen at train time. Thus, it may be impossible to match the distribution as in standard unsupervised domain adaptation. We therefore propose to replace (4.5) with a support constraint, enforcing that the distribution of  $h_{\mathrm{test}} \circ \mathcal{D}_{\mathrm{test}}$  is contained in the support of  $h_{\mathrm{train}} \circ \mathcal{D}_{\mathrm{train}}$ . We consider the following idealized objective:
$$
\min  _ {\tau (\cdot) \geq 0, h _ {\text {t e s t}} (\cdot)} \mathrm {D} \left(\tau \cdot \mathcal {P} _ {\text {t r a i n}} \| \mathcal {P} _ {\text {t e s t}}\right) \text {s . t .} \mathbb {E} _ {x \sim \mathcal {P} _ {\text {t r a i n}}} [ \tau (x) ] = 1. \tag {4.6}
$$
Here, by  $\tau \cdot \mathcal{P}_{\mathrm{train}}$ , we mean the re-weighted density of  $\mathcal{P}_{\mathrm{train}} = h_{\mathrm{train}} \circ \mathcal{D}_{\mathrm{train}}$  by a function  $\tau(x)$ . The constraints  $\mathbb{E}_{\mathcal{P}_{\mathrm{train}}}[\tau(x)] = 1$  and  $\tau(\cdot) \geq 0$  ensures this reweighted distribution is also a valid probability distribution. The reweighting operation  $\tau \cdot \mathcal{P}_{\mathrm{train}}$  seems intractable at first, but we show that if we take  $\mathrm{D}(\cdot, \cdot) = \chi^2(\cdot, \cdot)$  to be the  $\chi^2$  divergence, then Eq. (4.6) admits the following tractable Lagrangian formulation (we refer readers to [65] and Appendix B for a thorough derivation)
$$
\min  _ {\tau (\cdot) \geq 0, h _ {\text {t e s t}} (\cdot) f (\cdot), \lambda} \max  _ {f (\cdot), \lambda} \mathbb {E} _ {\mathcal {P} _ {\text {t r a i n}}} [ \tau (x) \cdot f (x) ] - \mathbb {E} _ {\mathcal {P} _ {\text {t e s t}}} \left[ f (x) + \frac {1}{4} f (x) ^ {2} \right] + \lambda \left(\mathbb {E} _ {\mathcal {P} _ {\text {t r a i n}}} [ \tau (x) ] - 1\right), \tag {4.7}
$$
where above,  $\lambda \in \mathbb{R}$ ,  $f: \mathcal{X} \to \mathbb{R}$ , and the objective depends on  $h_{\mathrm{test}}$  through the definition  $\mathcal{P}_{\mathrm{test}} = h_{\mathrm{test}} \circ \mathcal{D}_{\mathrm{test}}$ . This objective is now a tractable saddle point optimization, which can be solved with standard stochastic optimization techniques. The optimization alternates between optimizing the reweighting  $\tau$  and the visual encoder  $h_{\mathrm{test}}$ , and the dual variables  $f, \lambda$ . 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  $\tau(\cdot)$  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  $\ell_2$  loss between the training and testing encoding of each pair of observations. We defer the details of the complete optimization to Appendix C.
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?
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 a variant of DeepMind Control Suite where the static background is replaced with natural videos (Fig. 4). For adaptation experiments, we train agents on static undistracted backgrounds and adapt them to distracted variants. (2) Realistic Maniskill is a benchmark we constructed based on the Maniskill2 benchmark [15], but with realistic backgrounds from [3] to simulate learning in a diverse range of human homes. We solve three tasks - LiftCube, PushCube, and TurnFaucet in a variety of background settings. (3) Lazy TurtleBot is a real-world robotic setup where a TurtleBot has to reach some goal location from egocentric observations in a 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 a state-of-the-art 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 task-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) DeepMDP [12] is a model-free method that learns a 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. (1) calibrated distribution matching, a variant of the method proposed in Section 4.1, using a distribution matching constraint rather than a 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, a variant of the method proposed in Section 4.1, using a distribution matching constraint, but without using paired examples, (4) invariance through latent alignment ILA [61], a technique for test-time adaptation of representations with distribution matching and enforcing consistency in latent dynamics, (5) calibration, a baseline 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 DeepMind 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] see occasional success when
Figure 6: Results on adaptation from static environments to dynamic environments in Deepmind control. RePo with calibrated support constraints outperforms ablations and previous techniques for domain adaptation.
Figure 7: Results of training agents on varying static environments in Maniskill [15]. RePo is able to learn more 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 variations in the real world, we evaluate its performance on Lazy TurtleBot, where a mobile robot has to navigate around a furnished room to reach the goal from egocentric observations (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. 1, RePo is able to reach the goal with nontrivial success within 15K environment steps, whereas Dreamer fails to reach the goal. We provide details about the setup in Appendix. D.
Table 1: Results on Lazy TurtleBot at  ${15}\mathrm{\;K}$  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 effectiveness 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 in environments with static backgrounds, and adapt the representation to domains with natural video distractors (as shown in 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 infeasible 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 calibration 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 its ability to learn in a 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.
Figure 9: Probing representations learned by RePo shows that the background is largely ignored, while [17] tries to reconstruct it at the cost of the agent prediction.
Figure 10: Top two principle components of RePo and Dreamer's latent representations across different backgrounds. RePo's latent representation is more compact than Dreamer's, which enables data sharing.
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 is unable to reach optimal task performance. TIA [9] suffers from high variance and occasionally fails to dissociate task-relevant from task-irrelevant features. Denoised MDP [56], TD-MPC [22], and DBC [63] learn to ignore the background in 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 representations by post-hoc training a separate image decoder to reconstruct image observations from the latents. We visualize the results in Fig. 9 and compare them with Dreamer reconstructions [17]. Our representation contains little information about background but is capable of reconstructing the agent, implying that it contains only task-relevant information.
In addition to probing, we qualitatively compare the latent states of RePo and Dreamer by visualizing their top two principal components. We collect the same trajectory across all backgrounds in 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.
Figure 8: Ablating objectives showing the importance of information bottleneck and KL balancing described in Section 4.
Ablation experiments We conduct ablation experiments to determine the effect of hyperparameters in Fig. 8. As we can see, the performance is crucially dependent on the information bottleneck  $\epsilon$ , as well as KL balancing. We refer readers to Appendix E for a more thorough 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 is that the learned dynamics model is no longer task-agnostic, as it only captures task-relevant information. This can be potentially addressed by simultaneously predicting multiple reward objectives. Our framework opens up several interesting directions for future research, such as: can a multi-task variant of RePo allow for representations applicable to a 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-centric environments.
We would like to thank Marius Memmel, Max Balsells, and many other members of the WEIRD Lab at University of Washington for valuable feedback and discussions.
[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 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 2226-2240. PMLR, 2022.
[60] M. Yang, D. Schuurmans, P. Abbeel, and O. Nachum. Dichotomy of control: Separating what you can control from what you cannot. CoRR, abs/2210.13435, 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, J. Pineau, Y. Gal, and D. 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 representations 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, abs/1811.06032, 2018.
[65] R. Zhang, B. Dai, L. Li, and D. Schuurmans. Gendice: Generalized offline estimation of stationary values. In 8th International Conference on 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.06745, 2021.
[67] F. Zhuang, X. Cheng, P. Luo, S. J. Pan, and Q. He. Supervised representation learning with double encoding-layer autoencoder for transfer learning. ACM Trans. Intell. Syst. Technol., 9(2):16:1-16: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.
Supplementary Materials for
"RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability"
A Algorithm Pseudocode
Algorithm 1 Resilient Model-Based RL by Regularizing Posterior Predictability (RePo)
1: Initialize dataset  $\mathcal{D}_{\mathrm{buf}}$  with  $S$  random seed episodes.  
2: Initialize neural network parameters  $\theta, \phi, \psi$  randomly and set dual variable  $\beta = \beta_0$ .  
3: while not converged do  
4: for update step  $c = 1 \dots C$  do  
5: // Dynamics Learning  
6: Draw  $B$  data sequences  $\{(o_t, a_t, r_t)\}_{t=1}^T \sim \mathcal{D}_{\mathrm{buf}}$ .  
7: Encode images  $x_{1:T} = h_\theta(o_{1:T})$ .  
8: Infer prior and posterior distributions  $q_\theta(z_t \mid z_{t-1}, a_{t-1})$ ,  $p_\theta(z_t \mid z_{t-1}, a_{t-1}, x_t)$ .  
9: Sample latent states  $z_t \sim p_\theta(z_t \mid z_{t-1}, a_{t-1}, x_t)$  and infer reward distributions  $q_\theta(r_t \mid z_t)$ .  
10: Compute the Lagrangian objective in (4.4)  
11: Update  $\theta$  and  $\beta$  using dual gradient descent.  
12: // Behavior Learning  
13: Imagine trajectories  $\{(z_\tau, a_\tau)\}_{\tau=t}^{t+H}$  from each  $z_t$ .  
14: Predict rewards  $r_\tau \sim q_\theta(r_\tau \mid z_\tau)$ , values  $v_\tau \sim V_\psi(v_\tau \mid z_\tau)$ , actions  $a_\tau \sim \pi_\phi(a_\tau \mid z_\tau)$ .  
15: Update  $\phi$  and  $\psi$  using actor-critic learning.  
16: // Data Collection  
17: for time step  $t = 1 \dots N$  do  
18: Compute  $z_t \sim p_\theta(z_t, z_{t-1}, a_{t-1}, h_\theta(o_t))$  from history and current observation.  
19: Compute  $a_t \sim \pi_\phi(a_t \mid z_t)$  from policy.  
20: Execute  $a_t$  in environment and collect  $o_{t+1}, a_{t+1}$ .  
21: Add experience to dataset  $\mathcal{D}_{\mathrm{buf}} \gets \mathcal{D}_{\mathrm{buf}} \cup \{o_t, a_t, r_t\}_{t=1}^N$ .