WAVE-0: 257M Latent World Model Nearly Doubles π₀.₅ Real-Robot Success Rate

Background: From VLA to World Action Models

Vision-Language-Action (VLA) models like π₀, π₀.₅, and OpenVLA are the mainstream approach for general robot manipulation. They combine a pre-trained vision-language model (VLM) for semantic understanding with robot demonstration data for action generation. However, VLAs typically lack explicit prediction of action consequences, making multi-step reasoning difficult. World Action Models (WAMs) aim to add this capability, but video-generation-based WAMs require costly pixel-level prediction. Latent-space methods avoid pixel generation but still face the challenge of how to integrate future predictions into decision-making.

Existing latent-space approaches have limitations: VLA-JEPA and JEPA-WAM use future prediction only for representation learning, not for explicit action generation; LaWAM predicts visual sub-goals to guide policy but infers latent actions from video changes via an inverse dynamics model, relying on ~4,500 hours of video pre-training plus subsequent policy training.

The authors ask: as annotated robot trajectories (real and simulated) grow, is it still necessary to indirectly learn environment dynamics from unlabeled video?

WAVE-0 Design: World-model-Augmented VLA Execution

WAVE-0 (World-model-Augmented VLA Execution) builds on a key data observation: every robot demonstration trajectory records both the action sequence and the subsequent observation frames. This provides paired supervision for "execute this action, how does the environment change?" Training a world model directly on this data yields future predictions that strictly correspond to executable actions, avoiding inverse-dynamics misalignment and eliminating the need for large-scale video pre-training.

A second pillar is reusing the pre-trained VLA's capabilities. π₀.₅'s visual encoder already provides a semantically rich visual representation space; the world model only needs to learn "how representations evolve given an action in this space." Predicting representations instead of pixels allows the world model to be very small: the world model plus its encoder totals 257M trainable parameters, about 6% of the full WAVE-0 (≈4.3B parameters).

At inference, the VLM generates a latent action plan from the current observation and language instruction. The world model predicts future visual representations from this plan. A lightweight resampling module feeds these predictions into the Action Expert, which jointly conditions on the original visual-language context to produce the final action. This process operates at both training and inference time, letting predicted futures directly guide action decisions without generating future images, adding only the forward passes of the world model and resampling module.

Two-Stage Training

Stage 1: Learning How Actions Change the Environment

Robot demonstrations provide correspondences between current observation, true action sequence, and future observation. WAVE-0 uses a frozen visual encoder to extract current and future visual representations. The true action sequence and robot state are encoded and fed to the latent world model, which predicts the representation change from current to future. Training supervises both the future representation and its delta relative to the current state, teaching the model action-conditioned environment dynamics.

To increase sensitivity to action conditioning, WAVE-0 introduces Counterfactual Ranking supervision: keeping the current observation and robot state fixed, it feeds the true action sequence and a mismatched action sequence, requiring the true action's prediction to be closer to the actual future than the mismatched action's prediction. This encourages the model to use action information to distinguish different outcomes without collecting observations for alternative actions.

Stage 2: Aligning Latent Action Plans with the World Model

After freezing the Stage-1 world model and its encoder, WAVE-0 plugs them into pre-trained π₀.₅ and jointly trains the VLM, resampling module, and Action Expert. The policy learns to generate latent action plans and to use predicted futures to guide action generation.

Training is driven by three losses: (1) action generation loss supervises the final action; (2) action representation alignment loss pulls the VLM-generated plan toward the representation of the true action encoding; (3) future prediction loss constrains the plan's predicted future representation and its delta. The latter two auxiliary losses respectively constrain "how the plan is expressed" and "what consequences the plan brings," helping the VLA effectively exploit the environment dynamics learned in Stage 1.

Experimental Results

Standard Simulation Benchmarks

On the four LIBERO task suites, WAVE-0 achieves an average success rate of 98.7% , up 1.8 percentage points from π₀.₅'s 96.9%. LIBERO-Object reaches 100.0%; LIBERO-Long improves from 92.4% to 97.6% (+5.2 pp), showing stronger gains on long-horizon tasks where future prediction helps multi-step reasoning.

On RoboTwin 2.0 (50 dual-arm tasks), WAVE-0 raises the average success rate from π₀.₅'s 80.38% to 91.47% (+11.09 pp). Both clean and domain-randomized settings improve similarly (91.46% and 91.48%), matching LaWAM's 91.22% without any extra large-scale video pre-training.

Out-of-Distribution Generalization

Using the LIBERO-trained model directly (no fine-tuning), WAVE-0 scores 88.0% on LIBERO-Plus (+3.5 pp over π₀.₅) and 61.4% on LIBERO-Pro (+8.4 pp). On LIBERO-Para (instruction paraphrase), success rises from 71.4% to 75.0%. WAVE-0 outperforms all compared world-model methods on LIBERO-Plus and LIBERO-Pro, including LaWAM (which uses extra video pre-training) and ImageWAM (which does not), demonstrating robustness to visual, spatial, and language distribution shifts.

Real-Robot Evaluation

Three tasks on an AgileX Piper dual-arm platform: Book Organization, USB Storage, and Keyboard Pushback. Each task uses 200 demonstrations; all methods train on the same data and are tested 20 times per task.

WAVE-0 achieves 85.0% , 65.0% , and 70.0% success rates respectively — the highest among compared methods. The average success rate jumps from π₀.₅'s 38.3% to 73.3% (+35 pp, nearly double). Compared to LaWAM (43.3%) and ImageWAM (40.0%), WAVE-0 leads by 30.0 and 33.3 pp.

Data Efficiency Comparison

WAVE-0 requires no extra large-scale video pre-training. LaWAM uses ~4,500 hours of robot and human egocentric video; VLA-JEPA uses SSv2 and other video plus robot data. WAVE-0, Fast-WAM, JEPA-WAM, and ImageWAM all avoid independent video pre-training, but differ in how they learn world models and fuse them with policy.

Training sample count (including repeated sampling): WAVE-0 processes 33.28M samples total (Stage 1 shared pre-training on three tasks' demonstrations plus per-task policy training), versus LaWAM's 345.60M — roughly 1/10 the data — while achieving higher average real-robot success (73.3% vs. 43.3%).

Ablation Studies: How Future Prediction Contributes

On RoboTwin 2.0, two integration modes are compared: (a) using the world model only for auxiliary supervision during training, and (b) WAVE-0's future-conditioned approach where predicted futures directly participate in action generation at inference. Mode (a) lifts π₀.₅ from 80.38% to 89.51%; mode (b) further reaches 91.47%, showing that world model predictions add value at inference time, not just during training.

Accurate prediction alone is insufficient if the model ignores action differences. Counterfactual Ranking supervision forces the model to distinguish outcomes of different actions. Adding it raises LIBERO-Plus from 85.4% to 88.0% and LIBERO-Pro from 58.4% to 61.4%.

VLM-generated latent action plans must be understandable to the world model. WAVE-0 constrains plans with both future prediction loss and action representation alignment. Using only action generation loss yields 84.6% (Plus) and 58.0% (Pro); adding future prediction loss improves to 85.3% and 59.7%; adding both reaches 88.0% and 61.4%. This shows that aligning plans to true action representations while constraining their predicted consequences further improves policy performance.

Resources

Paper: https://arxiv.org/abs/2610.18568 Project page: https://hmwang2002.github.io/release/wave0/ Code:

https://github.com/AutoLab-SAI-SJTU/WAVE-0