UniW(A)M: A Unified Mixture-of-Experts Framework for World and World-Action Models

Jingbo Xu; Jianjie Fang; Zhuohang Li; Zihao Song; Zhaolu Wang; Xvyuan Liu; Xin Zhang; Haisheng Su

cs.ROOct 13, 2026

UniW(A)M: A Unified Mixture-of-Experts Framework for World and World-Action Models

Authors: Jingbo Xu, Jianjie Fang, Zhuohang Li, Zihao Song, Zhaolu Wang, Xvyuan Liu, Xin Zhang, Haisheng Su,

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

Figure 1: Overview of UniW(A)M. The model unifies AC-WM and four WAM paradigms over shared embodied context. AC-WM predicts future video, while action execution adaptively switches among the four WAM paradigms across successive chunks, illustrated by a cloth-folding task.

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

Table 1: Atomic capability composition and visibility characteristics of different paradigms.

Figure 2: Architecture of UniW(A)M. (a) Stage I: Unified learning. Five interfaces share a diffusion Transformer with CATR expert routing. (b) Stage II: PAIR adaptation. With the backbone frozen, PAIR learns interface selection from offline MAE-based supervision.

Figure 3: Real-world tasks: cloth folding, box folding, and tabletop cleaning.

Panel of Figure 4 — Figure 4: RoboTwin 2.0 success rates with clean-only training. (a) C2C evaluates on clean scenes; (b) C2R evaluates on randomized scenes. Scores are averaged over 50 tasks separately for each setting. Gray bars show official leaderboard baselines ( RoboTwin Team, 2026 ) : $\pi_{0}$ uses single-task fine-tuning, while the others use co-training. Blue bars show UniW(A)M.

Method	Folding Clothes	Folding Box	Tabletop Clean.	Average
Method	(Deformable)	(Contact-Rich)	(Long-Horizon)	Average
Baselines
$\pi_{0.5}$ ( 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

Table 2: Success rates (%) on real-world robotic tasks under the standard evaluation configuration. The best results within each task are shown in bold .

Model	JEPA Similarity $\uparrow$	Motion Smoothness $\uparrow$	Interaction Quality $\uparrow$	Trajectory Accuracy $\uparrow$	Instruction Following $\uparrow$	EWM Score-P $\uparrow$
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

Table 3: Evaluation of action-driven and interaction metrics of WorldArena on three real-world tasks. All metrics are reported on a

0

–

100

scale. The best and second-best results within each task are shown in bold and underlined , respectively.

Interface	Folding Clothes	Folding Box	Tabletop Clean.	Average
Interface	(Deformable)	(Contact-Rich)	(Long-Horizon)	Average
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

Table 4: OOD success rates (%) for UniW(A)M interface adaptation on real-world tasks. The best results within each task are shown in bold .

Panel of Figure 5 — Figure 5: Inference efficiency and task-level Pareto analysis of UniW(A)M. (a) Task success rate versus average completion time over successful real-robot trials, where UniW(A)M + PAIR occupies the upper-left Pareto frontier. (b) Full chunk inference time divided by the configured step count of 16 across the four standalone WAM interfaces; IDM-WAM uses 16 steps per stage (Appendix C ). Chunk-boundary routing is excluded from these bars: ensemble prediction takes approximately 2 ms, while routing including shared-expert feature computation averages approximately 63 ms under the measurement scope in Appendix B.2 .

Appendix figures & tables25 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Joint-WAM visibility masks. Future video and action tokens attend to both modalities, coupling their predictions during joint denoising. Training and inference use the same bidirectional visibility pattern.

Figure 7: IDM-WAM visibility masks. Action queries can access future video, while video queries cannot access future actions. At inference, video is generated first and then held fixed as the condition for action generation.

Figure 8: Action-Centric-WAM visibility masks. During training, future video queries access actions, but action queries cannot access future video. At deployment, future video tokens are omitted and WM is inactive, while observed-context tokens remain available to action generation.

Figure 9: Action-Only-WAM visibility masks. Neither future modality attends to the other during training, although both access the observed context. Inference retains only the context and action tokens, omitting auxiliary video prediction.

Figure 10: AC-WM visibility masks. Future video queries attend to the supplied action sequence during both training and inference. Actions remain clean conditioning inputs, so the interface predicts video without generating actions.

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

Appendix

Table 5: Training trajectory organization. All rows use the same four-block trajectory window. The initial video frame is conditioning only and excluded from video loss. Clean copies serve teacher forcing rather than additional supervision.

Figure 11: Expert activation across five video-action interfaces. Shared experts are active for every interface, while atomic experts specialize according to the conditional pathway. In Action-Centric-WAM, the grey dashed WM circle indicates that the expert is inactive when future-video prediction is omitted. In Action-Only-WAM, the coloured dashed VGM circle indicates that the expert remains active for the initial conditioning frame, without predicting future video. For these two interfaces, future action tokens are processed only by shared and VLA experts.

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)

