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

OSWAM learns compact imagination tokens from a video DiT as structured start points for a one-step action DiT denoising, achieving competitive planning and manipulation results with 1 action denoising step and substantially lower latency than multi-step world action models.

World action models that follow the imagine-then-execute paradigm have prohibitive inference latency, and few-step distillation reduces this latency at the cost of a substantial drop in control accuracy. Direct-policy WAMs avoid test-time video rollout and are more than 3x faster, but still need tens of action denoising steps, and performance degrades when steps are reduced, dropping by 20% on NAVSIM and 40% on RoboTwin in the one-step setting. The authors attribute this to the lack of action structure in random noise: multiple denoising steps can gradually impose temporal and physical consistency, whereas one-step generation must recover a coherent action sequence in a single update. In addition, direct-policy WAMs compress future dynamics into a first-frame-only representation, which may provide insufficient temporal context for future-aware control.

OSWAM keeps the direct-policy interface (video modeling is only a training signal) and augments the Video DiT input with two groups of noise-initialized latent tokens, Action Imagination Tokens (AIT) and Scene Imagination Tokens (SIT), adapted via LoRA. AIT attends to first-frame visual tokens, SIT, and other AIT, and an auxiliary decoder maps AIT representations to an action velocity producing a structured action prior that is detached and re-noised to an intermediate timestep. The Action DiT is initialized from this structured state and performs a single denoising step toward the ground-truth action chunk. Role-aware 3D RoPE assigns video tokens their original spatiotemporal coordinates, AIT spatial coordinates outside the video grid, and SIT the grid center, while AIT and SIT share temporal coordinates that convert action positions to the video latent timeline. Privileged-information distillation trains a frozen teacher Video DiT whose AIT attends to the full future clip and aligns OSWAM's AIT features with teacher features at selected Video DiT blocks; the teacher is discarded at inference, preserving the first-frame-only interface. The overall objective combines video flow matching, the AIT action-velocity loss, the distillation loss, and the one-step action loss.

• On NAVSIM, OSWAM matches SimWAM's PDMS while reducing action denoising from ten steps to one, outperforms Metis by 1.2 PDMS with a 3.24x speedup, and achieves a 4.36x speedup over DriveVA with only a 0.2 lower PDMS. §4.3 Main Results
• On RoboTwin 2.0, OSWAM reaches an average success rate with one action denoising step, reducing action denoising from ten steps versus Fast-WAM with only a small success decrease, and improving over the one-step Flash-WAM variant. §4.3 Main Results
• With cached language context, OSWAM reduces action GPU time versus Fast-WAM on RoboTwin 2.0 while imagination tokens add a small amount of visual GPU time, and compiled inference supports rates above 10 Hz. §4.3 Main Results
• Component ablations starting from a one-step Fast-WAM baseline show that AIT, then privileged distillation, then SIT each improve NAVSIM PDMS and RoboTwin success rates, with SIT giving the best performance on both benchmarks. §4.4 Ablation Studies
• Ablations on distillation weight, action initialization timestep, and teacher video timestep select the default settings used in OSWAM. §4.4 Ablation Studies
• In real-world PiPER manipulation, OSWAM achieves higher success rates than Fast-WAM on both tasks and reduces mean video-estimated completion time among successful episodes. §4.5 Real-World Evaluation
• On nuScenes open-loop planning without auxiliary supervision, OSWAM matches PWM's average L2 error while reducing the average collision rate. §A.2 More Results

• The paper states that its study focuses on learning action DiT structured initialization, and that the potential of SIT and privileged-information distillation to directly enhance the Action DiT remains unexplored.
• Real-world evaluation covers only two tabletop manipulation tasks with 50 training demonstrations each and 20 trials per task, so real-world generality beyond these tasks is left open by the experiments.
• Evaluation is confined to NAVSIM, NAVSIM-v2, nuScenes, and RoboTwin 2.0 simulation plus the two real-world tasks; the paper reports no other embodiments or datasets.

• Direct-policy WAMs are more than 3x faster than the imagine-then-execute paradigm but still require tens of denoising steps, which remains too slow for real-time applications such as autonomous driving. §1 Introduction
• Reducing denoising steps degrades performance, dropping by 20% on NAVSIM and 40% on RoboTwin in the one-step setting. §1 Introduction
• AIT and SIT are initialized from noise in the Video DiT and comprise far fewer tokens than the first-frame visual input, so they introduce only a small computational overhead during inference. §3.2 Overall Architecture
• AIT is constructed from an initial action noise sample at fixed timestep t_max, and its decoded prior is detached and re-noised with the same noise before one-step Action DiT refinement; the Action DiT target is (a_init - a_gt)/sigma_r. §3.3 Imagination Tokens
• Role-aware 3D RoPE assigns AIT spatial coordinates outside the video grid and SIT the grid center, while both share temporal coordinates t_n = n s_a / (s_v d_t). §3.3 Imagination Tokens
• The frozen teacher receives ground-truth future video latents corrupted at a fixed diffusion timestep and OSWAM's AIT features are aligned to teacher features at a set of Video DiT blocks; the teacher is discarded at inference. §3.4 Privileged-Information Distillation
• The training objective sums video flow matching, the AIT action-velocity loss, the privileged feature distillation loss, and the one-step action loss. §3.5 Training Objective
• On NAVSIM-v2, OSWAM achieves the highest EPDMS among the compared planners with only one action denoising step. §A.2 More Results