PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yunfan Lou, Xiaowei Chi, Xiaojie Zhang, Zezhong Qian, Chengxuan Li, Rongyu Zhang, Yaoxu Lyu, Guoyu Song, Chuyao Fu, Haoxuan Xu, Pengwei Wang, Shanghang Zhang

arXiv:2604.19683v2cs.RO

TL;DR

RGB-focused world models can spend capacity on visual variation that is weakly related to action selection, limiting robust robot control under appearance shifts. MWM predicts future semantic masks with a diffusion world model and uses its mask-centric features to condition a diffusion policy. Across simulation and real-robot evaluations, it consistently outperforms strong RGB-centric baselines and shows improved robustness and generalization.

  • Problem

    RGB video prediction can entangle control-relevant dynamics with nuisance appearance changes, making reliable generalist manipulation under visual variability difficult.

  • Method

    MWM forecasts future semantic masks with conditional diffusion and conditions a diffusion policy on the resulting mask-centric predictive features, using semantic supervision only during training.

  • Results

    MWM consistently outperforms strong RGB-centric baselines across LIBERO, RLBench, and real-robot evaluations, achieving 98.3% average success on LIBERO, 68.3% on RLBench, and 67.5% on a real Franka robot.

  • Takeaways & Limitations

    Semantic prediction provides a practical information bottleneck for learning more decision-relevant dynamics representations for generalist robot manipulation.

  • Takeaways & Limitations

    Semantic masks supervise the backbone during training and must be available for task-relevant entities, including the robot arm, gripper, and manipulated objects.

Abstract

from arXiv · show

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.

1. Introduction

MWM addresses the mismatch between photometric RGB prediction and robust control by forecasting semantic masks and conditioning a diffusion policy on mask-centric predictive features. It reports gains across simulation and real-robot evaluations while using semantic supervision only during training.

  • Motivation: RGB prediction entangles action-relevant dynamics with nuisance variation such as texture, lighting, reflections, and dynamic backgrounds, causing predictive drift under distribution shifts.These appearance-driven errors can compound during closed-loop execution, producing brittle policies.
  • Approach: MWM shifts prediction from future RGB frames to future semantic masks, imposing a geometric bottleneck that preserves object identity, spatial layout, and interaction structure.Semantic labels are used offline during training, while deployment requires only raw multi-view RGB.
  • Approach: MWM first forecasts future semantic masks with conditional diffusion, then trains a diffusion policy conditioned on intermediate mask-centric features to generate actions.The two-stage design explicitly couples predictive representations to control rather than using masks only for visualization.
  • Results: 98.3% average success rate on LIBERO, 68.3% on RLBench, and 67.5% average success on a real Franka robot across four tasks outperform strong RGB-centric baselines.The evaluations also include visual appearance shifts and random visual token pruning.
  • Results: Mask-centric designs consistently outperform future RGB prediction, indicating that gains primarily arise from changing the representation and prediction objectives rather than relying on one architecture.This comparison supports the semantic shift as the central design factor.

2. Related Work

Prior robotics world models and structured visual representations seek compact or semantically useful features, but many predictive systems still target photometric video. MWM applies semantic structure directly as the prediction space for a diffusion world model in manipulation.

  • Predictive world models: Recent robotics systems repurpose diffusion and transformer video generators as predictive backbones, often coupling multi-view video prediction with action heads for manipulation.These approaches build on earlier latent-space world models that avoid direct pixel prediction.
  • Semantic prediction: Despite strong results, related predictive methods primarily retain photometric targets, whereas MWM forecasts future semantic mask dynamics and needs no external segmenter at test time.Semantic supervision is used only during training.
  • Semantic representations: VLA policies can benefit from representations that expose object state and interaction geometry while reducing coupling to photometric variation in precise contact-sensitive tasks.This motivates injecting semantics for perceptual redundancy reduction and grounding.
  • Structured representations: Object-centric representations, scene factorization, masked modeling, and token dropping support the broader intuition that compact, structure-biased features can improve stability over raw pixels.MWM instantiates this intuition by making semantic structure the predictive space of a diffusion world model for manipulation.

3. Method

