AnyStep-WAM: Budget-Aligned Distillation and Adaptive Inference for World Action Models

AI summary

Overview

Research area: Robotics — efficient inference for world-action models (WAMs) that combine predictive visual modeling with action generation for manipulation.

Technical level: Advanced (assumes familiarity with diffusion/flow matching, denoising trajectories, flow maps, and low-rank adaptation).

Scope: The paper introduces AnyStep-WAM, a framework that trains a single WAM to predict denoising transitions at multiple budgets and then adaptively picks the smallest budget per action chunk, validated on three WAM backbones in simulation (RoboTwin 2.0) and on six real-world manipulation tasks.

What This Paper Is About

World-action models generate robot action chunks by iteratively denoising through large diffusion transformers, and they typically use a fixed number of denoising steps for every action chunk. The authors argue that this wastes computation on easy motions while under-refining contact-sensitive ones, so the goal is a single model that can generate accurately at any step budget and a lightweight scheduler that chooses the cheapest budget sufficient for the current scene.

Key Contributions

1. Budget-aligned WAM distillation: A method that samples intervals from candidate inference schedules and supervises each finite transition with the exact displacement along a frozen teacher trajectory, avoiding derivative- or finite-difference-based targets that depend on the evolving student. It also adds corrective continuation training from one-step student preview states, and uses shared LoRA adapters so one adapted WAM supports all candidate budgets.

2. Risk-benefit adaptive inference: A formulation of budget selection as a decision problem that separately predicts teacher-trajectory difficulty (a budget-independent risk proxy) and budget-specific student–teacher fidelity, then picks the smallest budget whose predicted fidelity meets a difficulty-dependent threshold.

3. Efficient single-preview execution: A one-step preview both informs budget selection and is reused to initialize the selected denoising trajectory, so the method requires exactly K* denoising evaluations with no online teacher access and no candidate rollouts.

4. Experimental validation: Evaluations on Motus, FastWAM, and LingBotVA using RoboTwin 2.0 (50 bimanual tasks, clean and randomized settings) and six real-world tasks on a Unitree G1D dual-arm robot, showing large reductions in denoising steps and latency while maintaining comparable or higher task success.

Main Findings

• Step reductions on RoboTwin 2.0: Adaptive inference cut mean denoising steps by 60.2% on Motus, 49.8% on FastWAM, and 85.28% on LingBotVA, with the paper also reporting 92.64% for LingBotVA's action branch. Mean selected budgets were 3.98, 5.03, and 3.68 steps respectively (defaults: 10 steps for Motus and FastWAM; 25 video steps and 50 action steps for LingBotVA).

• Success preserved: Average success rate stayed within 0.24 percentage points of the corresponding full-budget baselines. On Motus, adaptive results were 88.12% clean / 87.80% randomized (87.96% average), versus 88.66 / 87.02 / 87.84 for Base. On FastWAM: 92.12 / 91.06 / 91.59 versus 91.88 / 91.78 / 91.83. On LingBotVA: 92.74 / 91.70 / 92.22 versus 92.93 / 91.55 / 92.24.

• Latency on RoboTwin (single NVIDIA A100): Per-inference-call latency fell from 1.932 s to 0.859 s (Motus), 0.496 s to 0.295 s (FastWAM), and 9.732 s to 3.247 s (LingBotVA). Budget-selection overhead was 2.881 ms, 2.4996 ms, and 2.9772 ms per prediction.

• One-step budget gains: Under a matched one-step denoising budget, AnyStep improved average success over Base one-step by 7.07, 12.08, and 8.94 percentage points on Motus, FastWAM, and LingBotVA, and beat Flash-WAM by 3.19, 7.58, and 1.30 percentage points respectively. AnyStep one-step average success rates were 82.12% (Motus), 89.28% (FastWAM), and 82.71% (LingBotVA).

• Comparison to training-free gating: An AnyStep+Gate variant using percentile gating on action-coordination and visual-temporal scores reached 87.26 / 90.16 / 90.34 average success with 5.09 / 6.66 / 4.81 mean steps. The learned scheduler improved average success over gating by 0.70, 1.43, and 1.88 percentage points while reducing average denoising steps by 21.8%, 24.5%, and 23.5% on Motus, FastWAM, and LingBotVA.

• Training stability: Finite-difference supervision produced wide per-sample target RMS distributions and frequent gradient spikes, with 45.9% of steps exceeding the clipping threshold of 1.0; frozen-teacher interval targets remained concentrated with substantially more stable gradients. AnyStep was best at 1, 2, and 10 steps in the fixed-budget comparison.

