AnyStep-WAM: Budget-Aligned Distillation and Adaptive Inference for World Action Models
Abstract Overview
This paper introduces AnyStep-WAM, a framework for world action models that supports tunable denoising budgets and adaptive allocation of computation across action chunks. The method distills a pretrained teacher into a single student that can operate from one-step prediction to multi-step refinement using teacher-anchored interval supervision and shared LoRA adapters. It also adds a lightweight scheduler that uses a one-step preview to estimate scene difficulty and budget-specific student-teacher fidelity, then selects the smallest budget predicted to satisfy a risk-conditioned threshold. Experiments are reported on three world action model backbones in RoboTwin 2.0 and on six real-world manipulation tasks.
Novelty
The distinctive contribution is the combination of budget-aligned integral flow-map distillation with risk-conditioned adaptive inference for world action models. Rather than relying on derivative- or finite-difference-based supervision, the method trains finite-time transitions directly from frozen teacher trajectories and couples this with a single-preview scheduler that predicts both denoising difficulty and budget-dependent fidelity.
Results
On RoboTwin 2.0, the method reduces average denoising steps by 60.2% on Motus, 49.8% on FastWAM, and 85.28% on LingBotVA while keeping average success within 0.24 percentage points of the corresponding full-budget baselines. Under a one-step budget, it improves task success over the base models by 7.07, 12.08, and 8.94 percentage points on Motus, FastWAM, and LingBotVA, respectively. Real-world tests on six tasks also show reduced latency and mean denoising steps with comparable or higher average success rates.
Key Points
- AnyStep-WAM trains a single world action model to handle multiple denoising budgets using teacher-trajectory interval supervision and shared LoRA-based adaptation.
- A lightweight scheduler predicts risk and budget-specific fidelity from a one-step preview, enabling adaptive step selection without online teacher access or candidate rollouts.
- Empirically, the method cuts denoising steps and inference latency across simulation and real-robot tasks while preserving, and sometimes improving, task success.