# Enhancing Policy Learning with World-Action Model

arXiv: https://arxiv.org/abs/2603.28955 · HTML: https://mlsanity.com/paper/2603.28955/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Embodied and robot world models, one of 11 directions in World models.
- Builds on 18 of its 19 references that are in the map.
- First posted on arXiv 30 Mar 2026.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)
- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 18 of 18)

- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Towards Accurate Generative Models of Video: A New Metric & Challenges](https://mlsanity.com/paper/1812.01717/) (2018)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [CALVIN: A Benchmark for Language-Conditioned Policy Learning for Long-Horizon Robot Manipulation Tasks](https://mlsanity.com/paper/2112.03227/) (2021)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [WorldVLA: Towards Autoregressive Action World Model](https://mlsanity.com/paper/2506.21539/) (2025)
- [Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos](https://mlsanity.com/paper/2206.11795/) (2022)
- [Diffusion Policy Policy Optimization](https://mlsanity.com/paper/2409.00588/) (2024)
- [WMPO: World Model-based Policy Optimization for Vision-Language-Action Models](https://mlsanity.com/paper/2511.09515/) (2025)
- [World4RL: Diffusion World Models for Policy Refinement with Reinforcement Learning for Robotic Manipulation](https://mlsanity.com/paper/2509.19080/) (2025)
- [DiWA: Diffusion Policy Adaptation with World Models](https://mlsanity.com/paper/2508.03645/) (2025)
- [DreamingV2: Reinforcement Learning with Discrete World Models without Reconstruction](https://mlsanity.com/paper/2203.00494/) (2022)
- [DITTO: Offline Imitation Learning with World Models](https://mlsanity.com/paper/2302.03086/) (2023)
- [LUMOS: Language-Conditioned Imitation Learning with World Models](https://mlsanity.com/paper/2503.10370/) (2025)

## Built on by (most cited first, 6 of 6)

- [Sensorimotor World Models: Perception for Action via Inverse Dynamics](https://mlsanity.com/paper/2606.20104/) (2026)
- [What Makes Video World Model Latents Action-Relevant: Prediction over Reconstruction](https://mlsanity.com/paper/2606.07687/) (2026)
- [ATM: Action-Consistency Transfer Matrix for Diagnosing and Improving Latent World Models](https://mlsanity.com/paper/2606.09028/) (2026)
- [ActSWM: Action-Sensitive World Models for Long-Horizon Planning in Open-World Games](https://mlsanity.com/paper/2607.26712/) (2026)
- [TacWAM: Anchor-Guided World Action Model with Mechanics-Aware Tactile Prediction](https://mlsanity.com/paper/2607.28391/) (2026)
- [verdi: retrieval is not transfer for continual world model optimization](https://mlsanity.com/paper/2608.09537/) (2026)

## Abstract (arXiv)

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.