PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Guowen Zhang, Liyi Chen, Xu Hu, Pengfei Wang, Jingfan Chen, Chenhang He, Lei Zhang

arXiv:2610.35850v1cs.CVcs.RO

TL;DR

Direct-policy WAMs avoid explicit future visual prediction at inference but still require many denoising steps, limiting latency-sensitive use. OSWAM initializes action diffusion with imagination tokens and distills future context into them during training. It achieves competitive results across autonomous driving and embodied control while reducing inference latency.

  • Problem

    Direct-policy WAMs require tens of denoising steps for competitive action quality, while reducing steps alone degrades performance and latency remains high for real-time applications.

    Source: §1 · pp. 1–2

  • Method

    OSWAM uses imagination tokens from the Video DiT as structured action-diffusion starting points and distills full-future teacher representations into tokens used by a first-frame-only student at inference.

    Source: §1 · pp. 2–3

  • Results

    OSWAM is competitive with multi-step WAMs in autonomous driving and embodied control, with a 3.39× overall speedup on RoboTwin 2.0 over Fast-WAM.

    Source: §1 · p. 3, §4.3 · p. 9

  • Takeaways & Limitations

    Compiled inference reaches 40.81ms on RoboTwin 2.0 and 52.64ms on NAVSIM, supporting inference rates above 10Hz.

    Source: §4.3 · p. 9

    The study focuses on structured initialization for the Action DiT; directly enhancing the Action DiT with scene-context conditioning and privileged supervision remains unexplored.

    Source: §5 · p. 11

Abstract

from arXiv · show

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/

1 Introduction

Direct-policy WAMs avoid test-time video prediction but still need tens of action-denoising steps, and reducing steps alone degrades action quality. OSWAM uses imagination tokens for structured action initialization and privileged-information distillation for future-aware control, achieving competitive results with one-step generation.

1 Introduction

  • Direct-policy WAMs need tens of denoising steps for competitive action quality, and one-step generation drops performance by 20% on NAVSIM and 40% on RoboTwin.The paper attributes the degradation to random-noise initialization lacking action structure, which iterative denoising can gradually impose.Source: §1 · p. 2
  • Direct-policy WAMs use a first-frame-only interface, but compressing future information into that representation may leave insufficient temporal context for effective control.Adding future features at test time incurs additional inference overhead.Source: §1 · p. 2
  • Imagination tokens learned from the Video DiT provide structured action starting points in one forward pass with negligible computation, replacing unstructured initialization for one-step decoding.The tokens interact with visual features to accumulate physically grounded, action-relevant context.Source: §1 · pp. 2–3
  • Privileged-information distillation transfers long-horizon context into imagination tokens: the teacher attends to the full future horizon, while the student uses only the first frame at test time.This preserves the direct-policy WAM interface without future-dependent inference features.Source: §1 · p. 3
  • OSWAM achieves competitive results with multi-step WAMs in autonomous driving and embodied control while generating actions in one denoising step.The approach uses imagination tokens as structured starting points and privileged-information distillation for future-aware control.Source: §1 · p. 3

2 Related Work

Prior world action models either incur test-time visual rollout or retain iterative action denoising, limiting inference efficiency. OSWAM addresses the remaining denoising bottleneck with one-step action generation while retaining video-generative priors.

2.1 World Action Models for Autonomous Driving

  • For autonomous driving, OSWAM retains video-generative priors without test-time video rollout and reduces action generation to one denoising step.BEV approaches are computationally efficient but do not exploit large-scale video generative priors; image-domain future rollout remains inefficient.Source: §2.1 · p. 3

2.2 World Action Models for Embodied Control

  • For embodied control, direct-policy WAMs avoid test-time visual rollout, but iterative action denoising remains a bottleneck for high-frequency closed-loop control.OSWAM preserves video-generative priors without test-time rollout and reduces generation to one denoising step.Source: §2.2 · p. 3

2.3 Effective World Action Models

  • OSWAM targets WAM latency by replacing multi-step action denoising with one step, while its Video DiT produces action and scene imagination tokens for structured initialization.Flash-WAM still synthesizes future video at inference; Fast-WAM and SimWAM retain iterative action denoising, while Light-WAM underperforms on challenging benchmarks such as RoboTwin.Source: §2.3 · p. 4, Figure 3 · p. 4

3 Method

OSWAM preserves the direct-policy interface while replacing iterative action denoising with a single update from structured imagination-token initialization. It also distills future-aware representations into tokens that use only first-frame visual context at inference.

3.1 Preliminary

  • OSWAM replaces direct-policy WAMs’ iterative action denoising with one update from a structured intermediate state, rather than pure noise.Direct-policy WAMs avoid future-video generation at inference, but high-quality actions typically require multiple velocity evaluations.Source: §3.1 · p. 5

