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 ℓ 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.
π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.
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 ↑
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 ↑
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 ↓
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 ↓
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 ↓
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 (%) ↑
94.95
95.95
87.80
59.27
56.98
38.01
PSNR ↑
22.95
22.44
23.12
21.05
20.57
20.90
SSIM ↑
0.8423
0.8347
0.8443
0.8049
0.7943
0.7993
LPIPS ↓
0.1458
0.1517
0.1444
0.1935
0.2013
0.1966
FID ↓
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.
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%.
Jingbo Xu, Jianjie Fang, Zhuohang Li +9
Tsinghua University · Zhongguancun Academy · East China Normal University +1
World models are becoming core infrastructure for embodied intelligence, with action-conditioned video generation providing controllable predictions of how scenes evolve after agent interventions. Yet existing models are commonly trained with space-time-uniform mean squared error, allowing abundant background tokens to dominate the gradient while sparse interaction dynamics remain under-optimized; such uniform fitting rewards reconstructing appearance rather than learning how actions change the world. We introduce Causal Action Effect Reweighting (CAER), a general training paradigm that redistributes supervision toward the tokens whose predicted future is causally affected by the action. CAER contrasts the model's own predictions with and without action conditioning to localize these tokens online, then normalizes the resulting effect map into a weight that preserves the total coefficient mass and changes only where it is spent. This online signal requires no external annotations or offline preprocessing, avoids additional data-processing time, and scales naturally with model and dataset size. Experiments across heterogeneous action-conditioned world-model tasks show that CAER converges to better solutions than uniform MSE training, with consistent improvements in the physical consistency, controllability, and visual quality of generated videos.
Jianjie Fang, Xvyuan Liu, Ziyou Wang +9
1Tsinghua University · University of Science and Technology of China · 2Manifold AI
Actions that condition world models are highly heterogeneous, and thus existing action-conditioned world models are typically built for individual action types, fragmenting already limited action-annotated data across separate models. We introduce UniAction, a unified action-conditioning framework that represents any action as a sparse set of action elements, each with a discrete semantic identity and a continuous spatiotemporal grounding coordinate. By mapping these coordinates directly into the pretrained video model's native positional space, UniAction allows heterogeneous actions to condition a shared world model within the same spatiotemporal structure. We further introduce cardinality-calibrated attention to address the bias from differing action representation granularities. We evaluate UniAction on four representative action forms: force, camera control, hand pose, and robot pose. UniAction outperforms conditioning alternatives in action following, and a single jointly trained model surpasses the action-specific experts, showing that heterogeneous action data benefit from unified model and joint training.