Worldscape-MoE

A Unified Mixture-of-Experts World Model for Scalable Heterogeneous Action Control

Jianjie Fang^1,*, Yongyan Xu^1,*, Ziyou Wang^1,*, Chen Gao^1,*,†, Yuchao Huang¹, Zhaolu Wang¹, Rongze Tang¹, Mingyuan Jia¹, Baining Zhao¹, Weichen Zhang¹, Xin Zhang², Haisheng Su², Yu Shang¹, Wei Wu², Xinlei Chen¹, Yong Li^1,†
1Tsinghua University    ²Manifold AI
_{^*Equal contribution.  ^†Corresponding authors.}

Overview

Worldscape-MoE is a 5B-parameter video world model built on Wan2.2 for scalable,heterogeneous action control. It unifies camera trajectories, dual-arm robot actions, dense action maps, and LIBERO single-arm actions within a shared world-dynamics backbone, rather than treating each control modality as a separate task.

The model combines modality-aware control injection with a shared feed-forward expert and modality-specific experts. A control router activates the relevant expert while preserving knowledge of shared environment dynamics across modalities. This design supports heterogeneous supervision, out-of-distribution generalization, and continual extension to new action interfaces without requiring a separate world model for each control type.

The full demonstration is available on YouTube.

Release Scope

Important

The current release covers a subset of the full training collection. It contains 5,000 samples for each of the four modalities: camera control, dual-arm control, dense action-map control, and LIBERO single-arm control.The 5,000 camera-control samples come from RealEstate10K.

The repository includes the standard video-based training pipeline,inference entry points, dataset preparation and validation tools, and an optional offline VAE encoding pipeline. Model weights and dataset media are distributed separately under their respective terms.

Installation

Worldscape-MoE requires Python 3.10 or later. Before installing the project,install a CUDA-enabled PyTorch build compatible with your local CUDA toolchain.

git clone https://github.com/EmbodiedCity/Worldscape-MoE.code.git
cd Worldscape-MoE.code

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

# Install the appropriate PyTorch build first:
# https://pytorch.org/get-started/locally/
python -m pip install -e .

The reference environment uses PyTorch 2.9.1, Torchvision 0.24.1,Accelerate 1.12.0, Diffusers 0.36.0, and Transformers 4.57.3.

Data and Model Setup

Download the released dataset subset to the default project location:

hf download EmbodiedCity/Worldscape-MoE-Dataset \
  --repo-type dataset \
  --local-dir datasets/OpenSource_MOE

Extract the four modality archives before training or inference:

