OSWAM: An Effective World Action Model with One-Step Action Generation

AI summary

Overview

• Research area: Computer Vision / robotics — specifically World Action Models (WAMs) for autonomous driving and embodied (robotic) control.
• Technical level: Advanced. The paper assumes familiarity with diffusion models, flow matching, Diffusion Transformers (DiTs), and vision-language-action policies.
• Scope: This paper proposes OSWAM, a direct-policy world action model that replaces multi-step action denoising with a single step by initializing the action generator from learned "imagination tokens" rather than random Gaussian noise, and evaluates it on NAVSIM (driving) and RoboTwin 2.0 (bimanual manipulation) plus a real robot arm.

What This Paper Is About

World Action Models predict future visual states alongside actions, but that "imagine-then-execute" loop is far too slow for real-time control. Direct-policy WAMs avoid generating future video at test time, yet they still need tens of action denoising steps, and the paper reports that simply cutting to one step degrades performance by 20% on NAVSIM and 40% on RoboTwin. The goal is to make high-quality action generation possible in a single denoising step so these models become deployable in latency-critical settings.

Key Contributions

1. One-step action generation via imagination tokens. OSWAM learns two compact groups of tokens inside the Video DiT — Action Imagination Tokens (AIT) and Scene Imagination Tokens (SIT) — that are produced in a single forward pass and serve as a structured, physically grounded starting point for the Action DiT instead of unstructured Gaussian noise.
2. Privileged-information distillation with future conditioning. A frozen teacher Video DiT attends to video tokens over the full future horizon during training; its AIT features are distilled into the student's AIT at multiple Video DiT depths. At inference the teacher is discarded, so the model keeps a first-frame-only interface with no future-dependent features.
3. Role-aware 3D RoPE. Video tokens keep their original spatiotemporal coordinates, AIT are placed at spatial coordinates outside the video grid, and SIT are anchored at the grid center while sharing temporal coordinates with AIT — explicitly disentangling the three token roles in one positional system.
4. Efficiency and generality. The design is validated on both autonomous driving (NAVSIM) and embodied control (RoboTwin 2.0) plus a real PiPER arm, delivering roughly 3× speedups over direct-policy WAM baselines and compiled latency as low as 40 ms on an NVIDIA RTX 5090 GPU.

Main Findings

• NAVSIM parity at one step: OSWAM reaches a PDMS of 90.3 at 136 ms latency, matching SimWAM's 90.3 PDMS while cutting action denoising from 10 steps to 1 — a 3.01× speedup. Sub-scores: NC 98.6, DAC 98.0, EP 84.3, TTC 95.6, C 100.0.
• Beats Metis: 1.2 higher PDMS than Metis (89.1) with a 3.24× speedup (Metis: 440 ms).
• Close to DriveVA, much faster: DriveVA scores 90.5 PDMS at 593 ms (latency reproduced by the authors); OSWAM is only 0.2 PDMS lower with a 4.36× speedup.
• RoboTwin 2.0: OSWAM averages 90.00% success (Clean 90.04%, Randomized 89.96%) at 142 ms, versus Fast-WAM's 91.83% average at 397 ms — a 2.80× speedup for a 1.83% drop. Against the one-step Flash-WAM variant (81.41% average), OSWAM gains 8.59% with a 4.29× speedup. Light-WAM is faster (107 ms) but scores only 76.40% average.
• Latency breakdown: On RoboTwin 2.0 (RTX 5090, cached language context), action GPU time falls from 67.61 ms (Fast-WAM) to 6.78 ms, while imagination tokens add 5.48 ms to visual GPU time (18.45 → 23.93 ms). Overall latency drops from 323.67 ms to 95.35 ms (3.39×) with peak memory nearly unchanged (12.70 → 12.75 GB). On NAVSIM the analogous speedup over SimWAM is 3.32×, with overall latency at 99.70 ms.
• Compiled latency: With torch.compile and CUDA Graph, OSWAM reaches 40.81 ms on RoboTwin 2.0 and 52.64 ms on NAVSIM, supporting rates above 10 Hz.
• Ablations show AIT is the dominant ingredient: Starting from one-step Fast-WAM (NAVSIM PDMS 71.04, ADE 1.049; RoboTwin 57.38% Clean / 55.80% Rand), adding AIT lifts NAVSIM to 89.81 and RoboTwin to 87.42% / 88.16%. Distillation adds a further small gain (89.96; 89.42% / 89.62%) at no latency cost, and SIT gives the best result (90.28; 90.04% / 89.96%).
• Distillation strength matters: λ_dist = 0.01 gives PDMS 89.71 (below the 89.81 no-distillation baseline), while λ_dist = 10 gives the best 90.28 PDMS and 0.587 ADE.
• Re-noising beats rectified-flow refinement: All re-noising variants beat rectified-flow refinement (89.97 PDMS / 0.609 ADE); τ_r = 500 is best (90.28 / 0.587), with 400 giving 90.17 and 600 giving 90.08.
• Teacher robustness: Varying the teacher video timestep from 0 to 800 changes PDMS only between 90.16 and 90.28; even at timestep 800, where future video inputs are pure noise, PDMS is 90.16 — close to the clean-future setting (90.20).
• Real-world manipulation: On a PiPER arm with 50 demonstrations per task and 20 trials each, OSWAM succeeds 85% (bread into bowl) and 80% (eraser into drawer) versus 80% and 75% for Fast-WAM, and reduces video-estimated completion time by 34.2% and 21.8% respectively (bread 24.6 s → 16.2 s; eraser 17.6 s → 13.8 s).

Methodology in Plain English

The starting point is a "direct-policy" world action model that uses a large pretrained video generator (Wan2.2-5B) as its backbone, but never actually generates future video at test time — video prediction is only a training signal. Actions are produced by a smaller Action DiT whose weights are derived from the video model and which shares attention with it in a Mixture-of-Transformer layout.

The authors' diagnosis is that one-step action generation fails not because one step is inherently insufficient, but because starting from pure random noise gives the model no structure to work with. So they add two small sets of learnable tokens inside the video model, adapted in with LoRA (rank 8, α = 4):

• Action Imagination Tokens (AIT) are created by encoding a sample of action noise at a fixed timestep. Inside the video model they look at the first-frame visual tokens, at the scene tokens, and at each other, absorbing visual evidence and task conditioning. A small auxiliary decoder turns them into a rough action guess.
• Scene Imagination Tokens (SIT) act as extra slots that gather scene context for the AIT to consult, placed at the center of the video token grid.

That rough action guess is deliberately treated as a fixed starting point: the model adds back a controlled amount of noise (τ_r = 500 of t_max = 1000, i.e. σ_r = 0.5), and the Action DiT then performs exactly one denoising update to turn it into the final action chunk. Training targets are scaled by 1/σ_r so that a single Euler step lands on the ground-truth action.

To give the tokens knowledge of what happens next without ever showing the model the future at test time, the authors train a separate teacher model that is allowed to see the whole future clip (its future frames noised at a fixed timestep, 200 by default). The student's AIT features are pulled toward the teacher's at four depths (blocks 8, 16, 24, 30) with weight λ_dist = 10, starting after 500 steps and ramping in over the next 500. At inference the teacher is gone, and only the current frame is needed.

Training combines four losses: video flow matching, the AIT action-velocity loss, the distillation loss, and the one-step action loss. Optimization uses AdamW at 1e-4 with 0.01 weight decay and cosine decay with 5% warmup — 40,000 steps at batch 256 on NAVSIM, 29,335 steps at batch 1024 on RoboTwin.

Why This Matters

The paper targets the gap between "a world model that scores well on a benchmark" and "a world model that can run on a vehicle or a robot at control frequency." Its central claim — that the missing ingredient for one-step action generation is a structured initialization rather than a better distillation recipe — is a useful reframing for anyone working on few-step generative policies.

Real-world applications:

• Onboard autonomous driving, where the reported 52.64 ms compiled NAVSIM latency and >10 Hz inference rate fit within control loops that multi-step WAMs at 330–410 ms cannot.
• High-frequency closed-loop robotic manipulation, where a 95.35 ms overall latency (40.81 ms compiled) supports fast reactive correction during grasping and placement.
• Bimanual and dexterous robot platforms operating across many task variants — the RoboTwin 2.0 setting covers more than 50 tasks with heavy scene randomization.
• Cost-constrained deployment, since the model keeps nearly unchanged peak memory (12.75 GB vs 12.70 GB) while cutting action-generation cost by roughly an order of magnitude.

Industry relevance: The speedups are reported on commodity GPUs (H20 for the main tables, RTX 5090 for the efficiency table), and the approach preserves the first-frame-only inference interface of prior direct-policy models — meaning it is a drop-in acceleration rather than a change in the sensor or rollout requirements of a deployed system.

Future Directions

• Scaling and task diversity. Evaluation covers NAVSIM (1,192 training / 136 test scenarios) and RoboTwin 2.0 plus two real-world tabletop tasks with 50 demonstrations each; whether the one-step prior holds on broader real-world task distributions is untested here.
• Pushing latency lower. The paper reports compiled latencies of 40.81 ms and 52.64 ms but does not explore combining the structured initialization with step-distillation or caching techniques that prior work used for further compression.
• What the prior actually encodes. The teacher-timestep ablation shows nearly flat performance (PDMS 90.16–90.28) even when future frames are pure noise, which raises an open question about how much genuinely future-derived information the distilled tokens carry versus general scene structure.
• Remaining accuracy gap on manipulation. OSWAM's 90.00% average on RoboTwin trails Fast-WAM's 91.83%; closing that gap while keeping one denoising step is an explicit open problem the results leave on the table.

Target Audience

Researchers and engineers working on generative policies, world models, and vision-language-action systems for robotics and autonomous driving; practitioners who need diffusion-based planners to hit real-time control budgets; and readers already familiar with diffusion/flow-matching transformers who want a concrete recipe for collapsing multi-step action denoising into one step without retraining a backbone from scratch.

Authors’ abstract

Direct-policy World Action Models (WAM) have demonstrated the benefits of decoupling action generation from explicit future visual prediction for efficient inference in autonomous driving and embodied control. However, they still require tens of denoising steps to achieve competitive action quality, incurring substantial latency and hindering real-time deployment. Simply reducing the number of denoising steps will collapse the nonlinear multi-step denoising trajectory into coarse updates, substantially degrading action quality. To reduce the denoising steps while enabling real-time action generation, we present an effective world action model with one-step action generation, namely OSWAM, for autonomous driving and embodied control. We argue that it is the absence of structured action priors in random noise initialization that forces direct-policy WAMs to recover coherent action sequences through iterative denoising. Hence, we construct imagination tokens to initialize the diffusion process, providing physical priors and action-conditioned representations. We further introduce privileged information distillation with future conditioning to enhance future-aware action generation. As a result, OSWAM can efficiently and effectively generate high quality actions in just one diffusion step. Our experiments demonstrate that OSWAM achieves comparable or even better results in both autonomous driving and embodied control than direct-policy WAM and Joint-WAM that require tens of steps. It obtains approximately $3\times$ speedup over direct-policy WAM baselines and supports real-time compiled inference with latency as low as 40ms on an NVIDIA RTX 5090 GPU. The project page is available at https://gwenzhang.github.io/OSWAM/