SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chuning Zhu
Seattle, WA 98105
zchuning@cs.washington.edu
Boston, MA 02139
msimchow@mit.edu
Siri Gadipudi
sg06@uw.edu
Abhishek Gupta
abhgupta@cs.washington.edu
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 resilient to spurious variations, 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. We show its effectiveness in simulation benchmarks with significant spurious variations as well as a real-world egocentric navigation task with noisy TVs in the background.
Videos and code: https://zchuning.github.io/repo-website/.
Consider the difference between training a single robot arm against a plain background with reinforcement learning (RL), and learning to operate the same arm 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 RL agent - resilience to which is 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 task-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 of the demands of downstream
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 Posterior Predictability (RePo) – an algorithm for learning lossy latent representations resilient to spurious variations. A representation is satisfactory if it (a) predicts its own dynamics and (b) accurately predicts the reward. To satisfy these criteria, RePo jointly learns (i) a visual encoder mapping high-dimensional observations to intermediate image “encodings” (ii) a latent encoder which compresses histories of intermediate image encodings into compressed latent representations (iii) a 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] is a new desideratum of predictability: that, conditioned on past latents and actions, future latent dynamics should look as deterministic 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 a novel loss which encourages predictability, thereby discarding a broad range of spurious variations in aspects of the environment
which are 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 learning in environments with spurious variations - including dynamic elements like humans, changes in lighting and training across a range of visual appearances.
RePo implements a deceptively simple modification to recurrent state-space models for model-based RL [17, 59, 46]. 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, we optimize a 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 as 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, we 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 (e.g. a 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 a simple representation learning algorithm RePo for learning representations that are informative of rewards, while being as predictable as 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 to a 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.
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 dynamics models from high dimensional visual feedback (typically via recurrent state-space models). Perhaps most relevant to RePo is DREAMER [17]. Section 4 explains the salient differences between DREAMER and RePo; notably, we eschew a 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 exogenous distractors - parts of the state space whose dynamics is independent of the agent's actions. [56] introduces a 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 different 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 MMD [5, 33, 53, 24], KL divergence [67, 35] or Jensen-Shannon divergence [11, 52]. In [61], distribution matching was extended to sequential decision making. While domain adaptation settings typically assume that the source and target share an underlying marginal or joint distribution in a 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.
MDPs. A (discounted) MDP  $\mathcal{M} = (\mathcal{S},\mathcal{A},\gamma ,P,P_0,r)$  consists of a state-space  $\mathcal{S}$ , action space  $\mathcal{A}$ , discount factor,  $\gamma \in (0,1)$ , transition and  $P(\cdot ,\cdot):S\times \mathcal{A}\to \triangle (\mathcal{S})$ , initial state distribution  $P_{0}\in \triangle (\mathcal{S})$ , and reward function  $r(\cdot ,\cdot):S\times \mathcal{A}\rightarrow [0,1]$  (assumed deterministic for simplicity). A policy  $\pi :S\to \triangle (\mathcal{A})$  is a mapping from states to distributions over actions. We let  $\mathbb{E}_{\mathcal{M}}^{\pi}$  denote expectations under  $s_0\sim P_0$ ,  $a_{t}\sim \pi (s_{t})$ , and  $s_{t + 1}\sim P(s_t,a_t)$ ; the value is  $V_{\mathcal{M}}^{\pi}(s)\coloneqq \mathbb{E}_{\mathcal{M}}^{\pi}\left[\sum_{t = 0}^{\infty}\gamma^{h}r(s_{t},a_{t})\mid s_{0} = s\right]$ , and  $V_{\mathcal{M}}^{\pi} = \mathbb{E}_{s_{0}\sim P_{0}}[V_{\mathcal{M}}^{\pi}(s_{0})]$ . The goal is to learn a policy  $\pi$  that maximizes the sum of expected returns  $\mathbb{E}_{\mathcal{M}}^{\pi}\left[\sum_{t = 0}^{\infty}\gamma^{h}r(s_{t},a_{t})\mid s_{0} = s\right]$ , 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  $s_t$  to be visual observations  $s_t \equiv o_t \in \mathcal{O}$ ; for simplicity, we avoid explicitly describing a 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_t$  may be high-dimensional, so we learn encoders  $h: \mathcal{O} \to \mathcal{X}$  to an encoding space  $\mathcal{X}$ . We compress these encodings  $x_t$  further into latent states  $z_t$ , described at length in our method in Section 4.
Spurious variation. By spurious variation, we informally mean the presence of features of the states  $s_t$  which are irrelevant to 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 so easy to disentangle with the state: for example, lighting conditions will affect all observations, and hence will affect the appearance of transition dynamics.
Consider the following canonical example: an MDP with state space  $S_1 \times S_2$ , where for  $s = (s^{(1)}, s^{(2)}) \in S_1 \times S_2$ , the reward  $r(s, a)$  is a function  $\bar{r}(s^{(1)}, a)$  only of the projection onto  $S^1$ . Moreover, suppose that  $\mathbb{P}[(s^+)^{(1)} \in \cdot \mid s, a]$ , where  $s^+ \sim P(s, a)$ , is a distribution  $\bar{P}(s^{(1)}, a)$  again only depending on  $s^{(1)}$ . Then, the states  $s^{(2)}$  can be viewed as spuriously various. For example, if
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.
$s^{(1)}$  is a Lagrangian state and  $s^{(2)}$  is a static background, then it is clear that transitions of Lagrangian state and reward do not depend on  $s^{(2)}$ . Our template also encompasses dynamic distractors; e.g. a 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  $s^{(1)}$  (e.g. Lagrangian state) affect  $s^{(2)}$  (e.g. certain visual features), but not the other way round. In all cases, "spurious" implies that states  $(s_t^{(2)})_{t\geq 0}$  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 to 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  $\mathcal{O}$  be the space of image observations,  $\mathcal{X}$  the space of encoded observations, where  $h: \mathcal{O} \to \mathcal{X}$  represents the encoding function from images observations to encoded observations, and  $\mathcal{Z}$  the space of latent representations. Note that  $x_{t+1}$  is simply the instantaneous encoding of the image  $o_{t+1}$  as  $x_{t+1} = h(o_{t+1})$ , but the latent representation  $z_{t+1}$  at time step  $t+1$  is an aggregation of the current encoding  $x_{t+1}$  and previous latent  $z_t$  and action  $a_t$ . Let  $\mathcal{P}_{\mathrm{post}}$  denote the space of "posteriors" on latent dynamics  $z$  of the form  $p(z_{t+1} \in \cdot \mid z_t, a_t, x_{t+1})$ , where  $z_t, z_{t+1} \in \mathcal{Z}, a_t \in \mathcal{A}, x_{t+1} \in \mathcal{X}$ , and where and  $z_0 \sim p_0$  has some initial distribution  $p_0$ . 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 on this latent representation  $z_{t+1}$ , and not simply the image encoding  $x_{t+1}$  so as to incorporate historical information.
Let  $\mathcal{D}_{\mathrm{buf}}$  denote the distribution over experienced actions, observations and rewards from the environment  $((a_{1:T},o_{1:T},r_{1:T})\sim \mathcal{D}_{\mathrm{buf}})$ . For  $p\in \mathcal{P}_{\mathrm{post}}$ , let  $\mathbb{E}_{p,h}$  denote expectation of  $(a_{1:T},o_{1:T},r_{1:T})\sim \mathcal{D}_{\mathrm{buf}}$ ,  $x_{t} = h(o_{t})$  and the latents  $z_{t + 1}\sim p(\cdot \mid z_t,a_t,x_{t + 1})$  drawn from the latent posterior, with the initial latent  $z_0\sim p_0$ . 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  _ {p, h} \mathrm {I} _ {p, h} \left(z _ {1: T}; r _ {1: T} \mid a _ {1: T}\right) \text {s . t .} \mathrm {I} _ {p, h} \left(z _ {1: T}; o _ {1: T} \mid a _ {1: T}\right) <   \epsilon . \tag {4.1}
$$
Above,  $\mathrm{I}_{p,h}(z_{1:T};r_{1:T}\mid a_{1:T})$  denotes mutual information between latents and rewards conditioned actions under the  $\mathbb{E}_{p,h}$  distribution, and distribution  $\mathrm{I}_{p,h}(z_{1:T};o_{1:T}\mid a_{1:T})$  measures information
between latents and observations under  $\mathbb{E}_{p,h}$  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)
$$
\mathrm {I} _ {p, h} \left(z _ {1: T}; r _ {1: T} \mid a _ {1: T}\right) \geq \mathbb {E} _ {p, h} \left[ \sum_ {t = 1} ^ {T} \log q _ {\mathrm {r}} \left(r _ {t} \mid z _ {t}\right) \right] \tag {4.2}
$$
$$
\mathrm {I} _ {p, h} \left(z _ {1: T}; o _ {1: T} \mid a _ {1: T}\right) \leq \mathbb {E} _ {p, h} \left[ \sum_ {t = 0} ^ {T - 1} \mathrm {D} _ {\mathrm {K L}} \left(p (\cdot \mid z _ {t}, a _ {t}, x _ {t + 1}) \| q _ {\mathrm {z}} (\cdot \mid z _ {t}, a _ {t})\right) \right], \tag {4.3}
$$
where  $q_{\mathrm{r}}$  and  $q_{\mathrm{z}}$  are variational families representing beliefs over rewards  $r_t$  and latent representations  $z_{t+1}$ , respectively. We refer to  $z_{t+1} \sim p(\cdot \mid z_t, a_t, x_{t+1})$  as the latent posterior, because it conditions on the latest encoded observation  $x_{t+1} = h(o_{t+1})$ . We call the variational approximation  $q_{\mathrm{z}}(\cdot \mid z_t, a_t)$  the latent prior because it does not use the current observation  $o_{t+1}$  (or it's encoding  $x_{t+1}$ ) to determine  $z_{t+1}$ . Note that the right hand side of Eq. (4.3) depends on  $h$  through  $x_{t+1} = h(o_{t+1})$ , and thus gradients of this expression incorporate gradients through  $h$ .
The magic of Eq. (4.3). The upper bound in (4.3) reveals a striking feature which is at the core of our method: that, in order to reduce extraneous information in the latents  $z_{t}$  about observations  $o_{t}$ , it is enough to match the latent posterior  $z_{t + 1} \sim p(\cdot \mid z_{t},a_{t},x_{t + 1})$  to our latent prior  $q_{z}(\cdot \mid z_{t},a_{t})$  that does not condition on current  $x_{t + 1}$ . Elements that are spurious variations can be captured by  $p(\cdot \mid z_{t},a_{t},x_{t + 1})$ , but not by  $q_{z}(\cdot \mid z_{t},a_{t})$ , since  $q_{z}$  is not informed by the latest observation encoding  $x_{t + 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-prediction. We make these substitutions, and move to a penalty-formulation amenable to constrained optimization methods like dual-gradient descent [2]. The resulting objective we optimize to learn the latent posterior  $p$ , latent prior  $q_{\mathrm{z}}$ , reward predictor  $q_{\mathrm{r}}$  and observation encoder  $h$  jointly is:
$$
\max  _ {p, q _ {\mathrm {r}}, q _ {\mathrm {z}}, h} \min  _ {\beta} \mathbb {E} _ {p, h} \left[ \sum_ {t = 1} ^ {T} \log q _ {\mathrm {r}} \left(r _ {t} \mid z _ {t}\right) \right] + \beta \left(\mathbb {E} _ {p, h} \left[ \sum_ {t = 0} ^ {T - 1} \mathrm {D} _ {\mathrm {K L}} \left(p (\cdot \mid z _ {t}, a _ {t}, x _ {t + 1}) \| q _ {\mathrm {z}} (\cdot \mid z _ {t}, a _ {t})\right) \right] - \epsilon\right). \tag {4.4}
$$
Implementation details. We parameterize  $p$  and  $q$  using a recurrent state-space model (RSSM) [17]. The RSSM consists of an encoder  $h_{\theta}(x_t \mid o_t)$ , a latent dynamics model  $q_{\theta}(z_{t+1} \mid z_t, a_t)$  corresponding to the prior, a representation model  $p_{\theta}(z_{t+1} \mid z_t, a_t, x_{t+1})$  corresponding to the posterior, and a reward predictor  $q_{\theta}(r_t \mid z_t)$ . 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  $\mathrm{D}_{\mathrm{KL}}(p \parallel q) = \alpha \mathrm{D}_{\mathrm{KL}}(\lfloor p \rfloor \parallel q) + (1 - \alpha) \mathrm{D}_{\mathrm{KL}}(p \parallel \lfloor q \rfloor)$ , where  $\lfloor \cdot \rfloor$  denotes the stop gradient operator and  $\alpha \in [0,1]$  is the balancing parameter. With the removal of reconstruction, the KL balancing parameters become 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  $z$  by solving the optimization problem outlined in Eq. (4.4) to infer a latent posterior  $p(z_{t+1} \mid z_t, a_t, x_{t+1})$ , a latent prior  $q_z(z_{t+1} \mid z_t, a_t)$ , an encoder  $x_t = h(o_t)$  and a reward predictor  $q_r(r_t \mid z_t)$ , and (2) Policy Learning: using the inferred representation, dynamics model and reward predictor to learn a policy  $\pi_\phi(a_t \mid z_t)$  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  $V_\psi(z)$  is trained to predict the discounted cumulative reward given a latent state, and the actor  $\pi_\phi(a \mid z)$  is trained to take the action that maximizes the critic'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(z_{t+1} \mid z_t, a_t, x_{t+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.
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 a reweighted distribution constraint.