cd datasets/OpenSource_MOE
for archive in archives/*.tar.zst; do
  tar --zstd -xf "$archive"
done
cd ../..

Download the Wan2.2 TI2V 5B base model:

hf download Wan-AI/Wan2.2-TI2V-5B \
  --local-dir pretrained/Wan2.2-TI2V-5B

Place a compatible Wan2.2 5B control base model at:

pretrained/Wan2.2-Fun-5B-Control-Camera/

The dataset path and training model path can be overridden with DATA_ROOT and MODEL_NAME, respectively. See Dataset format and preparation for the manifest schema, directory layout, provenance, and validation workflow.

Note

Our trained Worldscape-MoE model weights will be open-sourced soon.

Worldscape-MoE Three-Modality Training

The default launcher provides an 8-GPU recipe for machines with 80 GB of GPU memory per card. To train the three-modality model with camera trajectories, dual-arm robot actions, and dense action maps, set the local model path and run:

MODEL_NAME=/path/to/Wan2.2-Fun-5B-Control-Camera \
bash scripts/wan2.2_fun/train_worldscape_moe_3modal_5b_8gpu.sh

64-GPU Multi-Node Training

For larger clusters, use the 64-GPU launcherscripts/wan2.2_fun/train_worldscape_moe_3modal_5b_64gpu.sh, configured for eight nodes with eight GPUs per node. It follows the original VideoX-FunMLP_* queue variables; for other cluster environments, adjust the mapping block and the NNODES, NPROC_PER_NODE, EXPECTED_WORLD_SIZE, NODE_RANK,MASTER_ADDR, and MASTER_PORT settings at the top of the launcher.

Worldscape-MoE Tuning

Worldscape-MoE Tuning progressively extends a trained three-modality model to additional control modalities. To add a fourth modality, we use LIBERO single-arm control as an example: introduce its control branch and modality-specific expert, initialize the new expert from a compatible existing expert, and tune the shared components conservatively while the new components specialize to the added control interface.

Starting from a three-modality checkpoint, the current four-modality launcher can be run as follows:

TRANSFORMER_PATH=/path/to/three-modal/checkpoint-250/diffusion_pytorch_model.safetensors \
MODEL_NAME=/path/to/Wan2.2-Fun-5B-Control-Camera \
bash scripts/wan2.2_fun/train_worldscape_moe_4modal_libero_5b_8gpu.sh

LIBERO is only one tuning instance, not a limit of the paradigm. The same procedure can be repeated to extend the model from three modalities to a growing set of control modalities. More broadly, Worldscape-MoE Tuning supports both modality expansion with newly composed modalities and continued tuning on new combinations of existing data sources, such as paired arm-and-camera data. This gives the model strong multimodal composition ability beyond the single LIBERO example. See MoE architecture for implementation details.

VAE-Encoded Training

Worldscape-MoE also supports offline VAE encoding for the released camera,dual-arm, and action-map data. Moving the frozen VAE out of the training loop avoids repeated video decoding and VAE computation during optimization.

Precompute latents for the three modalities with the default eight-GPU configuration:

bash worldscape_moe_vae/scripts/precompute_3modal.sh

The command processes 5,000 released samples for each supported modality and writes portable latent tensors and manifests to worldscape_moe_vae/outputs/latents. Interrupted runs resume from validated outputs by default. Use MAX_SAMPLES_PER_MODALITY=8 for a quick preprocessing check.

Train directly from the generated latents:

bash worldscape_moe_vae/scripts/train_3modal.sh

The latent-training launcher does not load the VAE or decode source videos during training. Override the default locations with DATA_ROOT, MODEL_NAME, LATENT_ROOT, and OUTPUT_DIR as needed. Validation runs automatically before training begins.

In the reference configuration on NVIDIA H20 GPUs, precomputed latents reduced per-step time by approximately 50%, excluding the one-time preparation cost. Actual throughput depends on the hardware, storage, and training configuration.

Training input	Relative step time	Relative throughput
Online VAE encoding	1.00x	1.00x
Precomputed VAE latents	~0.50x	~2.00x

Inference

Use a dedicated entry point for each control type.

Camera Trajectory Control

python scripts/wan2.2_fun/infer_camera_case.py \
  --model_name /path/to/Wan2.2-Fun-5B-Control-Camera \
  --checkpoint_path /path/to/checkpoint \
  --start_image /path/to/start.png \
  --camera_pose /path/to/camera_pose.txt \
  --prompt "a camera moving through the scene" \
  --output_path outputs/camera_case.mp4

Dual-Arm Action Control

python scripts/wan2.2_fun/infer_arm_case.py \
  --model_name /path/to/Wan2.2-Fun-5B-Control-Camera \
  --checkpoint_path /path/to/checkpoint \
  --start_image /path/to/start.png \
  --action_path /path/to/dual_arm_actions.json \
  --control_type arm \
  --prompt "move the blocks into the requested order" \
  --output_path outputs/arm_case.mp4

LIBERO Single-Arm Control

python scripts/wan2.2_fun/infer_arm_case.py \
  --model_name /path/to/Wan2.2-Fun-5B-Control-Camera \
  --checkpoint_path /path/to/checkpoint \
  --start_image /path/to/start.png \
  --action_path /path/to/libero_actions.json \
  --stat_path /path/to/libero_action_stats.json \
  --control_type libero \
  --prompt "put the object in the drawer" \
  --output_path outputs/libero_rollout.mp4

The LIBERO action file must provide a 7D state array for each frame. Thelibero control type enables the LIBERO expert automatically.

Dense Action-Map Control

python scripts/wan2.2_fun/infer_actionmap_case.py \
  --model_name /path/to/Wan2.2-Fun-5B-Control-Camera \
  --checkpoint_path /path/to/checkpoint \
  --start_image /path/to/start.png \
  --action_map_video /path/to/action_map.mp4 \
  --prompt "interact with the object" \
  --output_path outputs/actionmap_case.mp4

For checkpoints produced by the LIBERO extension above, append --enable_libero_expert to the camera, dual-arm, and action-map commands.

For all input formats, optional arguments, and inference for extended models,see Inference.

Documentation

• Training
• Inference
• Dataset format and preparation
• MoE architecture
• Contributing
• Security policy

Citation

If you find Worldscape-MoE useful, please cite:

@misc{fang2026worldscapemoe,
  title         = {Worldscape-MoE: A Unified Mixture-of-Experts World Model for Scalable Heterogeneous Action Control},
  author        = {Fang, Jianjie and Xu, Yongyan and Wang, Ziyou and Gao, Chen and Huang, Yuchao and Wang, Zhaolu and Tang, Rongze and Jia, Mingyuan and Zhao, Baining and Zhang, Weichen and Zhang, Xin and Su, Haisheng and Shang, Yu and Wu, Wei and Chen, Xinlei and Li, Yong},
  year          = {2026},
  eprint        = {2607.03964},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2607.03964}
}

License

The source code is released under the Apache License 2.0. Dataset media and model weights are distributed separately under their applicable terms; third-party notices are listed in NOTICE.