Reinforcement Learning with Action-Free Pre-Training from Videos

Introduction

Deep reinforcement learning (RL) has made significant advance in solving various sequential decision-making problems (Mnih et al., 2015; Levine et al., 2016; Silver et al., 2017; Vinyals et al., 2019; Berner et al., 2019; Akkaya et al., 2019; Kalashnikov et al., 2021). However, existing RL methods often start learning tabula rasa without any prior knowledge of the world, therefore requiring a large amount of environment interaction for learning meaningful behaviors. By contrast, within the computer vision (CV) and natural language processing (NLP) domains, recent unsupervised pre-training approaches have shown to be effective by leveraging the pre-trained representations for fine-tuning in downstream tasks (Mikolov et al., 2013; Pennington et al., 2014; Noroozi & Favaro, 2016; Gidaris et al., 2018; Devlin et al., 2019; Radford et al., 2018; He et al., 2020).

Recent works have shown promise in adopting such pre-training and fine-tuning paradigm to vision-based RL, by demonstrating that representations pre-trained with various unsupervised representation learning schemes can be effective in downstream tasks (Finn et al., 2016b; Dwibedi et al., 2018; Zhan et al., 2020; Laskin et al., 2020b; Stooke et al., 2021; Schwarzer et al., 2021b). Notably, Stooke et al. (2021) show that weight initialization with contrastively pre-trained representations leads to performance improvement. These works, however, mostly focus on the setup where pre-training datasets are collected in the target domains, or in different domains but with very similar visuals. Instead, we would like to leverage videos from diverse domains for pre-training, and transfer the pre-trained representations for solving newly encountered domains.

In this paper, we present APV: Action-Free Pre-training from Videos, a novel framework that performs generative pre-training on videos for improving the sample-efficiency of vision-based RL. Since our goal is to learn the representations that can be transferred to various downstream tasks from readily available videos, our framework do not require the videos to be collected in the same domain of the downstream tasks, and also do not assume the datasets contain action information. Summarized in Figure 1, our framework comprises two phases: we first pre-train an action-free latent video prediction model to learn useful representations from videos, then fine-tune the pre-trained model for learning action-conditional world models on downstream tasks. To leverage the fine-tuned world models for behavior learning, we build APV on top of DreamerV2 (Hafner et al., 2021).

The key ingredients of APV are as follows:

Action-free pre-training from videos: To capture rich dynamics information from diverse videos, we pre-train an action-free latent video prediction model. We find that the representations from the pre-trained model can be transferred to various downstream tasks.

Stacked latent prediction model: To incorporate additional action inputs during fine-tuning, we introduce a new architecture that stacks an action-conditional latent dynamics model on top of the action-free model.

Video-based intrinsic bonus: For better exploration, we propose an intrinsic bonus that utilizes video representations from the action-free model. Since the pre-trained representations contain information useful for understanding dynamics of environments, our intrinsic bonus effectively encourages agents to learn diverse behaviors.

In our experiments, we pre-train the action-free prediction model using 4950 videos collected on 99 manipulation tasks from RLBench (James et al., 2020) and fine-tune the pre-trained model on a range of manipulation tasks from Meta-world (Yu et al., 2020). Despite a big domain gap between RLBench and Meta-world, we demonstrate that APV significantly outperforms DreamerV2. For example, APV achieves the aggregate success rate of 95.4% on six manipulation tasks, while DreamerV2 achieves 67.9%. Moreover, we show that RLBench pre-trained representations can also be effective in learning locomotion tasks from DeepMind Control Suite (Tassa et al., 2020), where both the visuals and objectives significantly differ from RLBench videos.

Related Work

Recently, unsupervised representation learning methods have been actively studied in the domain of CV. Various representation learning methods, including reconstruction (He et al., 2021), rotation (Gidaris et al., 2018), solving zigsaw puzzles (Noroozi & Favaro, 2016), and contrastive learning (He et al., 2020; Chen et al., 2020), have reduced the gap with supervised pre-training with labels. In the domain of NLP, unsupervised representation learning has been successfully applied to training language models with generalization ability (Devlin et al., 2019; Radford et al., 2018; Yang et al., 2019). Notably, Devlin et al. (2019) and Radford et al. (2018) trained large transformer networks (Vaswani et al., 2017) with masked token prediction and generative pre-training, respectively, and showed that pre-trained models can be effectively fine-tuned on downstream tasks. In this work, we demonstrate that unsupervised pre-training can also be effective for vision-based RL.

Unsupervised representation learning for RL.

