Reinforcement Learning with Action-Free Pre-Training from Videos

TL;DR

APV tackles sample inefficiency in vision-based RL by decoupling representation learning from control through action-free pre-training on diverse videos, followed by a stacked latent dynamics model that integrates actions during fine-tuning. A novel video-based intrinsic bonus leverages pre-trained dynamics representations to encourage exploration. Empirical results show substantial gains in Meta-world and transfer to DM Control Suite, with RLBench-based pretraining yielding near-impressive performance and outmatching DreamerV2 in multiple tasks. The work demonstrates meaningful cross-domain transfer from unlabeled video data to downstream control problems and provides a scalable blueprint for future unsupervised pre-training in RL.

Abstract

Recent unsupervised pre-training methods have shown to be effective on language and vision domains by learning useful representations for multiple downstream tasks. In this paper, we investigate if such unsupervised pre-training methods can also be effective for vision-based reinforcement learning (RL). To this end, we introduce a framework that learns representations useful for understanding the dynamics via generative pre-training on videos. Our framework consists of two phases: we pre-train an action-free latent video prediction model, and then utilize the pre-trained representations for efficiently learning action-conditional world models on unseen environments. To incorporate additional action inputs during fine-tuning, we introduce a new architecture that stacks an action-conditional latent prediction model on top of the pre-trained action-free prediction model. Moreover, for better exploration, we propose a video-based intrinsic bonus that leverages pre-trained representations. We demonstrate that our framework significantly improves both final performances and sample-efficiency of vision-based RL in a variety of manipulation and locomotion tasks. Code is available at https://github.com/younggyoseo/apv.

Figures (14)

• Figure 1: We pre-train an action-free latent video prediction model using videos from different domains (left), and then fine-tune the pre-trained model on target domains (right).
• Figure 2: Illustration of action-free latent video prediction model. The model is trained to capture visual and dynamics information from action-free videos by reconstructing image observations. At inference time, the transition model is used to predict future states in the latent space without conditioning on predicted frames.
• Figure 3: Illustration of our framework. (a) We stack an action-conditional prediction model on top of the pre-trained action-free prediction model. At inference time, the transition model in the action-conditional model is used to predict future states in the latent space conditioned on future potential actions. (b) To compute the intrinsic bonus, we first average pool a sequence of model states from the action-free prediction model, and apply random projection to reduce the dimension of representations while preserving distances. The intrinsic bonus for each observation is computed as the distance in the representation space to its $k$-nearest neighbor in samples from a replay buffer.
• Figure 4: Illustration of experimental setups in our experiments with examples of image observations from environments. One can see that visuals in pre-training videos are notably different from the visuals in downstream manipulation and locomotion tasks.
• Figure 5: Learning curves on manipulation tasks from Meta-world as measured on the success rate. APV with generative pre-training on videos collected in manipulation tasks from RLBench consistently outperforms DreamerV2 in terms of sample-efficiency. The solid line and shaded regions represent the interquartile mean and bootstrap confidence intervals, respectively, across eight runs.