MWM: Mask World Model

June 8, 2026 · View on GitHub

Predicting What Matters for Robust Robot Policy Learning

arXiv ICML 2026 License

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 ICML 2026 · arXiv:2604.19683

TL;DR. Most video world models for robot control predict RGB pixels, which forces them to model nuisance variation (texture, lighting, background) and leads to brittle, appearance-driven policies. MWM instead predicts the evolution of semantic masks, imposing a geometric information bottleneck that captures contact-relevant dynamics while filtering out visual noise. Semantic supervision is used only offline during training — at test time MWM runs purely on raw multi-view RGB, with no external segmentation model. A mask-centric dynamics backbone is coupled to a diffusion policy head for robust end-to-end control. MWM achieves 98.3% average success on LIBERO and 68.3% on RLBench, substantially outperforming strong RGB-based world models while showing markedly better robustness under visual token pruning.

About this repository

This is the official PyTorch implementation of MWM, built on the Genie Envisioner stack. Training and inference share the same pipeline as Genie-Envisioner (LTX-Video backbone, action expert, flow matching): download GE-Base-style weights, prepare a LeRobot dataset, run Stage-1 mask-dynamics pretraining, then Stage-2 action post-training. This release is a trimmed fork: it keeps the LIBERO experiments and core models / runner / data paths, and omits the heavier upstream extras (GE-Sim / Cosmos demos, web deployment server, generic multi-domain policy configs, etc.). For full platform features and long-form tutorials, see the upstream README.

1. Install

git clone https://github.com/LYFCLOUDFAN/mask-world-model.git
cd mask-world-model

conda create -n mwm python=3.10.4 -y
conda activate mwm
pip install -r requirements.txt

2. What you need on disk

ItemNotes
GE-Base (fast)Pretrained world / video backbone, e.g. from Hugging Face · Genie-Envisioner.
LTX tokenizer + text encoder + VAEFrom Lightricks/LTX-Video: tokenizer, text_encoder, vae, and model_index.json in one directory. You do not need the full LTX UNet weights if you only load GE-Base into diffusion_model.model_path.
LeRobot dataset with masksLIBERO-style layout (parquet + videos + meta/) containing both RGB and semantic-mask fields (image / image_mask, wrist_image / wrist_image_mask). See Section 3 for how to generate the masks.

Point these paths in the YAML files under configs/ltx_model/ (see below).

3. Getting the mask data (required for Stage-1 training)

MWM is trained on semantic masks, not RGB pixels. Masks are used only offline as Stage-1 supervision — at deployment the policy runs on raw RGB. You therefore need a dataset where each RGB frame is paired with a per-pixel semantic mask for the robot arm / gripper and the task-relevant objects. The masks are then rendered into a fixed color palette (one color per semantic class) and encoded by the same video VAE as the RGB frames.

3.1 Render masks from the LIBERO simulator

The simulator can produce ground-truth instance segmentation for free, so no external segmentation model is needed.

  1. Render LIBERO demos with segmentation. Re-collect / re-render the LIBERO demonstrations with instance segmentation enabled. LIBERO (MuJoCo) supports this through the camera_segmentations rendering option and the get_segmentation_of_interest() helper in LIBERO/libero/libero/envs/. The resulting per-task HDF5 file should contain, under each data/demo_*/obs group:

    • agentview_rgb_2, `eye_in_hand_rgb_2$ — \text{RGB} \text{observations} (256 \times 256),
    • $agentview_mask, eye_in_hand_mask` — per-pixel instance-id masks,
    • mask_object_id — a dict string mapping object name → instance id,
    • object_of_interest — list of task-relevant object names.

    You can inspect any HDF5 to confirm these keys exist with python scripts/read_hdf5.py (edit the path inside).

  2. Convert HDF5 → LeRobot with colorized semantic masks. scripts/convert_libero_hdf5_to_lerobot.py reads the instance masks, maps each semantic class to a fixed color (robot arm → cyan, gripper → orange, task objects → yellow, background → black; see build_id2color / colorize_mask), and writes a LeRobot dataset with the extra image_mask / wrist_image_mask fields alongside image / wrist_image:

python scripts/convert_libero_hdf5_to_lerobot.py \
    --data_dir /path/to/libero/hdf5/files \
    --output_dir /path/to/output/lerobot_dataset \
    --repo_name libero/libero_object \
    --fps 10

To convert all four LIBERO suites at once, edit the paths in scripts/batch_convert_libero.sh and run bash scripts/batch_convert_libero.sh.

Note. scripts/convert_libero_data_to_lerobot.py converts the RGB-only RLDS release and does not produce masks; use convert_libero_hdf5_to_lerobot.py (above) whenever you need mask supervision.

3.2 How masks are consumed during training