Appendix

Table 6: Eligible experts for future target groups during training. Weights are learned within each listed set; experts outside the set do not contribute to that group’s output.

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

Appendix

Table 7: Expert sets and mixing rules for conditioning and action groups in the two action-only inference pathways. Shared-only groups have fixed weight one; learned mixtures use Equations 8 and 9 .

Figure 12: Phase-conditioned PAIR interface selection across three real-world tasks. Each panel reports 100%-stacked proportions for the Reach, Contact, and Core stages, revealing how the router shifts between fast action pathways and video-aware WAMs as interaction demands change.

Standard evaluation: baselines
Method / interface	Folding Clothes	Folding Box	Tabletop Cleaning
$\pi_{0.5}$ ( 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]

Appendix

Table 8: Per-task uncertainty for standard and OOD real-world evaluations. Each entry gives successes out of 20 trials, followed by the 95% Wilson interval in percent. Standard and OOD rows correspond to separate evaluation settings.

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

Appendix

Table 9: Per-task uncertainty for the dense, sample-level routing, and four-WAM-only ablations. Entries give successes out of 20 trials and 95% Wilson intervals in percent. Counts and intervals for the full five-paradigm CATR model appear in the standard UniW(A)M block of Table 8 .

Figure 13: Real-world visualizations of Folding Clothes for the four WAM interfaces. The sequences illustrate garment grasping, repositioning, and folding through successive stages of manipulation.

Figure 14: Real-world visualizations of Folding Box for the four WAM interfaces. Selected frames show the progression from a flat cardboard box to folding its walls and manipulating its flaps.

Figure 15: Real-world visualizations of Tabletop Cleaning for the four WAM interfaces. The sequences show object pickup, transfer, and placement into a collection basket, with complementary views of reaching and grasping.

Figure 16: Case study of critical operations in Folding Clothes, Folding Box, and Tabletop Cleaning. Within each task group, the upper row shows successful operations and the lower row shows corresponding failure cases, with red circles highlighting the affected interactions. Columns represent distinct cases, read from left to right. The examples illustrate local execution errors that motivate adapting the WAM interface at chunk boundaries.

Vision-Language-Action Models
Model	Spatial	Object	Goal	Long	Average
OpenVLA ( Kim et al., 2024 )	84.7	88.4	79.2	53.7	76.5
$\pi_{0}$ ( Black et al., 2024 )	96.8	98.8	95.8	85.2	94.1
$\pi_{0.5}$ ( 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

Appendix

Table 10: Success rates (%) on the four LIBERO suites and their arithmetic mean. The best and second-best results in each column are shown in bold and underlined , respectively.

Model	C2C	C2R
$\pi_{0}$ ( Black et al., 2024 ) (single-task)	46.42	16.34
$\pi_{0.5}$ ( 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

Appendix

Table 11: RoboTwin 2.0 success rates (%) with clean-only training, averaged across 50 tasks separately for C2C and C2R. Official leaderboard baselines ( RoboTwin Team, 2026 ) use co-training except

\pi_{0}

, which uses single-task fine-tuning. Best and second-best scores are bold and underlined .

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)

Appendix

Table 12: WorldArena component scores for visual quality, motion quality, and content consistency. The highest score for each metric within each task is shown in bold , including ties.

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)

Appendix

Table 13: WorldArena component scores for physics adherence, 3D accuracy, and controllability. The highest score for each metric within each task is shown in bold , including ties.

