Enhancing Policy Learning with World-Action Model

Yuci Han; Alper Yilmaz

cs.AIMar 30, 2026

Enhancing Policy Learning with World-Action Model

Authors: Yuci Han, Alper Yilmaz

Organizations: Photogrammetry and Computer Vision Lab, The Ohio State University, Columbus, OH 43210, USA.

Abstract

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 $x_{t}$ are encoded and produce posterior $z_{t}$ , which is regularized toward the prior $\hat{z}_{t}$ via KL divergence. The inverse dynamics head predicts actions $\hat{a}_{t}$ from consecutive encoder embeddings, cascading action-aware structure through the posterior to the prior. The decoder reconstructs observations $\hat{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 $\uparrow$	22.10 $\pm$ 2.22	21.66 $\pm$ 2.20
SSIM $\uparrow$	0.814 $\pm$ 0.061	0.807 $\pm$ 0.067
LPIPS $\downarrow$	0.144 $\pm$ 0.072	0.149 $\pm$ 0.073
FVD $\downarrow$	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\pm 4.2$	$\mathbf{89.7\pm 3.1}$
open_drawer	$53.3\pm 5.1$	$\mathbf{73.3\pm 4.8}$
move_slider_left	$50.0\pm 3.7$	$\mathbf{68.8\pm 5.2}$
move_slider_right	$51.7\pm 4.5$	$\mathbf{82.8\pm 3.9}$
turn_on_lightbulb	$42.4\pm 3.3$	$\mathbf{51.5\pm 4.6}$
turn_off_lightbulb	$3.4\pm 1.8$	$\mathbf{17.2\pm 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
Task	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 $\pm$ 4.8	74.44 $\pm$ 1.92	96.7 $\pm$ 2.4	117,600 $\pm$ 23,758	134,400 $\pm$ 26,508
close drawer	89.7 $\pm$ 3.1	91.95 $\pm$ 1.99	96.6 $\pm$ 1.8	600,600 $\pm$ 27,651	1,545,600 $\pm$ 261,346
move slider left	68.8 $\pm$ 5.2	83.33 $\pm$ 1.80	87.5 $\pm$ 3.7	270,933 $\pm$ 28,780	1,377,600 $\pm$ 251,439
move slider right	82.8 $\pm$ 3.9	82.76 $\pm$ 3.45	89.7 $\pm$ 3.2	249,600 $\pm$ 09,050	537,600 $\pm$ 23,758
turn on lightbulb	51.5 $\pm$ 4.6	91.92 $\pm$ 1.75	100.0 $\pm$ 0.0	302,933 $\pm$ 15,964	588,000 $\pm$ 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.

Explore similar work

World2Act: Latent Action Post-Training from World Model Dynamics
Mar 11, 2026An Dinh Vuong, Tuan Van Vo, Abdullah Sohail +6Vision-Language-Action Models VLMs for Robotics
MBZUAI
DeltaWAM: Delta World Action Models for Bimanual Manipulation
Sep 23, 2026Han Yan, Zishang Xiang, Haokai Jiang +6Video Diffusion Models Bimanual Robotic Manipulation
School of Computer Science, Peking University · AI2 Robotics
$R^2$ -WAM: Repair-and-Reject Post-Training for World Action Models
Oct 4, 2026Ruiyan Xu, Haisheng Su, Sixu Lin +4World Action Models Action-Conditioned World Models
The Chinese University of Hong Kong, Shenzhen · Manifold AI · Southeast University +1

Mar 11, 2026J/K move · Enter open · S save

World2Act: Latent Action Post-Training from World Model Dynamics

An Dinh Vuong, Tuan Van Vo, Abdullah Sohail, Haoran Ding+5

MBZUAI

Vision-Language-Action Models VLMs for Robotics Latent Action Learning

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

World2Act: Latent Action Post-Training from World Model Dynamics

DeltaWAM: Delta World Action Models for Bimanual Manipulation

R2R^2R2-WAM: Repair-and-Reject Post-Training for World Action Models

$R^2$ -WAM: Repair-and-Reject Post-Training for World Action Models