PaperlayerView exact revision on arXiv

Source-linked AI summary

Masked World Models for Visual Control

Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, Pieter Abbeel

arXiv:2206.14244v3cs.ROcs.AIcs.CVcs.LG

TL;DR

Visual model-based RL struggles with complex visual interactions, especially involving small objects. MWM decouples representation learning from latent dynamics learning using masked convolutional features and auxiliary reward prediction, and solves visual control tasks across several benchmarks.

  • Problem

    Visual world models must accurately model complex visual observations, including interactions with small objects, but end-to-end approaches trade off representation and dynamics learning.

  • Method

    MWM alternates autoencoder training with convolutional feature masking and reward prediction, latent dynamics learning, and online environment interaction.

  • Results

    MWM solves a variety of visual control tasks from Meta-world, RLBench, and DeepMind Control Suite, and significantly outperforms a state-of-the-art model-based baseline on Meta-world and RLBench.

  • Takeaways & Limitations

    Convolutional feature masking and auxiliary reward prediction provide a visual model-based RL framework that captures fine-grained and task-relevant information.

  • Takeaways & Limitations

    Performance heavily depends on auxiliary reward prediction, and the model uses only RGB pixels from a single camera viewpoint.

Abstract

from arXiv · show

Visual model-based reinforcement learning (RL) has the potential to enable sample-efficient robot learning from visual observations. Yet the current approaches typically train a single model end-to-end for learning both visual representations and dynamics, making it difficult to accurately model the interaction between robots and small objects. In this work, we introduce a visual model-based RL framework that decouples visual representation learning and dynamics learning. Specifically, we train an autoencoder with convolutional layers and vision transformers (ViT) to reconstruct pixels given masked convolutional features, and learn a latent dynamics model that operates on the representations from the autoencoder. Moreover, to encode task-relevant information, we introduce an auxiliary reward prediction objective for the autoencoder. We continually update both autoencoder and dynamics model using online samples collected from environment interaction. We demonstrate that our decoupling approach achieves state-of-the-art performance on a variety of visual robotic tasks from Meta-world and RLBench, e.g., we achieve 81.7% success rate on 50 visual robotic manipulation tasks from Meta-world, while the baseline achieves 67.9%. Code is available on the project website: https://sites.google.com/view/mwm-rl.

1 Introduction

Visual model-based RL seeks sample-efficient robot learning but struggles to model fine-grained interactions with small objects. MWM addresses this by decoupling representation and dynamics learning, using convolutional feature masking and reward prediction to retain visual and task-relevant information.

  • Visual world models remain challenged by complex observations, particularly interactions with small objects.
  • End-to-end training creates a trade-off between representation quality and dynamics learning, while separately trained representations may miss dynamics-suitable or task-relevant details.
  • Masked autoencoders motivate scalable ViT representations but can struggle to capture fine-grained details within patches needed for visual control.
  • MWM alternates autoencoder training with convolutional feature masking and reward prediction, latent dynamics learning, and online environment sampling.
  • MWM significantly outperforms a state-of-the-art model-based baseline across visual control tasks from Meta-world and RLBench.
  • Auxiliary reward prediction improves representations by encoding task-relevant information beyond pixel reconstruction.

2 Related Work

Prior work applies visual representations, latent dynamics, and self-supervised ViTs to model-based control. MWM builds on these directions while specifically decoupling representation learning from dynamics learning.

  • Visual model-based methods learn representations through image reconstruction, video prediction, latent dynamics, bisimulation, or contrastive objectives.
  • The experiments cover visual robot control tasks from Meta-world, RLBench, and DeepMind Control Suite.
  • The paper positions MWM as a framework that decouples visual representation learning and dynamics learning.
  • Self-supervised ViTs have been trained with contrastive learning, self-distillation, and masked image modeling.

3 Preliminaries

The preliminaries formulate visual control as a POMDP and review Dreamer and masked autoencoders. These foundations describe latent state prediction, pixel reconstruction, reward prediction, and patch-masked visual representation learning.

  • Problem formulation: A visual control task is modeled as a POMDP with observation space O, action space A, transition dynamics p, reward function r, and discount factor γ.
  • Dreamer: Dreamer learns a world model from pixels and trains an actor-critic through latent imagination.
  • Dreamer: Dreamer’s representation model uses prior state, previous action, and current observation, while its transition model predicts future state without the current observation.
  • Dreamer: The image decoder reconstructs pixels for learning, and the reward predictor computes rewards from future model states without decoding future frames.
  • Dreamer: Dreamer jointly trains model parameters to learn visual representations and environment dynamics, with β controlling their trade-off.
  • Masked autoencoder: Masked autoencoders reconstruct raw pixels from randomly masked patches using a ViT-based encoder-decoder.

