SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jialong Wu\* Haoyu Ma\* Chaoyi Deng, Mingsheng Long
wujialong0229@gmail.com, {mhy22,dengcy23}@mails.tsinghua.edu.cn
mingsheng@tsinghua.edu.cn
Unsupervised pre-training methods utilizing large and diverse datasets have achieved tremendous success across a range of domains. Recent work has investigated such unsupervised pre-training methods for model-based reinforcement learning (MBRL) but is limited to domain-specific or simulated data. In this paper, we study the problem of pre-training world models with abundant in-the-wild videos for efficient learning of downstream visual control tasks. However, in-the-wild videos are complicated with various contextual factors, such as intricate backgrounds and textured appearance, which precludes a world model from extracting shared world knowledge to generalize better. To tackle this issue, we introduce Contextualized World Models (ContextWM) that explicitly separate context and dynamics modeling to overcome the complexity and diversity of in-the-wild videos and facilitate knowledge transfer between distinct scenes. Specifically, a contextualized extension of the latent dynamics model is elaborately realized by incorporating a context encoder to retain contextual information and empower the image decoder, which encourages the latent dynamics model to concentrate on essential temporal variations. Our experiments show that in-the-wild video pre-training equipped with ContextWM can significantly improve the sample efficiency of MBRL in various domains, including robotic manipulation, locomotion, and autonomous driving. Code is available at this repository: https://github.com/thuml/ContextWM.
Model-based reinforcement learning (MBRL) holds the promise of sample-efficient learning for visual control. Typically, a world model [20, 40] is learned to approximate state transitions and control signals of the environment to generate imaginary trajectories for planning [22] or behavior learning [71]. In the wake of revolutionary advances in deep learning, world models have been realized as action-conditional video prediction models [34] or latent dynamics models [22, 21]. However, given the expressiveness and complexity of deep neural networks, the sample efficiency of MBRL can still be limited by the failure to learn an accurate and generalizable world model efficiently.
Pre-training and fine-tuning paradigm has been highly successful in computer vision [82, 26] and natural language processing [57, 13] to fast adapt pre-trained representations for downstream tasks, while learning tabula rasa is still dominant in MBRL. Recent work has taken the first step towards pre-training a world model, named Action-free Pre-training from Videos (APV) [68]. However, it has been conducted by pre-training on domain-specific and carefully simulated video datasets, rather than abundant in-the-wild video data [17, 18, 30, 83]. Prior attempts of leveraging real-world video data result in either underfitting for video pre-training or negligible benefits for downstream visual control tasks [68]. Against this backdrop, we naturally ask the following question:
Can world models pre-trained on diverse in-the-wild videos benefit sample-efficient learning of downstream visual control tasks?
Figure 1: Illustration of In-the-wild Pre-training from Videos (IPV). In various domains, we pre-train world models with in-the-wild videos by action-free video prediction (left) and then fine-tune the pre-trained one on downstream visual control tasks with model-based reinforcement learning (right).
It opens up a way to utilize the vast amount of videos available on the Internet and thereby fully release the potential of this methodology by applying it at scale. Large and diverse data spanning various scenes and tasks can provide world knowledge that is widely generalizable and applicable to a variety of downstream tasks. For instance, as depicted in Figure 1, world models for robotic manipulation tasks can probably benefit not only from videos of humans interacting with target objects in diverse ways but also from motion videos that embody rigorous physical rules. Analogous to APV [68], we brand our paradigm as In-the-wild Pre-training from Videos (IPV).
However, in-the-wild videos are of inherent high complexity with various factors, such as intricate backgrounds, appearance, and shapes, as well as complicated dynamics. Coarse and entangled modeling of context and dynamics can waste a significant amount of model capacity on modeling low-level visual details of what is there and prevent world models from capturing essential shared knowledge of what is happening. Biological studies reveal that in natural visual systems,  $\sim 80\%$  retinal ganglion cells are P-type operating on spatial detail and color, while  $\sim 20\%$  are M-type operating on temporal changes [28, 44]. Partially inspired by this, we suggest that leveraging IPV for world models calls for not only appropriate data and model scaling [35] but more essentially, dedicated design for context and dynamics modeling.
In this paper, we present the Contextualized World Model (ContextWM), a latent dynamics model with visual observations that explicitly separates context and dynamics modeling to facilitate knowledge transfer of semantically similar dynamics between visually diverse scenes. Concretely, a contextualized latent dynamics model is derived to learn with variational lower bound [33, 38] and elaborately realized by incorporating a context encoder to retain contextual information and a parsimonious latent dynamics model to concentrate on essential temporal variations. Moreover, dual reward predictors are introduced to enable simultaneously enhancing task-relevant representation learning as well as learning an exploratory behavior. The main contributions of this work are three-fold:
- From a data-centric view, we systematically study the paradigm of pre-training world models with in-the-wild videos for sample-efficient learning of downstream tasks, which is much more accessible and thus makes it possible to build a general-purpose world model.
- We propose Contextualized World Models (ContextWM), which explicitly model both the context and dynamics to handle complicated in-the-wild videos during pre-training and also encourage task-relevant representation learning when fine-tuning with MBRL.
- Our experiments show that equipped with our ContextWM, in-the-wild video pre-training can significantly improve the sample efficiency of MBRL on various domains.
Pre-training in RL. Three categories of pre-training exist in RL: unsupervised online pre-training, offline pre-training, and visual representation pre-training. The first two are both domain-specific, which learn initial behaviors, primitive skills, or domain-specific representations by either online environment interaction [55, 39] or offline demonstrations [41, 3]. Recent work has explored general-
purpose visual pre-training using in-the-wild datasets [69, 36, 77, 53, 86, 49, 45] to accelerate policy learning in model-free manners. Plan2Explore [66], as a model-based method, pre-trains its world model and exploratory policy simultaneously by online interaction. Ye et al. [81] and Xu et al. [78] build upon EfficientZero [80] and pre-train their world models using offline experience datasets across Atari games. To the best of our knowledge, APV [68] is the first to pre-train world models with datasets from different domains, though the datasets are manually simulated demonstrations by scripted policies across tasks from RLBench [32], which still lacks diversity and scale.
Visual control with in-the-wild videos. While several existing works leverage demonstration videos for learning visual control [43, 84, 62], rare efforts have been made to utilize off-the-shelf video datasets from the Internet. Shao et al. [70] and Chen et al. [9] learn reward functions using videos from the Something-Something dataset [17]. R3M [49] leverages diverse Ego4D dataset [18] to learn reusable visual representations by time contrastive learning. Other applications with such videos include video games [7, 6], visual navigation [8], and autonomous driving [88]. We instead pre-train world models by video prediction on in-the-wild video datasets.
World models for visual RL. World models that explicitly model the state transitions and reward signals are widely utilized to boost sample efficiency in visual RL. A straightforward method is to learn an action-conditioned video prediction model [50, 34] to generate imaginary trajectories. Ha and Schmidhuber [20] propose to first use a variational autoencoder to compress visual observation into latent vectors and then use a recurrent neural network to predict transition on compressed representations. Dreamer [22, 21, 23], a series of such latent dynamics models learning via image reconstruction, have demonstrated their effectiveness for both video games and visual robot control. There also have been other methods that learn latent representations by contrastive learning [51, 11] or value prediction [63, 25]. Recent works also use Transformers [75] as visual encoders [67] or dynamics models [10, 47, 58] that make world models much more scalable and can be complementary with our effort to pre-train world models with a vast amount of diverse video data.
Concurrent to our work, SWIM [46] also pre-trains a world model from large-scale human videos and then transfers it to robotic manipulation tasks in the real world. However, it designs a structured action space via visual affordances to connect human videos with robot tasks, incorporating a strong inductive bias and thus remarkably limiting the range of pre-training videos and downstream tasks.
Problem formulation. We formulate a visual control task as a partially observable Markov decision process (POMDP), defined as a tuple  $(\mathcal{O},\mathcal{A},p,r,\gamma)$ .  $\mathcal{O}$  is the observation space,  $\mathcal{A}$  is the action space,  $p(o_{t} \mid o_{<t},a_{<t})$  is the transition dynamics,  $r(o_{\leq t},a_{<t})$  is the reward function, and  $\gamma$  is the discount factor. The goal of MBRL is to learn an agent  $\pi$  that maximizes the expected cumulative rewards  $\mathbb{E}_{p,\pi}[\sum_{t=1}^{T}\gamma^{t-1}r_{t}]$ , with a learned world model  $(\hat{p},\hat{r})$  approximating the unknown environment. In our pre-training and fine-tuning paradigm, we also have access to an in-the-wild video dataset  $\mathcal{D} = \{(o_{t})_{t=1}^{T}\}$  without actions and rewards to pre-train the world model.
Dreamer. Dreamer [21, 23, 24] is a visual model-based RL method where the world model is formulated as a latent dynamics model [16, 22] with the following four components:
Representation model:  $z_{t} \sim q_{\theta}(z_{t} \mid z_{t-1}, a_{t-1}, o_{t})$  Image decoder:  $\hat{o}_{t} \sim p_{\theta}(\hat{o}_{t} \mid z_{t})$  Transition model:  $\hat{\hat{x}}_t \sim p_t(\hat{\hat{x}}_t \mid \hat{z}_t, a_{t-1})$  Reward predictor:  $\hat{\hat{x}}_t \sim p_t(\hat{\hat{x}}_t \mid \hat{z}_t)$  (1)
Transition model:  $\hat{z}_t\sim p_\theta (\hat{z}_t|z_{t - 1},a_{t - 1})$  Reward predictor:  $\hat{r}_t\sim p_\theta (\hat{r}_t|z_t)$
The representation model, also known as posterior of  $z_{t}$ , approximates latent state  $z_{t}$  from previous state  $z_{t-1}$ , previous action  $a_{t-1}$  and current observation  $o_{t}$ , while the transition model, also known as prior of  $z_{t}$ , predicts it directly from  $z_{t-1}$  and  $a_{t-1}$ . The overall models are jointly learned by minimizing the negative variational lower bound (ELBO) [33, 38]:
$$
\begin{array}{l} \mathcal {L} (\theta) \doteq \mathbb {E} _ {q _ {\theta} \left(z _ {1: T} \mid a _ {1: T}, o _ {1: T}\right)} \left[ \sum_ {t = 1} ^ {T} \left(- \ln p _ {\theta} \left(o _ {t} \mid z _ {t}\right) - \ln p _ {\theta} \left(r _ {t} \mid z _ {t}\right) \right. \right. \tag {2} \\ \left. \left. + \beta_ {z} \operatorname {K L} \left[ q _ {\theta} \left(z _ {t} \mid z _ {t - 1}, a _ {t - 1}, o _ {t}\right) \| p _ {\theta} \left(\hat {z} _ {t} \mid z _ {t - 1}, a _ {t - 1}\right) \right]\right) \right]. \\ \end{array}
$$
Behavior learning of Dreamer can be conducted by actor-critic learning purely on imaginary latent trajectories, for which we refer readers to Hafner et al. [21] for details.
Action-free Pre-training from Videos (APV). To leverage action-free video data, APV [68] first pre-trains an action-free variant of the latent dynamics model (with representation  $q_{\theta}(z_t \mid z_{t-1}, o_t)$ , transition  $p_{\theta}(\hat{z}_t \mid z_{t-1})$  and image decoder  $p_{\theta}(\hat{o}_t \mid z_t)$ ) as a video prediction model, which drops action conditions and the reward predictor from Eq. (1) and Eq. (2). When fine-tuned for MBRL, APV stacks an action-conditional model (with representation  $q_{\phi}(s_t \mid s_{t-1}, a_{t-1}, z_t)$  and transition  $p_{\phi}(\hat{s}_t \mid s_{t-1}, a_{t-1})$ ) on top of the action-free one. Action-free dynamics models and image decoders are initialized with pre-trained weights  $\theta$  from video pre-training. Furthermore, to utilize pre-trained representations for better exploration, APV introduces a video-based intrinsic bonus  $r_t^{\mathrm{int}}$  measuring the diversity of visited trajectories by random projection and nearest neighbors. The overall models during fine-tuning are optimized by minimizing the following objective:
$$
\mathcal {L} (\phi , \theta) \doteq \mathbb {E} _ {q _ {\phi} (s _ {1: T} | a _ {1: T}, z _ {1: T}), q _ {\theta} (z _ {1: T} \mid o _ {1: T})} \left[ \sum_ {t = 1} ^ {T} \left(\underbrace {- \ln p _ {\theta} \left(o _ {t} \mid s _ {t}\right)} _ {\text {i m a g e l o g l o s s}} \underbrace {- \ln p _ {\phi} \left(r _ {t} + \lambda r _ {t} ^ {\text {i n t}} \mid s _ {t}\right)} _ {\text {r e w a r d l o g l o s s}}\right) \right. \tag {3}
$$
$$
\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}} \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 ].
$$
4 Contextualized World Models
In this section, we present Contextualized World Models (ContextWM), a framework for both action-free video prediction and visual model-based RL with explicit modeling of context and dynamics from observations. Our method introduces: (i) a contextualized extension of the latent dynamics model (see Section 4.1), (ii) a concrete implementation of contextualized latent dynamics models tailored for visual control (see Section 4.2). We provide the overview and pseudo-code of our ContextWM in Figure 3 and Appendix A, respectively.
4.1 Contextualized Latent Dynamics Models
While latent dynamics models have been successfully applied for simple and synthetic scenes, in-the-wild images and videos are naturalistic and complex. Our intuition is that there are two groups of information in the sequence of observations, namely time-invariant context and time-dependent dynamics [76, 12, 42, 15]. The context encodes static information about objects and concepts in the scene, such as texture, shapes, and colors, while the dynamics encode the temporal transitions of the concepts, such as positions, layouts, and motions. To overcome the complexity of in-the-wild videos, it is necessary to explicitly represent the complicated contextual information, thus shared knowledge with respect to dynamics can be learned and generalized across distinct scenes.
Figure 2: Probabilistic graphical model of a contextualized latent dynamics model, assuming time-invariant context  $c$  and time-dependent dynamics  $z_{t}$ . Solid lines show the generative model and dotted lines show variational inference.
We thus propose a contextualized extension of latent dynamics models [16, 22], extending Eq. (1). As shown in
Figure 2, the probabilistic model generates the observation  $o_{t}$  conditioned not only on the current latent dynamics  $z_{t}$  but also on a context variable  $c$  that can include rich information (solid lines in Figure 2), while the variational inference model approximates the posterior of the latent dynamics  $z_{t}$  conditioned on observations  $o_{t}$  and previous states  $z_{t-1}$ , without referring to the context (dotted lines in Figure 2). The proposed design allows contextualized image decoders to reconstruct diverse and complex observations using rich contextual information, expanding their ability beyond the expressiveness of latent dynamics variables  $z_{t}$ . On the other hand, through the variational information bottleneck [5, 1], latent dynamics variables  $z_{t}$  are encouraged to selectively exclude contextual information and only capture essential temporal variations that are not included in the context  $c$ . Furthermore, the design of context-unaware latent dynamics inference encourages representation and transition models to learn this temporal information within a high-level semantic space, in contrast to merely differentiating with the context frame in the low-level visual space, facilitating the acquisition of representations more transferable and robust to intricate contexts.
(a) Action-free video pre-training
(b) Action-conditioned fine-tuning with MBRL
Figure 3: Architecture of Contextualized World Models. Building upon the stacked latent model from APV [68], we empower the image decoder by incorporating a context encoder that operates in parallel with the latent dynamics model. The context encoder captures rich contextual information from a randomly selected context frame and enhances the decoder features of each timestep with a multi-scale cross-attention mechanism that enables flexible information shortcuts across spatial positions. This design encourages the latent dynamics model to focus only on essential temporal variations, while also allowing the decoder to reconstruct complex observations more effectively.
As derived in Appendix B, the overall model can be learned with ELBO of conditional log probability  $\ln p_{\theta}(o_{1:T}, r_{1:T} \mid a_{1:T}, c)$ , without the need to model the complex distribution of contexts  $p(c)$ :
$$
\begin{array}{l} \mathcal {L} (\theta) \doteq \underbrace {\mathbb {E} _ {q _ {\theta} \left(z _ {1 : T} \mid a _ {1 : T} , 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}} - \ln p _ {\theta} \left(r _ {t} \mid z _ {t}\right) \right. \right. \tag {4} \\ \left. \left. + \beta_ {z} \operatorname {K L} \left[ q _ {\theta} \left(z _ {t} \mid z _ {t - 1}, a _ {t - 1}, o _ {t}\right) \parallel p _ {\theta} \left(\hat {z} _ {t} \mid z _ {t - 1}, a _ {t - 1}\right) \right]\right) \right]. \\ \end{array}
$$
4.2 Contextualized World Model Architectures
We then introduce a concrete implementation of contextualized latent dynamics models, specifically tailored for visual control. Our approach builds upon Dreamer [21], which utilizes a Recurrent State Space Model (RSSM) [22]. To enable action-free pre-training from videos and action-conditioned fine-tuning for MBRL, we also incorporate the stacked latent model from APV [68] into our design. Notably, our work distinguishes itself from previous research through the elaborate design of contextualized components and a dual reward predictor structure that enhances task-relevant representation learning. An overview of the overall architecture is illustrated in Figure 3.
Context formulation. There exist various choices of contextual information to be conditioned on, including text descriptions, pre-trained representations, semantic maps, or more sophisticated structured data [59]. As our focus is on end-to-end visual control, we choose the simplest one solely based on visual observations: we consider the context  $c$  as a single frame of observation randomly sampled from the trajectory segment  $o_{1:T}$ :  $c \doteq o_{\tilde{t}}$ ,  $\tilde{t} \sim \text{Uniform}\{1,2,\dots,T\}$ . It is assumed that contextual information lies equally in each frame, and by randomly selecting a context frame, our context encoder should learn to be robust to temporal variations.
Multi-scale cross-attention conditioning. Given the context in the form of a frame of the same size as the reconstruction  $\hat{o}_t$ , a U-Net [60] architecture would be one of the most appropriate choices to empower a contextualized image decoder  $p_{\theta}(\hat{o}_t \mid z_t, c)$ , as its multi-scale shortcuts help to propagate the information directly from context to reconstruction. Moreover, in our case, the image decoder learns by also conditioning on latent dynamics  $z_t$ , as shown in Figure 3. While a conventional U-Net architecture incorporates shortcut features directly by concatenation or summation and thus forces a spatial alignment between them, temporal variations such as motions or deformations cannot be neglected in our case. Inspired by recent advances in generative models [59, 4], we augment the decoder feature  $X \in \mathbb{R}^{c \times h \times w}$  with the context feature  $Z \in \mathbb{R}^{c \times h \times w}$  by a cross-attention mechanism
[75], shown as follows (we refer to Appendix C.3 for details):
$$
\text {F l a t t e n :} Q = \operatorname {R e s h a p e} (X) \in \mathbb {R} ^ {h w \times c}, K = V = \operatorname {R e s h a p e} (Z) \in \mathbb {R} ^ {h w \times c}
$$
$$
\text {C r o s s - A t t e n t i o n :} R = \text {A t t e n t i o n} \left(Q W ^ {Q}, K W ^ {K}, V W ^ {V}\right) \in \mathbb {R} ^ {h w \times c} \tag {5}
$$
$$
\text {R e s i d u a l - C o n n e c t i o n :} X = \operatorname {R e L U} (X + \operatorname {B a t c h N o r m} (\operatorname {R e s h a p e} (R))) \in \mathbb {R} ^ {c \times h \times w}.
$$
Dual reward predictors. For simplicity of behavior learning, we only feed the latent variable  $s_t$  into the actor  $\pi(a_t|s_t)$  and critic  $v(s_t)$ , thus introducing no extra computational costs when determining actions. However, this raises the question of whether the latent variable  $s_t$  contains sufficient task-relevant information for visual control and avoids taking shortcuts through the U-Net skip connections. For example, in robotics manipulation tasks, the positions of static target objects may not be captured by the time-dependent  $s_t$ , but by the time-invariant context  $c$ . We remark that reward predictors  $p_{\phi}(r_t|s_t)$  play a crucial role in compelling the latent variable  $s_t$  to encode task-relevant information, as the reward itself defines the task to be accomplished. Nevertheless, for hard-exploration tasks such as Meta-world [85], the video-based intrinsic bonus [68] may distort the exact reward regressed for behavioral learning and constantly drift during the training process [61], making it difficult for the latent dynamics model to capture task-relevant signals. Therefore, we propose a dual reward predictor structure comprising a behavioral reward predictor that regresses the exploratory reward  $r_t + \lambda r_t^{\mathrm{int}}$  for behavior learning, and additionally, a representative reward predictor that regresses the pure reward  $r_t$  to enhance task-relevant representation learning [31].
Overall objective. The model parameters of ContextWM can be jointly optimized as follows: