Enhancing Policy Learning with World-Action Model

TL;DR

The paper introduces the World-Action Model (WAM), a world model that jointly reasons over visual observations and actions, improving policy learning efficiency and success rates in manipulation tasks.

Contribution

WAM incorporates an inverse dynamics objective into DreamerV2, capturing action-relevant structure and enhancing downstream control performance.

Findings

WAM improves behavioral cloning success from 59.4% to 71.2%.

WAM achieves 92.8% success after PPO fine-tuning, outperforming baselines.

WAM reaches 100% success in two tasks with 8.7x fewer training steps.

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.