MWM learns mask-centric predictive dynamics from multi-view RGB and language, then uses those predictive features for diffusion-based action generation. Its geometric bottleneck preserves task-relevant structure while reducing reliance on nuisance appearance variation.

  • Two-stage training: Training uses two stages: mask-dynamics pretraining with offline semantic supervision, followed by action-policy learning that updates the backbone through the action loss.The second stage freezes the VAE and jointly trains the DiT backbone and action expert without video or mask reconstruction loss.
  • Mask-centric prediction: MWM forecasts future semantic mask latents rather than RGB pixels, preserving object identity, spatial layout, and interaction-relevant structure while discarding redundant appearance.Semantic masks provide the predictive representation, while RGB observations remain the deployment input.
  • Latent representation: RGB observations and rendered semantic masks share a pretrained VAE interface, allowing mask targets to be encoded as continuous latents without modifying the VAE.Masks are rendered with a fixed color palette before encoding, while deployment uses raw RGB only.
  • Token preparation: Latents are channel-normalized, temporally interpolated, stacked across views, and flattened into a fixed-length token sequence for the transformer backbone.This processing supports varying temporal subsampling rates and VAE compression artifacts.
  • Mask dynamics: The backbone predicts future mask latents from fixed RGB memory slots and language using conditional flow matching, with conditioning injected through cross-attention into the predicted velocity field.The setup uses n=4 memory frames, τ=5 future latent frames, and applies the mask loss only to future slots.
  • Policy learning: A multi-level predictive feature bank is exposed to an action diffusion head, which predicts 15-dimensional action-state chunks through conditional denoising and receding-horizon execution.The action head attends to backbone features via cross-attention; action chunks have Ha=36, and inference iteratively denoises each chunk before execution.

4. Experiments

MWM is evaluated across simulated and real-world manipulation benchmarks, using semantic-mask prediction and multi-view observations to test performance, robustness, and generalization. It consistently outperforms RGB-centric baselines across standard success-rate evaluations and remains more reliable under appearance shifts and token pruning.

  • Simulation benchmarks: LIBERO evaluation uses four suites, while RLBench evaluation covers standardized tabletop manipulation tasks with language-conditioned goals.LIBERO results are evaluated over 500 episodes per suite; RLBench checkpoints are evaluated over 20 randomized episodes.
  • LIBERO results: 0.810→0.918 average SR: MWM-C2 outperforms MWM-C1, while MWM-C1 improves over Cosmos w/ IDM from 0.675 to 0.810.The largest MWM-C1 gain occurs on long-horizon LIBERO-10, rising from 0.488 to 0.704; MWM-C2 improves over Cosmos w/ Latent IDM from 0.873 to 0.918.
  • RLBench results: 68.3% average SR: MWM outperforms GE-ACT at 30.8% and FiS-VLA at 50.0% on RLBench.The largest gains occur on tasks requiring precise object interaction and goal satisfaction.
  • Real-world results: 67.5% average SR: MWM exceeds GE-ACT at 23.8% and π0 at 38.8% across four Franka robot tasks.Each task uses 50 demonstrations for post-training, and each method is evaluated over 20 randomized trials per task.
  • Token-pruning robustness: 0.648 vs. 0.629 nPAUC: MWM shows slightly higher robustness than GE-ACT under random visual-token pruning.The metric averages success retention across pruning ratios from 0.1 to 0.9 and the four LIBERO suites.
  • Visual generalization: 42.1% OOD-SR: MWM achieves the highest average performance under background, lighting, and object-color shifts, versus 12.5% for GE-ACT and 19.2% for π0.MWM reaches 27.5% BG-SR, 56.3% under lighting shifts, 42.5% under color shifts, and 0.62 retention of in-distribution performance.

5. Conclusion

MWM forecasts future semantic masks instead of pixels to align world-model prediction with control utility. It uses offline semantic supervision during training, raw multi-view RGB at test time, and a diffusion policy head, with consistent gains across evaluations.

  • MWM forecasts future semantic masks to reduce the mismatch between photometric RGB prediction and control utility.
  • Semantic supervision is used only offline, while inference operates directly on raw multi-view RGB observations.
  • MWM couples mask-centric predictive features with a diffusion policy head and consistently outperforms strong RGB-centric baselines across simulation and real-robot evaluations.

A. Implementation Details

The implementation overview provides a structured account of MWM's architecture, training protocol, and reproducibility settings. Detailed hyperparameters are collected in Table 6.

  • The implementation details cover MWM's architecture.
  • The overview documents the training protocol.
  • Detailed hyperparameters are summarized in Table 6.

A.1. Network Architecture

MWM uses a pretrained 3D VAE for video compression, a multi-view Diffusion Transformer mask backbone, and a hierarchical action expert that consumes backbone features.

  • Video VAE & Tokenization: A fixed pretrained 3D VAE compresses 256×256 RGB frames into 8×8 latent patches with channel dimension D=128.The spatial and temporal compression ratios are fs=32 and ft=8.
  • Mask Dynamics Backbone: The mask dynamics backbone is a 28-layer DiT with hidden dimension 2048 and 32 attention heads, processing third-person and wrist views with shared spatiotemporal attention.Cross-view attention is inserted at every third block.
  • Action Expert Head: The 28-layer action transformer uses Predictive Feature Bank cross-attention to consume spatial-temporal features from corresponding DiT layers.The action head has 512 hidden dimensions and 16 attention heads.