3.2 Overall Architecture

  • The model uses a Video DiT and Action DiT in a shared-attention architecture, with future video used for training but not rolled out at inference.For driving, an ego encoder supplies ego state and command information; the architecture follows the direct-policy setting.Source: §3.2 · p. 5

3.3 Imagination Tokens

  • Role-aware 3D RoPE gives video tokens, AIT, and SIT distinct spatial anchors while aligning AIT and SIT on the action-position timeline.AIT uses spatial coordinates outside the video grid, while SIT is anchored at its center.Source: Figure 4 · p. 6, §3.3 · p. 6
  • AIT representations replace Gaussian action initialization with a structured prior, which the Action DiT refines in one denoising step.AIT attends to first-frame visual tokens, SIT, and other AIT; an auxiliary decoder predicts an action velocity to form the prior.Source: §3.3 · pp. 5–6

3.4 Privileged-Information Distillation

  • Privileged-information distillation transfers teacher representations learned with full-future video context into OSWAM’s AIT, which uses first-frame context at inference.The teacher’s AIT attends to the full video clip; OSWAM’s AIT attends to first-frame visual tokens, SIT, and other AIT.Source: §3.4 · pp. 6–7

3.5 Training Objective

  • The training objective combines video flow matching, AIT action and privileged-feature supervision, and one-step Action DiT prediction of demonstrated action chunks.The Action DiT is trained to predict action chunks from the AIT-derived initialization in one denoising step.Source: §3.5 · p. 7

4 Experiments

Experiments evaluate OSWAM on autonomous-driving and embodied-control benchmarks, comparing one-step action generation with multi-step baselines. Ablations examine its components and settings, while PiPER trials test real-world manipulation.

4.1 Datasets and Evaluation Metrics

  • The evaluation uses NAVSIM’s PDMS for driving and average success over 100 trials per task under clean and randomized conditions on RoboTwin 2.0.RoboTwin 2.0 covers more than 50 bimanual tasks; NAVSIM data include 1,192 training and 136 test scenarios.Source: §4.1 · p. 7

4.2 Implementation Details

  • The implementation uses one Action DiT denoising step at inference, with benchmark-specific observations, action horizons, and training schedules.Unless otherwise specified, latency is measured on one NVIDIA H20 GPU and an Intel Xeon Max 9468 CPU.Source: §4.2 · p. 8

4.3 Main Results

  • OSWAM matches SimWAM’s NAVSIM PDMS while using one action-denoising step, and remains competitive with multi-step WAMs on RoboTwin 2.0.It also outperforms Metis by 1.2 PDMS; on RoboTwin, it achieves a 90.00% average success rate.Source: §4.3 · pp. 8–9, Table 1 · p. 8, Table 2 · p. 9
  • Compiled inference reaches 40.81 ms on RoboTwin 2.0 and 52.64 ms on NAVSIM; OSWAM also reports 3.39× and 3.32× overall speedups, respectively.The 3.39× comparison is against Fast-WAM on RoboTwin; the 3.32× comparison is against SimWAM on NAVSIM.Source: §4.3 · p. 9, Table 3 · p. 9

4.4 Ablation Studies

  • Adding AIT raises one-step Fast-WAM’s NAVSIM PDMS from 71.04 to 89.81 and improves RoboTwin success; distillation and SIT provide further gains.Distillation raises NAVSIM PDMS to 89.96 without increasing inference latency, while SIT yields the best reported performance on both benchmarks.Source: §4.4 · pp. 9–10, Table 4 · p. 10
  • Distillation weight 10 yields 90.28 PDMS and 0.587 ADE, while action re-noising at τr=500 outperforms the tested alternatives.The teacher-timestep sweep changes PDMS only slightly, from 90.16 to 90.28, including when future video inputs are pure noise.Source: §4.4 · p. 10, Table 5 · p. 10, Table 6 · p. 10

4.5 Real-World Evaluation

  • On PiPER, OSWAM succeeds on 85% of bread placements and 80% of eraser placements, compared with Fast-WAM’s 80% and 75%.Among successful episodes, OSWAM’s video-estimated completion times are lower by 34.2% and 21.8%, respectively.Source: Table 8 · p. 10, §4.5 · p. 11, Figure 5 · p. 11

5 Conclusion

OSWAM combines structured action initialization, scene-context tokens, and privileged-information distillation for one-step action generation. Its potential to directly enhance the Action DiT with scene-context conditioning and privileged supervision remains unexplored.