• Real-world results (six tasks, Unitree G1D with Unitree Dex1-1 grippers): Average success rose to 69.17% (AnyStep-Motus, from 65.83%), 71.67% (AnyStep-FastWAM, from 70.00%), and 70.00% (AnyStep-LingBotVA, from 68.33%), while mean denoising steps decreased by 59.4%, 63.9%, and 85.84% (mean selected steps 4.06, 3.61, 3.54). Latency on a single RTX 5090 dropped from 1.868 s to 1.117 s (Motus, 1.67x), 0.280 s to 0.122 s (FastWAM, 2.30x), and 4.654 s to 0.758 s (LingBotVA, 6.14x).

• Loss ablations (Motus): Removing interval supervision (w/o L_int) lowered one-step average success by 1.40 percentage points, dropped two-step success from 86.21% to 78.94%, and cut adaptive success from 87.96% to 78.09% while raising the average budget from 3.98 to 5.27 steps. Removing endpoint supervision (w/o L_end) lowered one-step success from 82.12% to 75.53% and raised the adaptive budget to 4.74 steps. Removing continuation supervision (w/o L_rec) reduced adaptive success from 87.96% to 86.96%.

• Scheduler generalization and threshold sensitivity: A scheduler trained on clean-only data reached 87.56% success with 4.05 steps under randomized conditions, close to Full (Ours) at 87.80% with 3.94 steps. Lowering all thresholds by five percentile levels gave 87.60% at 3.08 steps; raising them gave 87.93% at 5.85 steps — both within 0.36 percentage points of Full.

• Risk correlates with contact: Contact-intensive stages tended to show higher predicted risk and were allocated larger budgets; in Bowl Pouring, more steps were assigned to grasping and pouring.

Methodology in Plain English

The authors start from a pretrained WAM whose denoising network is reformulated to predict finite-time transitions — a flow map conditioned on both the start and end time of an interval — rather than instantaneous velocities. Instead of deriving training targets from the student's own changing predictions, they run the frozen teacher model along a fine-grained reference trajectory and compute the exact average velocity over pairs of teacher states. Because both endpoints come from the frozen teacher, the targets are stable and independent of student parameters.

Crucially, the intervals used for training are drawn from the actual schedules the model will use at inference: they sample a candidate budget K and then a specific interval within that budget's schedule. Alongside this interval loss, the objective includes a local flow-matching term (to preserve instantaneous velocity prediction), an endpoint term with dedicated (t, r) = (1, 0) supervision for the one-step preview, and a recovery term that restarts the frozen teacher from a detached first recycled student state so the model learns to continue correctly after preview reuse. Adaptation uses LoRA in selected attention, feed-forward, and time-conditioning projections, with all pretrained weights frozen and adapters shared across budgets.

For scheduling, the paper defines risk from teacher velocity variation — average directional disagreement (cosine) and relative velocity change (RMS-based) along the teacher trajectory — converted to percentile ranks on the training set. Fidelity is defined as budget-specific agreement between the executed preview-recycled student path and the corresponding teacher reference, combined with exponentially decaying interval weights (2^-i normalized) that emphasize earlier intervals. A lightweight predictor takes features from a single one-step preview, plus contextual features, and outputs sigmoid estimates of risk and per-budget fidelity, trained with Smooth L1 losses against offline targets.

At inference, the preview produces a velocity, a predicted clean sample, and scheduler features. Predicted risk interpolates between easy and hard thresholds per modality, and the system selects the smallest budget whose predicted fidelity clears the threshold for all active modalities (falling back to the largest summed benefit if none qualifies). If K* = 1, the preview output is returned directly; otherwise the predicted clean sample is re-noised to the first intermediate time and K* − 1 further interval updates are run, for exactly K* denoising evaluations total.

Training used four NVIDIA H100 GPUs for 8,000 steps with a global batch size of 32. RoboTwin evaluation used a single A100 with 100 trials per task per method; real-world evaluation used 50 training episodes per task and 20 trials per task on average on an RTX 5090. Default thresholds used the 10th/20th percentiles of offline training-set video similarity and the 20th/80th percentiles of action similarity.

Why This Matters

