UniW(A)M: A Unified Mixture-of-Experts Framework for World and World-Action Models
Organizations: Tsinghua University · Zhongguancun Academy · East China Normal University · Manifold AI
Abstract
World models (WMs) and world-action models (WAMs) focus on future world prediction and executable action generation, respectively. Although these models learn shared physical dynamics and interaction semantics from embodied trajectories, they are typically trained in isolation, which restricts knowledge reuse between prediction and control. The central challenge is therefore to develop a unified learning framework that shares these dynamics while preserving the distinct modeling capabilities required by each paradigm. We propose UniW(A)M, a mixture-of-experts world-(action) model built upon diffusion Transformers that unifies action-conditioned world modeling and diverse WAM variants. Shared experts capture common dynamics, while atomic capability experts are reused across interfaces so that each interface keeps its own conditional pathway. At deployment, a plug-in adaptive interface router (PAIR) learns from a small subset of target-task chunk scenarios to select a suitable WAM interface for each execution chunk, balancing task success and execution efficiency while keeping the backbone frozen. We evaluate UniW(A)M on LIBERO, RoboTwin 2.0, and three real-world manipulation tasks, where it achieves 88.3% average real-world success versus 76.7% for the strongest baseline. On WorldArena, it leads action-conditioned world modeling on all three tasks. AC-WM co-training contributes to action generation on every action interface and PAIR raises OOD success from 63.3% to 76.7%.
Figures & tables
| Paradigm | Atomic Capabilities | Visibility / Role |
|---|---|---|
| AC-WM | WM | Action-conditioned world modeling |
| Joint-WAM | VGM + IDM | Mutual visibility |
| IDM-WAM | VGM + IDM | Sequential vision-to-action |
| Action-Centric-WAM | VLA + WM | Action-centric with video supervision |
| Action-Only-WAM | VLA + VGM | Action generation with auxiliary video |
| Method | Folding Clothes | Folding Box | Tabletop Clean. | Average |
| (Deformable) | (Contact-Rich) | (Long-Horizon) | ||
| Baselines | ||||
| ( Physical Intelligence et al., 2025 ) | 60 | 65 | 70 | 65.0 |
| DreamZero ( Ye et al., 2026b ) | 45 | 55 | 60 | 53.3 |
| WorldScape Policy 2.0 ( Su et al., 2026 ) | 75 | 75 | 80 | 76.7 |
| UniW(A)M | ||||
| Model | JEPA Similarity | Motion Smoothness | Interaction Quality | Trajectory Accuracy | Instruction Following | EWM Score-P |
|---|---|---|---|---|---|---|
| Task 1: Folding Clothes (Deformable) | ||||||
| Cosmos-Predict 2.5 ( NVIDIA et al., 2025 ) | 21.86 | 60.70 | 64.00 | 7.20 | 56.50 | 48.08 |
| CtrlWorld ( Guo et al., 2025 ) | 17.57 | 78.21 | 63.50 | 18.56 | 51.50 | 53.09 |
| UniW(A)M (Ours) | 44.39 | 87.77 | 80.00 | 22.87 | 93.50 | 61.51 |
| Task 2: Folding Box (Contact-Rich) | ||||||
| Cosmos-Predict 2.5 ( NVIDIA et al., 2025 ) | 40.27 | 52.47 | 76.50 | 10.91 | 95.00 | 51.07 |
| Interface | Folding Clothes | Folding Box | Tabletop Clean. | Average |
|---|---|---|---|---|
| (Deformable) | (Contact-Rich) | (Long-Horizon) | ||
| Standalone Interfaces | ||||
| Joint-WAM | 65 | 55 | 70 | 63.3 |
| IDM-WAM | 65 | 50 | 75 | 63.3 |
| Action-Centric-WAM | 50 | 40 | 55 | 48.3 |
| Action-Only-WAM | 50 | 45 | 55 | 50.0 |
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
| Interface | Conditioning copies and inputs | Denoising targets |
|---|---|---|
| Joint-WAM | Clean video history, current state, and language | Future video and actions, with joint current-block visibility |
| IDM-WAM | Video and action conditioning copies, states, and language under interleaved teacher forcing | Future video and actions, ordered video before action |
| Action-Centric-WAM | Clean video history, current state, and language | Future video and actions; video additionally reads current noisy actions |
| Action-Only-WAM | Clean video history, current state, and language | Future video and actions without current-block cross-modal attention |
| AC-WM | Clean video history, supplied clean actions, and language | Future video only |
| Interface | Future video targets | Action targets |
|---|---|---|
| Joint-WAM / IDM-WAM | Shared + VGM | Shared + IDM |
| Action-Centric-WAM | Shared + WM | Shared + VLA |
| Action-Only-WAM | Shared + VGM | Shared + VLA |
| AC-WM | Shared + WM | None (clean conditions) |
| Token group | Action-Centric | Action-Only | Mixing |
|---|---|---|---|
| Clean history copies | Shared | Shared | Fixed |
| Initial conditioning frame (seed) | Shared | Shared + VGM | Fixed / learned |
| Future action targets | Shared + VLA | Shared + VLA | Learned |
| State tokens | Shared | Shared | Fixed |
| Method / interface | Folding Clothes | Folding Box | Tabletop Cleaning |
|---|---|---|---|
| Standard evaluation: baselines | |||
| ( Physical Intelligence et al., 2025 ) | 12/20; [38.7, 78.1] | 13/20; [43.3, 81.9] | 14/20; [48.1, 85.5] |
| DreamZero ( Ye et al., 2026b ) | 9/20; [25.8, 65.8] | 11/20; [34.2, 74.2] | 12/20; [38.7, 78.1] |
| WorldScape Policy 2.0 ( Su et al., 2026 ) | 15/20; [53.1, 88.8] | 15/20; [53.1, 88.8] | 16/20; [58.4, 91.9] |
| Standard evaluation: UniW(A)M with CATR | |||
| Joint-WAM | 18/20; [69.9, 97.2] | 16/20; [58.4, 91.9] | 18/20; [69.9, 97.2] |
| Method / interface | Folding Clothes | Folding Box | Tabletop Cleaning |
|---|---|---|---|
| Single DiT (Dense) | |||
| Joint-WAM | 2/20; [2.8, 30.1] | 2/20; [2.8, 30.1] | 4/20; [8.1, 41.6] |
| IDM-WAM | 4/20; [8.1, 41.6] | 0/20; [0.0, 16.1] | 6/20; [14.5, 51.9] |
| Action-Centric-WAM | 0/20; [0.0, 16.1] | 0/20; [0.0, 16.1] | 2/20; [2.8, 30.1] |
| Action-Only-WAM | 0/20; [0.0, 16.1] | 0/20; [0.0, 16.1] | 1/20; [0.9, 23.6] |
| MoE with Sample-Level Routing | |||
| Model | Spatial | Object | Goal | Long | Average |
|---|---|---|---|---|---|
| Vision-Language-Action Models | |||||
| OpenVLA ( Kim et al., 2024 ) | 84.7 | 88.4 | 79.2 | 53.7 | 76.5 |
| ( Black et al., 2024 ) | 96.8 | 98.8 | 95.8 | 85.2 | 94.1 |
| ( Physical Intelligence et al., 2025 ) | 98.8 | 98.2 | 98.0 | 92.4 | 96.9 |
| World Action Models | |||||
| Motus ( Bi et al., 2025 ) | 96.8 | 99.8 | 96.6 | 97.6 | 97.7 |
| Model | C2C | C2R |
|---|---|---|
| ( Black et al., 2024 ) (single-task) | 46.42 | 16.34 |
| ( Physical Intelligence et al., 2025 ) | 70.7 | 46.0 |
| Fast-WAM ( Yuan et al., 2026 ) | 77.8 | 1.9 |
| WorldScape Policy 2.0 ( Su et al., 2026 ) | 84.92 | 38.5 |
| UniW(A)M | 84.38 | 66.60 |
| Model | Visual Quality | Motion Quality | Content Consistency | ||||||
|---|---|---|---|---|---|---|---|---|---|
| Image Quality | Aesthetic Quality | JEPA Similarity | Dynamic Degree | Flow Score | Motion Smoothness | Subject Consistency | Backgrd. Consistency | Photom. Consistency | |
| Task 1: Folding Clothes (Deformable) | |||||||||
| Cosmos-Predict 2.5 ( NVIDIA et al., 2025 ) | 0.3623 | 0.2389 | 0.2186 | 0.3057 | 0.1940 | 0.6070 | 0.7122 | 0.8976 | 0.2339 |
| CtrlWorld ( Guo et al., 2025 ) | 0.3237 | 0.2549 | 0.1757 | 0.6430 | 0.6827 | 0.7821 | 0.6917 | 0.8932 | 0.0469 |
| UniW(A)M (Ours) | 0.3158 | 0.2189 | 0.4439 | 0.6778 | 0.8025 | 0.8777 | 0.6924 | 0.8772 | 0.0177 |
| Task 2: Folding Box (Contact-Rich) | |||||||||
| Model | Physics Adherence | 3D Accuracy | Controllability | |||
|---|---|---|---|---|---|---|
| Interaction Quality | Trajectory Accuracy | Depth Accuracy | Perspectivity | Instruction Following | Semantic Alignment | |
| Task 1: Folding Clothes (Deformable) | ||||||
| Cosmos-Predict 2.5 ( NVIDIA et al., 2025 ) | 0.6400 | 0.0720 | 0.5474 | 0.8000 | 0.5650 | 0.8175 |
| CtrlWorld ( Guo et al., 2025 ) | 0.6350 | 0.1856 | 0.4754 | 0.8000 | 0.5150 | 0.8593 |
| UniW(A)M (Ours) | 0.8000 | 0.2287 | 0.6460 | 0.8000 | 0.9350 | 0.8929 |
| Task 2: Folding Box (Contact-Rich) | ||||||
| Configuration | Backbone | PAIR classification head |
|---|---|---|
| Optimizer | AdamW | AdamW |
| Learning rate | ||
| Learning-rate schedule | Constant | Constant, no warm-up |
| Weight decay | ||
| Global batch size | 192 | 128 per head |
| Interface | Folding Clothes | Folding Box | Tabletop Clean. | Average |
| (Deformable) | (Contact-Rich) | (Long-Horizon) | ||
| Single DiT (Dense) | ||||
| Joint-WAM | 10 | 10 | 20 | 13.3 |
| IDM-WAM | 20 | 0 | 30 | 16.7 |
| Action-Centric-WAM | 0 | 0 | 10 | 3.3 |
| Action-Only-WAM | 0 | 0 | 5 | 1.7 |
World Models for Robotics
World models for robotics predict how scenes and robot states change under actions to support planning, control, and learning.
See all 379 papers