SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathrm {K L} \left[ q (z | a, o) \| p (z | a) \right] = \int_ {z} q (z | a, o) \ln \prod_ {t = 1} ^ {T} \frac {q (z _ {t} | z _ {t - 1} , a _ {t - 1} , o _ {t})}{p (z _ {t} | z _ {t - 1} , a _ {t - 1})} \mathrm {d} z \\ = \sum_ {t = 1} ^ {T} \left[ \int_ {z _ {1: t - 1}} \mathrm {d} z _ {1: t - 1} q \left(z _ {1: t - 1} \mid o, a\right) \int_ {z _ {t}} \mathrm {d} z _ {t} q \left(z _ {t} \mid z _ {t - 1}, a _ {t - 1}, o _ {t}\right) \ln \frac {q \left(z _ {t} \mid z _ {t - 1} , a _ {t - 1} , o _ {t}\right)}{p \left(z _ {t} \mid z _ {t - 1} , a _ {t - 1}\right)} \right. \\ \left. \int_ {z _ {t + 1: T}} \mathrm {d} z _ {t + 1: T} q (z _ {t + 1: T} | a, o) \right] \\ = \sum_ {t = 1} ^ {T} \mathbb {E} _ {q (z _ {t - 1} | a, o)} \mathrm {K L} [ q (z _ {t} | z _ {t - 1}, a _ {t - 1}, o _ {t}) \| p (z _ {t} | z _ {t - 1}, a _ {t - 1}) ]. \\ \end{array}
$$
3For brevity, we denote  $z_{1:T},a_{1:T},o_{1:T},r_{1:T}$  as  $z,a,o,r$  , respectively.
Algorithm 1 Contextualized World Models with In-the-wild Pre-training from Videos
1: // Action-free pre-training with in-the-wild videos  
2: Initialize parameters  $\theta$  of action-free dynamics model, image encoder, and decoder randomly  
3: Load in-the-wild video dataset  $\mathcal{D}$   
4: for each training step do  
5: // WORLD MODEL PRE-TRAINING  
6: Sample random minibatch  $\{o_{1:T}\} \sim \mathcal{D}$   
7: Get context  $c \gets o_{\tilde{t}}$ ,  $\tilde{t} \sim$  Uniform  $\{1,T\}$   
8: Update world model by minimizing  $\mathcal{L}^{\mathrm{CWM - pt}}(\theta)$  in Eq. (7)  
9: end for  
10: // Action-conditioned fine-tuning with MBRL  
11: Load pre-trained parameters  $\theta$  of action-free dynamics model, image encoder, and decoder  
12: Initialize parameters  $\phi, \varphi$  of action-conditioned dynamics model and reward predictors randomly  
13: Initialize parameters  $\psi, \xi$  of actor  $\pi_{\psi}(a|s)$  and critic  $v_{\xi}(s)$   
14: Initialize replay buffer  $\mathcal{B}$  with random seed episodes  
15: for each timestep  $t$  do  
16: // COLLECT TRANSITIONS  
17: Get representation  $z_t \sim q_\theta(z_t|z_{t-1}, a_{t-1}, o_t)$ ,  $s_t \sim q_\phi(s_t|s_{t-1}, a_{t-1}, z_t)$   
18: Get action  $a_t \sim \pi_\psi(a_t|s_t)$   
19: Add transition  $\{(o_t, a_t, r_t)\}$  to replay buffer  $\mathcal{B}$   
20: // WORLD MODEL FINE-TUNING  
21: Sample random minibatch  $\{(o_{1:T}, a_{1:T}, r_{1:T})\} \sim \mathcal{B}$   
22: Get context  $c \gets o_{\tilde{t}}$ ,  $\tilde{t} \sim$  Uniform  $\{1,T\}$   
23: Update world model by minimizing  $\mathcal{L}^{\mathrm{CWM}}(\phi, \varphi, \theta)$  in Eq. (6)  
24: // BEHAVIOR LEARNING  
25: Imagine future rollouts  $\{\hat{s}_{\tau}, \hat{a}_{\tau}, \hat{r}_{\tau} + \lambda \hat{r}_{\tau}^{\mathrm{int}}\}$  using world model and actor  
26: Update actor and critic by minimizing objectives in Eq. (8) and Eq. (10)  
27: end for
Overall, we have a lower bound on the log-likelihood of the data:
$$
\ln p (o, r | a, c) \geq \mathbb {E} _ {q (z | a, o)} \left[ \sum_ {t = 1} ^ {T} \big (\ln p (o _ {t} | z _ {t}, c) + \ln p (r _ {t} | z _ {t}) - \mathrm {K L} \big [ q (z _ {t} | z _ {t - 1}, a _ {t - 1}, o _ {t}) \| p (z _ {t} | z _ {t - 1}, a _ {t - 1}) ] \big) \right].
$$
We can also calculate that the tightness of the bound is KL  $[q(z|a,o)\| p(z|a,o,r,c)]$
C.1 Data Preprocessing
We use three common video datasets for in-the-wild video pre-training. All inputs are resized to  $64 \times 64$  pixels.
Something-Something-v2. The Something-Something-v2 dataset [17] shows footage of humans interacting with common everyday objects. We use all videos in the training set but filter out videos less than 25 frames long, resulting in a total of 162K videos for pre-training.
Human3.6M. The Human3.6M dataset [30] contains videos of various activities performed by 11 human subjects. We use the particular dataset processed by Pavlakos et al. [56], which includes a total of 840 videos with 210 scenarios and 4 different viewpoints for each scenario. Each video has around 500 frames. To make sure the images used for pre-training better fit our downstream tasks, we place the subject at the center of each frame by cutting a bounding box calculated from the 3D joint positions provided by the dataset. In addition, a padding of 200 pixels is performed around
<sup>4</sup>https://developer=qualcomm.com/software/ai-datasets/something-something
<sup>5</sup>http://vision.imar.ro/human3.6m/description.php
the bounding box so that the subject is properly positioned. The image clipped using the padded bounding box is then resized to  $64 \times 64$  pixels.
YouTube Driving. The YouTube Driving dataset [88] consists of real-world driving videos with different road and weather conditions. As videos in this dataset usually have a short opening, the first 100 frames of each video are skipped. The original video is downsampled 5 times. We employ all 134 videos in our pre-training process, with each video containing roughly 10000 to 20000 frames.
Assembled Dataset. We compose the three datasets mentioned above into an assembled dataset. On each sample, we uniformly choose a dataset from the three choices and sample videos from that dataset accordingly.
C.2 Benchmark Environments
Meta-world. Meta-world [85] is a benchmark of 50 distinct robotic manipulation tasks. Following APV [68], we use a subset of six tasks, namely Lever Pull, Printer Open, Door Lock, Button Press Topdown Wall, Reach, and Dial Turn. In all tasks, the episode length is 500 steps without any action repeat. The action dimension is 4, and the reward ranges from 0 to 10. All methods are trained over 250K environment steps, which is consistent with the setting of APV [68].
DMC Remastered. The DMC Remastered (DMCR) Suite [19] is a variant of the DeepMind Control Suite [73] with randomly generated graphics emphasizing visual diversity. On initialization of each episode for both training and evaluation, the environment makes a random sample for 7 factors affecting visual conditions, including floor texture, background, robot body color, target color, reflectance, camera position, and lighting. Our agents are trained and evaluated on three tasks: Cheetah Run, Hopper Stand, and Walker Run. All variation factors are used except camera position in the Hopper Stand task, where we find it too difficult for the agent to learn when the camera is randomly positioned and rotated. Following the common setup of DeepMind Control Suite [21, 79], we set the episode length to 1000 steps with an action repeat of 2, and the reward ranges from 0 to 1. All methods are trained over 1M environment steps.
CARLA. CARLA [14] is an autonomous driving simulator. We conduct a task using a similar setting as Zhang et al. [87], where the agent's goal is to progress along a highway in 1000 time-steps and avoid collision with 20 other vehicles moving along. The maximum episode length is set to 1000 without action repeat. Possible reasons to end an episode early include deviation from the planned route, driving off the road, and being stuck in traffic due to collisions. Each action in the CARLA environment is made up of a steering dimension and an accelerating dimension (or braking for negative acceleration). In addition to the reward function in Zhang et al., we add a lane-keeping reward to prevent our agent from driving on the shoulder. Our reward function is as follows:
$$
r _ {t} = v _ {\mathrm {e g o}} ^ {T} \hat {u} _ {h} \cdot \Delta t \cdot | \mathrm {c e n t e r} | - \lambda_ {1} \cdot \mathrm {i m p l u s e} - \lambda_ {2} \cdot | \mathrm {s t e e r} |
$$
The first term represents the effective distance traveled on the highway, where  $v_{\mathrm{ego}}$  is the velocity vector projected onto the highway's unit vector  $\hat{u}_h$ , multiplied by a discretized time-step  $\Delta t = 0.05$  and a lane-keeping term |center| to penalize driving on the shoulder. The impulse term represents collisions during driving, while the steer  $\in [-1,1]$  term prevents excessive steering. The hyperparameters  $\lambda_1$  and  $\lambda_2$  are set to  $10^{-3}$  and 1, respectively. The central camera of the agent, observed as  $64 \times 64$  pixels, is used as our observation. We use three different weather and sunlight conditions offered by CARLA for training and evaluation: ClearNoon, WetSunset, and dynamic weather. In the dynamic weather setting, a random weather and sunlight condition is sampled at the start of each episode and changes realistically during the episode. To aggregate per-weather scores, we normalize raw episode returns of each weather to comparable ranges. Namely, we linearly rescale episode returns from  $[-400,1200]$  under ClearNoon,  $[-400,800]$  under WetSunset, and  $[-400,600]$  under dynamic weather to a unified range of  $[0,1]$ . All methods are trained over 150K environment steps.
$^{6}$ https://github.com/metadriverse/ACO
stage name	output size	image encoder
conv_in	32×32	3×3, 48, stride 2
stage1	16×16	[3×3, 48/3×3, 48]×2
average pool 2×2, stride 2
stage2	8×8	[3×3, 96/3×3, 96]×2
average pool 2×2, stride 2
stage3	4×4	[3×3, 192/3×3, 192]×2
average pool 2×2, stride 2
Table 1: Architecture for image encoders in world models. Building residual blocks [27] are shown in brackets, with the number of blocks stacked. Numbers inside each bracket indicate the kernel size and output channel number of the corresponding building block, e.g.  $3 \times 3$ , 192 means one of the convolutions in the block uses a  $3 \times 3$  kernel and the number of output channels is 192. We use Batch Normalization [29] and ReLU activation function.
C.3 Model Details
Overall model. We adopt the stack latent model introduced by APV [68] and extend it into a Contextualized World Model (ContextWM) with the following components:
Action-free
$\int$  Representation model:  $z_{t}\sim q_{\theta}(z_{t}|z_{t - 1},o_{t})$
Transition model:  $\hat{z}_t\sim p_\theta (\hat{z}_t|z_{t - 1})$
Action-conditional
$\int$  Representation model:  $s_t \sim q_\phi(s_t | s_{t-1}, a_{t-1}, z_t)$
Transition model:  $\hat{s}_t\sim p_\phi (\hat{s}_t|s_{t - 1},a_{t - 1})$
Dual reward predictors
$\int$  Representative reward predictor:  $\hat{r}_t \sim p_\varphi(\hat{r}_t \mid s_t)$
$\lfloor$  Behavioral reward predictor:  $\hat{r}_t + \lambda \hat{r}_t^{\mathrm{int}}\sim p_\phi (\hat{r}_t + \lambda \hat{r}_t^{\mathrm{int}}\mid s_t)$
Contextualized image decoder:  $\hat{\sigma}_t\sim p_\theta (\hat{\sigma}_t|s_t,c)$  (11)
Specifically, the action-free representation model utilizes a ResNet-style image encoder to process observation input  $o_{t}$ . Moreover, the contextualized image decoder also incorporates a similar context encoder to extract contextual information from the context frame  $c$  and augment the decoder features with a cross-attention mechanism. We elaborate on the details below.
Image encoders and decoder. To ensure a sufficient capacity for both plain WM and ContextWM pre-trained on diverse in-the-wild videos, we implement image encoders and decoders as 13-layer ResNets [27]. For the sake of simplicity, we only show the architecture of the image encoder in Table 1, as the architecture of the context encoder is exactly the same, and the architecture of the image decoder is symmetric to the encoders. In the decoder, we use nearest-neighbor upsampling for unpooling layers. The outputs of the last residual block of two stages in the context encoder (stage2 and stage3) before average pooling (thus in the shape of  $16 \times 16$  and  $8 \times 8$  respectively) are passed to the corresponding residual block of the image decoder and used to augment the incoming decoder features with cross-attention.
Context cross-attention. The input features  $X \in \mathbb{R}^{c \times h \times w}$  of residual blocks in the image decoder are augmented with shortcut features  $Z \in \mathbb{R}^{c \times h \times w}$  of the same shape from the context encoder, by a cross-attention mechanism as Eq. (5). Here we elaborate on the details. After  $Z$  is flattened into a sequence of tokens  $\in \mathbb{R}^{hw \times c}$ , to mitigate the quadratic complexity of cross-attention, we random
mask  $75\%$  of the tokens, resulting tokens of the length  $\lfloor \frac{1}{4} hw \rfloor$ . Moreover, we use learned position embedding  $\in \mathbb{R}^{hw \times c}$  initialized with zeros for both  $Q$  and  $K$ ,  $V$ , respectively.
Context augmentation. We randomly sample a frame of observation from the trajectory segment as the input to the context encoder. To prevent the image decoder from taking a shortcut by trivially copying the context frame and thus impeding the training process, we apply random Cutout [89] as image augmentation on the context frame for all Meta-world tasks (except drawer open), as Meta-World observations have a static background and a fixed camera position. Note that we do not utilize any augmentation for DMCR and CARLA tasks. Our comparison with plain WM is fair since we do not introduce any augmentation into the image encoder of the latent dynamics models, thus the agent determines its actions solely based on the original observation.
Latent dynamics model. The stacked latent model from APV [68] is adopted to enable action-free pre-training and action-conditioned fine-tuning, where both the action-free and action-conditioned models are built upon the discrete latent dynamics model introduced in DreamerV2 [23], where the latent state consists of a deterministic part and a discrete stochastic part. Inputs to the latent dynamics model are representations from the image encoder. Following APV [68], we use 1024 as the hidden size of dense layers and the deterministic model state dimension.
Reward predictors, actor, and critic. Both of the dual reward predictors, as well as actor and critic, are all implemented as 4-layer MLPs with the hidden size 400 and ELU activation function, following DreamerV2 [23] and APV [68].
For the newly introduced hyperparameter, we use  $\beta_r = 1.0$  for all tasks. We set  $\beta_z = 1.0$  for pre-training, and  $\beta_z = 0.0$ ,  $\beta_s = 1.0$  for fine-tuning, following APV [68]. Despite we adopt a deeper architecture for both image encoders and decoders, we find that hyperparameters are relatively robust to architectures, and the new deeper architecture under the same hyperparameters can match or even slightly improve the performance originally reported by APV. Thus, we use the same hyperparameters with APV. For completeness, we report important hyperparameters in Table 2.
C.5 Qualitative Analysis Details
We explain the visualization scheme for Figure 8. For video prediction in Figure 8a, we sample a video clip of length 25 and let the model observe the first 15 frames and then predict the future 10 frames open-loop by the latent dynamics model. The context frame is selected as the last of the observed 15 frames. To visualize context attention, we compute the cross-attention weight across the  $16 \times 16$  feature maps without random masking and plot a heatmap of the weights averaged over attention heads and over attention targets of a small region (for example,  $3 \times 3$ ) on the decoder feature map. For video representations in Figure 8b, we sample video clips of length 25 and compute the averaged model states  $\mathrm{avg}(z_{1:T})$  of dimension 2048 for each video clip, which are then visualized by t-SNE [74].
C.6 Computational Resources
We implement all the methods based on PyTorch [54] and train with automatic mixed precision. In terms of parameter counts, ContextWM consists of 26M and 47M parameters for video pre-training and MBRL, respectively, while a plain WM consists of 24M and 44M parameters. In terms of training time, it takes  $\sim 37$  hours for pre-training of ContextWM over 600K iterations, and fine-tuning of ContextWM with MBRL requires  $\sim 24$  hours for each run of Meta-world experiments over 250K environment steps,  $\sim 23$  hours for each run of DMC Remastered experiments over 1M environment steps and  $\sim 16$  hours for each run of CARLA experiments over 150K environment steps, respectively. Although ContextWM introduces new components of the model, we find it does not significantly impact the training time, since the context encoder only needs to operate on one frame of observations, and the cross-attention between features also accounts for a relatively small amount of computation. In terms of memory usage, Meta-world experiments require  $\sim 18\mathrm{GB}$  GPU memory, and DMC Remastered and CARLA experiments require  $\sim 7\mathrm{GB}$  GPU memory, thus the experiments can be done using typical 24GB and 12GB GPUs, respectively.
Table 2: Hyperparameters in our experiments. We use the same hyperparameters as APV [68].
Hyperparameter	Value
Pre-training from Videos	Image size	64 × 64 × 3
Image preprocess	Linearly rescale from [0, 255] to [-0.5, 0.5]
Video segment length T	25
KL weight βz	1.0
Optimizer	Adam [37]
Learning rate	3 × 10-4
Batch size	16
Training iterations	6 × 105for SSv2 / Human / YouTubeDriving
1.2 × 106for Assembled dataset
Fine-tuning with MBRL	Observation size	64 × 64 × 3
Observation preprocess	Linearly rescale from [0, 255] to [-0.5, 0.5]
Trajectory segment length T	50
Random exploration	5000 environment steps for Meta-world
1000 environment steps for DMCR and CARLA
Replay buffer capacity	106
Training frequency	Every 5 environment steps
Action-free KL weight βz	0.0
Action-conditional KL weight βs	1.0
Representative reward predictor weight βr	1.0
Intrinsic reward weight λ	1.0 for Meta-world
0.1 for DMCR and CARLA
Imagination horizon H	15
Discount γ	0.99
λ-target discount	0.95
Entropy regularization η	1 × 10-4
Batch size	50 for Meta-world
16 for DMCR and CARLA
World model optimizer	Adam
World model learning rate	3 × 10-4
Actor optimizer	Adam
Actor learning rate	8 × 10-5
Critic optimizer	Adam
Critic learning rate	8 × 10-5
Evaluation episodes	10 for Meta-world and DMCR
5 for CARLA
D.1 Comparison with Additional Baselines
We compare our approach with a state-of-the-art model-free RL method DrQ-v2 [79] and a model-based RL method Iso-Dream [52]. We adapt their official implementations $^{78}$  to all the tasks in our experiments and present the results in Figure 9. As shown, our method consistently outperforms state-of-the-art baselines, which demonstrates the merit of pre-training from in-the-wild videos. Model-free methods such as DrQ-v2 have been shown to perform worse than model-based methods [68, 79], and our results are consistent with this. We also remark that our method differs from Iso-Dream [52] since we separate modeling of context and dynamics at the semantic level while Iso-Dream isolates controllable and noncontrollable parts at the pixel level. Thus, contextual information such as body color has to be modeled by the dynamics branch in Iso-dream, which probably results in worse decoupling of context and dynamics as well as inferior performance for challenging tasks. Moreover, compared to Iso-Dream, our model makes a simpler adjustment to Dreamer architecture.
<sup>8</sup>https://github.com/panmt/Iso-Dream
(a) Aggregate results
(b) Meta-world
(c) DMC Remastered
Figure 9: Comparison with additional baselines, DrQ-v2 [79] and Iso-Dream [52]. (a) Aggregate performance of different methods across all tasks on each domain. (b)(c)(d) Learning curves of different methods on each task from Meta-world, DMC Remastered, and CARLA, respectively. We report aggregate results across eight runs for each task.
(d) CARLA
Figure 10: Learning curves of IPV with ContextWM on Meta-world manipulation tasks when only the pre-trained parameters of the visual encoders and decoder are transferred, i.e., without transferring dynamics knowledge obtained by pre-training (denoted by Enc/Dec Only). We report success rates, aggregated across eight runs for each task.
D.2 Effects of Dynamics Knowledge
To investigate the importance of dynamics knowledge captured by pre-trained ContextWM, we report the performance of IPV with ContextWM when only the pre-trained parameters of the visual encoders and decoder are transferred. Figure 10 indicates that only transferring visual encoders and decoder performs worse on four of six tasks. These results align with the observation of APV [68] and highlight that pre-trained world models are capable to capture essential dynamics knowledge beyond learning visual representations solely [49, 53], further boosting visual control learning.
Particularly strong performance gains on DMCR tasks. Our method obtains considerable performance gains on DMCR tasks. The main reason is that DMCR is a purposefully designed benchmark, which measures visual generalization and requires the agent to extract task-relevant information as well as ignore visual distractors [19]. As demonstrated in Figure 8c, our ContextWM has the advantage of separately modeling contexts (task-irrelevant in DMCR) and dynamics (task-relevant in DMCR), which avoids wasting the capacity of dynamics models in modeling low-level visual details. Furthermore, pre-training with in-the-wild videos enables our models to eliminate diverse distractors and capture shared motions, which is essential for visual generalization in RL. In contrast, a plain WM needs to model complicated contexts and dynamics in an entangled manner, which adds difficulty to dynamics learning and behavior learning on these features. We also emphasize that, motivated by separating contexts and enhancing temporal dynamics modeling, our proposed IPV w/ ContextWM is a general-purpose framework and, as shown in our experiments, can obtain adequate performance gains on various benchmarks that have more complicated entangling of contexts and dynamics beyond DMCR.
Discussion with Transformer-based world models. While a cross-attention mechanism is leveraged for context information conditioning, we still use an RNN-based latent dynamics model from Dreamer [21] to model history observations. More powerful sequential backbones such as Transformers are orthogonal to our primary technical contribution of explicit context modeling and conditioning. Nevertheless, exploring Transformer-based world models [10, 47, 58] is valuable, not only because of its favorable scalability but also of its flexibility to incorporate pre-training on action-free videos by simply changing conditioning variables and to condition the contextualized decoder in a more native way. Overall, the combination of our pre-training framework with a Transformer architecture holds great potential, and we will delve deeper into this aspect in future work.
F Additional Visualizations
Video prediction. We present additional showcases of video prediction by plain WM and ContextWM, respectively, on three video datasets, in Figure 11. On the challenging Something-Something-v2 and YouTube Driving datasets, our model provides better prediction quality in comparison to a plain WM and well captures dynamics information (e.g., how objects are moving) with concentrated cross-attention. However, on the Human3.6M dataset, although our model better predicts human motions, we find its context cross-attention is relatively divergent. This result supports our conjecture that our model pre-trained on the Human3.6M dataset lacking in diversity may suffer from overfitting, which results in insufficient modeling of context and dynamics as well as inferior performance when fine-tuned with MBRL.
(a) Something-Something-V2
(b) Human3.6M
(c) YouTube Driving
Figure 11: Future frames predicted by plain WM and ContextWM, respectively, on three in-the-wild video datasets. We also visualize context attention weight from a target region of prediction (red box) to the context frame. We refer to Appendix C.5 for details.
Video representation. We present in Figure 12 additional t-SNE visualization similar to Figure 8. Our ContextWM learns well-distributed representations that clearly separate motion directions while plain WM fails.
Compositional decoding. Figure 13 includes additional compositional decoding analysis on the DMCR Walker Run task.
Figure 12: Additional t-SNE visualization of average pooled representations of SSv2 videos with two distinct labels, from a pre-trained plain WM and ContextWM, respectively.
Figure 13: Additional compositional decoding analysis on the DMCR Walker Run task.