The dataset loader (data/libero_dataset.py) pairs each RGB camera key with its mask key via valid_cam / valid_cam_mask (by default it appends _mask, e.g. image → image_mask). Make sure your YAML data.*.valid_cam / valid_cam_mask keys match the field names written by the converter.

4. Config checklist (same idea as GE, fewer files)

Before any train.sh run, edit the chosen config and set at least:

  • pretrained_model_name_or_path → directory with LTX tokenizer / text encoder / VAE (+ model_index.json).
  • output_dir → where checkpoints and logs go.
  • data.train / data.val → data_roots, domains, valid_cam, stat_file, etc.
  • Stage 1 (mask dynamics): diffusion_model.model_path → GE-Base (or compatible) checkpoint.
  • Stage 2 (action policy): diffusion_model.model_path → your Stage-1 mask-dynamics checkpoint.

Statistic JSON for normalization: run scripts/get_statistics.py, or reuse configs/ltx_model/libero/libero_all.json and point stat_file in the YAML to that path.

5. Training (two stages — mirrors GE, minimal text)

MWM follows a two-stage protocol: (1) pretrain the mask-dynamics backbone to forecast future semantic mask latents with a flow-matching objective, then (2) train the diffusion policy head end-to-end (gradients from the action loss also refine the backbone so its predictive features become control-aligned). Semantic masks are used only as Stage-1 supervision targets — deployment uses raw RGB only.

Trains the mask-centric world model on your robot / cameras so the latent dynamics match the domain. A memory window of 4 RGB frames is used to predict 5 future mask latents.

bash scripts/train.sh main.py configs/ltx_model/libero/video_model_libero.yaml

Stage 2 — Action post-training (mask-guided diffusion policy)

Starting from the Stage-1 checkpoint, enables action prediction (return_action: true, train_mode: action_full, action_expert: true in the diffusion config).

bash scripts/train.sh main.py configs/ltx_model/libero/action_model_libero.yaml

For LIBERO-Goal style setups, see action_model_libero_goal.yaml if applicable.

scripts/train.sh wraps torchrun on all visible GPUs (same pattern as Genie-Envisioner).

6. Open-loop inference (validation plots)

bash scripts/infer.sh main.py \
  configs/ltx_model/libero/action_model_libero.yaml \
  path/to/checkpoint.safetensors \
  path/to/save/outputs \
  DOMAIN_NAME

DOMAIN_NAME should match the domain key used in your dataset config (e.g. as in the YAML data.*.domains).

7. Closed-loop evaluation (simulation)

Scripts and benchmark numbers are documented in experiments/RUN.md.

Edit checkpoint paths in experiments/eval_libero.sh, align configs/ltx_model/libero/action_model_libero.yaml, then run:

bash experiments/eval_libero.sh

The LIBERO/ subtree provides the simulation environment code expected by the eval scripts.

8. Repository layout (high level)

configs/ltx_model/     # LIBERO YAMLs (video / action)
data/                  # Dataset adapters (e.g. LeRobot LIBERO)
experiments/           # eval_*.py, eval_*.sh, RUN.md
models/                # LTX transformer, pipeline, VAE, etc.
runner/                # Trainer / inferencer
scripts/               # train.sh, infer.sh, convert + stats helpers
utils/
main.py

9. Citation

If you find this work useful, please cite our paper and the Genie Envisioner report:

@inproceedings{lou2026mwm,
  title={Mask World Model: Predicting What Matters for Robust Robot Policy Learning},
  author={Lou, Yunfan and Chi, Xiaowei and Zhang, Xiaojie and Qian, Zezhong and Li, Chengxuan and Zhang, Rongyu and Lyu, Yaoxu and Song, Guoyu and Fu, Chuyao and Xu, Haoxuan and Wang, Pengwei and Zhang, Shanghang},
  booktitle={Proceedings of the 43rd International Conference on Machine Learning (ICML)},
  year={2026}
}

@article{liao2025genie,
  title={Genie Envisioner: A Unified World Foundation Platform for Robotic Manipulation},
  author={Liao, Yue and Zhou, Pengfei and Huang, Siyuan and Yang, Donglin and Chen, Shengcong and Jiang, Yuxin and Hu, Yue and Cai, Jingbin and Liu, Si and Luo, Jianlan and Chen, Liliang and Yan, Shuicheng and Yao, Maoqing and Ren, Guanghui},
  journal={arXiv preprint arXiv:2508.05635},
  year={2025}
}

10. Acknowledgments

  • Genie-Envisioner — upstream training/inference design, GE-Base weights, and documentation.
  • LTX-Video — tokenizer, VAE, and architecture lineage.
  • Diffusers and components noted in the upstream License section.

Licensing follows the upstream Genie-Envisioner policy: Apache 2.0–licensed third-party subtrees (e.g. parts of models/, web_infer_utils if present upstream) vs. CC BY-NC-SA 4.0 for other original materials — confirm LICENSE / notices in this repo before redistributing.