The work shows that the fixed denoising budget assumed by most WAM pipelines is unnecessary: a single adapted model can serve many budgets, and choosing among them per scene recovers baseline success at a fraction of the computation. This decouples the accuracy-versus-latency trade-off that currently constrains deployment of diffusion-based robot policies, and it does so without online teacher access, candidate rollouts, or extra denoising passes — the scheduler adds only a few milliseconds.

Impact on research: the paper offers an alternative to derivative-based flow-map supervision, arguing and empirically demonstrating that frozen-teacher finite-interval targets avoid the noisy targets and gradient spikes seen with finite-difference supervision (45.9% of steps exceeded the clipping threshold of 1.0 in the finite-difference setting). It also frames adaptive inference for WAMs as a two-quantity problem — teacher-side difficulty and student-side budget fidelity — rather than a single uncertainty signal, and reports that no prior work specifically addressed adaptive budget selection for flow-map-trained WAMs.

Real-world applications:

• Warehouse and logistics manipulation, where pick-and-place and stacking tasks could run on faster inference budgets during simple reaches and spend more compute on contact-rich grasps.
• Household service robots, including pouring and table clearing, where the paper's real-world tasks (Bowl Pouring, Clean Table, Put Block, Put Cup, Stack Blocks, Stack Bowls) are directly representative.
• Manufacturing and assembly, where contact-sensitive steps can be allocated stricter fidelity thresholds while free-space motion is generated cheaply.
• Onboard/edge robotics, since the reported 1.67x to 6.14x per-call speedups on a single RTX 5090 target local deployment rather than datacenter inference.

Industry relevance: the method is architecture-agnostic across three different WAM backbones (joint video-action, serial, and action-only denoising) and is implemented with LoRA adapters, which makes it compatible with existing pretrained checkpoints and existing fine-tuning workflows rather than requiring models to be trained from scratch.

Future Directions

• Broadening backbone coverage: Validation covers Motus, FastWAM, and LingBotVA; extending budget-aligned distillation to other WAM formulations — including those that generate video and actions in different orders or use visual prediction only as a training auxiliary — remains open, though Appendix A argues the framework accommodates these distinctions.

• Improving and characterizing fidelity estimation: The paper states that S_{K,m} "remains a proxy rather than a direct measure of terminal action error or task success," so better-calibrated benefit estimates that track task outcomes more directly are a natural next step.

• Threshold selection and automation: The easy/hard thresholds were fixed at specific percentiles (10th/20th for video similarity, 20th/80th for action similarity), with sensitivity tested by shifting all levels by five. Automatically adapting these thresholds per deployment or per task is untested.

• Generalization beyond the evaluated settings: The scheduler was tested for clean-only training generalization, but behavior across longer horizons, different robot embodiments, larger task suites, and non-manipulation domains is not reported.

Target Audience

Researchers and engineers working on efficient generative models for robotics — particularly those building or deploying world-action models, diffusion policies, or flow-matching-based visuomotor systems. It is also relevant to practitioners interested in distillation, adaptive computation, and low-rank adaptation, and to roboticists needing lower-latency inference on local hardware. Readers without a background in diffusion or flow matching will find the preliminaries and appendices necessary to follow the method.

Authors’ abstract

World-action models (WAMs) couple predictive visual modeling with action generation, typically relying on iterative denoising with a fixed denoising steps. However, manipulation tasks contain actions chunks with varying sensitivity to generation errors: critical actions require precision, while less sensitive actions allow faster generation with fewer denoising steps. Here we introduce AnyStep World Action Model, a general framework for tunable-budget prediction and scene-dependent computation allocation. Our budget-aligned teacher-trajectory distillation trains interval-conditioned flow maps using explicit frozen-teacher transitions and shared low-rank adapters, supporting action generation from one-step prediction to multi-step refinement. Building on this capability, a lightweight risk-benefit scheduler predicts teacher-curvature-based difficulty and budget-specific student-teacher fidelity from a single one-step preview, selecting the smallest budget predicted to satisfy risk-adaptive fidelity requirements. We evaluate our framework on three widely used WAMs Motus, FastWAM, and LingBotVA using RoboTwin 2.0. Our method reduces average denoising steps by 60.2%, 49.8%, and 85.28%, respectively, while maintaining baseline task success rates. In particular, our AnyStep training substantially improves model performance under a one-step denoising budget, increasing task success rates by 7.07%, 12.08%, and 8.94% on Motus, FastWAM, and LingBotVA, respectively. Experiments on six real-world manipulation tasks further validate its effectiveness.