AutoLab-SAI-SJTU / WAVE-0

WAVE-0: Bridging VLA Policies and WAMs with Lightweight Latent Future Prediction 
## 📚 Introduction
We introduce **WAVE-0 (World-model-Augmented VLA Execution)**, a framework that bridges **vision–language–action (VLA) policies** and **world action models (WAMs)** through action-grounded latent future prediction. Built on π0.5, WAVE-0 unifies latent planning, future prediction, and action generation, allowing anticipated action consequences to directly guide robotic manipulation.
- **🌊 Lightweight future-conditioned action generation** — A **257M-parameter latent world model** transforms VLM-generated latent action plans into predicted future visual representations that condition the action expert during both training and inference, without reconstructing future images.
- **🎯 Action-grounded two-stage training** — We first learn action-dependent latent dynamics directly from ground-truth robot actions with a counterfactual ranking objective, then integrate the frozen world model into a pretrained VLA to guide action generation with predicted future representations, without requiring additional large-scale video pretraining.
- **🤖 Strong performance and generalization** — WAVE-0 achieves **91.47%** success on **RoboTwin 2.0** and **98.7%** on **LIBERO**, with improved generalization under visual, spatial, and language distribution shifts on **LIBERO-Plus**, **LIBERO-Pro**, and **LIBERO-Para**. Across three real-world manipulation tasks, it improves the average success rate of π0.5 from **38.3% to 73.3%**.
## 🔥 News
- **[2026-10-14]** 👋 Upload paper and init project. Read
## 📝 Open-Source Plan
- [x] Paper
- [x] Model weights
- [x] Training scripts
## 📌 Quick Start
### ⚙️ Installation
WAVE-0 is built upon the π0.5 implementation in LeRobot v0.4.4.
 [code] 
### 🚀 Train
#### LIBERO
##### 1. Download the model and dataset
Download the π0.5 base model and the standard 40-task LIBERO dataset. Both stages use the mean/std statistics in the dataset's `meta/stats.json`; no quantile computation is needed. Run the following commands from the repository root.
 [code] 
##### 2. Train Stage 1
Train the latent world model for 20,000 steps:
 [code] 
##### 3. Train Stage 2
Initialize the policy from π0.5 and load the Stage-1 checkpoint to train for 15,000 steps:
 [code] 
Both scripts default to 8 GPUs with a per-GPU batch size of 64 (global batch 512). The action chunk length is 32, with 10 actions executed per prediction and a future prediction horizon of 9 frames. Override `NUM_GPUS` and `BATCH_SIZE` as needed. We use the Stage-2 checkpoint at 10,000 steps for the LIBERO results reported in our paper.
#### RoboTwin2.0
##### 1. Download the model and dataset
Before training, we need to download the π0.5 base model and the RoboTwin dataset.
 [code] 
##### 2. Prepare quantile statistics
Both stages use quantile normalization for actions and robot states, requiring the 1st and 99th percentiles (`q01` and `q99`). We provide precomputed statistics for the full RoboTwin