Unsupervised representation learning for RL has also been studied to improve the sample-efficiency of RL algorithms. Notably, Jaderberg et al. (2017) showed that optimizing auxiliary unsupervised losses can improve the performance of RL agents. This has been followed by a series of works which demonstrated the effectiveness of various unsupervised learning objectives, including world-model learning (Hafner et al., 2019, 2021), reconstruction (Yarats et al., 2021c), future representation prediction (Gelada et al., 2019; Schwarzer et al., 2021a), bisimulation (Castro, 2020; Zhang et al., 2021), and contrastive learning (Oord et al., 2018; Anand et al., 2019; Mazoure et al., 2020; Srinivas et al., 2020). While these works optimize auxiliary unsupervised objectives to accelerate the training of RL agents, we instead aim to pre-train representations as in CV and NLP domains.

There have been several approaches to perform unsupervised pre-training for RL (Finn et al., 2016b; Dwibedi et al., 2018; Zhan et al., 2020; Srinivas et al., 2020; Stooke et al., 2021; Schwarzer et al., 2021b). In particular, Schwarzer et al. (2021b) proposed several self-supervised learning objectives that rely on actions, but assume access to action information of downstream tasks which may not be available in practice. The work closest to ours is Stooke et al. (2021), which demonstrated that the representations contrastively pre-trained without actions and rewards can be effective on unseen downstream tasks but with very similar visuals. In this work, we instead develop a framework that leverages action-free videos from diverse domains with different visuals and embodiments for pre-training. Concurrent to our work, Xiao et al. (2022) showed that pre-training a visual encoder on videos can be effective for downstream RL tasks. We instead pre-train a video prediction model instead of a visual encoder that operates on a single image.

Behavior learning with videos.

Video datasets have also been utilized for behavior learning in various ways (Peng et al., 2018; Torabi et al., 2018; Aytar et al., 2018; Liu et al., 2018; Sermanet et al., 2018; Edwards et al., 2019; Schmeckpeper et al., 2020b, a; Chang et al., 2020; Chen et al., 2021; Zakka et al., 2022). Aytar et al. (2018) solved hard exploration tasks on Atari benchmark by designing an imitation reward based on YouTube videos, and Peng et al. (2018) proposed to learn physical skills from human demonstration videos by extracting reference motions and training an RL agent that imitates the extracted motions. Our work differs in that we utilize videos for pre-training representations, instead of learning behaviors from videos.

We provide more discussion on related fields in Appendix C.

Method

In order to pre-train representations from action-free videos, we first learn a latent video prediction model, which is an action-free variant of a latent dynamics model (Hafner et al., 2019). Unlike autoregressive video prediction models that predict a next frame and utilize it as an input for the following prediction, the model instead operates on the latent space (Zhang et al., 2019; Hafner et al., 2019; Franceschi et al., 2020). Specifically, the model consists of three main components: (i) the representation model that encodes observations $o_{t}$ to a model state $z_{t}$ with Markovian transitions, (ii) the transition model that predicts future model states $\hat{z}_{t}$ without access to the observation, and (iii) the image decoder that reconstructs image observations $\hat{o}_{t}$. The model can be summarized as follow (see Figure 2):

We train the model to reconstruct image observations, and to make the prediction from the representation model and transition model be close to each other. All model parameters $\phi$ are jointly optimized by minimizing the negative variational lower bound (ELBO; Kingma & Welling 2014):

where $\beta_{z}$ is a scale hyperparameter and $T$ is the length of training sequences in a minibatch. Since the transition model does not condition on observations, it allows us to efficiently predict future states in the latent space without needing to predict future images using the image decoder at inference time. We implement the transition model as an action-free recurrent state-space model (RSSM; Hafner et al. 2019), which consists of both deterministic and stochastic components, and the representation model by combining the action-free RSSM with an image encoder. We refer to Appendix B for a more detailed formulation.

2 Stacked Latent Prediction Model

Once we pre-train the action-free prediction model, we fine-tune it into an action-conditional prediction model that can be used for solving various visual control tasks. Since actions and rewards, which provide more information about target tasks, are available during fine-tuning, it motivates incorporating them into the model. One naïve approach would be to initialize the action-conditional prediction model with the action-free model, and learn a reward predictor on top of it. But we find this fine-tuning scheme rapidly erases the useful knowledge in pre-trained models (see Figure 6(a) for supporting results). To effectively utilize the pre-trained representations, we introduce a new architecture that stacks an action-conditional prediction model on top of the action-free model as below (see Figure 3(a)):

which is optimized by minimizing the following objective:

