AnyStep-WAM: Budget-Aligned Distillation and Adaptive Inference for World Action Models
Rui Wang, Xiangyu Wang, Donglin Yang, Yibo Li, Canyang Chen, Zhongrui Wang, Xiaojuan Qi
cs.RO
2026-09-28
Teacher-trajectory distillation plus preview scheduling. RoboTwin: three WAMs cut steps 60.2%/49.8%/85.28% with success held; one-step SR +7.07/+12.08/+8.94 pp.
World-action models (WAMs) pair future-visual prediction with action generation. At inference they still walk a large diffusion transformer through a fixed denoising schedule. Motus and FastWAM use 10 steps; LingBotVA is serial, 25 steps on the video DiT and 50 on the action DiT. A simple reach and a contact-heavy grasp pay the same compute. Teacher trajectories also bend and fan out more in contact, so coarse integration error grows there, while near-straight reaches barely change if you add extra steps.
Few-step generation already exists. Flash-WAM ports consistency distillation onto WAMs for one-step prediction. AnyFlow trains variable-step video generators with student-side finite differences. Porting AnyFlow onto WAMs makes the target move with the student: 45.9% of training steps exceed a gradient-norm clip of 1.0. One-step success collapses to 13.75% on Motus, 32.01% on FastWAM, and 49.96% on LingBotVA. Manipulation fails on small contact errors, so a bad long jump is expensive.
Two gaps remain. One model has to be accurate from a single jump through multi-step refinement. Step choice for the current action chunk has to track both how curved the teacher trajectory is and how faithfully the student matches that teacher at each budget.
AnyStep WAM is two pieces: a shared model that can run any candidate budget, then a scheduler that picks the smallest budget that looks sufficient.
Distillation rewrites the pretrained denoising branch as an interval-conditioned flow map. Given a noisy state and endpoints (t, r), the network predicts the average velocity over that interval; the displacement is that velocity times the time gap. Targets come from a frozen teacher integrated on a fine grid, not from the student's own derivatives. The teacher's mean velocity on an interval is the endpoint displacement divided by the time gap, so the target does not drift as student weights change.
Training samples a budget K from {1,...,10}, then one interval on that schedule, so the student sees the spans it will actually run. Four regression terms share a dimension-normalized velocity loss:
Weights are λFM=0.20, λend=0.50, λint=0.30, λrec=0.20. Backbone weights stay frozen. Shared LoRA adapters (rank 128) are the only trainable piece: 286.39M, 104.40M, and 232.91M parameters on Motus, FastWAM, and LingBotVA. One adapter set covers every budget, so switching K does not switch weights. Lrec matches deployment: inference recycles the one-step preview into a multi-step schedule, and skipping student-generated states would train a different path. FastWAM does not denoise video at inference, so only action losses remain.
The scheduler treats step selection as a risk-benefit call. Risk is a percentile of teacher velocity turning and relative speed change; it scores denoising difficulty, not failure chance. Benefit is student-teacher velocity agreement along the same preview-recycle path used at deployment, a proxy rather than terminal action error. Both heads read summaries from a single one-step preview: a four-layer Transformer with one risk query and a set of budget queries, sigmoid outputs, Smooth L1 losses.
Inference is short. One evaluation at (t,r)=(1,0) yields the preview. Predicted risk interpolates between easy and hard fidelity thresholds; the smallest K that clears every active modality is kept. If K=1 the preview is the output. Otherwise it is re-noised with the original noise to the first intermediate time, then the remaining K-1 updates run. Total network calls equal K. The teacher is not queried online, and candidate rollouts are not run. Scheduler cost is 2.5–3.0 ms. Contact-heavy stages such as grasp and pour tend to draw larger budgets; simple reaches often land at 1–2 steps.
RoboTwin 2.0, 50 bimanual tasks, 100 trials per task in clean and randomized settings.
| Backbone | Full-budget SR / steps | AnyStep adaptive SR / steps | One-step SR: Base → AnyStep |
| Motus | 87.84% / 10 | 87.96% / 3.98 (-60.2%) | 75.05% → 82.12% (+7.07) |
| FastWAM | 91.83% / 10 | 91.59% / 5.03 (-49.8%) | 77.20% → 89.28% (+12.08) |
| LingBotVA | 92.24% / 25+50 | 92.22% / 3.68 (-85.28% video / -92.64% action) | 73.77% → 82.71% (+8.94) |
Adaptive mean success stays within 0.24 points of the full-budget baselines. Per-call latency on one A100 falls from 1.932s to 0.859s (Motus), 0.496s to 0.295s (FastWAM), and 9.732s to 3.247s (LingBotVA).
Under a matched one-step budget, AnyStep beats Flash-WAM by 3.19, 7.58, and 1.30 points. A training-free percentile gate on the same distilled models, AnyStep+Gate, uses more steps and wins less often: the learned scheduler is +0.70/+1.43/+1.88 SR and -21.8%/-24.5%/-23.5% steps.
Real robot: Unitree G1D dual-arm, six tasks, about 20 trials each.
| Backbone | Base SR | AnyStep SR | Steps | RTX 5090 latency |
| Motus | 65.83% | 69.17% | 10→4.06 (-59.4%) | 1.868s→1.117s (1.67×) |
| FastWAM | 70.00% | 71.67% | 10→3.61 (-63.9%) | 0.280s→0.122s (2.30×) |
| LingBotVA | 68.33% | 70.00% | 25/50→3.54 (-85.84%) | 4.654s→0.758s (6.14×) |
Motus ablations: drop Lint and adaptive SR falls from 87.96% to 78.09%; drop Lend and one-step SR falls from 82.12% to 75.53%; drop Lrec and adaptive SR falls by 1.00 point. Shifting all threshold percentiles by 5 keeps SR in 87.60%–87.93% while mean steps move from 3.08 to 5.85. A scheduler trained on clean demos only still reaches 87.56% SR / 4.05 steps under randomization, near the full 87.80% / 3.94.
This is LoRA distillation plus a tiny scheduler on frozen public WAMs. The world model itself does not change. Deployment saves DiT denoising evaluations; the scheduler itself is milliseconds. Even the fixed one-step setting beats both the original one-step baseline and Flash-WAM, which is the practical evidence that integral teacher displacements hold up better than consistency or student finite differences when the jump is large.
If a stack already runs Motus, FastWAM, or LingBotVA and is waiting on a fixed denoising schedule, this is closer to a drop-in budget layer than a research-only recipe. The real-robot numbers are directional: six tasks, about 20 trials each.
There is no standalone Limitations section. The authors state that risk ρ is relative denoising difficulty, not failure probability, and that fidelity S is student-teacher velocity agreement, not terminal action error or task success. Thresholds are the 10th/20th percentiles of training-set video similarity and the 20th/80th of action similarity; a new domain needs retuning.
Training still rolls a frozen teacher on a fine grid, then sweeps preview-recycle paths offline to label the scheduler. Inference drops the teacher; distillation does not. Candidate budgets are locked to 1–10. LingBotVA’s 85% step cut partly reflects a 25/50 baseline, not the same 10-step starting point as Motus. Real-world data are 50 demos per task and about 20 trials; Stack Blocks sits at 50–55% base success, so a couple of extra wins move the mean. The adaptive control is AnyStep+Gate, a homemade training-free gate, not a published adaptive policy. FastWAM adaptive SR is 91.59% against a 91.83% baseline: held success on that backbone is a small loss.