# Mask World Model: Predicting What Matters for Robust Robot Policy Learning

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

## Key facts

- Cited by 17 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 22 of its 35 references that are in the map.
- First posted on arXiv 21 Apr 2026.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)
- Vision-language-action models / [World model-based VLA](https://mlsanity.com/vla/world-model-based-vla/)

## Builds on (most cited first, 22 of 22)

- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [WorldSimBench: Towards Video Generation Models as World Simulators](https://mlsanity.com/paper/2410.18072/) (2024)
- [V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning](https://mlsanity.com/paper/2506.09985/) (2025)
- [CogACT: A Foundational Vision-Language-Action Model for Synergizing Cognition and Action in Robotic Manipulation](https://mlsanity.com/paper/2411.19650/) (2024)
- [Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets](https://mlsanity.com/paper/2504.02792/) (2025)
- [Masked World Models for Visual Control](https://mlsanity.com/paper/2206.14244/) (2022)
- [mimic-video: Video-Action Models for Generalizable Robot Control Beyond VLAs](https://mlsanity.com/paper/2512.15692/) (2025)
- [Training Agents Inside of Scalable World Models](https://mlsanity.com/paper/2509.24527/) (2025)
- [ReconVLA: Reconstructive Vision-Language-Action Model as Effective Robot Perceiver](https://mlsanity.com/paper/2508.10333/) (2025)
- [MoLe-VLA: Dynamic Layer-skipping Vision Language Action Model via Mixture-of-Layers for Efficient Robot Manipulation](https://mlsanity.com/paper/2503.20384/) (2025)
- [WoW: Towards a World omniscient World model Through Embodied Interaction](https://mlsanity.com/paper/2509.22642/) (2025)
- [EVA: An Embodied World Model for Future Video Anticipation](https://mlsanity.com/paper/2410.15461/) (2024)
- [WristWorld: Generating Wrist-Views via 4D World Models for Robotic Manipulation](https://mlsanity.com/paper/2510.07313/) (2025)
- [SemanticVLA: Semantic-Aligned Sparsification and Enhancement for Efficient Robotic Manipulation](https://mlsanity.com/paper/2511.10518/) (2025)
- [MinD: Learning A Dual-System World Model for Real-Time Planning and Implicit Risk Analysis](https://mlsanity.com/paper/2506.18897/) (2025)
- [Video2Act: A Dual-System Video Diffusion Policy with Robotic Spatio-Motional Modeling](https://mlsanity.com/paper/2512.03044/) (2025)
- [TTF-VLA: Temporal Token Fusion via Pixel-Attention Integration for Vision-Language-Action Models](https://mlsanity.com/paper/2508.19257/) (2025)
- [Can World Models Benefit VLMs for World Dynamics?](https://mlsanity.com/paper/2510.00855/) (2025)
- [TC-IDM: Grounding Video Generation for Executable Zero-shot Robot Motion](https://mlsanity.com/paper/2601.18323/) (2026)
- [Disentangled World Models: Learning to Transfer Semantic Knowledge from Distracting Videos for Reinforcement Learning](https://mlsanity.com/paper/2503.08751/) (2025)
- [MoWM: Mixture-of-World-Models for Embodied Planning via Latent-to-Pixel Feature Modulation](https://mlsanity.com/paper/2509.21797/) (2025)

## Built on by (most cited first, 17 of 17)

- [World Action Models: The Next Frontier in Embodied AI](https://mlsanity.com/paper/2605.12090/) (2026)
- [Dream-Tac: A Unified Tactile World Action Model for Contact-Rich Robot Manipulation](https://mlsanity.com/paper/2606.08737/) (2026)
- [MaskWAM: Unifying Mask Prompting and Prediction for World-Action Models](https://mlsanity.com/paper/2606.13515/) (2026)
- [World Action Models: A Survey](https://mlsanity.com/paper/2606.20781/) (2026)
- [ST-WAM: Semantic-Temporal World Action Model for Robust Manipulation under Visual Distribution Shifts](https://mlsanity.com/paper/2607.28993/) (2026)
- [Kairos: A Regret-Aware Native World-Action Model Stack for Physical AI](https://mlsanity.com/paper/2606.16533/) (2026)
- [Data Pyramid for Embodied Manipulation: A Survey](https://mlsanity.com/paper/2607.24744/) (2026)
- [DynamicWAM: Dual-Path Motion Conditioning for World-Action Models in Dynamic Manipulation](https://mlsanity.com/paper/2608.00793/) (2026)
- [Foresight Without Seeing: Latent Futures for World Action Models](https://mlsanity.com/paper/2608.11605/) (2026)
- [SG-WAM: Text-Grounded and Spatial-aware Semantic Guidance for World-Action Models](https://mlsanity.com/paper/2608.08839/) (2026)
- [Mask2Real-WM: Segmentation Masks as a Sim-to-Real Bridge for Controllable Dexterous World Models](https://mlsanity.com/paper/2607.04546/) (2026)
- [DC-WAM: Dynamic-Centric Visual Supervision and Reasoning for World-Action Models](https://mlsanity.com/paper/2607.25918/) (2026)
- [AtlasVLA: Persistent World-Ego State Modeling for Vision-Language-Action Models](https://mlsanity.com/paper/2608.06729/) (2026)
- [CAER: Causal Action Effect Reweighting for World Model Training](https://mlsanity.com/paper/2608.30897/) (2026)
- [Token-World: World Modeling in Vision-Language Model Token Space for Robot Manipulation](https://mlsanity.com/paper/2610.00575/) (2026)
- [Native Action-Prior Learning from Videos for World Action Models](https://mlsanity.com/paper/2610.03391/) (2026)
- [Box-WAM: Box Trajectories as Compact State Representations for World Action Models](https://mlsanity.com/paper/2610.17526/) (2026)

## Abstract (arXiv)

World models derived from large-scale video generative pre-training have emerged as a promising paradigm for generalist robot policy learning. However, standard approaches often focus on high-fidelity RGB video prediction, this can result in overfitting to irrelevant factors, such as dynamic backgrounds and illumination changes. These distractions reduce the model's ability to generalize, ultimately leading to unreliable and fragile control policies. To address this, we introduce the Mask World Model (MWM), which leverages video diffusion architectures to predict the evolution of semantic masks instead of pixels. This shift imposes a geometric information bottleneck, forcing the model to capture essential physical dynamics and contact relations while filtering out visual noise. We seamlessly integrate this mask dynamics backbone with a diffusion-based policy head to enable robust end-to-end control. Extensive evaluations demonstrate the superiority of MWM on the LIBERO and RLBench simulation benchmarks, significantly outperforming the state-of-the-art RGB-based world models. Furthermore, real-world experiments and robustness evaluation (via random token pruning) reveal that MWM exhibits superior generalization capabilities and robust resilience to texture information loss.