Figure 17: Action-conditioned future-world prediction on Folding Clothes. UniW(A)M in AC-WM mode better preserves garment visibility and appearance in this early-sequence comparison. Red circles highlight missing garment regions and inconsistent fabric appearance in the baselines.

Figure 18: Action-conditioned future-world prediction on Folding Box. The comparison emphasizes cardboard structure and wrist-view contact details. Red circles identify localized blur and poorly resolved boundaries in baseline predictions.

Figure 19: Action-conditioned future-world prediction on Tabletop Cleaning. UniW(A)M retains more recognizable bottle-label details and scene appearance. Red circles highlight lost label texture, color drift, and distorted object or background regions in the baselines.

Configuration	Backbone	PAIR classification head
Optimizer	AdamW	AdamW
Learning rate	$1\times 10^{-5}$	$1\times 10^{-3}$
Learning-rate schedule	Constant	Constant, no warm-up
Weight decay	$1\times 10^{-5}$	$1\times 10^{-2}$
Global batch size	192	128 per head

Appendix

Table 14: Optimization settings for backbone training and PAIR classification-head training.

Interface	Folding Clothes	Folding Box	Tabletop Clean.	Average
Interface	(Deformable)	(Contact-Rich)	(Long-Horizon)	Average
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

Appendix

Table 15: Architecture, routing, and training-paradigm ablations on three real-world tasks. The dense, sample-level routing, and full five-paradigm CATR configurations use the same training steps, learning-rate settings, and training data volume; the latter two have identical expert architectures and parameter counts. The four-WAM-only CATR variant shares the full model’s initialization, training steps, learning-rate settings, batch size, and dataset, but samples only the four WAMs uniformly. Success rates are reported in percentage (%).

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

UniWAM: Unified World-Action Model
Oct 1, 2026Wenxuan Song, Jiayi Chen, Jingbo Wang +13Vision-Language-Action Models World Action Models
Can World and Action Modeling Benefit Each Other under Unified Causal Training?
Oct 15, 2026Yue Su, Bokai Lin, Xinkai Wang +9Action-Conditioned World Models World Model Learning
IAAI Team · SII · THU +1
World-Action Models for Robot Learning and Control: A Survey
Sep 13, 2026Zuxing Lu, Hongjia Zhai, Guanzhi Wang +13World Action Models Action-Conditioned World Models
MBZUAI · Caltech · Amazon FAR +4

Oct 1, 2026J/K move · Enter open · S save

UniWAM: Unified World-Action Model

Wenxuan Song, Jiayi Chen, Jingbo Wang, Shuai Zhou+12

Vision-Language-Action Models World Action Models Reasoning for VLA Models Embodied World Models Robot Learning

Open paper AskPDF arXiv

Vision-language-action models benefit from the understanding and reasoning capabilities of pretrained vision-language models, but action-only supervision provides limited grounding in world dynamics. Conversely, world-action models inherit spatiotemporal priors from video generation models, yet remain limited in semantic understanding and reasoning under distribution shifts. We introduce UniWAM, a unified architecture that integrates a physical reasoner, a world generator, and an action predictor to jointly learn semantic understanding of the physical world, visual generation, and action prediction. To ensure the quality of the training data, we developed a rigorous data cleaning and annotation pipeline for both human egocentric data and robot data. To adapt the vision-language component to embodied tasks while preserving its inherited language capabilities, we represent low-level actions in natural language and introduce a pre-training recipe that assigns complementary supervision from visual question answering (VQA) data, human egocentric data, and robot demonstrations to the appropriate model components. During post-training, future visual noise augmentation reduces reliance on precise future predictions, while history-conditioned flow matching uses encoded action history to initialize action generation. Together, these designs significantly reduce denoising steps while maintaining performance. UniWAM achieves state-of-the-art (SOTA) performance across multiple evaluations, including in-distribution performance, robustness, generalization, instruction following, and long-horizon task execution. Furthermore, we uncover a log-linear scaling law of unified human-robot co-training, demonstrating the effectiveness of large-scale pre-training on a mixture of human and robot data.

Abstract

World Models for Robotics

Explore similar work

UniWAM: Unified World-Action Model

Can World and Action Modeling Benefit Each Other under Unified Causal Training?

World-Action Models for Robot Learning and Control: A Survey