Skip to content

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

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

ICLR 2026; world-model RL policy optimization for VLA; goes beyond demonstrations

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

WMPO (World Model-based Policy Optimization) is a framework that enables on-policy reinforcement learning for Vision-Language-Action (VLA) robotic manipulation models entirely within a learned pixel-space video world model, eliminating the need for costly real-world rollouts. The world model is pre-trained on large-scale robot trajectories (OXE) and aligned to the current policy's behavior distribution, then used as a synthetic environment for GRPO-based policy optimization. This produces emergent self-correction behaviors not present in the demonstration data.

Problem

VLA models trained via imitation learning are brittle under distribution shift and cannot recover from failures. Direct RL on real robots suffers from prohibitive sample complexity. Existing model-based RL methods use latent-space world models (e.g., RSSM/Dreamer), which are misaligned with VLAs pretrained on web-scale images — the pretrained visual representations cannot operate on abstract latent dynamics, preventing effective use of rich pretrained knowledge.

Method

WMPO grounds VLA RL in an action-conditioned, pixel-space video diffusion world model (modified OpenSora backbone with SDXL 2D VAE replacing the 3D VAE for better motion fidelity). Three key mechanisms:

  1. Policy Behavior Alignment: The world model is pre-trained on OXE trajectories then fine-tuned on real rollouts from the current policy \(\pi_\theta\), adapting the imagined distribution to include failure cases.

  2. Long-horizon generation: Clip-level autoregressive generation produces complete trials. Noisy-frame conditioning perturbs conditioning frames with mild diffusion noise (step 50 of 1000) during training to prevent error accumulation. Frame-level action control injects action signals via an extended AdaLN block — for frame \(i\), modulation coefficients \((\gamma_1^i, \beta_1^i, \alpha_1^i)\) are produced by an MLP from action \(a_i\), yielding \(x^i = x^i + (1 + \alpha_1^i) \cdot \text{Block}(\gamma_1^i \cdot \text{LayerNorm}(x^i) + \beta_1^i)\).

  3. Reward model: A VideoMAE encoder with a linear head trained with binary cross-entropy on real trajectories to predict success/failure from clip segments. Sparse binary reward \(R_\psi(\tau) \in \{0,1\}\) is assigned per imagined trajectory.

Policy optimization uses GRPO with the objective: $\(\mathcal{J}(\theta) = \mathbb{E}\left[\frac{1}{G}\sum_{i=1}^G \frac{1}{T}\sum_{t=0}^T \min\left(r_{i,t}(\theta)\hat{A}_i,\ \text{clip}(r_{i,t}(\theta), 1-\epsilon_{\text{low}}, 1+\epsilon_{\text{high}})\hat{A}_i\right)\right]\)$ where \(r_{i,t}(\theta) = \pi_\theta(a_{i,t}|s_{i,t}) / \pi_{\theta_{\text{old}}}(a_{i,t}|s_{i,t})\) and \(\hat{A}_i = (R_i - \text{mean}(\{R_i\})) / \text{std}(\{R_i\})\). KL divergence regularization is removed (following DAPO). Dynamic Sampling discards all-success or all-failure groups to prevent vanishing gradients.

Key Contributions

  • First framework for on-policy GRPO training of VLA models entirely within a pixel-space video world model, no real-world interaction during RL
  • Policy Behavior Alignment technique to adapt the world model's imagination to cover the current policy's failure distribution
  • Noisy-frame conditioning and frame-level AdaLN action control for stable long-horizon generation
  • Lightweight VideoMAE-based reward model achieving F1 > 0.95 for success/failure classification
  • Demonstration of emergent self-correction behavior not present in demonstration data

Results

  • On Mimicgen simulation (4 tasks: Coffee D0, StackThree D0, ThreePieceAssembly D0, Square D0), base policy OpenVLA-OFT fine-tuned on 300 expert demos/task:
  • P=128 real rollouts: WMPO 47.1% vs GRPO 33.2% (+13.9 pts), DPO 37.3% (+9.8 pts), Base 33.6%
  • P=1280 real rollouts: WMPO 57.6% vs GRPO 37.1% (+20.5 pts), DPO 42.4% (+15.2 pts)
  • Reward model achieves F1 > 0.95 across all tasks
  • WMPO outperforms GRPO and DPO on all four individual tasks at both budget levels
  • Generalization experiments (position, background, and texture disruption) and lifelong learning (alternating VLA/world model updates) validated qualitatively in real-robot settings

Limitations

  • Assumes fully observable MDP (image-only state); POMDPs explicitly left for future work
  • World model quality degrades if the policy behavior alignment data is too limited or non-representative
  • Pixel-space generation is computationally heavier than latent-space alternatives
  • Reward model trained only on real trajectories may misclassify unusual failure modes in imagination
  • No quantitative real-robot success rates reported; real-robot evaluation remains largely qualitative

Relevance to Vision-Language Models

WMPO is directly relevant to the VLM-for-robotics thread: it builds on VLMs (OpenVLA-OFT, which inherits a VLM backbone) and addresses the key gap between passive VLM pretraining and active, environment-interactive policy improvement. The pixel-space world model is specifically motivated by the need to preserve compatibility with VLM visual representations pretrained on web-scale images — a fundamental alignment insight for anyone extending VLMs to embodied agents. The GRPO adaptation from LLM post-training to continuous robotic action spaces demonstrates a generalizable recipe for bringing language-model-style RL (DeepSeek-style GRPO) into the VLA domain, a direction of growing interest in multimodal foundation model research.