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

Yue Su; Bokai Lin; Xinkai Wang; Chubin Zhang; Xiangnan Wu; Kuan Zhang; Weiyang Jin; Jiwen Yu

cs.CVOct 15, 2026

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

Authors: Yue Su, Bokai Lin, Xinkai Wang, Chubin Zhang, Xiangnan Wu, Kuan Zhang, Weiyang Jin, Jiwen Yu,

Organizations: IAAI Team · SII · THU · UCAS

Abstract

World and action modeling are closely related, yet their reciprocal benefits under unified causal training remain unclear. We introduce OneWM, a simple framework for unified causal training that treats actions as both conditions for world prediction and targets for policy learning within a shared video-action architecture. Using this framework, we investigate whether learning forward dynamics conditioned on specified actions improves action learning, and whether learning to generate actions in turn improves action-conditioned world modeling. We use LIBERO, LIBERO-Plus, and RoboCasa365 as embodied testbeds to evaluate both policy learning and world modeling, complemented by real-world policy experiments. Learning the visual consequences of specified actions improves policy performance and generalization across changing scenes. In turn, action learning improves multi-view world-prediction fidelity under the same architecture. Together, these results answer both questions positively, supporting reciprocal benefits between world modeling and action generation under unified causal training.

Figures & tables

Figure 1: \CaptionLead Overview of OneWM. A VAE encodes frames into latents; the observed prefix remains clean, and the video expert denoises future latents. Action chunks are embedded for the action expert. Causal joint attention couples the two experts, each receiving instruction $\ell$ through its own text cross-attention (Sec. 3.2 ). The lower panels illustrate all-clean, clean-prefix/noisy-suffix, and all-noisy action construction, unifying world prediction and action generation (Sec. 3.3 ).

Figure 2: \CaptionLead Evaluation tasks and environments. We evaluate OneWM on LIBERO, LIBERO-Plus, and RoboCasa365 in simulation, and on three real-world tasks: block stacking, block insertion, and breakfast arrangement.

\sdtablehead LIBERO
Method	Spatial	Object	Goal	Long	Avg.
$\pi_{0}$ [ 51 ]	98.00	96.80	94.40	88.40	94.40
WorldVLA [ 46 ]	87.60	96.20	83.40	60.00	81.80
Fast-WAM [ 7 ]	98.20	100.00	97.00	95.20	97.60
WAM mask	93.80	96.20	96.80	83.00	92.45
WAM noisy	94.00	95.80	95.60	85.00	92.60

Table 1: \CaptionLead Policy performance in simulation. Success rates (%) on LIBERO, LIBERO-Plus, and RoboCasa365. Bold highlights the best result among internal variants. Bkgd., Lang., and Init. denote background, language, and robot initialization, respectively; S/U and Comp. denote seen/unseen and composite tasks in RoboCasa365.

Panel of Figure 3 — Figure 3: \CaptionLead Real-world policy performance. Success rates on three manipulation tasks; Overall aggregates the tasks.

	\sdtablehead LIBERO				\sdtablehead LIBERO-Plus				\sdtablehead RoboCasa365
\sdtablehead Metric	Wan	WM AdaLN	WM	OneWM	Wan	WM AdaLN	WM	OneWM	Wan	WM AdaLN	WM	OneWM
PSNR $\uparrow$	20.52	20.52	22.33	22.95	19.08	19.58	20.54	21.05	23.27	22.22	24.72	25.21
SSIM $\uparrow$	0.8045	0.8057	0.8325	0.8423	0.7575	0.7813	0.7937	0.8049	0.8176	0.7873	0.8433	0.8515
LPIPS $\downarrow$	0.1817	0.1793	0.1529	0.1458	0.2448	0.2185	0.2014	0.1935	0.1292	0.1564	0.1087	0.1015
FID $\downarrow$	14.00	12.88	11.89	12.45	24.01	15.91	18.06	17.01	9.04	10.28	9.20	9.16
FVD $\downarrow$	15.83	14.73	13.68	13.58	49.73	29.53	32.48	31.66	20.23	15.98	18.55	17.74

Table 2: \CaptionLead World-prediction performance. Arrows indicate the preferred direction. Bold and underline mark the best and second-best values within each benchmark, respectively.

	\sdtablehead LIBERO			\sdtablehead LIBERO-Plus
\sdtablehead Metric	60/25/15	50/35/15	30/45/25	60/25/15	50/35/15	30/45/25
SR (%) $\uparrow$	94.95	95.95	87.80	59.27	56.98	38.01
PSNR $\uparrow$	22.95	22.44	23.12	21.05	20.57	20.90
SSIM $\uparrow$	0.8423	0.8347	0.8443	0.8049	0.7943	0.7993
LPIPS $\downarrow$	0.1458	0.1517	0.1444	0.1935	0.2013	0.1966
FID $\downarrow$	12.45	12.57	12.36	17.01	17.32	17.18

Table 3: \CaptionLead Conditioning mixtures affect both tasks. Columns specify all-noisy / partial-clean / all-clean percentages; 60/25/15 is the default. SR is success rate (%). Bold marks the best value within each benchmark.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: \CaptionLead Main-camera world predictions. Six cases from LIBERO and LIBERO-Plus compare GT, OneWM, and the indicated world-model baseline. Orange boxes mark the regions enlarged at $t+32$ in each panel’s final column.

Figure 5: \CaptionLead Wrist-camera world predictions. Six matched comparisons emphasize object position, shape, and pose near the gripper. Each final column enlarges the same marked region at $t+32$ for all three rows.

Figure 6: RoboCasa365 world predictions. Six matched cases compare GT, OneWM, WM , and WM AdaLN from left (a,b), wrist (c,d), and right (e,f) cameras. The final column enlarges the same marked region at $t+16$ for all four rows.

Figure 7: \CaptionLead Endless dream on LIBERO. Six rollouts generated by alternating action prediction and action-conditioned world prediction. Each panel shows main- and wrist-camera views; blue outlines mark the initial observations, and every later frame is generated. Columns indicate modeled time, with the final column showing the end of round 16.

Action-Conditioned World Models

Action-conditioned world models predict how environments change under actions to support planning, control, and simulation.

See all 337 papers

UniW(A)M: A Unified Mixture-of-Experts Framework for World and World-Action Models
Oct 13, 2026Jingbo Xu, Jianjie Fang, Zhuohang Li +9World Models for Robotics World Action Models
Tsinghua University · Zhongguancun Academy · East China Normal University +1
CAER: Causal Action Effect Reweighting for World Model Training
Aug 31, 2026Jianjie Fang, Xvyuan Liu, Ziyou Wang +9Action-Conditioned Video Generation World Model Learning
1Tsinghua University · University of Science and Technology of China · 2Manifold AI
UniAction: Unified Action Conditioning for World Models
Oct 23, 2026Bowen Xue, Zizhang Li, Hong-Xing Yu +1Action-Conditioned World Models Action-Conditioned Video Generation
Stanford University

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

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

Jingbo Xu, Jianjie Fang, Zhuohang Li, Zihao Song+8

Tsinghua University · Zhongguancun Academy · East China Normal University+1

World Models for Robotics World Action Models Embodied World Models Action-Conditioned World Models Mixture-of-Experts Models

Open paper AskPDF arXiv

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%.

Abstract

Action-Conditioned World Models

Explore similar work

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

CAER: Causal Action Effect Reweighting for World Model Training

UniAction: Unified Action Conditioning for World Models