Source-linked AI summary
WMPO: World Model-based Policy Optimization for Vision-Language-Action Models
Fangqi Zhu, Zhengyang Yan, Zicong Hong, Quanxin Shou, Xiao Ma, Song Guo
TL;DR
VLA policies are limited by dependence on demonstrations, while real-robot reinforcement learning is sample-intensive. WMPO uses a pixel-based, action-conditioned world model to support on-policy GRPO without real-world interaction, yielding stronger sample efficiency and performance alongside self-correction, generalization, and lifelong learning. The paper’s scope includes an image-observation state assumption and the broader challenge of accurate simulation for diverse scenarios.
Problem
VLA imitation learning is brittle on unseen states and lacks learning from failures, while real-robot reinforcement learning requires costly interaction data.
Method
WMPO pretrains and policy-aligns an action-conditioned pixel-space video world model, then performs on-policy GRPO on imagined trajectories with outcome-based rewards.
Results
WMPO consistently outperforms model-free or real-trajectory baselines while improving sample efficiency and exhibiting self-correction, generalization, and lifelong learning in simulation and real-world experiments.
Takeaways & Limitations
WMPO provides a scalable learned-environment paradigm for VLA reinforcement learning without costly real-world policy-optimization interactions.
Takeaways & Limitations
The formulation assumes robot states are fully defined by image observations and leaves partially observable settings for future work.
Abstract
from arXiv · showhide
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.
1 Introduction
VLA imitation learning is brittle outside demonstrations, while real-world RL is costly and difficult to run on-policy. WMPO addresses these constraints with a pixel-space world model, policy-aligned imagined rollouts, and on-policy optimization without real-world interaction.
- Motivation: Imitation learning struggles with failures and self-correction, whereas real-world RL requires impractical, unsafe, and time-consuming interaction data.Real-robot RL may require millions of interactions, while imitation-trained policies can compound errors on unseen states.
- WMPO: WMPO grounds VLA reinforcement learning in an action-conditioned pixel-space video world model pretrained on large-scale robotic trajectories.The model creates a scalable training environment for downstream tasks using realistic visual dynamics.
- Method: WMPO aligns the world model with policy behavior and generates complete trials for more reliable outcome-based reward assignment.These designs address state-distribution mismatch and the reward-hacking risks of short-horizon prediction, while introducing visual distortion and action–frame misalignment challenges.
- Results: Experiments in simulation and real-world environments show substantially higher sample efficiency and consistent improvements over VLA RL methods using real trajectories.The reported evaluation covers both Mimicgen simulation and real-world settings.
- Results: WMPO exhibits emergent self-correction, stronger generalization, and lifelong learning through alternating updates between the VLA policy and world model.The self-correction strategies were not present in the demonstration data, and tasks were often completed faster and more smoothly.
2 Related Work
Prior VLA reinforcement-learning methods use human supervision, simulation, or real-world interaction, but remain difficult to scale because of labor, sample inefficiency, or simulator construction costs. World-model research has progressed from abstract latent rollouts toward pixel-level fidelity, motivating WMPO’s learned-environment approach.
- Reinforcement Learning for VLA Models: VLA reinforcement learning complements imitation learning by improving robustness and recovery, but on-policy optimization remains sample-inefficient and systemically complex.Existing approaches include human-guided learning, simulation, and real-world optimization.
- Reinforcement Learning for VLA Models: Human-guided methods reduce exploration cost but require continuous supervision, while simulation and real-world methods face poor efficiency and prohibitive scenario-specific engineering.Prior methods include trajectory-level DPO, PPO, and GRPO applied in simulation or real environments.
- World Models: WMPO differs by performing policy optimization entirely within a learned world model, substantially improving sample efficiency over direct reinforcement learning.This approach avoids reliance on costly real-world interactions during policy optimization.
- World Models: World models enable policy learning from imagined trajectories, progressing from efficient but abstract latent dynamics toward diffusion-based pixel-level predictions.Pixel-level fidelity has been identified as important for reinforcement learning with Gaussian policies.
3 World Model-Based Policy Optimization
WMPO performs VLA policy optimization entirely within a learned pixel-space world model, using imagined trajectories instead of costly real-world interactions. Its pipeline combines trajectory generation, reward-based sampling, and on-policy policy updates, with alignment and conditioning techniques supporting long-horizon simulation.
- The training loop alternates imagined trajectory generation, reward-based trajectory sampling, and policy updates, repeating these stages throughout training.The policy and world model generate trajectories, the reward model evaluates them, and policy parameters are optimized iteratively.
- 3.1 Problem Formulation: WMPO formulates manipulation as an MDP with image-and-language states, action chunks, a learned world-model transition function, and binary trajectory-level rewards.The framework assumes robot states can be defined solely by image observations and leaves partially observable settings for future study.
- WMPO decouples VLA reinforcement learning from real-world interactions by optimizing policies within a learned generative world model.The framework operates directly in pixel space to better bridge pretrained VLA features and imagined trajectories.
- 3.2 Generative World Model: The world model generates full imagined trajectories autoregressively from frames, actions, and language-conditioned policy outputs, which the reward model labels for optimization.Each trajectory is represented as a trajectory-label pair before policy optimization.
- 3.2 Generative World Model: Noisy-frame conditioning and frame-level action control mitigate long-horizon visual degradation and action–frame misalignment in autoregressive video prediction.The reward model uses clips from full trajectories to classify task success, enabling outcome-based evaluation of imagined rollouts.
- 3.2 Generative World Model: Policy Behavior Alignment fine-tunes the world model with policy-collected trajectories so it can simulate diverse scenarios, including failures.The model is pretrained on diverse robot demonstrations, while downstream policy data addresses the underrepresentation of failure scenarios.
- 3.4 On-Policy Reinforcement Learning with WMPO: On-policy GRPO uses imagined trajectories, trajectory rewards, probability ratios, and normalized advantages to optimize the VLA policy.The policy objective uses the ratio between new and old policies and trajectory-level normalized advantages.
4 Experiments
Experiments evaluate WMPO against online and offline RL baselines in simulation, testing efficiency, emergent behavior, generalization, lifelong learning, and real-robot performance. Across these settings, WMPO improves success, supports self-correction, generalizes across visual disruptions, and achieves the strongest reported real-world success rate.
- 4.2 Comparison Experiments: WMPO consistently outperforms GRPO and DPO across four simulation tasks, with its advantage growing from +9.8 points at P=128 to +15.2 points on average at P=1280.The rollout budget P denotes the number of full real trajectories available for optimization.
- 4.3 Emergent Behavior of WMPO: WMPO learns self-correction behaviors that recover from collisions by lifting, realigning, and successfully inserting the square, unlike the base policy.The base policy continues pushing after collision and reaches the time limit, whereas WMPO uses imagined trajectories to learn recovery.
- 4.3 Emergent Behavior of WMPO: WMPO-trained policies produce significantly shorter successful trajectories and rarely become stuck in suboptimal states.The authors associate shorter trajectories with discouraging stuck behaviors that often lead to timeout failures.
- 4.4 Generalization to Novel Tasks: WMPO achieves the best performance across spatial, background, and texture disruptions, while DPO degrades under visual changes and GRPO remains similar to the base policy.The disruption scenarios vary object position, tabletop background, and base texture.
- 4.5 Lifelong Learning: WMPO achieves stable and substantial iterative improvements on StackThree, whereas DPO fails to improve iteratively because of unstable training.WMPO alternates policy optimization with collection of additional real trajectories, while using policy-collected data rather than more human demonstrations.
- 4.6 Real-world Experiments: On the real-world square-insertion task, WMPO reaches a 70% success rate, compared with 60% for DPO and 53% for the base policy.All models are evaluated under identical conditions over 30 trials.
5 Conclusion
WMPO enables on-policy VLA reinforcement learning entirely within a video-generative world model, avoiding costly real-world interactions while maintaining consistency with pretrained VLA representations. Experiments show improved performance, self-correction, generalization, and lifelong learning.
- WMPO consistently outperforms state-of-the-art model-free baselines in simulation and real-world experiments.
- WMPO grounds policy optimization in a video-generative world model, eliminating costly real-world interactions while preserving consistency with pretrained VLA representations.
- WMPO-trained policies exhibit emergent self-correction behavior not present in the demonstration data.
- WMPO generalizes reliably to unseen scenarios and supports iterative lifelong learning.
A Training Details
WMPO training uses a pretrained and behavior-aligned world model together with GRPO policy optimization, using generated trajectories and outcome rewards to update the VLA policy.
- Training excludes robot proprioceptive states and wrist-camera inputs, while using parallel prediction of discrete action tokens.Supervised fine-tuning uses 8 H100 GPUs, and world-model training plus policy optimization use 32 H100 GPUs.
- The world model is pretrained on OXE data and then fine-tuned on downstream policy behavior data.
- Algorithm 1 takes a world model, reward model, VLA policy, group size, batch size, mini-batch size, epochs, and horizon as inputs.
- The algorithm samples initial states, generates grouped trajectories with the world model, assigns rewards, and updates the VLA policy using GRPO.The procedure maintains an old policy, forms trajectory batches, filters degenerate reward groups, and updates policy parameters over mini-batches.
B Baseline Details
The baseline implementations compare online GRPO and offline DPO with WMPO under matched rollout budgets and evaluation protocols, while exposing the real-trajectory cost of online GRPO.
- Online GRPO requires at least 512 real trajectories for one update with batch size 64 and group size 8.Dynamic sampling removes groups with success rates of 0 or 1, making large batches infeasible at low rollout budgets.
- At rollout budget P = 1280, online GRPO can be updated only once or twice with the larger batch configuration.
- The study also evaluates online GRPO with batch size 8 and proportionally reduced learning rate, reporting the best result across both batch configurations.
- The DPO baseline uses preference data from supervised-fine-tuned OpenVLA-OFT trajectories and matches GRPO architecture, optimizer settings, rollout budgets, and evaluation protocols.
C Real World Cases
Real-world world-model predictions usually capture successful and failed outcomes, but rare failures remain near the final moment of contact.
- The world model correctly predicts that a misaligned square and stick cannot produce a successful insertion.
- A rare failure occurs when subtle perturbations prevent the model from capturing the moment a square becomes stuck, despite accurate earlier frames.
- Validation results indicate that the world model reliably predicts successful and failed outcomes in most scenarios.This reliability is presented as important for stable policy optimization.
D Limitation
WMPO currently focuses on discretized action representations, while flow-based policy classes remain future work.
- WMPO's current implementation focuses on discretized action representations.
- Flow-based policies are supported in principle but are not evaluated in this work.
- The authors plan to extend WMPO to flow-matching policies and investigate FlowGRPO for broader action-space applicability.
E LLM Usage Statement
GPT-5 was used only for manuscript language refinement and stylistic polishing, with no role in the study's scientific work.
- GPT-5 assisted solely with language refinement and stylistic polishing of the manuscript.
- The authors state that GPT-5 did not generate research concepts, design the study, analyze data, or interpret findings.
- The supplied materials also include world-model and GRPO hyperparameter tables and real-world failure-case figures.