PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforcement Learning with Action-Free Pre-Training from Videos

Younggyo Seo, Kimin Lee, Stephen James, Pieter Abbeel

arXiv:2203.13880v2cs.CVcs.AI

TL;DR

The paper asks whether action-free video pre-training can provide transferable representations for vision-based RL beyond the pre-training domain. APV pre-trains an action-free latent video model, adds action-conditional fine-tuning and a video-based exploration bonus, and improves performance and sample-efficiency across manipulation and locomotion tasks.

  • Problem

    Existing vision-based RL methods often learn tabula rasa, while prior pre-training work mostly uses target-domain data or visually similar domains.

  • Method

    APV generatively pre-trains an action-free latent video prediction model, stacks an action-conditional latent dynamics model for fine-tuning, and uses video representations for an intrinsic exploration bonus.

  • Results

    APV improves final performance and sample-efficiency across manipulation and locomotion tasks, achieving 95.4% aggregate success on six manipulation tasks versus DreamerV2's 67.9%.

  • Takeaways & Limitations

    Pre-trained representations can transfer from unseen video domains to downstream vision-based RL tasks, even when visuals and objectives differ substantially.

  • Takeaways & Limitations

    Pre-training is limited to simulated robotic videos, and real-world pre-training did not outperform the baseline without pre-training.

Abstract

from arXiv · show

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.

1. Introduction

APV addresses the high interaction cost of vision-based RL by pre-training action-free video representations on diverse domains, then transferring them to downstream tasks. Its stacked latent model and video-based intrinsic bonus support action-conditional learning and diverse exploration, with strong results across manipulation and locomotion.

  • Vision-based RL methods often learn tabula rasa, requiring substantial environment interaction before acquiring meaningful behaviors.
  • Prior vision-based RL pre-training usually uses target-domain data or domains with similar visuals, motivating transfer from diverse videos to newly encountered domains.
  • APV pre-trains an action-free latent video prediction model, then fine-tunes it to learn action-conditional world models on downstream tasks.
  • The stacked latent prediction model adds action-conditional dynamics on top of the action-free model during fine-tuning.
  • APV uses pre-trained video representations for an intrinsic bonus that encourages agents to learn diverse behaviors through better exploration.
  • 95.4% aggregate success rate on six Meta-world manipulation tasks compares with 67.9% for DreamerV2, despite pre-training on different RLBench domains.

2. Related Work

Related work spans unsupervised representation learning in vision, language, and RL, as well as video-based behavior learning. APV builds on this literature by applying action-free generative video pre-training to vision-based RL.

  • Unsupervised CV and NLP methods use objectives such as reconstruction, rotation, masked prediction, contrastive learning, and generative pre-training for downstream fine-tuning.
  • APV demonstrates that unsupervised pre-training can also be effective for vision-based RL.
  • RL representation learning has explored auxiliary losses, world-model learning, reconstruction, and other unsupervised objectives to improve sample-efficiency.
  • Some RL pre-training methods rely on action information from downstream tasks, which may not be available in practice.
  • Video datasets have been used for behavior learning through imitation rewards and related approaches across multiple domains.

3. Method

The method pre-trains an action-free latent video model, fine-tunes a stacked action-conditional model for visual control, and adds a trajectory-diversity intrinsic bonus for exploration.

  • The task is formulated as a partially observable Markov decision process with high-dimensional observations, actions, transition dynamics, rewards, and discounting.
  • Action-free pre-training: The action-free latent video model encodes observations into model states, predicts future states without observations, and reconstructs image observations.Training jointly encourages reconstructed observations and agreement between representation-model and transition-model predictions.
  • Action-free pre-training: Inference predicts future states directly in latent space without decoding future images.
  • Stacked latent prediction model: Fine-tuning stacks an action-conditional prediction model on the pre-trained action-free model to incorporate task actions and rewards.The action-conditional model is used for future imagination after fine-tuning.
  • Video-based intrinsic bonus: The intrinsic bonus averages action-free model states across trajectories, applies random projection, and rewards distance to a replay-buffer nearest neighbor.This trajectory-level diversity signal is intended to encourage longer-term exploration and more diverse behaviors.

4. Experiments

The experiments evaluate APV's transfer across manipulation and locomotion tasks, compare it with DreamerV2 and naïve fine-tuning, and analyze the contributions of pre-training, dynamics representations, intrinsic bonuses, and dataset choice.

  • Meta-world experiments: APV consistently outperforms DreamerV2 in sample-efficiency across six Meta-world manipulation tasks using RLBench videos.APV exceeds 60% success on Lever Pull, while DreamerV2 fails to solve it.
  • Fine-tuning comparison: APV without intrinsic bonus achieves > 10% higher success rate than DreamerV2 from the beginning, whereas naïve fine-tuning provides little gain.The comparison uses the same pre-trained model, indicating that the stacked architecture preserves useful pre-trained representations more effectively.
  • Ablations: Generative pre-training and the intrinsic bonus each improve performance, with their combination achieving the best results.The stacked architecture alone performs similarly to DreamerV2, so the gain depends on how it is used for fine-tuning.
  • Ablations: Using τ = 5 future model states yields better performance than τ ∈ {1, 3}, while τ = 10 performs worse than τ = 5.The authors attribute the τ = 5 advantage to contextual information from sequences of observations.
  • Representation analysis: Dynamics information is crucial: transferring only the image encoder and decoder performs worse than transferring APV representations.Pre-trained representations cluster by task in t-SNE visualizations, unlike randomly initialized representations.
  • Pre-training datasets: Additional in-domain Meta-world videos offer little benefit when all parameters are fine-tuned but significantly help when the representation model is frozen.RLBench-only pre-training already provides useful representations for Meta-world tasks.
  • Representation analysis: RLBench pre-training gives low prediction error for proprioceptive states and rewards from the beginning, while manipulation videos outperform locomotion videos on manipulation tasks.The results support transfer of task-relevant information from action-free videos.