5 Conclusion

  • OSWAM combines Action Imagination Tokens, Scene Imagination Tokens, and privileged-information distillation to support one-step action generation.The conclusion describes a single Video DiT pass followed by one Action DiT denoising step.Source: §5 · p. 11
  • Experiments on NAVSIM and RoboTwin 2.0 show competitive planning and manipulation performance with substantial latency reductions over direct-policy baselines.The paper characterizes OSWAM as a practical approach to efficient world-action modeling for latency-critical control.Source: §5 · p. 11
  • Whether SIT and privileged-information distillation can directly enhance the Action DiT remains unexplored.The study focuses on learning structured initialization for the Action DiT.Source: §5 · p. 11

Appendix

The appendix lists additional experiment results and qualitative visualizations.

Appendix

  • The appendix introduces supplementary materials.It identifies further results and visualizations in later sections.Source: Appendix · p. 16
  • Section A contains more experiment results.Source: Appendix · p. 16
  • Section B contains qualitative visualizations.Source: Appendix · p. 16

A.1 More Implementation Details

The nuScenes appendix describes the dataset, action representation, and training and distillation setup used for an additional autonomous-driving evaluation.

A.1 More Implementation Details

  • nuScenes contains 1,000 real-world driving scenes, split into 700 training, 150 validation, and 150 test scenes, with six cameras covering 360 degrees horizontally.The evaluation uses L2 Displacement Error and Collision Rate.Source: §A.1 · p. 16
  • Each nuScenes training sample uses the current frame and eight future frames, with eight action steps matching the four-second video horizon.Supervision and evaluation are limited to the first three seconds; actions are represented in the current LiDAR coordinate frame.Source: §A.1 · p. 16
  • Training runs for 10,000 steps, and distillation uses a frozen teacher at video timestep 200 with feature alignment after Video DiT blocks 8, 16, 24, and 30.The AIT diffusion timestep is 1,000 and the Action DiT re-noising timestep is 500.Source: §A.1 · p. 16

A.2 More Results

Across autonomous driving and embodied control, OSWAM delivers strong planning and task performance with one-step action generation, including results competitive with multi-step methods.

NAVSIM-v2

  • On NAVSIM-v2, OSWAM achieves the highest EPDMS among compared planners at 90.3 with one action denoising step, slightly surpassing SimWAM at 90.2 and outperforming Metis at 89.5.Source: §A.2 · p. 16
  • OSWAM's results under EPDMS are evaluated across driving direction, traffic-light and lane-keeping compliance, and comfort criteria.EPDMS extends the original comfort metric with History Comfort.Source: §A.2 · pp. 16–17

nuScenes

  • OSWAM matches PWM's average 0.41 m L2 error on nuScenes while reducing average collision rate from 0.04% to 0.03%.It also has the lowest L2 errors at 1 and 2 seconds and the lowest collision rates across evaluated horizons.Source: §A.2 · p. 17, Table A.2 · p. 17

RoboTwin per-task results

  • OSWAM matches or exceeds Light-WAM on 49 clean-setting tasks and 46 randomized-setting tasks out of 50 RoboTwin 2.0 manipulation tasks.It matches or exceeds Fast-WAM on 22 clean and 15 randomized tasks; the reported inference speedup over Fast-WAM is 2.80×.Source: §A.2 · p. 17

One-step action generation

  • On NAVSIM, OSWAM's one-step PDMS of 90.3 exceeds SimWAM, Metis, and DriveVA by 21.4, 3.3, and 77.1 points, respectively.On RoboTwin 2.0, it raises average success from 56.59% to 90.00% over the one-step Fast-WAM baseline, with gains in both Clean and Randomized settings.Source: §A.2 · p. 17, Table A.3 · p. 17, Table A.4 · p. 17

B Qualitative visualizations

Figure B.1 qualitatively compares one-step SimWAM and OSWAM trajectories on Navsim v1, highlighting differences in smoothness and driving-rule compliance.

Qualitative comparison

  • OSWAM produces smoother trajectories than SimWAM in difficult turns and interactions, while satisfying no-at-fault-collision and drivable-area criteria in one denoising step.SimWAM struggles in these examples, with drivable-area violations and at-fault collisions.Source: B Qualitative visualizations · p. 17, Figure B.1 · p. 19
  • The comparison focuses on challenging turns and interactions with surrounding agents, where SimWAM's trajectories can be unsmooth.Figure B.1 compares the two methods under one-step action generation on Navsim v1.Source: B Qualitative visualizations · p. 17, Figure B.1 · p. 19
  • The stated criteria for the OSWAM examples are no at-fault collision and drivable-area compliance.The passage does not provide quantitative results for this visual comparison.Source: B Qualitative visualizations · p. 17

Related papers

Loading 2610.35850v1…