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

• URL: http://arxiv.org/abs/2604.19683v2
• Date: Wed, 22 Apr 2026 17:44:56 GMT
• Title: Mask World Model: Predicting What Matters for Robust Robot Policy Learning
• Abstract summary: Mask World Model (MWM) is a generalization of large-scale video generative pre-training world models.<n>We show that MWM exhibits superior generalization capabilities and robust resilience to texture information loss.
• Score: 31.96162737409163
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• WM-Craftnet: World Synesthesia Model for Generalizable and Robust Dexterous In-Hand Manipulation [7.577131730633025]
  Generalizable and robust in-hand manipulation requires a policy to infer object pose, geometry, contact, and potential slip.<n>We propose WM-Craftnet, a world-model-conditioned framework that learns compact action-conditioned latent dynamics.
  arXiv  Detail & Related papers  (2026-09-07T03:48:06Z)
• Sword: Style-Robust World Models as Simulators via Dynamic Latent Bootstrapping for VLA Policy Post-Training [19.013714390803052]
  One approach treats learned World Models as generative simulators, enabling policy optimization entirely within "imagination"<n>However, when deployed as simulators for specific environments such as the LIBERO benchmark, existing World Models often suffer from poor generalization and long-horizon error accumulation.<n>We propose Sword, a robust World Model framework, to mitigate these problems.
  arXiv  Detail & Related papers  (2026-05-08T05:54:33Z)
• Opening the Sim-to-Real Door for Humanoid Pixel-to-Action Policy Transfer [59.02729900344616]
  GPU-accelerated, photorealistic simulation has opened a scalable data-generation path for robot learning.<n>We develop a teacher-student-bootstrap learning framework for vision-based humanoid loco-manipulation.<n>This represents the first humanoid sim-to-real policy capable of diverse articulated loco-manipulation using pure RGB perception.
  arXiv  Detail & Related papers  (2025-11-30T20:07:13Z)
• DyMoDreamer: World Modeling with Dynamic Modulation [52.27044216359359]
  A critical bottleneck in deep reinforcement learning (DRL) is sample inefficiency, as training high-performance agents often demands extensive environmental interactions.<n>We introduce DyMoDreamer, a novel algorithm that incorporates a dynamic modulation mechanism to improve the extraction of dynamic features and enrich the temporal information.<n>Experiments demonstrate that DyMoDreamer sets a new state-of-the-art on the Atari $100$k benchmark with a $156.6$% mean human-normalized score.
  arXiv  Detail & Related papers  (2025-09-29T13:54:42Z)
• GWM: Towards Scalable Gaussian World Models for Robotic Manipulation [53.51622803589185]
  We propose a novel branch of world model named Gaussian World Model (GWM) for robotic manipulation.<n>At its core is a latent Diffusion Transformer (DiT) combined with a 3D variational autoencoder, enabling fine-grained scene-level future state reconstruction.<n>Both simulated and real-world experiments depict that GWM can precisely predict future scenes conditioned on diverse robot actions.
  arXiv  Detail & Related papers  (2025-08-25T02:01:09Z)
• MaskGWM: A Generalizable Driving World Model with Video Mask Reconstruction [8.503246256880612]
  We propose MaskGWM: a Generalizable driving World Model embodied with Video Mask reconstruction.<n>Our model contains two variants: MaskGWM-long, focusing on long-horizon prediction, and MaskGWM-mview, dedicated to multi-view generation.
  arXiv  Detail & Related papers  (2025-02-17T10:53:56Z)
• Learning Latent Dynamic Robust Representations for World Models [9.806852421730165]
  Visual Model-Based Reinforcement Learning (MBL) promises to agent's knowledge about the underlying dynamics of the environment. Top-temporal agents such as Dreamer often struggle with visual pixel-based inputs in the presence of irrelevant noise in the observation space. We apply a-temporal masking strategy, combined with latent reconstruction, to capture endogenous task-specific aspects of the environment for world models.
  arXiv  Detail & Related papers  (2024-05-10T06:28:42Z)
• Transferring Foundation Models for Generalizable Robotic Manipulation [82.12754319808197]
  We propose a novel paradigm that effectively leverages language-reasoning segmentation mask generated by internet-scale foundation models.<n>Our approach can effectively and robustly perceive object pose and enable sample-efficient generalization learning.<n>Demos can be found in our submitted video, and more comprehensive ones can be found in link1 or link2.
  arXiv  Detail & Related papers  (2023-06-09T07:22:12Z)
• Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learning [54.67880602409801]
  In this paper, we study the problem of pre-training world models with abundant in-the-wild videos for efficient learning of visual control tasks. We introduce Contextualized World Models (ContextWM) that explicitly separate context and dynamics modeling. Our experiments show that in-the-wild video pre-training equipped with ContextWM can significantly improve the sample efficiency of model-based reinforcement learning.
  arXiv  Detail & Related papers  (2023-05-29T14:29:12Z)
• Virtual Data Augmentation: A Robust and General Framework for Fine-tuning Pre-trained Models [51.46732511844122]
  Powerful pre-trained language models (PLM) can be fooled by small perturbations or intentional attacks. We present Virtual Data Augmentation (VDA), a general framework for robustly fine-tuning PLMs. Our approach is able to improve the robustness of PLMs and alleviate the performance degradation under adversarial attacks.
  arXiv  Detail & Related papers  (2021-09-13T09:15:28Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.