where $\beta$ is a scale hyperparameter. Here, we note that we initialize the image decoder $p_{\theta}(o|s_{t})$ with the pre-trained image decoder $p_{\phi}(o|z_{t})$. We implement the transition model of action-conditional prediction model as RSSM, and the representation model as RSSM with dense layers that receive the model states of the action-free model as inputs. We refer to Appendix B for a more detailed formulation. In our experiments, we use $\beta_{z}=0$ during fine-tuning, and only utilize the action-conditional RSSM for future imagination.

3 Video-based Intrinsic Bonus

It has been observed that good representations are crucial for efficient exploration in environments with high-dimensional observations (Laskin et al., 2021). To utilize useful information captured in the pre-trained representations for exploration, we propose a video-based intrinsic bonus. Our main idea is to increase the diversity of visited trajectories by utilizing it as an intrinsic bonus. Specifically, given a sequence of model states from the action-free prediction model $z_{t:t+\tau}$, we apply average pooling across the sequence dimension to obtain a trajectory representation $y_{t}=\text{Avg}(z_{t:t+\tau})$. Then, we utilize the distance of $y_{t}$ to its $k$-nearest neighbor in samples from a replay buffer as a metric for measuring the diversity of trajectories. To summarize, our intrinsic bonus is defined as below (see Figure 3(b) for illustration):

where $\psi$ is a random projection (Bingham & Mannila, 2001) that maps the model state to a low-dimensional representation for compute-efficient distance computation, and $y_{t}^{k}$ is a $k$-nearest neighbor of $y_{t}$ in a minibatch. By explicitly encouraging the agents to visit more diverse trajectories instead of single states (Pathak et al., 2017; Burda et al., 2019; Pathak et al., 2019; Liu & Abbeel, 2021), it effectively encourages the agents to explore environments in a more long term manner, and thus learn more diverse behaviors. Then, the reward predictor is trained to predict the sum of $r_{t}$ and $r_{t}^{\tt{int}}$ as below:

where $\lambda$ is a hyperparameter that adjusts the tradeoff between exploitation and exploration. We find the intrinsic bonus provides large gains when combined with pre-training, as the pre-trained representations already contain useful representation from the beginning of the fine-tuning (see Figure 7(a) for supporting results). In our experiments, we utilize a sliding window of size $\tau$ for constructing a set of $\{y_{t}\}$ from trajectories in a minibatch, then compute the intrinsic bonus using them. For behavior learning, we utilize the actor-critic learning scheme of DreamerV2 (Hafner et al., 2021) that learns values with imagined rewards from future imaginary states and a policy that maximizes the values (see Appendix A for details). We also summarize the difference of APV to DreamerV2 in Appendix D.

Experiments

We designed our experiments to investigate the following:

Can APV improve the sample-efficiency of vision-based RL in robotic manipulation tasks by performing action-free pre-training on videos from different domains?

Can representations pre-trained on videos from manipulation tasks transfer to locomotion tasks?

How does APV compare to a naïve fine-tuning scheme?

What is the contribution of each of the proposed techniques in APV?

How does pre-trained representations qualitatively differ from the randomly initialized representations?

How does APV perform when additional in-domain videos or real-world natural videos are available?

Following Agarwal et al. (2021), we report the interquartile mean with bootstrap confidence interval (CI) and stratified bootstrap CI for results on individual tasks and aggregate results, respectively, across 8 runs for each task. Source codes and other resources are available at https://github.com/younggyoseo/apv.

We first evaluate APV on various vision-based robotic manipulation tasks from Meta-world (Yu et al., 2020). In all manipulation tasks, the episode length is 500 steps without any action repeat, action dimension is 4, and reward ranges from 0 to 10. To evaluate the ability of APV to learn useful representations from different domains, we use videos collected in robotic manipulation tasks from RLBench (James et al., 2020) as pre-training data (see Figure 4).In this work, we do not consider a setup where we perform pre-training on Meta-world videos and fine-tuning for solving RLBench manipulation tasks, as existing RL algorithms struggle to solve challenging, sparsely-rewarded RLBench tasks. Specifically, we collect 10 demonstrations rendered with 5 camera views in 99 tasks from RLBench; giving a total of 4950 videos. We then train the action-free video prediction model by minimizing the objective in Equation 2 for 600K gradient steps. For downstream tasks, we fine-tune the model by minimizing the objective in Equation 6 for 250K environment steps, i.e., 500 episodes.

DeepMind Control Suite experiments.

