Source note
Mask World Model: Predicting What Matters for Robust Robot Policy Learning
Summary
Mask World Model replaces RGB video prediction with future semantic mask prediction so a robot world model learns object geometry and contact dynamics instead of texture and lighting. The paper claims this makes policy learning more robust and improves success rates in simulation and on a real robot while still using raw RGB at test time.
Problem
- Standard robot world models often predict future RGB frames, which pushes capacity toward texture, lighting, reflections, and background changes that do not help action selection.
- Those appearance factors can cause predictive drift in closed-loop control, which hurts generalization and makes policies fragile under visual shifts.
- A useful robot policy needs features tied to object identity, spatial layout, and contact-relevant motion, especially for long-horizon manipulation.
Approach
- MWM trains a world model to predict future semantic masks instead of future pixels. The masks cover task-relevant objects and the robot, creating a geometric bottleneck that drops photometric noise.
- Training has two stages. Stage 1 learns mask dynamics with a diffusion / flow-matching backbone conditioned on past multi-view RGB and language. Stage 2 trains a diffusion policy head that reads the backbone's predictive features and outputs actions.
- Semantic masks are used only as offline supervision during training. At inference, the system takes raw multi-view RGB plus language and does not require an external segmentation model.
- The model reuses a shared video VAE for both RGB frames and rendered mask images, then uses transformer features from the mask-prediction backbone as the policy context.
Results
- On LIBERO, MWM reports 98.3% average success rate, beating GE-ACT at 96.5%, pi0 at 94.2%, CogACT at 93.6%, Cosmos w/ Latent IDM at 91.9%, and OpenVLA at 76.5%.
- On the harder LIBERO-10 subset, MWM reaches 96.0%, compared with 94.4% for GE-ACT, 85.2% for pi0, 84.2% for Cosmos w/ Latent IDM, and 48.8% for Cosmos w/ IDM.
- The mask-based ablations support the main claim: MWM-C1 improves over Cosmos w/ IDM from 67.5% to 81.0% average SR; MWM-C2 improves over Cosmos w/ Latent IDM to 91.8% average SR.
- On RLBench, MWM reports 68.3% average success rate, ahead of FiS-VLA at 50.0%, CogACT at 42.5%, GE-ACT at 30.8%, pi0 at 33.3%, and OpenVLA at 23.3%.
- The paper also states 67.5% average success on a real Franka robot across four tasks and says MWM handles background, lighting, object-color shifts, and random token pruning better than RGB-based baselines.
- Quantitative robustness details for the real-world OOD tests are only partially visible in the provided excerpt, so the strongest complete claims are the benchmark success-rate gains above.