This paper presents the World-Action Model (WAM), an action-regularized world model that jointly reasons over future visual observations and the actions that drive state transitions. Unlike conventional world models trained solely via image prediction, WAM incorporates an inverse dynamics objective into DreamerV2 that predicts actions from latent state transitions, encouraging the learned representations to capture action-relevant structure critical for downstream control. We evaluate WAM on enhancing policy learning across eight manipulation tasks from the CALVIN benchmark. We first pretrain a diffusion policy via behavioral cloning on world model latents, then refine it with model-based PPO inside the frozen world model. Without modifying the policy architecture or training procedure, WAM improves average behavioral cloning success from 59.4% to 71.2% over DreamerV2 and DiWA baselines. After PPO fine-tuning, WAM achieves 92.8% average success versus 79.8% for the baseline, with two tasks reaching 100%, using 8.7x fewer training steps.
Figures & tables
Fig. 1 : (a) Standard world models predict only future observations, treating actions solely as conditioning inputs. (b) Our World-Action Model adds an inverse dynamics head that jointly predicts observations and actions during training. (c) The action-aware world model serves as a learned simulator for offline policy fine-tuning via PPO.
Fig. 2 : WAM architecture. Observations xt are encoded and produce posterior zt , which is regularized toward the prior z^t via KL divergence. The inverse dynamics head predicts actions a^t from consecutive encoder embeddings, cascading action-aware structure through the posterior to the prior. The decoder reconstructs observations x^t and a reward estimator provides task-completion signals for policy fine-tuning.
Fig. 3 : Qualitative comparison of imagined rollouts on the CALVIN benchmark. We visualize predicted frames at selected timesteps from both static and gripper cameras. Compared to DreamerV2, our WAM produces more realistic future state predictions across the entire rollout horizon.
Metric
Ours (WAM)
Baseline (DreamerV2)
PSNR ↑
22.10 ± 2.22
21.66 ± 2.20
SSIM ↑
0.814 ± 0.061
0.807 ± 0.067
LPIPS ↓
0.144 ± 0.072
0.149 ± 0.073
FVD ↓
10.82
12.13
TABLE I : Quantitative evaluation of imagination quality on the CALVIN benchmark. We report standard video prediction metrics between our World-Action Model (WAM) and the baseline (DreamerV2).
DiWA
Ours
World Model
DreamerV2
WAM
Diffusion Policy
DiffusionMLP
DiffusionMLP
TABLE II : Baseline comparison. For downstream policy learning, we share the same policy architecture and training pipeline as DiWA, differing only in the world model.
Task
DiWA
WAM (Ours)
close_drawer
58.6±4.2
89.7±3.1
open_drawer
53.3±5.1
73.3±4.8
move_slider_left
50.0±3.7
68.8±5.2
move_slider_right
51.7±4.5
82.8±3.9
turn_on_lightbulb
42.4±3.3
51.5±4.6
turn_off_lightbulb
3.4±1.8
17.2±3.4
TABLE III : Behavioral Cloning (BC) success rates (%) on 8 CALVIN manipulation tasks. Both methods use identical diffusion policy architectures and training procedures; the only difference is the world model used for feature extraction.
Task
Diffusion Policy
DiWA
WAM (Ours)
DPPO
Base
Offline Fine-Tuning
Online Fine-Tuning
Vision WM Encoder Env Steps to Match DiWA
Vision Env Steps to Match DiWA
open drawer
73.3 ± 4.8
74.44 ± 1.92
96.7 ± 2.4
117,600 ± 23,758
134,400 ± 26,508
close drawer
89.7 ± 3.1
91.95 ± 1.99
96.6 ± 1.8
600,600 ± 27,651
1,545,600 ± 261,346
move slider left
68.8 ± 5.2
83.33 ± 1.80
87.5 ± 3.7
270,933 ± 28,780
1,377,600 ± 251,439
move slider right
82.8 ± 3.9
82.76 ± 3.45
89.7 ± 3.2
249,600 ± 09,050
537,600 ± 23,758
turn on lightbulb
51.5 ± 4.6
91.92 ± 1.75
100.0 ± 0.0
302,933 ± 15,964
588,000 ± 62,859
TABLE IV : Diffusion policy success rates (%) on 8 CALVIN tasks after 800 iterations of model-based PPO fine-tuning. DPPO columns report environment steps required to match DiWA performance (lower = more sample-efficient).
Fig. 4 : Behavioral cloning evaluation curves across all eight CALVIN tasks. WAM (orange) consistently reaches higher success rates than DiWA (blue) and converges faster on most tasks.
World Models (WMs) offer a promising mechanism for post-training Vision-Language-Action (VLA) policies by providing dynamics priors that improve generalization under task and scene variation. However, most WM-based post-training methods rely on pixel-space supervision, making policies sensitive to visual artifacts introduced by imperfect WM rollouts. We present World2Act, a latent-space post-training framework that transfers WM dynamics to the VLA policy without pixel-space supervision. World2Act operates in two stages: 1) it induces a shared video-action latent space by contrastively aligning WM-dynamics latents with action embeddings, and 2) it post-trains the VLA by guiding policy action representations toward WM-imagined dynamics rather than decoded pixels. Built on GR00T-N1.6, World2Act delivers absolute success-rate gains of up to +2.5% on simulation benchmarks (RoboCasa, LIBERO, Bridge-SIMPLER) and +6.7% on a real robot over finetuned VLA baselines. Notably, it outperforms pixel-space WM supervision by up to +6.0%, including on LIBERO where pixel supervision degrades the baseline, suggesting that latent WM dynamics offer a more stable WM-based post-training alternative to pixel-space transfer.
World-action models (WAMs) transfer visual and motion priors from pretrained video generators to robot control by jointly modeling visual dynamics and actions. Existing WAMs, however, predict dense future frames during training, repeatedly modeling largely unchanged content and coupling action-conditioned dynamics to nuisance appearance variations. At inference, processing each complete observation with the heavy video expert bottlenecks few-step action generation. Accordingly, we propose DeltaWAM, which jointly predicts visual deltas and actions using dense-anchor, sparse-delta, and action streams, with three architectures that differ in representation and computation sharing. We further develop Streaming Delta Memory (SDM), which updates cached anchor context with compact observed deltas, reducing heavy video-expert processing. On RoboTwin, DeltaWAM with SDM improves average success over Fast-WAM from 81.3% to 85.4% in the clean setting and from 75.8% to 83.9% under visual randomization. The three architectures reduce training FLOPs by 17.78-23.77%, while SDM reduces one-step inference latency and FLOPs by 36.57% and 31.55%, respectively; real-world evaluations further show the highest overall success rate and normalized progress among the evaluated policies. Code: https://github.com/AIGeeksGroup/DeltaWAM. Website: https://aigeeksgroup.github.io/DeltaWAM.
Han Yan, Zishang Xiang, Haokai Jiang +6
School of Computer Science, Peking University · AI2 Robotics
World Action Models (WAMs) emerge as a promising foundation for policy refinement by predicting the consequences of sampled actions. However, visually plausible predictions can mislead policy refinement if they fail to reflect the input actions. To address this mismatch, we introduce R2-WAM, a two-stage repair-and-reject post-training framework that first improves the consistency of predicted futures with input actions, then uses these futures to select inferior action samples for negative fine-tuning. The repair stage grounds imagination in observed robot behavior through a kinematic alignment score that measures agreement between predicted and demonstrated motion, enabling the predicted video to faithfully reflect its input actions. Using the repaired video model, the rejection stage compares imagined outcomes of sampled and demonstrated actions, selectively applying negative fine-tuning to samples whose predicted task progress falls below the demonstrated reference by a prescribed margin. Together, the two stages extend video prediction from representation learning to consequence-based policy refinement without additional environment interaction or changes to the inference procedure. R2-WAM achieves 93.8% average success on RoboTwin 2.0 across clean and randomized settings. On the long-horizon real-world Fold Shirt task, it achieves 87.5% average success, compared with 0% for Fast-WAM. Our project page is available at https://r2-wam.github.io/.
Ruiyan Xu, Haisheng Su, Sixu Lin +4
The Chinese University of Hong Kong, Shenzhen · Manifold AI · Southeast University +1