We also consider widely used robotic locomotion tasks from DeepMind Control Suite (Tassa et al., 2020). Following the common setup in this benchmark (Hafner et al., 2020), the episode length is 1000 steps with the action repeat of 2, and reward ranges from 0 to 1. For pre-training, we consider two datasets: (i) 1000 videos collected from Triped Walk (see Figure 10) and (ii) manipulation videos from RLBench. The former one is for evaluating the performance of APV on in-domain transfer setup similar to the setup in Stooke et al. (2021), while the latter one is for investigating whether the pre-trained representations can be transferred to extremely different domains, i.e., out-of-domain transfer. Specifically, we collect 1000 videos encountered during the training of DreamerV2 agent in Triped Walk and use these videos for pre-training. For downstream tasks, we fine-tune the model for 1M environment steps. See Appendix E for more details.

Hyperparameters.

For newly introduced hyperparameters, we use $\beta_{z}=1.0$ for pre-training, and $\beta_{z}=0,\beta=1.0$ for fine-tuning. We use $\tau=5$ for computing the intrinsic bonus. To make the scale of intrinsic bonus be 10% of extrinsic reward, we normalize the intrinsic reward and use $\lambda=0.1,1.0$ for manipulation and locomotion tasks, respectively. We find that increasing the hidden size of dense layers and the model state dimension from 200 to 1024 improves the performance of both APV and DreamerV2. We use $T=25,50$ for manipulation and locomotion tasks, respectively, during pre-training. Unless otherwise specified, we use the default hyperparameters of DreamerV2.

2 Meta-world Experiments

Figure 5 shows the learning curves of APV pre-trained using the RLBench videos on six robotic manipulation tasks from Meta-world. We find that APV consistently outperforms DreamerV2 in terms of sample-efficiency in all considered tasks. In particular, our framework achieves success rate above 60% on Lever Pull task while DreamerV2 completely fails to solve the task. These results show that APV can leverage action-free videos for learning useful representations that improve the sample-efficiency of vision-based RL.

Comparison with DrQ-v2.

We also compare APV with a state-of-the-art model-free RL method DrQ-v2 (Yarats et al., 2021a) in Appendix F. We find that APV also significantly outperforms DrQ-v2 on most tasks, while DrQ-v2 struggles to achieve strong performance. While this aligns with the observation of Yarats et al. (2021a) where DreamerV2 outperformed DrQ-v2 on DeepMind Control Suite, investigating why DrQ-v2 fails is an interesting future direction.

Comparison with naïve fine-tuning.

To verify the necessity of the proposed architecture for fine-tuning, we compare APV to a naïve fine-tuning scheme that initializes the action-conditional latent dynamics model with the pre-trained parameters of the action-free model (see Appendix E for the details). For a fair comparison, we do not utilize the intrinsic bonus for APV. Figure 6(a) shows that DreamerV2 with this naïve fine-tuning scheme (DreamerV2 w/ Naïve FT) does not provide large gains over DreamerV2, which implies that naïve fine-tuning quickly loses pre-trained representations. By contrast, we find that APV without intrinsic bonus consistently outperforms DreamerV2 by achieving $>10\%$ higher success rate from the beginning of the fine-tuning, even though the same pre-trained model is used for fine-tuning. This shows the proposed architecture is crucial for effective fine-tuning.

Ablation study.

To evaluate the contribution of the proposed techniques in APV, we report the performance of our framework with or without generative pre-training and intrinsic bonus in Figure 6(b). First, we observe that APV (Pre: X / Int: X), whose difference to DreamerV2 is the usage of stacked latent prediction model, achieves similar performance to DreamerV2. This implies that our performance gain is not from the architecture itself, but the way we utilize it for fine-tuning is important. We also find that our intrinsic bonus can improve the performance with or without pre-training, and generative pre-training can also improve the performance with or without intrinsic bonus. Importantly, one can see that the best performance is achieved when both components are combined. This implies that our proposed techniques synergistically contribute to the performance improvement.

Effects of video-based intrinsic bonus.

We investigate the effect of considering multiple model states of length $\tau$ in Equation 5 instead of a single model state. Figure 6(c) shows that APV with $\tau=5$ achieves better performance than $\tau\in\{1,3\}$. We think this is because considering a sequence of observations, i.e., videos, enables us to utilize contextual information for encouraging agents to perform diverse behaviors. But we also find that APV with $\tau=10$ performs worse than $\tau=5$, which might be due to the complexity from average pooling over longer videos.

Qualitative analysis.

We visually investigate why the pre-trained representations can be useful for unseen Meta-world tasks. Specifically, we sample video clips of length 25 from the 10 videos of six manipulation task, and visualize the averaged model states from the sampled videos using t-SNE (Van der Maaten & Hinton, 2008) in Figure 7(a), where colors indicate the tasks. We find that the pre-trained representations from each task are clustered, while randomly initialized ones are entangled. This shows that the pre-trained representations capture information about the tasks without access to Meta-world videos during pre-training.