4 Masked World Models

Masked World Models decouples visual representation learning from latent dynamics learning. Its autoencoder reconstructs pixels from masked convolutional features, predicts rewards to encode task-relevant information, and supplies representations for a separate dynamics model.

  • Framework: MWM alternates autoencoder updates, latent dynamics learning, and online environment interaction.The autoencoder uses convolutional feature masking and auxiliary reward prediction, while the dynamics model operates in its latent space.
  • Visual Representation Learning: Convolutional feature masking preserves fine-grained visual details by reconstructing pixels from nearby unmasked features.Early convolutional layers mix low-level details, enabling reconstruction of details within patches while retaining masked-modeling stability and compute efficiency.
  • Visual Representation Learning: Auxiliary reward prediction encourages the autoencoder to encode task-relevant information alongside raw pixels.The reward is predicted from the masked latent representation using an additional decoder mask token and output head.
  • Latent Dynamics Learning: The latent dynamics model reconstructs rich autoencoder representations rather than raw pixels, allowing it to focus on dynamics learning.MWM uses frozen autoencoder representations as inputs and targets, including spatially informative elements rather than only a class token.

5 Experiments

MWM is evaluated across Meta-world, RLBench, and DeepMind Control Suite, with ablations and visualizations examining its design. It outperforms DreamerV2 on challenging small-object manipulation, exceeds 80% success on RLBench Reach Target, and remains competitive on visual locomotion.

  • RLBench Experiments: > 80% success rate on RLBench Reach Target for MWM, versus < 20% success rate for DreamerV2.The authors report that these RLBench results are preliminary because they remain too sample-inefficient for real-world scenarios.
  • DeepMind Control Suite Experiments: MWM achieves competitive performance to DreamerV2 on visual locomotion tasks and outperforms it on Reach Duplo.The locomotion results concern Quadruped tasks in DeepMind Control Suite; Reach Duplo is one of the benchmark’s manipulation tasks.
  • Ablation Study: Convolutional feature masking significantly outperforms pixel masking in the ablation study.The comparison uses reward prediction for both conditions and supports the importance of capturing fine-grained details within patches.
  • Ablation Study: Performance significantly degrades without reward prediction, indicating reconstruction alone may not learn task-relevant information.The study identifies reward-free task-relevant representation learning as a future direction because rewards may be unavailable in practice.
  • Qualitative Analysis: MWM visualizations capture changing red-block target positions while ignoring non-task-relevant blue and orange blocks in latent predictions.The autoencoder reconstructions capture all details, whereas the latent dynamics model selectively models task-relevant components.

6 Discussion

The discussion presents MWM as a framework that separates visual representation learning from dynamics learning, while identifying dependencies on reward prediction, limited modalities, and sample efficiency as areas for improvement.

  • Discussion: MWM decouples visual representation learning from dynamics learning for visual model-based reinforcement learning.Its latent dynamics model operates on representations learned by a self-supervised ViT.
  • Limitations: Performance depends heavily on the auxiliary reward-prediction task, possibly because the autoencoder does not learn temporal information.The authors propose video representation learning and auxiliary objectives for other modalities as directions for improvement.
  • Limitations: The model uses only RGB pixels from one camera viewpoint and trains behaviors from scratch, limiting multimodal scope and real-world sample efficiency.The paper suggests demonstrations, path planning, and video pretraining as possible future directions.
  • Behavior learning: The behavior learner uses imagined trajectories, a critic trained on λ-targets, and an actor trained by backpropagating through learned world models.Entropy regularization is used to encourage exploration.
  • Representation learning: Convolutional feature masking is intended to preserve fine-grained visual details that pixel-patch masking may miss, such as object positions.The paper motivates this design with reconstructions from Coffee Pull and Peg Insert Side, where MAE struggles with object positions inside masked patches.
  • Method: The method repeatedly updates a masked autoencoder, a latent dynamics model, and behavior-learning components using environment interaction.The autoencoder uses convolutional feature masking and auxiliary reward prediction before dynamics and actor-critic updates.

F Experiments Details

