Enhancing Policy Learning with World-Action Model

2026/03/30 by Yuci Han, Alper Yilmaz · 1 voice
Computer Science · #Cloning (programming) #Generative Adversarial Networks and Image Synthesis #Human Pose and Action Recognition #Key (lock) #Perspective (graphical) #Policy learning #Reinforcement Learning in Robotics #State (computer science) #Training (meteorology) #cs.AI

paper · pdf · doi:10.48550/arxiv.2603.28955

openalex publication_date 2026/03/30 · arxiv published 2026/03/30 · arxiv updated 2026/03/30 · openalex created_date 2026/04/02 · openalex updated_date 2026/07/28

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.

Citations

• WMPO: World Model-based Policy Optimization for Vision-Language-Action Models
• World4RL: Diffusion World Models for Policy Refinement with Reinforcement Learning for Robotic Manipulation
• DiWA: Diffusion Policy Adaptation with World Models
• WorldVLA: Towards Autoregressive Action World Model
• LUMOS: Language-Conditioned Imitation Learning with World Models
• Diffusion Policy Policy Optimization
• DITTO: Offline Imitation Learning with World Models
• DayDreamer: World Models for Physical Robot Learning
• Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos
• DreamingV2: Reinforcement Learning with Discrete World Models without Reconstruction
• CALVIN: A Benchmark for Language-Conditioned Policy Learning for Long-Horizon Robot Manipulation Tasks
• Mastering Atari with Discrete World Models
• Dream to Control: Learning Behaviors by Latent Imagination
• Towards Accurate Generative Models of Video: A New Metric & Challenges
• Learning Latent Dynamics for Planning from Pixels
• Recurrent World Models Facilitate Policy Evolution
• The Unreasonable Effectiveness of Deep Features as a Perceptual Metric
• Curiosity-driven Exploration by Self-supervised Prediction

Discussions

• How can AI policy learning improve by integrating world-action models? New research suggests combining these models enhances decision-making and adaptability in complex environments. 🔗 https://arxiv. [bsky, 0 points, 0 comments]

Related

• WMPO: World Model-based Policy Optimization for Vision-Language-Action Models
• World Action Models: The Next Frontier in Embodied AI
• AdaWorld: Learning Adaptable World Models with Latent Actions
• DWM: Separating World Effects from Actions in Latent World Models
• Self-Improving World Modelling with Latent Actions