Importance of dynamics information.

To investigate whether the performance gain comes from utilizing the dynamics information captured in pre-trained representations or visual information captured in the image encoder and decoder, we report the performance of APV when only the pre-trained parameters of the convolutional image encoder and decoder are transferred. Figure 7(b) shows that APV (Encoder / Decoder Only) performs worse than APV, which demonstrates that utilizing the dynamics information learned from diverse videos is crucial for performance.

Pre-training with additional in-domain videos.

We also consider an experimental setup where we have an access to additional videos with similar visuals collected on the target domain. Specifically, we collect 10 videos on each task from ML-10 training tasks in Meta-world, i.e., total of 100 videos, and utilize these videos for pre-training in conjunction with RLBench videos. Then, we fine-tune the pre-trained model for solving 4 manipulation tasks that were not seen during pre-training. To evaluate how additional videos affect representation learning, we report the performance of APV when the representation model of the action-free prediction model is frozen, and use it as a proxy for evaluating the quality of representations.

Figure 7(c) shows that pre-training with additional Meta-world videos (APV (RLB + MW)) achieves almost similar performance to pre-training with only RLBench videos (APV (RLB)). This shows that pre-trained representations with RLBench videos already learns useful representations so that they can be quickly fine-tuned for solving Meta-world tasks. However, we observe that pre-training with in-domain videos significantly improves the performance when the pre-trained representation model is frozen. This implies that additional in-domain videos help for addressing the domain gap between pre-training and fine-tuning.

State and reward regression analysis.

Our hypothesis for how pre-training from videos is useful for RL is that pre-trained representations can capture useful information (e.g., rewards and proprioceptive states) for solving RL tasks. To test this hypothesis, we train a regression model that predicts proprioceptive states and rewards on the pre-collected Triped dataset. Figure 8(a) and Figure 8(b) show the regression performances with and without RLBench pre-trained representations. We find that the RLBench pre-trained model not only has a small prediction error at the beginning, but also quickly converges compared to the model trained from scratch. This shows that pre-training helps the agent to understand environment.

Effects of pre-training datasets

To investigate the effect of pre-training domains, we evaluate the performance of APV on Meta-world when pre-trained with different datasets. To this end, we consider pre-training on manipulation videos from RLBench (i.e., APV w/ Manipulation Videos), locomotion videos from Triped dataset (i.e., APV w/ Locomotion Videos), and real-world natural videos of people performing diverse behaviors (i.e., APV w/ Real-world). Specifically, for real-world videos, we utilize Something-Something-V2 dataset (Goyal et al., 2017), which contains 159K videos of people performing actions. Because we find that our video prediction model suffers from severe underfitting to real-world video datasets (see Appendix G for the examples of blurry predicted future frames), we use subsampled 1.5K videos for pre-training. Moreover, it would be interesting to study whether pre-training on complex robotics video (e.g., RoboNet (Dasari et al., 2019)) can improve performance on more complex robotics tasks (e.g., RLBench).

In Figure 8(c), we observe that pre-training with videos from more similar domains can be helpful. For instance, APV w/ Manipulation outperforms APV w/ Locomotion, which shows that pre-training with manipulation videos can be more effective for manipulation tasks. However, APV w/ Real-world struggles to outperform the baseline without pre-training even though there is no underfitting issue, which might be due to the small number of training data and domain gap between egocentric real-world videos and third-person robotic videos. It would be an interesting direction to investigate how different aspects of pre-training datasets, e.g., size, domain, and point of view.

3 DeepMind Control Suite Experiments

Figure 9 shows the learning curves of APV and DreamerV2 on locomotion tasks. Interesingly, we find that APV pre-trained on manipulation videos from RLBench (pink curves) consistently achieves better performance than DreamerV2 (black curves). This demonstrates that representations pre-trained using manipulation videos, which have notably different visuals and objectives, effectively capture dynamics information useful for quickly learning the dynamics of locomotion environments. Also, by utilizing the videos from a similar domain, i.e., Triped environment, the performance of APV is further improved. We provide additional experimental results that evaluate the contribution of the proposed techniques on DeepMind Control Suite tasks in Appendix I.

Discussion

