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

In the field of robot policy learning, maintaining reliability under visual variability remains a central challenge. Traditional approaches often rely on high-fidelity RGB video prediction, which can lead to overfitting to irrelevant factors such as dynamic backgrounds and illumination changes, ultimately reducing the model's ability to generalize. Recently, with the development of video generative pre-training techniques, world models derived from large-scale video generative pre-training have emerged as a promising paradigm for generalist robot policy learning. However, these methods often focus on predicting RGB pixels, and this photometric objective is often misaligned with control tasks. RGB frames contain substantial nuisance variation, including texture, lighting, reflections, and dynamic backgrounds, which are weakly related to action selection. Pixel prediction compels a model to allocate capacity to these factors and to entangle appearance with dynamics, treating changes in illumination or background as comparable to contact-relevant motion. In closed-loop execution, this misallocation becomes more damaging: small appearance-driven errors can compound over time, causing predictive drift and brittle policies under modest distribution shifts.

Traditional RGB video prediction methods in robot policy learning face overfitting issues, particularly to irrelevant factors such as dynamic backgrounds and illumination changes. This reduces the model's ability to generalize, ultimately leading to unreliable and fragile control policies. To achieve robust robot policy learning, there is an urgent need for a method that can effectively filter out visual noise and capture decision-relevant geometric information.

The core innovation of the Mask World Model (MWM) lies in shifting the predictive space from RGB frames to future semantic masks. Semantic masks impose a geometric bottleneck that preserves object identity, spatial layout, and interaction-relevant structure while discarding redundant appearance. MWM does not require an external segmentation model at inference: semantic labels are used only offline during training, while deployment uses only raw multi-view RGB. The training pipeline adopts a two-stage strategy, first learning a mask-centric predictive model via conditional diffusion objectives, and then training a diffusion policy head based on mask-centric predictive features.

• �� MWM employs video diffusion architectures to predict the evolution of semantic masks rather than pixels.

• �� Training is conducted in two stages: first, a mask-centric predictive model is learned via conditional diffusion objectives, and then a diffusion policy head is trained based on mask-centric predictive features.

• �� The mask dynamics backbone is seamlessly integrated with a diffusion-based policy head to enable robust end-to-end control.

• �� By introducing a geometric information bottleneck, the model is forced to capture essential physical dynamics and contact relations while filtering out visual noise.

The experimental design includes evaluations on the LIBERO and RLBench benchmarks. LIBERO consists of 130 simulated manipulation tasks evaluated using templated language instructions. RLBench contains 100 tabletop manipulation tasks evaluated using standardized multi-view observations and natural language goals. In RLBench, 20 evaluation episodes per task are conducted using randomized seeds and initializations. Baseline models include OpenVLA, CogACT, π, Cosmos+IDM, Cosmos+LatentIDM, and GE-ACT.

In the LIBERO benchmark, MWM achieved a 98.3% average success rate, significantly outperforming RGB-based world models. In RLBench, MWM's average success rate was 68.3%, also surpassing existing RGB models. In real-world experiments, MWM achieved an average success rate of 67.5% across four tasks, far exceeding GE-ACT's 23.8% and π's 38.8%. Robustness evaluation via random visual token pruning revealed that MWM exhibits superior generalization capabilities and robust resilience to texture information loss.

MWM has broad application potential in robot policy learning, particularly in scenarios requiring handling visual variability and improving decision-relevant geometric information capture. Its application scenarios include automated manufacturing, autonomous vehicles, and smart home robots. MWM's robustness and generalization capabilities make it suitable for various complex real-world tasks.

Despite MWM's outstanding performance across multiple benchmarks, it may still face limitations in handling extreme lighting changes and complex backgrounds, especially when lacking sufficient training data. Additionally, the model's computational resource requirements are relatively high, potentially unsuitable for resource-constrained real-time applications. In certain specific tasks, the accuracy of mask prediction may not be as high as direct RGB prediction. Future research directions include optimizing MWM's computational efficiency for resource-constrained applications; exploring more semantic mask generation methods to improve model generalization; and validating MWM's effectiveness in more real-world tasks.