5. Discussion

APV transfers dynamics-relevant representations learned from action-free video pre-training to unseen vision-based RL domains. The authors report gains in sample-efficiency and final performance, while identifying simulated-video pre-training and domain differences as limitations and future directions.

  • APV improves sample-efficiency and final performance across manipulation and locomotion tasks by transferring representations from unseen domains.
  • Pre-training is limited to simulated robotic videos because of an underfitting issue reported in the analysis.
  • Scaling the architecture or using newer video-prediction architectures is proposed to investigate how prediction quality affects performance.
  • Generalization approaches are proposed to address differences between pre-training and fine-tuning domains.
  • Future work could examine masked prediction, latent reconstruction, contrastive learning, and applications requiring longer-term reasoning.

A. Behavior Learning

Behavior learning uses imagined future states from a learned world model, with actor-critic optimization of predicted returns. APV stacks action-free and action-conditional latent models so pre-trained representations can support action-conditioned prediction.

  • The actor maximizes values of imagined future states by propagating analytic gradients through the separately learned world model.
  • Future states are recursively predicted from initial model states using the stochastic actor and the world model's transition predictor.
  • The critic is trained to regress λ-targets, and the actor maximizes the resulting λ-return through a straight-through estimator.
  • Actor entropy is maximized to encourage exploration, with η controlling the strength of entropy regularization.
  • The action-free RSSM computes deterministic and stochastic representations, predicts latent states without image observations, and decodes future model states into frames.
  • The stacked latent prediction model places an action-conditional RSSM after the action-free RSSM.

C. Extended Related Work

APV relates to latent video prediction, intrinsic exploration, and world-model architectures. Its distinguishing design is stacked action-free and action-conditional latent prediction for action-conditioned world modeling.

  • Video prediction: Prior video-prediction methods condition future-frame prediction on images, text, or actions, whereas APV builds on latent state-space video prediction.
  • Video prediction: State-space models allow future states to be predicted without conditioning on predicted frames.
  • Exploration in RL: RL exploration methods include action entropy, prediction-error bonuses, count-based novelty, and state entropy.
  • Architecture: APV first processes observations with an action-free prediction model, then processes the resulting representations with an action-conditional model using additional actions.
  • Architecture: DreamerV2 directly processes observations with an action-conditional RSSM, whereas APV uses stacked latent prediction.
  • Behavior learning: DreamerV2 is a specialized case of APV with intrinsic-bonus scale λ = 0.

E. Experimental Details

Experiments use DreamerV2-based implementations with substantial pre-training and fine-tuning compute, multiple scripted-video datasets, and task-dependent modeling choices. Additional implementation details cover intrinsic exploration and naïve fine-tuning.

  • Implementation: APV pre-training takes 24 hours, while fine-tuning takes 4.75 hours for Meta-world and 6.25 hours for DeepMind Control Suite experiments.
  • Implementation: Experiments use one Nvidia RTX3090 GPU and 10 CPU cores per training run, with XLA optimization for fine-tuning.
  • Dataset details: Meta-world and RLBench videos use scripted policies, with RLBench collecting five camera viewpoints.
  • Model details: DeepMind Control Suite uses concatenated model-state and image-encoder representations, while Meta-world uses only the model state because concatenation showed no difference.
  • Video-based intrinsic bonus: The intrinsic bonus uses k = 16 nearest neighbors and a queue of 4096 recent action-free representations.
  • Naïve fine-tuning: Naïve fine-tuning masks actions during pre-training and reinitializes action-input parameters at fine-tuning, but gradient clipping and warm-up do not resolve the added-action difficulty.

F. Meta-world Experiments with DrQ-v2

DrQ-v2 struggles to achieve competitive performance on most considered Meta-world manipulation tasks under the reported vision-based setup.

  • DrQ-v2 struggles to achieve competitive performance on most considered Meta-world manipulation tasks.The experiments use 64 × 64 × 3 image observations; action repeat 1 and frame stacking 6 perform best.

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

The action-free video prediction model severely underfits on Something-Something-V2, producing blurry future frames. The authors identify lightweight, high-fidelity video prediction as a future direction.

  • The model severely suffers from underfitting and generates blurry future frames on Something-Something-V2.The authors suggest adopting a Transformer-based RSSM architecture as one possible future direction.

H. Video Prediction on RLBench and Meta-world

Predictions on RLBench and Meta-world capture robot dynamics information, although RLBench predictions are lower quality than Meta-world predictions.

  • Predicted frames capture robot dynamics information on both RLBench and Meta-world.The frames represent behaviors such as robots moving toward objects, unlike the blurry Something-Something-V2 predictions.
  • RLBench predictions are lower quality than Meta-world predictions because Meta-world is a simpler domain.

I. Ablation Study on DeepMind Control Suite

Ablations show that pre-training and intrinsic bonuses contribute differently across locomotion tasks. Their combination performs best, while exploration difficulty and representation quality affect the gains.

  • APV with both pre-training and intrinsic bonus performs best in the locomotion ablations.Pre-training representations provide useful information from the beginning of fine-tuning.
  • On Quadruped tasks, pre-training improves performance without intrinsic bonus, whereas intrinsic bonus alone does not significantly improve vanilla DreamerV2.The visual observations are complex, making bonuses based on randomly initialized representations less useful.
  • On Hopper Hop, APV without intrinsic reward struggles to outperform DreamerV2 because exploration is more difficult.
  • Without pre-training, APV cannot outperform DreamerV2 before 300K environment steps.The model requires many samples to learn representations useful for capturing environmental dynamics.
Loading 2203.13880v2…