In this work, we introduce a vision-based RL framework that learns representations useful for understanding the dynamics of downstream domains via action-free pre-training on videos, and utilizes the pre-trained representations for fine-tuning. Our experimental results demonstrate that APV can improve both sample-efficiency and final performances of vision-based RL on various manipulation and locomotion tasks, by effectively transferring the pre-trained representations from unseen domains. However, one limitation of our work is that pre-training is conducted only on simulated robotic videos, which is because of the underfitting issue reported in our analysis. Given that, one interesting direction would be to scale up the architecture or utilize recently developed video prediction architectures (Babaeizadeh et al., 2021; Yan et al., 2021), and investigate how the prediction quality affects the performance. Another interesting direction is to incorporate generalization approaches in RL (Tobin et al., 2017; Higgins et al., 2017; Laskin et al., 2020a) to deal with the difference between pre-training and fine-tuning domains. Moreover, while our work focuses on representation learning via generative pre-training, another interesting future direction would be to investigate the performance of representation learning schemes such as masked prediction (He et al., 2021; Xiao et al., 2022; Yu et al., 2022), latent reconstruction (Yu et al., 2021; Schwarzer et al., 2021a), and contrastive learning (Oord et al., 2018). It would also be interesting to utilize our pre-trained representations for tasks that require more long-term reasoning, such as reward learning from preferences (Park et al., 2022) or videos (Chen et al., 2021). By presenting a generic framework that can leverage videos with diverse visuals and embodiments for pre-training, we hope this work would facilitate future research on unsupervised pre-training for RL.

Acknowledgements

We would like to thank Danijar Hafner, Junsu Kim, Fangchen Liu, Jongjin Park, Jihoon Tack, Wilson Yan, and Sihyun Yu for helpful discussions. We also thank Cirrascale Cloud Serviceshttps://cirrascale.com for providing compute resources. This work was partially supported by Hong Kong Centre for Logistics Robotics, Center for Human Compatible AI (CHAI), Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST)), and Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.2022-0-00166, Self-directed AI Agents with Problem-solving Capability).

References

Appendix A Behavior Learning

For behavior learning, we utilize actor-critic learning scheme of Hafner et al. (2021) where the agent maximizes the values of imagined future states by propagating the analytic gradients back through the world model separately learned with Equation 3. Specifically, given a stochastic actor and a deterministic critic as below:

a sequence of H future states $\hat{s}_{1:H}$ is recursively predicted conditioned on initial states $\hat{s}_{0}$, which are model states encountered during world model training, using the stochastic actor and the transition predictor of world model. Then a deterministic critic is learned to regress the $\lambda$-target (Schulman et al., 2015; Sutton & Barto, 2018) as follows:

where sg is a stop gradient function, and $\lambda$-return $V_{t}^{\lambda}$ is defined using the future states as follows:

Then, the actor is trained to maximize the $\lambda$-return in Equation 8 by leveraging the straight-through estimator (Bengio et al., 2013) for backpropagating the value gradients through the discrete world model as follows:

where the entropy of actor $\,\text{H}\left[a_{t}|\hat{s}_{t}\right]$ is maximized to encourage exploration, and $\eta$ is a hyperparameter that adjusts the strength of entropy regularization.

Appendix B Formulation with Recurrent State-Space Model

While we follow the formulation of Hafner et al. (2020) for our main draft, we additionally provide a detailed formulation of the action-free latent video prediction model and stacked latent prediction model implemented with a recurrent state-space model (RSSM; Hafner et al. 2019) that consists of deterministic and stochastic components, for better understanding of our implementation.

The main component of our model is an action-free RSSM that consists of (i) a recurrent model that computes deterministic states $h^{\tt{AF}}_{t}$ conditioned on previous states for each time step $t$, (ii) a representation model that computes posterior stochastic states $z^{\tt{AF}}_{t}$ conditioned on determinsitic states $h^{\tt{AF}}_{t}$ and image observations $o_{t}$, (iii) a transition predictor that computes prior stochastic states $\hat{z}^{\tt{AF}}_{t}$ without access to image observations. We define the model state as the concatenation of $h^{\tt{AF}}_{t}$ and $z^{\tt{AF}}_{t}$, which is used as an input to the decoder. The model can be summarized as:

At inference time, the recurrent model and the transition predictor (i.e., learned prior) are used for predicting future model states, from which the image decoder reconstruct future frames. All the components of the model parameterized by $\phi$ are jointly optimized by minimizing the following loss:

which is a negative variational lower bound (ELBO; Kingma & Welling 2014) objective where $\beta^{\tt{AF}}$ is a scale hyperparameter and $T$ is the length of training sequences in a minibatch.

B.2 Stacked Latent Prediction Model

Our stacked latent prediction model consists of the action-free RSSM followed by the action-conditional RSSM as below:

All components parameterized by $\phi$ and $\theta$ are optimized jointly by minimizing the following:

where $\beta^{\tt{AC}}$ is a scale hyperparameter.

Appendix C Extended Related Work

A line of works close to our work is video prediction methods, that aims to predict the future frames conditioned on images (Michalski et al., 2014; Ranzato et al., 2014; Srivastava et al., 2015; Vondrick et al., 2016; Lotter et al., 2017), texts (Wu et al., 2021), and actions (Oh et al., 2015; Finn et al., 2016a), which would be useful for various applications, e.g., control with world models (Hafner et al., 2019; Kaiser et al., 2019; Rybkin et al., 2021), and simulator development (Kim et al., 2020, 2021). Recent successful approaches include generative adversarial networks (GANs; Goodfellow et al. 2014) known to generative a sequence of frames by introducing adversarial discriminators that makes prediction based on temporal information (Aigner & Körner, 2018; Jang et al., 2018; Kwon & Park, 2019; Clark et al., 2019; Luc et al., 2020), and autoregressive video prediction models (Babaeizadeh et al., 2018; Kalchbrenner et al., 2017; Reed et al., 2017; Denton & Fergus, 2018; Lee et al., 2018; Villegas et al., 2019; Weissenborn et al., 2020; Wu et al., 2021; Babaeizadeh et al., 2021; Yan et al., 2021; Seo et al., 2022; Yan et al., 2022). Our work builds on prior latent video prediction methods that utilize a state-space model to the domain of video prediction (Hafner et al., 2019; Franceschi et al., 2020), which enables us to predict future states without conditioning on predicted frames in an efficient manner. It would be an interesting direction to scale up current architecture or develop a new high-fidelity latent prediction model and utilize it for pre-training for RL.

Exploration in RL.

Exploration in RL has been studied for encouraging the agents to visit more diverse states by maximizing the entropy of the action space (Ziebart, 2010; Haarnoja et al., 2018) and the various form of intrinsic bonus, including prediction errors (Houthooft et al., 2016; Pathak et al., 2017; Burda et al., 2019; Sekar et al., 2020), count-based state novelty (Bellemare et al., 2016; Tang et al., 2017; Ostrovski et al., 2017), and state entropy (Hazan et al., 2019; Lee et al., 2019; Liu & Abbeel, 2021; Seo et al., 2021; Yarats et al., 2021b). Our framework differs in that we explicitly consider a sequence of future model states for computing the intrinsic bonus, while previous works only consider a single state. In this work, we show that this encourages the agents to explore the environments in a more long-term manner, thus perform more diverse behaviors.

Appendix D Difference to DreamerV2

The main difference lies in the architecture of world models used for imagining future model states. APV utilizes a stacked latent prediction model where the action-free prediction model first processes $o_{t}$ into representations $z_{t}$ and then the action-conditional prediction model processes them into $s_{t}$ conditioned on additional actions. In contrast, DreamerV2 utilizes a RSSM where the action-conditional prediction model directly processes $o_{t}$ into model states $s_{t}$ conditioned on actions. Namely, in terms of architecture, APV can be seen as a DreamerV2 that takes representations from the action-free prediction model as inputs instead of raw observations.

Behavior learning.

The behavior learning scheme of APV is same as DreamerV2 except that APV learns a reward predictor to predict the sum of both extrinsic reward and intrinsic reward. In terms of behavior learning, DreamerV2 can be seen as a specialized case of APV with a scale hyperparameter $\lambda=0$.

Appendix E Experimental Details

We build our framework on top of the official implementation of DreamerV2https://github.com/danijar/dreamerv2, which is based on TensorFlow (Abadi et al., 2016). We use a single Nvidia RTX3090 GPU and 10 CPU cores for each training run. The training time required for pre-training of APV is 24 hours, and fine-tuning of APV requires 4.75 hours for Meta-world experiments, and 6.25 hours for DeepMind Control Suite experiments, when used with XLA optimization. This takes longer than training of vanilla DreamerV2, which requires 3.5 hours for Meta-world experiments, and 4.25 hours for DeepMind Control Suite experiments.

Dataset details.

For data collection in Meta-world environment, we use the corner viewpoint for rendering. We use scripted policy availble in the official implementationhttps://github.com/rlworkgroup/metaworld/tree/master/metaworld. For RLBench environment, we use 5 camera viewpoints consisting of front_rgb, left_shoulder_rgb, overhead_rgb, right_shoulder_rgb, and wrist_rgb. We also utilize the scripted policy available in the official implementationhttps://github.com/stepjam/RLBench. For data collection in Triped Walk, we build the tasks by modifying the Quadruped Walk available in DeepMind Control Suite.

Model details.