The experiments use task-specific visual-control settings across Meta-world, RLBench, and DeepMind Control Suite, with MWM requiring more training time than DreamerV2 despite fewer parameters.

  • Meta-world: Meta-world experiments evaluate 50 tasks from a consistent modified corner2 camera viewpoint, with maximum episodes of 500 steps and action repeat of 2.Tasks are grouped by difficulty and trained for 500K to 3M environment steps.
  • RLBench: RLBench experiments use Reach Target and Push Button with dense rewards, delta-joint-position actions, front RGB observations, and maximum episodes of 200 steps.Push Button uses a shaped reward combining gripper-to-button L2 distance and button-push magnitude.
  • DeepMind Control Suite: DeepMind Control Suite experiments use action repeat of 2 and default camera configurations, but their results are not directly comparable with Hafner et al.The authors attribute this to using DreamerV2 with larger networks.
  • Computation: MWM has 25.9M parameters versus DreamerV2's 33.2M, but training takes 5.5 hours versus 3.5 hours over 500K environment steps.The higher training time results from processing observations through the ViT twice, with and without masking.

G Full Meta-world Experiments

Figure 9 reports learning curves for 50 visual robotic manipulation tasks from Meta-world using success rate, with means and bootstrap confidence intervals across five runs.

  • Full Meta-world experiments: Figure 9 measures success rate across 50 Meta-world visual robotic manipulation tasks.The curves show means and bootstrap confidence intervals across five runs.

H Additional DeepMind Control Suite Experiments

Figure 10 reports episode-return learning curves for eight visual robot tasks from the DeepMind Control Suite, summarizing eight runs with bootstrap confidence intervals.

  • Additional DeepMind Control Suite experiments: Figure 10 measures episode return across eight DeepMind Control Suite visual robot tasks.The solid lines and shaded regions represent means and bootstrap confidence intervals across eight runs.

I Extended Ablation Study and Analysis

The ablations examine representation inputs, autoencoder size, and encoder architecture. Using all convolutional and CLS representations helps, whereas DreamerV2 with ViT is unstable; autoencoder size has little effect on Meta-world tasks.

  • Representation inputs: Using all CLS + Conv representations for dynamics learning outperforms using only CLS representations.The broader representation encourages learning spatial information included across representations.
  • Model size: Three autoencoder sizes show no significant performance differences on the evaluated Meta-world tasks.The authors suggest this may reflect the relatively low visual complexity of Meta-world observations.
  • DreamerV2 with ViT: DreamerV2 with ViT exhibits severe training instability and sometimes becomes unable to solve the tasks.This comparison tests whether MWM’s gains arise solely from replacing CNNs with ViTs.
  • Extended ablation study: Figure 11 compares learning curves for three Meta-world manipulation tasks across world-model inputs, autoencoder sizes, and DreamerV2 encoder architectures.Curves show means and stratified bootstrap confidence intervals across 12 runs.
  • Extended ablation study: Figure 12 reports learning curves for individual tasks used in the ablation studies and analysis.Curves show means and bootstrap confidence intervals across four runs.

J Additional Meta-world Experiments with Longer Environment Steps

Longer-horizon and additional-baseline experiments test MWM’s sample efficiency, representation quality, and transfer. MWM remains effective against longer-trained DreamerV2 and alternative representation-learning approaches, while task-relevant reward prediction leaves broader generalization as an open direction.

  • Longer environment steps: DreamerV2 cannot outperform MWM even after 6M environment steps, supporting MWM’s sample-efficient learning and better asymptotic performance.The comparison extends DreamerV2 training substantially beyond the main experiments.
  • Comparison with VAE: A decoupled baseline using frozen VAE representations fails on most tasks, while MWM solves the target tasks.The authors identify representation learning as crucial to the decoupling approach’s performance improvement.
  • Comparison with contrastive baseline: DreamerV2 with contrastive representations outperforms DreamerV2 with reconstruction, but MWM still performs better, especially on Pick Place.The result indicates contrastive learning may capture fine-grained details from small objects more effectively than reconstruction within DreamerV2.
  • Generalizability of representations: Representation learning using only task-irrelevant information remains an open direction for improving applicability across diverse setups.The authors note that task-specific reward prediction can also encode shared robot-motion and object details across tasks.
  • Transfer to unseen tasks: Frozen representations trained on Push are evaluated on Push Back, Pick Place, and Drawer Open tasks with differing manipulation requirements.The transfer setup includes an unseen drawer object in the Drawer Open observations.
  • State regression analysis: State regression experiments compare representations learned with and without reward prediction on seen Push and unseen Pick-place tasks.Regression models predict simulator-provided proprioceptive states from frozen autoencoder representations.
  • State regression analysis: Figure 16 evaluates regression models predicting proprioceptive states from frozen Push-task representations on Push and Pick Place.The curves report means and bootstrap confidence intervals across four runs.
Loading 2206.14244v3…