Mask World Model: Predicting What Matters for Robust Robot Policy Learning
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2026
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866908986670841856 |
|---|---|
| author | Lou, Yunfan Chi, Xiaowei Zhang, Xiaojie Qian, Zezhong Li, Chengxuan Zhang, Rongyu Lyu, Yaoxu Song, Guoyu Fu, Chuyao Xu, Haoxuan Wang, Pengwei Zhang, Shanghang |
| author_facet | Lou, Yunfan Chi, Xiaowei Zhang, Xiaojie Qian, Zezhong Li, Chengxuan Zhang, Rongyu Lyu, Yaoxu Song, Guoyu Fu, Chuyao Xu, Haoxuan Wang, Pengwei Zhang, Shanghang |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2604_19683 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Mask World Model: Predicting What Matters for Robust Robot Policy Learning Lou, Yunfan Chi, Xiaowei Zhang, Xiaojie Qian, Zezhong Li, Chengxuan Zhang, Rongyu Lyu, Yaoxu Song, Guoyu Fu, Chuyao Xu, Haoxuan Wang, Pengwei Zhang, Shanghang Robotics 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. |
| title | Mask World Model: Predicting What Matters for Robust Robot Policy Learning |
| topic | Robotics |
| url | https://arxiv.org/abs/2604.19683 |