For DeepMind Control Suite experiments, we find that utilizing the concatenation of the model state $z_{t}$ and the representation from the image encoder as inputs to the action-conditional prediction model improves the performance, but in Meta-world experiments, we observe that there is no difference to only utilizing $z_{t}$. Therefore, we only use the concatenation as inputs for DeepMind Control Suite experiments. For Meta-world experiments, during fine-tuning, we find that not updating the pre-trained video prediction model at the initial phase of DreamerV2 (which is called pretrain in the original source code) leads to slightly better performance, so we use this for all Meta-world experiments. We also find that increasing the hidden size of linear layers and the dimension of deterministic states in RSSMs from 200 to 1024 improves the performance of both our framework and DreamerV2. In all experiments, we use increased model size for all experiments. Specifically, in our experiments, DreamerV2 has 31M parameters and APV has 45M parameters, while DreamerV2 with default hyperparameters has 15M parameters (see Figure 6(b) for the experimental results demonstrating that the performance gain from APV is not from the increased number of parameters). Unless otherwise specified, we follow the architectures and hyperparameters used in Hafner et al. (2021).

Video-based intrinsic bonus details.

We use $k=16$ for all experiments for computing intrinsic bonus by measuring the distance to $k$-NN representation. For computing the intrinsic bonus, we construct a queue of size 4096 that contains the recent representations from the action-free model, and use the representations in queue for finding $k$-NN state. Our preliminary internal results show that this technique of introducing the queue is not critical to the performance, but we leave the technique here, as it is used in our reported experimental results.

Naïve fine-tuning details.

We implement the naïve fine-tuning scheme in Figure 6(a) by (i) zero-masking the action inputs during pre-training, and (ii) removing the masks at fine-tuning phase and re-initializing the parameters whose inputs are actions. We also applied gradient clipping and warm up scheme to both this baseline and DreamerV2, and find that these techniques does not resolve the difficulty of dealing with additional action inputs.

Appendix F Meta-world Experiments with DrQ-v2

We report the performance of the state-of-the-art model-free RL method DrQ-v2 on vision-based robotic manipulation tasks from Meta-world. We use image observations of $64\times 64\times 3$, and conduct the hyperparameter search over action repeat of $\{1,2\}$ and the frame stacking of $\{3,6\}$. We find that action repeat of 1 with frame stacking of 6 performs the best. As shown in Figure 11, we find that DrQ-v2 struggles to achieve competitive performance on most of the considered tasks, which necessitates more analysis to understand the reason of a failure.

Appendix G Real-World Video Prediction on Something-Something-V2

We report the future frames predicted by our action-free video prediction model trained on real-world videos from Something-Something-V2 (Goyal et al., 2017). One can see that the model severely suffers from underfitting, and generates blurry frames. Developing a lightweight, high-fidelity video prediction model for RL would be an interesting future direction. For instance, we can consider adopting RSSM architecture based on Transformers (Chen et al., 2022).

Appendix H Video Prediction on RLBench and Meta-world

We report the future frames predicted by our action-free video prediction model trained on videos from RLBench (James et al., 2020) and Meta-world (Yu et al., 2020). One can see that predicted frames on both datasets capture dynamics information of robots (e.g., how robots are moving towards objects), in contrast to the prediction on Something-Something-V2 where predicted frames are so blurry that it is difficult to see which behavior is performed in the videos. We also observe that prediction quality in Meta-world is much better than that of RLBench, which is because Meta-world videos have more simple visuals and are collected in smaller number of tasks.

Appendix I Ablation Study on DeepMind Control Suite

We provide the results from ablation studies on locomotion tasks from DeepMind Control Suite in Figure 14. Interestingly, we find that pre-training can significantly improve the performance without intrinsic bonus on Quadruped tasks, but APV with only intrinsic bonus does not make significant difference over vanilla DreamerV2. This might be because the visual observations in Quadruped tasks are very complex, so intrinsic bonus based on randomly initialized representations becomes not particularly useful in the tasks. But APV with both pre-training and intrinsic bonus performs best, which shows that pre-training representations can provide useful information from the beginning of the fine-tuning. On the other hand, on Hopper Hop, we find that APV without intrinsic reward struggles to outperform DreamerV2, which is because Hopper Hop is more difficult in terms of exploration. We also observe that APV without pre-training cannot also outperform DreamerV2 before 300K environment steps, since the model needs a large amount of samples to learn representations that are useful for capturing the dynamics information of the environment. By utilizing the dynamics information captured in the pre-trained representations for exploration, APV with both pre-training and intrinsic bonus performs best from the initial phase of fine-tuning.