# Masked World Models for Visual Control

arXiv: https://arxiv.org/abs/2206.14244 · HTML: https://mlsanity.com/paper/2206.14244/

## Key facts

- Cited by 170 papers in the map (data through 14 Oct 2026).
- Placed in Embodied and robot world models, one of 11 directions in World models.
- Builds on 79 of its 88 references that are in the map.
- First posted on arXiv 28 Jun 2022.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)

## Builds on (most cited first, 30 of 79)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [ViViT: A Video Vision Transformer](https://mlsanity.com/paper/2103.15691/) (2021)
- [An Empirical Study of Training Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.02057/) (2021)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [VideoMAE: Masked Autoencoders are Data-Efficient Learners for Self-Supervised Video Pre-Training](https://mlsanity.com/paper/2203.12602/) (2022)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [SimMIM: A Simple Framework for Masked Image Modeling](https://mlsanity.com/paper/2111.09886/) (2021)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [CvT: Introducing Convolutions to Vision Transformers](https://mlsanity.com/paper/2103.15808/) (2021)
- [iBOT: Image BERT Pre-Training with Online Tokenizer](https://mlsanity.com/paper/2111.07832/) (2021)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [R3M: A Universal Visual Representation for Robot Manipulation](https://mlsanity.com/paper/2203.12601/) (2022)
- [RLBench: The Robot Learning Benchmark & Learning Environment](https://mlsanity.com/paper/1909.12271/) (2019)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Masked Feature Prediction for Self-Supervised Visual Pre-Training](https://mlsanity.com/paper/2112.09133/) (2021)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [LeViT: a Vision Transformer in ConvNet's Clothing for Faster Inference](https://mlsanity.com/paper/2104.01136/) (2021)

## Built on by (most cited first, 30 of 170)

- [SmolVLA: A Vision-Language-Action Model for Affordable and Efficient Robotics](https://mlsanity.com/paper/2506.01844/) (2025)
- [Unleashing Large-Scale Video Generative Pre-training for Visual Robot Manipulation](https://mlsanity.com/paper/2312.13139/) (2023)
- [TinyVLA: Towards Fast, Data-Efficient Vision-Language-Action Models for Robotic Manipulation](https://mlsanity.com/paper/2409.12514/) (2024)
- [Scaling Language-Image Pre-training via Masking](https://mlsanity.com/paper/2212.00794/) (2022)
- [GR-2: A Generative Video-Language-Action Model with Web-Scale Knowledge for Robot Manipulation](https://mlsanity.com/paper/2410.06158/) (2024)
- [A Survey on Vision-Language-Action Models for Embodied AI](https://mlsanity.com/paper/2405.14093/) (2024)
- [Navigation World Models](https://mlsanity.com/paper/2412.03572/) (2024)
- [DriveDreamer: Towards Real-world-driven World Models for Autonomous Driving](https://mlsanity.com/paper/2309.09777/) (2023)
- [DriveDreamer-2: LLM-Enhanced World Models for Diverse Driving Video Generation](https://mlsanity.com/paper/2403.06845/) (2024)
- [Foundation Models for Decision Making: Problems, Methods, and Opportunities](https://mlsanity.com/paper/2303.04129/) (2023)
- [Predictive Inverse Dynamics Models are Scalable Learners for Robotic Manipulation](https://mlsanity.com/paper/2412.15109/) (2024)
- [Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers](https://mlsanity.com/paper/2409.20537/) (2024)
- [3D Diffusion Policy: Generalizable Visuomotor Policy Learning via Simple 3D Representations](https://mlsanity.com/paper/2403.03954/) (2024)
- [Survey on Large Language Model-Enhanced Reinforcement Learning: Concept, Taxonomy, and Methods](https://mlsanity.com/paper/2404.00282/) (2024)
- [AdaWorld: Learning Adaptable World Models with Latent Actions](https://mlsanity.com/paper/2503.18938/) (2025)
- [ManiGaussian: Dynamic Gaussian Splatting for Multi-task Robotic Manipulation](https://mlsanity.com/paper/2403.08321/) (2024)
- [UniDexGrasp++: Improving Dexterous Grasping Policy Learning via Geometry-aware Curriculum and Iterative Generalist-Specialist Learning](https://mlsanity.com/paper/2304.00464/) (2023)
- [ADriver-I: A General World Model for Autonomous Driving](https://mlsanity.com/paper/2311.13549/) (2023)
- [WorldDreamer: Towards General World Models for Video Generation via Predicting Masked Tokens](https://mlsanity.com/paper/2401.09985/) (2024)
- [Dreamitate: Real-World Visuomotor Policy Learning via Video Generation](https://mlsanity.com/paper/2406.16862/) (2024)
- [Multi-View Masked World Models for Visual Robotic Manipulation](https://mlsanity.com/paper/2302.02408/) (2023)
- [OVRL-V2: A simple state-of-art baseline for ImageNav and ObjectNav](https://mlsanity.com/paper/2303.07798/) (2023)
- [Reinforcement Learning from Passive Data via Latent Intentions](https://mlsanity.com/paper/2304.04782/) (2023)
- [A Comprehensive Survey on World Models for Embodied AI](https://mlsanity.com/paper/2510.16732/) (2025)
- [Imitation Bootstrapped Reinforcement Learning](https://mlsanity.com/paper/2311.02198/) (2023)
- [InterDreamer: Zero-Shot Text to 3D Dynamic Human-Object Interaction](https://mlsanity.com/paper/2403.19652/) (2024)
- [MoDem: Accelerating Visual Model-Based Reinforcement Learning with Demonstrations](https://mlsanity.com/paper/2212.05698/) (2022)
- [A Survey on Masked Autoencoder for Self-supervised Learning in Vision and Beyond](https://mlsanity.com/paper/2208.00173/) (2022)
- [ManiCM: Real-time 3D Diffusion Policy via Consistency Model for Robotic Manipulation](https://mlsanity.com/paper/2406.01586/) (2024)
- [A Survey on Transformers in Reinforcement Learning](https://mlsanity.com/paper/2301.03044/) (2023)

## Abstract (arXiv)

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.