WMPO: World Model-based Policy Optimization for Vision-Language-Action Models

TL;DR

WMPO introduces a pixel-based world model for vision-language-action reinforcement learning, enabling efficient, on-policy training without real-world interactions, leading to improved performance and self-correction in robotic tasks.

Contribution

The paper presents WMPO, a novel framework that uses pixel-based world models for on-policy RL in VLA models, reducing sample complexity and enhancing robustness.

Findings

WMPO significantly improves sample efficiency in robotic learning.

WMPO achieves stronger performance than off-policy methods.

WMPO demonstrates emergent self-correction and lifelong learning capabilities.

Abstract

Vision-Language-Action (VLA) models have shown strong potential for general-purpose robotic manipulation, but their reliance on expert demonstrations limits their ability to learn from failures and perform self-corrections. Reinforcement learning (RL) addresses these through self-improving interactions with the physical environment, but suffers from high sample complexity on real robots. We introduce World-Model-based Policy Optimization (WMPO), a principled framework for on-policy VLA RL without interacting with the real environment. In contrast to widely used latent world models, WMPO focuses on pixel-based predictions that align the "imagined" trajectories with the VLA features pretrained with web-scale images. Crucially, WMPO enables the policy to perform on-policy GRPO that provides stronger performance than the often-used off-policy methods. Extensive experiments in both simulation and real-robot settings demonstrate that WMPO (i) substantially improves sample efficiency, (ii) achieves stronger overall performance, (iii) exhibits emergent behaviors such as self-correction, and (iv) demonstrates robust generalization and lifelong learning capabilities.