SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathcal {L} ^ {\mathrm {C W M}} (\phi , \varphi , \theta) \doteq \underbrace {\mathbb {E} _ {q _ {\phi} (s _ {1 : T} | a _ {1 : T} , z _ {1 : T}) , q _ {\theta} (z _ {1 : T} | o _ {1 : T})}} _ {\text {c o n t e x t - u n a w a r e l a t e n t i n f e r e n c e}} \left[ \sum_ {t = 1} ^ {T} \left(\underbrace {- \ln p _ {\theta} (o _ {t} | s _ {t} , c)} _ {\text {c o n t e x t u a l i z e d i m a g e l o s s}}\right) \right. \\ \underbrace {- \ln p _ {\phi} \left(r _ {t} + \lambda r _ {t} ^ {\text {i n t}} \mid s _ {t}\right)} _ {\text {b e h a v i o r a l r e w a r d l o s s}} \underbrace {- \beta_ {r} \ln p _ {\varphi} \left(r _ {t} \mid s _ {t}\right)} _ {\text {r e p r e s e n t a t i v e r e w a r d l o s s}} \underbrace {+ \beta_ {z} \mathbf {K L} \left[ q _ {\theta} \left(z _ {t} \mid z _ {t - 1} , o _ {t}\right) \| p _ {\theta} \left(\hat {z} _ {t} \mid z _ {t - 1}\right) \right]} _ {\text {a c t i o n - f r e e K L l o s s}} \tag {6} \\ \underbrace {+ \beta_ {s} \operatorname {K L} \left[ q _ {\phi} \left(s _ {t} \mid s _ {t - 1} , a _ {t - 1} , z _ {t}\right) \| p _ {\phi} (\hat {s} _ {t} \mid s _ {t - 1} , a _ {t - 1}) \right]} _ {\text {a c t i o n - c o n d i t i o n a l K L l o s s}} \Biggr) \Biggr ]. \\ \end{array}
$$
When pre-trained from in-the-wild videos, ContextWM can be optimized by minimizing an objective that drops action-conditioned dynamics and reward predictors from Eq. (6), similar to APV [68]. For behavior learning, we adopt the actor-critic learning scheme of DreamerV2 [23]. For a complete description of pre-training and fine-tuning of ContextWM for MBRL, we refer to Appendix A.
We conduct our experiments on various domains to evaluate In-the-wild Pre-training from Videos (IPV) with Contextualized World Models (ContextWM), in contrast to plain world models (WM) used by DreamerV2 and APV  $^{2}$ . Our experiments investigate the following questions:
- Can IPV improve the sample efficiency of MBRL?
- How does ContextWM compare to a plain WM quantitatively and qualitatively?
- What is the contribution of each of the proposed techniques in ContextWM?
- How do videos from different domains or of different amounts affect IPV with ContextWM?
Visual control tasks. As shown in Figure 4, we formulate our experiments on various visual control domains. Meta-world [85] is a benchmark of 50 distinct robotic manipulation tasks and we use the same six tasks as APV [68]. DMC Remastered [19] is a challenging extension of the widely used robotic locomotion benchmark, DeepMind Control Suite [73], by expanding a complicated graphical variety. We also conduct an autonomous driving task using the CARLA simulator [14], where an agent needs to drive as far as possible along Town04's highway without collision in 1000 timesteps. See Appendix C.2 for details on visual control tasks.
In our experiments, we rebrand the APV baseline [68] with the new name 'IPV w/ Plain WM' to note that it is pre-trained on In-the-wild data rather than only Action-free data.
Figure 4: Example image observations of our visual control benchmark tasks: Meta-world, DMControl Remastered, and CARLA (left to right).
Pre-training datasets. We utilize multiple in-the-wild video datasets which can potentially benefit visual control. Something-Something-v2 (SSv2) [17] dataset contains 193K videos of people interacting with objects. Human3.6M [30] dataset contains videos of human poses over 3M frames under 4 different viewpoints. YouTube Driving [88] dataset collects 134 real-world driving videos, over 120 hours, with various weather conditions and regions. We also merge the three datasets to construct an assembled dataset for the purpose of pre-training a general world model. See Figure 1 for example video frames and Appendix C.1 for details on data preprocessing.
Implementation details. For the newly introduced hyperparameter, we use  $\beta_r = 1.0$  in Eq. (6) for all tasks. To ensure a sufficient capacity for both plain WM and ContextWM pre-trained on diverse in-the-wild videos, we implement visual encoders and decoders as 13-layer ResNets [27]. Unless otherwise specified, we use the same hyperparameters with APV. See Appendix C.3 for more details.
Evaluation protocols. Following Agarwal et al. [2] and APV, we conduct 8 individual runs for each task and report interquartile mean with bootstrap confidence interval (CI) for individual tasks and with stratified bootstrap CI for aggregate results.
5.2 Meta-world Experiments
SSv2 pre-trained results. Figure 5a shows the learning curves on six robotic manipulation tasks from Meta-world. We observe that in-the-wild pre-training on videos from SSv2 dataset consistently improves the sample efficiency and final performance upon the DreamerV2 baseline. Moreover, our proposed ContextWM surpasses its plain counterpart in terms of sample efficiency on five of the six tasks. Notably, the dial-turn task proves challenging, as neither method is able to solve it. These results demonstrate that our method of separating context and dynamics modeling during pre-training and fine-tuning can facilitate knowledge transfer from in-the-wild videos to downstream tasks.
Ablation study. We first investigate the contribution of pre-training in our framework and report the aggregate performance with and without pre-training at the top of Figure 5b. Our results show that a plain WM seldom benefits from pre-training, indicating that the performance gain of IPV with plain WM is primarily due to the intrinsic exploration bonus. This supports our motivation that the complex contexts of in-the-wild videos can hinder knowledge transfer. In contrast, ContextWM significantly improves its performance with the aid of video pre-training. Additionally, we evaluate the contribution of the proposed techniques in ContextWM, as shown at the bottom of Figure 5b. We experiment with replacing the cross-attention mechanism (Eq. (5)) with simple concatenation or removing the dual reward predictor structure. Our results demonstrate that all these techniques contribute to the performance of ContextWM, as both variants outperform the plain WM.
Effects of dataset size. To investigate the effects of pre-training dataset size, we subsample 1.5k and 15k videos from SSv2 dataset to pre-train ContextWM. Figure 7a illustrates the performance of ContextWM with varying pre-training dataset sizes. We find that pre-training with only a small subset of in-the-wild videos can almost match the performance of pre-training with the full data. This is probably because, despite the diversity of contexts, SSv2 dataset still lacks diversity in dynamics patterns as there are only 174 classes of human-object interaction scenarios, which can be learned by proper models with only a small amount of data. It would be interesting for future work to explore whether there is a favorable scaling property w.r.t. pre-training dataset size when pre-training with more sophisticated video datasets with our in-the-wild video pre-training paradigm.
Effects of dataset domain. In Figure 7b, we assess pre-training on various video datasets, including RLBench [32] videos curated by APV [68] and our assembled data of three in-the-wild datasets. We observe that while pre-training always helps, pre-training from a more similar domain, RLBench, outperforms that from SSv2. Nevertheless, simulated videos lack diversity and scale and make it difficult to learn world models that are broadly applicable. Fortunately, in-the-wild video pre-training
(a) Learning curves
Figure 5: Meta-world results. (a) Learning curves with in-the-wild pre-training on SSv2 dataset, as measured on the success rate, aggregated across eight runs. (b) Performance of ContextWM and plain WM with or without pre-training (top) and performance of ContextWM and its variants that replace the cross-attention mechanism with naive concatenation or remove the dual reward predictor structure (bottom). We report aggregate results across a total of 48 runs over six tasks.
(b) Ablation study
can continuously improve its performance with a more diverse assembled dataset, suggesting a promising scalable alternative for domain-specific pre-training.
5.3 DMControl Remastered Experiments
In the top row of Figure 6a, we present the learning curves of IPV from SSv2 dataset and DreamerV2 on the DMC Remastered locomotion tasks. Remarkably, we observe that pre-training with the SSv2 dataset can significantly enhance the performance of ContextWM, even with a large domain gap between pre-training and fine-tuning. This finding suggests that ContextWM effectively transfers shared knowledge of separately modeling context and dynamics. In contrast, a plain WM also benefits from pre-training, but it still struggles to solve certain tasks, such as hopper stand. These results also suggest that our ContextWM could be a valuable practice for situations where visual generalization is critical, as ContextWM trained from scratch also presents a competitive performance in Figure 6b.
Effects of dataset domain. Figure 7b demonstrates the performance of pre-training on human motion videos from the Human3.6M dataset. However, we observe a negative transfer, as pre-training from Human3.6M leads to inferior performance compared to training from scratch. We argue that the Human3.6M dataset is collected in the laboratory environment, rather than truly in-the-wild. These results support our motivation that pre-training data lacking in diversity can hardly help learn world models that are generally beneficial. Additionally, we experiment with pre-training on the assembled data of three video datasets but find no significant improvement over SSv2 pre-training.
5.4 CARLA Experiments
The bottom row of Figure 6a displays the learning curves of IPV from the SSv2 dataset and DreamerV2 on the CARLA driving task under different weather and sunlight conditions. Similar to the DMC Remastered tasks, we find that MBRL benefits from pre-training on the SSv2 dataset, despite a significant domain gap. We observe that in almost all weather conditions, IPV with ContextWM learns faster than plain WM at the early stages of training and can also outperform plain WM in terms of the final performance. Minor superiority of ContextWM over plain WM also suggests that in autonomous driving scenarios, a single frame may not contain sufficient contextual information, and more sophisticated formulations of contextual information may further enhance performance.
Effects of dataset domain. We then assess the performance of ContextWM on CARLA by pretraining it from alternative domains. As shown in Figure 7b, we observe that pre-training from
Figure 6: DMC Remastered (top) and CARLA (bottom) results. (a) Learning curves with pre-training on SSv2 dataset, as measured on the episode return, aggregated across eight runs. (b) Performance of ContextWM and plain WM with or without pre-training, aggregated across 24 runs over three tasks. Episode returns of each task in CARLA are normalized to comparable ranges.
(b) Effects of pre-training
(a) Effects of dataset size
(b) Effects of dataset domain
Figure 7: Analysis on pre-training datasets. We report aggregated results over tasks. (a) Performance of ContextWM on Meta-world tasks, with varying pre-training dataset sizes. (b) Performance of ContextWM with pre-training on videos from various domains.
the YouTube Driving dataset or a combination of three in-the-wild datasets can both improve upon learning from scratch. However, neither can significantly surpass the performance achieved by pre-training from SSv2, despite a narrower domain gap between YouTube Driving and CARLA. We conjecture that this could be attributed to the higher complexity of the YouTube Driving dataset in comparison to SSv2. It would be promising to explore the potential of scaling ContextWM further to capture valuable knowledge from more complex videos.
Video prediction. We visually investigate the future frames predicted by a plain WM and ContextWM on SSv2 in Figure 8a. We find that our model effectively captures the shape and motion of the object, while the plain one fails. Moreover, we also observe that cross-attentions in our model (Eq. (5)) successfully attend to varying spatial positions of the context to facilitate the reconstruction of visual details. This shows our model works with better modeling of context and dynamics.
Video representations. We sample video clips of length 25 with two distinct labels (push something from right to left and from left to right) and visualize the averaged model states of the sampled videos using t-SNE [74] in Figure 8b. Note that we do not utilize any labels of the videos in pre-training. Video representations of plain WM may entangle with extra contextual information and thus are not sufficiently discriminative to object motions. However, ContextWM which separately models context and dynamics can provide representations well distributed according to different directions of motion.
(a) Video prediction
ContextWM Pre-trained
(b) Video representations
(c) Compositional decoding
Figure 8: Qualitative analysis. (a) Future frames predicted on SSv2. Predictions from our model well capture the shape and motion of the water cup. Moreover, our context cross-attentions from the left upper corner (red box) of different frames successfully attend to varying spatial positions of the context frame. (b) t-SNE visualization of average pooled representations of SSv2 videos with two distinct labels, from a pre-trained plain WM and ContextWM, respectively. (c) Compositional decoding analysis on the DMCR domain, where fine-tuned ContextWM is able to combine the new context with dynamics information from the original trajectory.
Compisitional decoding. We conduct a compositional decoding analysis by sampling a random frame from another trajectory to replace the original context and leaving dynamics the same. As shown in Figure 8c, ContextWM correctly combines the new context with the original dynamics information. These results show that our model, fine-tuned on the DMCR domain, has successfully learned disentangled representations of contexts and dynamics. In contrast, the plain WM suffers from learning entangled representations and thus makes poor predictions about the transitions.
This paper presents Contextualized World Models (ContextWM), a framework for both action-free video prediction and visual model-based RL. We apply ContextWM to the paradigm of In-the-wild Pre-training from Videos (IPV), followed by fine-tuning on downstream tasks to boost learning efficiency. Experiments demonstrate the effectiveness of our method in solving a variety of visual control tasks from Meta-world, DMC Remastered, and CARLA. Our work highlights not only the benefits of leveraging abundant in-the-wild videos but also the importance of innovative world model design that facilitates knowledge transfer and scalable learning.
Limitations and future work. One limitation of our current method is that a randomly selected single frame may not sufficiently capture complete contextual information of scenes in real-world applications, such as autonomous driving. Consequently, selecting and incorporating multiple context frames as well as multimodal information [59] for better context modeling need further investigation. Our work is also limited by medium-scale sizes in terms of both world models and pre-training data, which hinders learning more broadly applicable knowledge. An important direction is to systematically examine the scalability of our method by leveraging scalable architectures like Transformers [47, 67] and massive-scale video datasets [18, 48]. Lastly, our work focuses on pretraining world models via generative objectives, which utilize the model capacity inefficiently on image reconstruction to overcome intricate contexts. Exploring alternative pre-training objectives, such as contrastive learning [51, 11] or self-prediction [65, 80], could further release the potential of IPV by eliminating heavy components on context modeling and focusing on dynamics modeling.
We would like to thank many colleagues, in particular Haixu Wu, Baixu Chen, and Jincheng Zhong, for their valuable discussion. This work was supported by the National Key Research and Development Plan (2020AAA0109201), the National Natural Science Foundation of China (62022050 and 62021002), and the Beijing Nova Program (Z201100006820041).
[1] Alessandro Achille and Stefano Soatto. Emergence of invariance and disentanglement in deep representations. The Journal of Machine Learning Research, 19(1):1947-1980, 2018.
[2] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. In NeurIPS, 2021.
[3] Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. OPAL: Offline primitive discovery for accelerating offline reinforcement learning. In ICLR, 2021.
[4] Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. In NeurIPS, 2022.
[5] Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410, 2016.
[6] Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando De Freitas. Playing hard exploration games by watching youtube. In NeurIPS, 2018.
[7] Bowen Baker, Ilge Akkaya, Peter Zhokhov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. In NeurIPS, 2022.
[8] Matthew Chang, Arjun Gupta, and Saurabh Gupta. Semantic visual navigation by watching youtube videos. In NeurIPS, 2020.
[9] Annie S Chen, Suraj Nair, and Chelsea Finn. Learning generalizable robotic reward functions from "in-the-wild" human videos. In RSS, 2021.
[10] Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. TransDreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.09481, 2022.
[11] Fei Deng, Ingook Jang, and Sungjin Ahn. DreamerPro: Reconstruction-free model-based reinforcement learning with prototypical representations. In ICML, 2022.
[12] Emily Denton and Vighnesh Birodkar. Unsupervised learning of disentangled representations from video. In NeurIPS, 2017.
[13] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In NAACL, 2019.
[14] Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. CARLA: An open urban driving simulator. In CoRL, 2017.
[15] Jean-Yves Franceschi, Edouard Delasalles, Mickaël Chen, Sylvain Lamprier, and Patrick Gallinari. Stochastic latent residual video prediction. In ICML, 2020.
[16] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deep-MDP: Learning continuous latent space models for representation learning. In ICML, 2019.
[17] Raghav Goyal, Samira Ebrahimi Kahou, Vincent Michalski, Joanna Materzynska, Susanne Westphal, Heuna Kim, Valentin Haenel, Ingo Fruend, Peter Yianilos, Moritz Mueller-Freitag, et al. The "something something" video database for learning and evaluating visual common sense. In ICCV, 2017.
[79] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. In ICLR, 2022.
[80] Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. In NeurIPS, 2021.
[81] Weirui Ye, Yunsheng Zhang, Pieter Abbeel, and Yang Gao. Become a proficient player with limited data through watching pure videos. In ICLR, 2023.
[82] Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks? In NeurIPS, 2014.
[83] Fisher Yu, Haofeng Chen, Xin Wang, Wenqi Xian, Yingying Chen, Fangchen Liu, Vashisht Madhavan, and Trevor Darrell. BDD100K: A diverse driving dataset for heterogeneous multitask learning. In CVPR, 2020.
[84] Tianhe Yu, Chelsea Finn, Annie Xie, Sudeep Dasari, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot imitation from observing humans via domain-adaptive meta-learning. In RSS, 2018.
[85] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In CoRL, 2020.
[86] Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang, Yi Wu, Yang Gao, and Huazhe Xu. Pre-trained image encoder for generalizable visual reinforcement learning. In NeurIPS, 2022.
[87] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In ICLR, 2021.
[88] Qihang Zhang, Zhenghao Peng, and Bolei Zhou. Learning to drive by watching youtube videos: Action-conditioned contrastive policy pretraining. In ECCV, 2022.
[89] Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmentation. In AAAI, 2020.
For clarity, we first elaborate on the optimization objective for contextualized world models during action-free pre-training, which drops action-conditioned dynamics and reward predictors from the fine-tuning objective (Eq. (6)):
$$
\begin{array}{l} \mathcal {L} ^ {\mathrm {C W M - p t}} (\theta) \doteq \underbrace {\mathbb {E} _ {q _ {\theta} \left(z _ {1 : T} \mid o _ {1 : T}\right)}} _ {\text {c o n t e x t - u n a w a r e l a t e n t i n f e r e n c e}} \left[ \sum_ {t = 1} ^ {T} \left(\underbrace {- \ln p _ {\theta} \left(o _ {t} \mid z _ {t} , c\right)} _ {\text {c o n t e x t u a l i z e d i m a g e l o s s}}\right) \right. \tag {7} \\ \underbrace {+ \beta_ {z} \operatorname {K L} \left[ q _ {\theta} \left(z _ {t} \mid z _ {t - 1} , o _ {t}\right) \| p _ {\theta} (\hat {z} _ {t} \mid z _ {t - 1}) \right]} _ {\text {a c t i o n - f r e e K L l o s s}} \biggr) \biggr ]. \\ \end{array}
$$
For behavior learning, we adopt the same actor-critic learning scheme based purely on imaginary latent trajectories  $\hat{s}_{\tau},\hat{a}_{\tau},\hat{r}_{\tau}$  with horizon  $H$  as DreamerV2 [21, 23]. The critic  $v_{\xi}(s)$  is learned by regressing the  $\lambda$ -target [72, 64]:
$$
\mathcal {L} ^ {\text {c r i t i c}} (\xi) \doteq \mathbb {E} _ {p _ {\phi}, \pi_ {\psi}} \left[ \sum_ {\tau = t} ^ {t + H} \frac {1}{2} \left(v _ {\xi} \left(\hat {s} _ {\tau}\right) - \operatorname {s g} \left(V _ {\tau} ^ {\lambda}\right)\right) ^ {2} \right], \tag {8}
$$
$$
V _ {\tau} ^ {\lambda} \doteq \hat {r} _ {\tau} + \gamma \left\{ \begin{array}{l l} (1 - \lambda) v _ {\xi} (\hat {s} _ {\tau + 1}) + \lambda V _ {\tau + 1} ^ {\lambda} & \text {i f} \tau <   t + H \\ v _ {\xi} (\hat {s} _ {\tau + 1}) & \text {i f} \tau = t + H, \end{array} \right. \tag {9}
$$
where  $\mathrm{sg}$  is a stop gradient function. And the actor  $\pi_{\psi}(a|s)$  is learned by maximizing the imagined return by back-propagating the gradients through the learned world model, with an entropy regularizer:
$$
\mathcal {L} ^ {\text {a c t o r}} (\psi) \doteq \mathbb {E} _ {p _ {\phi}, \pi_ {\psi}} \left[ \sum_ {\tau = t} ^ {t + H} \left(- V _ {\tau} ^ {\lambda} - \eta \mathrm {H} \left[ \pi_ {\psi} \left(a _ {\tau} \mid \hat {s} _ {\tau}\right) \right]\right) \right], \tag {10}
$$
Overall, a complete description of pre-training and fine-tuning of ContextWM for MBRL is presented in Algorithm 1.
The variational bound for contextualized latent dynamics models  $p(o_{1:T}, r_{1:T} \mid a_{1:T}, c)$  and a variational posterior  $q(z_{1:T} \mid a_{1:T}, o_{1:T})$  follows from importance weighting and Jensen's Inequality as shown<sup>3</sup>,
$$
\begin{array}{l} \ln p (o, r | a, c) \geq \mathbb {E} _ {q (z | a, o)} \left[ \ln \frac {p (o , r , z | a , c)}{q (z | a , o)} \right] \\ = \mathbb {E} _ {q (z | a, o)} [ \ln p (o, r | z, a, c) ] - \operatorname {K L} [ q (z | a, o) \| p (z | a) ] \\ = \mathbb {E} _ {q (z | a, o)} \left[ \sum_ {t = 1} ^ {T} \ln p \left(o _ {t} \mid z _ {t}, c\right) + \ln p \left(r _ {t} \mid z _ {t}\right) \right] - \mathrm {K L} [ q (z | a, o) \| p (z | a) ], \\ \end{array}
$$