Predicting What Matters: Robust Generalist Robot Policy Learning via Future Semantic Mask

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.

Robots often struggle when the real world looks different from their training data. A robot may see a new background, different lighting, or a different object color and make worse decisions, even though the task itself has not changed. Many robot learning systems try to predict future camera images in detail, but this can make them spend effort on visual details that do not matter for action. In this work, we introduce Mask World Model, a method that teaches a robot to focus on what matters: which objects are present, where they are, and how they may move or interact. During training, the robot learns from simplified object maps, but at test time it only needs normal camera images. It then uses this learned understanding of future object motion to choose actions. Across simulated and real robot tasks, our method makes robots more reliable than strong image-based baselines, especially under changes in background, lighting, and object color. This suggests that focusing on object structure rather than visual appearance can help build robot policies that are more robust in everyday environments.