A.2. Infrastructure and Reproducibility

The infrastructure details specify multi-GPU training, preprocessing choices, fixed randomness, and a hyperparameter reference table for reproducibility.

  • Infrastructure: Training uses 8× NVIDIA A100 (80GB) GPUs with DeepSpeed ZeRO-2.
  • Infrastructure: Stage 1 takes approximately 3.5 days for 30k steps, while Stage 2 takes approximately 1.5 days for 18k steps per task suite.
  • Data Processing: Images are resized to 256×256, normalized to [−1, 1], and semantic masks are rendered with a consistent color palette.All experiments use fixed random seed 42.
  • Reproducibility: Detailed MWM hyperparameters are provided in Table 6.

B. Real-World Experimental Setup

The real-world evaluation uses a dual-view Franka Panda setup and semantic masks generated offline for MWM pre-training. The section also provides detailed training and inference procedures for both MWM variants.

  • Evaluation Tasks: The four real-world tasks involve placing food into a basket, opening a drawer to insert a pen, and two additional manipulation tasks shown in Figure 4.Figure 4 presents snapshots of the evaluated task environments.
  • Hardware Configuration: The physical testbed uses a 7-DoF Franka Emika Panda robot with a parallel gripper and two synchronized Intel RealSense D435i cameras.One camera provides a fixed third-person workspace view, while the wrist-mounted camera supplies fine-grained, contact-focused feedback.
  • Semantic Annotation Process: RoboEngine automatically generates pixel-wise semantic annotations for the robot arm, gripper, and task-relevant objects from demonstration videos.The masks are used as supervision targets during Stage 1 pre-training, while inference relies on raw multi-view observations.

D. Detailed Random Token Pruning Results

The Random Token Pruning stress test evaluates MWM and GE-ACT across pruning ratios from 0.1 to 0.9 on four LIBERO evaluation suites.

  • Pruning Setup: The experiment compares MWM (Ours) and GE-ACT under random token pruning ratios r ∈ {0.1, ..., 0.9}.The test examines robustness as increasing fractions of visual tokens are pruned.
  • Evaluation Suites: Success rates are reported for the LIBERO Spatial, Object, Goal, and Libero-10 evaluation suites.The section provides detailed success rates for each pruning ratio and evaluation suite.
  • Comparison: The stress test is designed to compare the robustness profiles of MWM and GE-ACT under visual token removal.The supplied section identifies the models, pruning range, and four evaluation suites but does not state the numerical outcomes.

E. Detailed Visual Generalization Results

The visual generalization evaluation measures GE-ACT, π0, and MWM across four real-world tasks under background, lighting, and object-color shifts. Detailed success-rate tables and visualizations define these stress tests, while additional tables describe token-pruning robustness.

  • Evaluation Scope: GE-ACT, π0, and MWM are evaluated across four real-world tasks under Background, Lighting, and Object Color shifts.The reported metric is task success rate for each method, task, and shift condition.
  • Purpose of Stress Tests: The stress tests examine whether policies overfit to nuisance variables such as table texture and illumination instead of learning manipulation geometry.Figure 5 visualizes the environmental variations used in these tests.
  • Token Pruning: Separate robustness-profile tables report success rates for MWM and GE-ACT across varying visual token-pruning ratios.Table 7 covers MWM, while Table 8 covers GE-ACT.
  • Reporting: Reported visual-generalization values are success rates (%) averaged over 20 trials per task.This reporting convention is stated for the detailed shift tables.
  • Model Procedures: The appendix includes algorithmic procedures for both MWM variants alongside the visual-generalization results.Algorithm 1 describes MWM-C1 with inverse dynamics, and Algorithm 2 describes MWM-C2 with a predictive feature policy.
  • Background Shift: Background shifts replace the table surface or tablecloth texture with unseen patterns while preserving the manipulation setting.Table 9 reports detailed success rates for this condition.
  • Lighting Shift: Lighting shifts significantly alter illumination intensity, including dimmer and brighter conditions.Table 10 reports detailed success rates under the lighting variation.
  • Object Color Shift: Object Color shifts replace task objects with unseen colors or instances while retaining their geometry.Table 11 reports detailed success rates for this condition.
Loading 2604.19683v2…