SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 2 Semi-Supervised Adaptation of Visual Encoder Using Support Constraint
1: Fix training encoder  $h_\theta$ , model  $p_\theta$ , policy  $\pi_\phi$  and replay buffer  $\mathcal{D}_{\mathrm{train}}$ .  
2: Initialize test-time replay buffer  $\mathcal{D}_{\mathrm{test}}$  with  $S$  random seed episodes.  
3: Initialize test-time encoder  $h_{\theta'}$  by setting  $\theta' = \theta$ .  
4: Initialize neural networks  $\tau_\rho$  and  $f_\omega$  randomly and set dual variable  $\lambda = \lambda_0$ .  
5: Collect  $N$  calibration trajectories with paired observations  $\mathcal{D}_{\mathrm{cal}} \gets \{(o_t^{\mathrm{train}}, o_t^{\mathrm{test}}, a_t, r_t)\}_{t=1}^T$  using an exploration policy.  
6: while not converged do  
7: for update step  $c = 1 \dots C$  do  
8: // Adaptation  
9: Draw  $B$  observations from training and test buffers  $o_{1:B}^{\mathrm{train}} \sim \mathcal{D}_{\mathrm{train}}, o_{1:B}^{\mathrm{test}} \sim \mathcal{D}_{\mathrm{test}}$ .  
10: Encode images using respective encoders  $x_{1:B}^{\mathrm{train}} = h_\theta(o_{1:B}^{\mathrm{train}})$ ,  $x_{1:B}^{\mathrm{test}} = h_{\theta'}(o_{1:B}^{\mathrm{test}})$ .  
11: Compute the Lagrangian objective in (4.7)  
12: Draw  $B$  pairs of observations from the calibration buffer  $\{(o_i^{\mathrm{cal\_train}}, o_i^{\mathrm{cal\_test}})\}_{i=1}^B \sim \mathcal{D}_{\mathrm{cal}}$ .  
13: Encode images using encoders  $x_{1:B}^{\mathrm{cal\_train}} = h_\theta(o_{1:B}^{\mathrm{train}})$ ,  $x_{1:B}^{\mathrm{cal\_test}} = h_{\theta'}(o_{1:B}^{\mathrm{cal\_test}})$ .  
14: Compute MSE loss between calibration encodings and add to the Lagrangian objective.  
15: Update  $\theta', \rho, \omega$  and  $\lambda$  using dual gradient descent.  
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 | 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{test}} \gets \mathcal{D}_{\mathrm{test}} \cup \{o_t, a_t, r_t\}_{t=1}^N$ .
B.1 RePo objective
Recall that  $\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}}$  (our class of posteriors), 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$ .
We derive a variational lower bound for the following information bottleneck objective from eq. (4.1).
$$
\max  _ {p, h} \mathrm {I} _ {p, h} (z _ {1: T}; r _ {1: T} \mid a _ {1: T}) \text {s . t .} \mathrm {I} _ {p, h} (z _ {1: T}; o _ {1: T} \mid a _ {1: T}) <   \epsilon
$$
where the mutual information is with respect to the distribution described by the expectation  $\mathbb{E}_{p,h}$  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,  $p$  is only a latent dynamics model. Thus, for the derivation we will introduce  $\tilde{p}$  as corresponding to the distribution of all variables under  $\mathbb{E}_{p,h}$ , which in particular takes into account the distribution from  $\mathcal{D}_{\mathrm{buf}}$ .
Definition B.1. We let  $\tilde{p}$  denote the joint distribution of  $(z_{1:T}, a_{1:T}, x_{1:T}, o_{1:T})$  under  $\mathbb{E}_{p,h}$ .
Lower bound on reward prediction. We first derive a variational lower bound for the objective  $\mathrm{I}_{p,h}(z_{1:T}; r_{1:T} \mid a_{1:T})$ . For our derivation, we make a simplifying assumption:
Assumption B.1. Assume that there is a function  $p_{\mathrm{r}}(r_t \mid z_t)$  such that  $\tilde{p}(r_t \mid z_{1:t}, a_{1:t}) = p_{\mathrm{r}}(r_t \mid z_t)$ .
The simplifying assumption enforces that  $z_{t}$  is sufficient for the reward given past latents, actions, and the current action  $a_{t}$ . 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:
$$
\begin{array}{l} \mathrm {I} _ {p, h} \left(z _ {1: T}; r _ {1: T} \mid a _ {1: T}\right) \\ = \mathbb {E} _ {p, h} \left[ \log \tilde {p} \left(r _ {1: T} \mid z _ {1: T}, a _ {1: T}\right) \right] + \mathrm {H} _ {p, h} \left(r _ {1: T} \mid a _ {1: T}\right) \\ \stackrel {\pm} {=} \mathbb {E} _ {p, h} \left[ \log \tilde {p} \left(r _ {1: T} \mid z _ {1: T}, a _ {1: T}\right) \right] \\ = \mathbb {E} _ {p, h} \left[ \sum_ {t = 1} ^ {T} \log p _ {\mathrm {r}} \left(r _ {t} \mid z _ {t}, a _ {t}\right)\right) \tag {Assumption B.1} \\ = \mathbb {E} _ {p, h} \left[ \sum_ {t = 1} ^ {T} \log q _ {\mathrm {r}} \left(r _ {t} \mid z _ {t}\right)\right) + \sum_ {t = 1} ^ {T} \mathrm {D} _ {\mathrm {K L}} \left(p _ {\mathrm {r}} \left(r _ {t} \mid z _ {t}\right) \| q _ {\mathrm {r}} \left(r _ {t} \mid z _ {t}\right)\right) \\ \geq \mathbb {E} _ {p, h} \left[ \sum_ {t = 1} ^ {T} \log q _ {\mathrm {r}} \left(r _ {t} \mid z _ {t}\right)\right), \\ \end{array}
$$
where above  $\mathrm{H}_{p,h}$  denotes an entropy term under distribution  $\mathbb{E}_{p,h}$ , and  $\stackrel{+}{\equiv}$  denotes equality up to a constant that does not depend on our choice of parametrization  $p$ , which only dictates latents  $z_{t}$ . In the fourth line, we add and subtract  $\log q_{\mathrm{r}}(r_t|z_t)$  in each term of the summation. The last step uses the nonnegativity of KL divergence.
Upper bound on dynamic compression. We proceed to derive an upper bound for the constraint  $\mathrm{I}_{p,h}(z_{1:T};o_{1:T}\mid a_{1:T})$ . We make a simplifying assumption analogous to Assumption B.1; again, this assumption can also be discarded, but illustrates that our lower bound is sharper when this assumption holds.
Assumption B.2. We assume that there is a function  $p_{\mathrm{z}}(z_{t + 1} \mid z_t, a_t)$  such that  $\tilde{p}(z_{t + 1} \mid z_{1:t}, a_{1:t}) = p_{\mathrm{z}}(z_{t + 1} \mid z_t, a_t)$ .
With Assumption B.2, we then invoke a variational upper bound which replaces  $p_{\mathrm{z}}(z_{t + 1} \mid z_t, a_t)$  with a variational approximation  $q_{\mathrm{z}}(z_{t + 1} \mid z_t, a_t)$ :
$$
\begin{array}{l} \mathrm {I} _ {p, h} \left(z _ {1: T}; o _ {1: T} \mid a _ {1: T}\right) \\ = \mathbb {E} _ {p, h} \left[ \log \tilde {p} \left(z _ {1: T} \mid o _ {1: T}, a _ {1: T}\right) - \log \tilde {p} \left(z _ {1: T} \mid a _ {1: T}\right) \right] \\ = \mathbb {E} _ {p, h} \left[ \left(\sum_ {t = 1} ^ {T} \log p \left(z _ {t + 1} \mid z _ {t}, a _ {t}, x _ {t + 1}\right)\right) - \log \tilde {p} \left(z _ {1: T} \mid a _ {1: T}\right) \right] (\text {d e f i n i t i o n o f l a t e n t d y n a m i c s}) \\ = \mathbb {E} _ {p, h} \left[ \sum_ {t = 1} ^ {T} \log p \left(z _ {t + 1} \mid z _ {t}, a _ {t}, x _ {t + 1}\right) - \sum_ {t = 1} ^ {T} \log p _ {\mathrm {z}} \left(z _ {t + 1} \mid z _ {t}, a _ {t}\right) \right] \tag {Assumption B.2} \\ = \mathbb {E} _ {p, h} \left[ \sum_ {t = 1} ^ {T} \log p (z _ {t + 1} \mid z _ {t}, a _ {t}, x _ {t + 1}) - \sum_ {t = 1} ^ {T} \log q _ {\mathrm {z}} (z _ {t + 1} \mid z _ {t}, a _ {t}) \right] - \sum_ {t = 1} ^ {T} \mathrm {D} _ {\mathrm {K L}} (p _ {\mathrm {z}} (z _ {t + 1} \mid z _ {t}, a _ {t}) \| q _ {\mathrm {z}} (z _ {t + 1} \mid z _ {t}, a _ {t})) \\ \leq \mathbb {E} _ {p, h} \left[ \sum_ {t = 1} ^ {T} \log p \left(z _ {t + 1} \mid z _ {t}, a _ {t}, x _ {t + 1}\right) - \sum_ {t = 1} ^ {T} \log q _ {z} \left(z _ {t + 1} \mid z _ {t}, a _ {t}\right) \right], \\ \end{array}
$$
where in the last two lines we add and subtract  $q_{z}(z_{t + 1} \mid z_{t}, a_{t})$  and apply the nonnegativity of KL divergence. In the third line, we use the fact that our transition dynamics model is Markov, with
$$
\tilde {p} \left(z _ {t + 1} \mid o _ {1: t}, a _ {1: t}, z _ {1: t}\right) = p \left(z _ {t + 1} \mid h \left(o _ {t + 1}\right), a _ {t}, z _ {t}\right) = p \left(z _ {t + 1} \mid x _ {t + 1}, a _ {t}, z _ {t}\right), \tag {B.1}
$$
where under  $\mathbb{E}_{p,h}$ ,  $x_{t} \equiv h(o_{t})$  holds by definition. Notice that the same identity would be true if instead we tried to lower bound the mutual information  $\mathrm{I}_{p,h}(z_{1:T}; o_{1:T} \mid a_{1:T})$ , showing that our objective is agnostic to whether we consider compressing information with respect to  $x_{1:T}$  or  $o_{1:T}$ . This may be seem surprising, but is a simple consequence of choosing a latent dynamics model which depends on  $o_{1:T}$  only through  $x_{1:T}$ .
Combining the derivations. Combining these results and writing in Lagrangian form, we arrive at the objective in eq. (4.4)
$$
\max_{p,q_{\mathrm{r}},q_{\mathrm{z}},h}\min_{\beta}\mathbb{E}_{p,h}\left[\sum_{t = 1}^{T}\log q_{\mathrm{r}}(r_{t}\mid z_{t})\right] + \beta \left(\mathbb{E}_{p,h}\left[\sum_{t = 0}^{T - 1}\mathrm{D}_{\mathrm{KL}}(p(\cdot \mid z_{t},a_{t},x_{t + 1}))\parallel q_{\mathrm{z}}(\cdot \mid z_{t},a_{t}))\right] - \epsilon\right).
$$
B.2 Tractable objective for support constraint
We derive a tractable variant for the support matching objective in eq. (4.6):
$$
\min  _ {\tau (\cdot) \geq 0, h _ {\mathrm {t e s t}} (\cdot)} \mathrm {D} (\tau \cdot \mathcal {P} _ {\mathrm {t r a i n}} \| \mathcal {P} _ {\mathrm {t e s t}}) \mathrm {s . t .} \mathbb {E} _ {x \sim \mathcal {P} _ {\mathrm {t r a i n}}} [ \tau (x) ] = 1,
$$
where  $\mathcal{P}_{\mathrm{train}}$  denotes the distribution of image encodings at training time and  $\mathcal{P}_{\mathrm{test}}$  denotes the distribution of image encodings at test time.  $\tau$  is a reweighting function and the constraint  $\mathbb{E}_{x\sim \mathcal{P}_{\mathrm{train}}}\big[\tau (x)\big] = 1$  ensures that  $\tau \cdot \mathcal{P}_{\mathrm{train}}$  is a valid distribution. Let  $\mathrm{D}_{\phi}$  be any  $f$ -divergence, we have
$$
\begin{array}{l} \min  _ {\tau (\cdot) \geq 0, h _ {\text {t e s t}} (\cdot)} \mathrm {D} _ {\phi} (\tau \cdot \mathcal {P} _ {\text {t r a i n}} \| \mathcal {P} _ {\text {t e s t}}) \\ = \min  _ {\tau (\cdot) \geq 0, h _ {\text {t e s t}} (\cdot)} \int_ {x} p _ {\text {t e s t}} (x) \phi \left(\frac {p _ {\text {t r a i n}} (x) \cdot \tau (x)}{p _ {\text {t e s t}} (x)}\right) \quad \text {(d e f i n i t i o n o f} f \text {- d i v e r g e n c e}) \\ = \min  _ {\tau (\cdot) \geq 0, h _ {\text {t e s t}} (\cdot)} \int_ {x} p _ {\text {t e s t}} (x) \left\{\max  _ {f (\cdot)} \left[ \frac {p _ {\text {t r a i n}} (x) \cdot \tau (x)}{p _ {\text {t e s t}} (x)} f (x) - \phi^ {*} (f (x)) \right] \right\} \quad (\text {b y}) \\ = \min  _ {\tau (\cdot) \geq 0, h _ {\text {t e s t}} (\cdot)} \max  _ {f (\cdot)} \int_ {x} p _ {\text {t e s t}} (x) \left[ \frac {p _ {\text {t r a i n}} (x) \cdot \tau (x)}{p _ {\text {t e s t}} (x)} f (x) - \phi^ {*} (f (x)) \right] \quad \left(\text {b y}\right) \\ = \min  _ {\tau (\cdot) \geq 0, h _ {\text {t e s t}} (\cdot)} \max  _ {f (\cdot)} \mathbb {E} _ {\mathcal {P} _ {\text {t r a i n}}} [ \tau (x) \cdot f (x) ] - \mathbb {E} _ {\mathcal {P} _ {\text {t e s t}}} [ \phi^ {*} (f (x)) ] \\ \end{array}
$$
We follow [4] and use the  $\chi^2$  divergence where  $\phi(x) = (x - 1)^2$  and  $\phi^*(y) = y + \frac{y^2}{4}$ . This gives
$$
\min _ {\tau (\cdot) \geq 0, h _ {\mathrm {t e s t}} (\cdot)} \max _ {f (\cdot)} \mathbb {E} _ {\mathcal {P} _ {\mathrm {t r a i n}}} [ \tau (x) \cdot f (x) ] - \mathbb {E} _ {\mathcal {P} _ {\mathrm {t e s t}}} \left[ f (x) + \frac {1}{4} f (x) ^ {2} \right].
$$
Incorporating the constraint as a Lagrange multiplier, we arrive at the objective in eq. (4.7).
$$
\min _ {\tau (\cdot) \geq 0, h _ {\mathrm {t e s t}} (\cdot) \atop f (\cdot), \lambda} \max _ {f (\cdot), \lambda} \mathbb {E} _ {\mathcal {P} _ {\mathrm {t r a i n}}} [ \tau (x) \cdot f (x) ] - \mathbb {E} _ {\mathcal {P} _ {\mathrm {t e s t}}} \left[ f (x) + \frac {1}{4} f (x) ^ {2} \right] + \lambda (\mathbb {E} _ {\mathcal {P} _ {\mathrm {t r a i n}}} [ \tau (x) ] - 1).
$$
Model architecture We base our implementation of RePo on Dreamer [17]. We fix the image size to  $64 \times 64$  and parameterize the image encoder using a 4-layer CNN with  $\{32, 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  $q_{\mathrm{z}}(z_{t+1} \mid z_t, a_t)$  as a Gaussian whose mean and standard deviation are predicted by passing the next belief through a 2-layer MLP. Similarly, we parametrize the posterior  $p(z_{t+1} \mid z_t, a_t, x_{t+1})$  as a Gaussian whose mean and standard deviation are predicted by passing the next belief along with the next image embedding through a 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 as 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 a 4-layer MLP outputting a squashed Gaussian distribution. We use 200 hidden units per layer and ELU activation for all MLPs.
Table 2: Hyperparameters for evaluation tasks
β0	ε	r
Walker	1e-5	3	5
Cheetah	1e-5	3	5
Hopper	1e-4	1	3
Cartpole	1e-4	3	4
Maniskill	1e-4	3	4
TurtleBot	1e-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 to predict the generalized value estimate [17] with  $\lambda = 0.95$  and discount factor  $\gamma = 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  $\beta = \beta_0$  and set its learning rate to 1e-4. We cast the KL balancing parameter  $\alpha$  to the ratio  $r$  between the number of prior training steps to posterior training steps, where  $r$  translates to  $\alpha = \frac{r}{r + 1}$ . We tune the initial Lagrange multiplier  $\beta_0$ , target KL  $\epsilon$ , and KL balancing ratio  $r$  on our evaluation tasks and report the best hyperparameters in Table 2.
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 use the official implementation along with their reported hyperparameters. We use a deterministic transition model for DBC as we find it to perform better than using a stochastic transition model.
Adaptation We parametrize the reweighting function  $\tau$  as 2-layer MLP with 256 hidden units in each layer and ReLU activation. The dual function  $f$  is analogous to the discriminator in a GAN architecture [14]. To prevent vanishing gradient, we parameterize  $f$  using a variational discriminator bottleneck [40] with 2 hidden layers of size 256 and a 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-time encoder with the weights of the training encoder and adapt it online, performing 100
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  $\ell_2$  loss between their encodings. All components are optimized using Adam [29] optimizer. We use a learning rate of 3e-4 for  $h_{\mathrm{test}}$ , 5e-5 for  $\tau$ , 1e-4 for  $f$ , and 5e-3 for the Lagrange multiplier  $\lambda$  initialized to 1e-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-Shannon 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 KL divergence between the prior and the posterior. We minimize dynamics violation by backpropagating the gradient through the model to the encodings.
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 a 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 manipulation tasks adapted from the ManiSkill 2 benchmark [15]:
- Push Cube: the goal is to push a cube to reach a position on the floor.
- Lift Cube: the goal is to lift a cube above a certain height.
- Turn Facet: 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 at the beginning of each episode. We use a time limit of 100 and an action repeat of 1. We evaluate all methods with 4 random seeds.
Lazy TurtleBot To evaluate RePo's resilience against spurious variations in the real world, we furnish a room to mimic a typical household setup and train a TurtleBot to reach a certain goal location from egocentric observations (Fig. 4). We introduce spurious variations by placing two TVs playing random YouTube 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)$ ,  $(-1,0)$ ,  $(0,-1)$  respectively, and continuous actions are mapped back to the closest discrete actions in L2 distance. The reward is the negative distance to goal, where the robot's state is estimated using an overhead camera mounted on the ceiling. We apply a time limit of 100 and an action repeat of 1. At the start of training, we prefetch the replay buffer with 5K offline transitions collected by sampling and reaching random goals using a state-based controller. This is followed by 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, we evaluate each method for 2 seeds.
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.
Figure 11: Comparison with Dreamer on standard DeepMind control environments. RePo matches Dreamer in asymptotic performance despite not being trained with a reconstruction objective.
E.2 Results on Distracted Walker Walk with scoreboard
We evaluate RePo on a variant of Distracted Walker Walk which consists of a scoreboard displaying the cumulative reward of the current episode. The goal is to investigate if RePo's reward-predictive nature causes its latent representation to collapse onto predicting the score. We present the probing 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 not only the current reward but also the dynamics and future rewards. The empirical performance with the scoreboard is  $872.86 \pm 43.71$  while without the scoreboard it is  $868.72 \pm 53.93$ .
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 on distracted DMC Walker Walk to analyze the effect of hyperparameters. All other hyperparameters are fixed to those reported in Table. 2. Fig. 13 illustrates the effect of  $\beta$  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 is too 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 to discard redundant information in the representation. Fig. 14 show the result of ablating the KL constraint value  $\epsilon$ . With too small a KL target, the representation becomes degenerate and fails to capture task-relevant information, as indicated by the increase in reward loss. A 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. In accordance with [19], we find that the dynamics model is generally harder to train than the representation model. Hence, a larger KL balancing parameter tends to result in better performance. However, if we train the prior too aggressively, then the posterior becomes poorly regularized and captures task-irrelevant information.
Figure 13: Ablating  $\beta$  learning rate. A larger  $\beta$  learning rate corresponds to faster convergence of the constraint, which can lead to ineffective exploration.
Figure 14: Ablating information bottleneck  $\epsilon$ . A tighter information bottleneck generally induces more parsimonious representations and more accurate dynamics models. But too tight a bottleneck can thwart reward learning.
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 a higher KL balancing ratio. However, training the prior too aggressively leads to poorly regularized posterior.
E.4 An illustrative example of support constraint
We provide a didactic example to illustrate the intuition behind using a support constraint for alignment. Consider two ground truth latent distributions shown in Fig. 16. The training distribution  $\mathcal{P}_{\mathrm{train}}$  is a uniform distribution spanning  $[0,6]\times [0,6]$ , and the test-time distribution  $\mathcal{P}_{\mathrm{test}}$  is a uniform distribution covering exactly half the support of  $\mathcal{P}_{\mathrm{train}}$ . This is to simulate insufficient online exploration at test time. We construct a nonlinear emission function  $f:\mathcal{X}\to \mathcal{O}:f(x) = 0.025e^{x}$  to generate the observations. For illustration purposes, we use the same emission function for both training