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

Introduction

Physical intelligence [1, 4], including autonomous driving [2] and embodied control [3], holds the promise of freeing human workers from physical tasks that are dangerous, laborious, or tedious. A general-purpose physical agent should not only act from visual observations, but also anticipate the physical consequences of their actions. Existing Vision-Language-Action (VLA) models, including those standard models [5, 7, 6] and specialized end-to-end ones [9, 8, 10], typically predict actions directly without explicitly modeling the future evolution of the environment. This has motivated a shift toward World Action Models (WAMs) [3, 16, 13, 15] for physical AI, which build on pretrained video generative models [32, 33] to jointly predict future visual states and their corresponding actions.

However, WAMs follow the imagine-then-execute paradigm [12, 13] and incur prohibitive inference latency, rendering them impractical for latency-critical applications such as autonomous driving. Borrowing few-step techniques from image and video generation, some methods apply step distillation [11] to reduce this latency, but often at the expense of a substantial drop in control accuracy, as illustrated in Figure 1. To address this issue, efforts [3, 16] propose to switch from the imagine-then-execute paradigm to direct-policy WAMs. These methods retain video prediction as a training-time objective that shapes physically grounded representations, without synthesizing future observations at test time. This direct-policy formulation is more than $3\times$ faster than the imagine-then-execute paradigm.

Despite this speedup, efficiently generating high-quality actions remains a key challenge for direct-policy WAMs. They require tens of denoising steps to achieve competitive action quality, resulting in latency that remains too high for real-time applications such as autonomous driving. This bottleneck cannot be resolved by reducing the number of denoising steps. Specifically, as shown in Figure 2, the performance degrades as the step count decreases, dropping by $20\%$ on NAVSIM (autonomous driving) and $40\%$ on RoboTwin (embodied control) in the one-step setting. We 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. Existing few-step accelerations, such as step distillation [11] and trajectory rectification [29], still start from Gaussian noise and, therefore, do not supply the missing action prior. This motivates introducing structured action priors at initialization, providing an informed starting point for high-quality action generation in a single denoising step.

In addition, direct-policy WAMs [3, 16] remain limited in their ability to exploit future dynamics for effective control. They use a consistent first-frame-only interface during training and inference, relying on video co-training to encode future dynamics into the current representation. However, compressing all future information into a single first-frame representation may provide insufficient temporal context for effective future-aware control. To address this limitation, subsequent methods reintroduce future features at inference and mitigate the associated rollout cost through caching, sparse interaction  [30], or compact latent representations [31]. Despite these efforts to reduce rollout cost, generating and incorporating future features at test time still incurs additional inference overhead.

To address the aforementioned issues, we propose OSWAM, a world action model with one-step action generation for more effective physical control. Our key observation is that previous one-step action generation methods fail mainly because the action diffusion processing is initialized from unstructured Gaussian noise. To supply the missing action structure, we learn a compact set of imagination tokens from the video Diffusion Transformer (DiT) and use them as structured start points for action denoising. These tokens are far fewer than video first-frame tokens and are obtained in a single forward pass, so they add negligible computation, yet they interact with visual features to accumulate physically grounded, action-relevant context. We then noise the imagination tokens to an intermediate timestep and decode them with the action DiT in one step, replacing the long denoising trajectory required by random initialization. To enhance future-aware control, we train these imagination tokens with privileged-information distillation. A teacher attends to visual tokens over the full future horizon, while the student imagination tokens, at test time, attend only to the first frame. By distilling the teacher’s privileged representations into the imagination tokens and subsequently into the action DiT, OSWAM transfers long-horizon context without reintroducing future-dependent features at inference. This design preserves the first-frame-only interface of direct-policy WAMs, while enabling high-quality action generation in a single denoising step.

Our contributions are summarized as follows:

• We propose OSWAM, a world action model that enables high-quality one-step action generation. It learns compact imagination tokens from the video DiT as structured start points for the action DiT. These tokens are obtained in a single forward pass with negligible overhead and are decoded in one denoising step.

• We introduce privileged-information distillation to transfer long-horizon context into the imagination tokens. A teacher attends to visual tokens over the full future horizon, while the student remains first-frame-only at test time, preserving the interface of direct-policy WAMs without future-dependent features at inference.

• Experiments on autonomous driving and embodied control show that OSWAM achieves competitive results with multi-step WAMs while being substantially faster.

Related Work

World Action Models for Autonomous Driving

Autonomous driving has recently shifted from modular pipelines with separate perception [53, 55, 56, 54] and planning [52] modules toward end-to-end frameworks [37]. To further integrate future dynamics modeling, world models [32, 58, 59, 57] for autonomous driving have recently been extended into world action models (WAMs). Existing methods fall into two categories: bird’s-eye view (BEV)-based and video-generation-based world models. BEV-based approaches [10, 21, 26] forecast occupancy evolution and ego motion, predict future latent scene features, or infer future BEV layouts for online trajectory evaluation. They are computationally efficient but do not exploit large-scale video generative priors. A complementary line of research leverages pretrained video generators and predicts the future in the image domain. These methods [14, 15, 27] transfer such priors to driving, but jointly rolling out video features at test time remains inefficient. SimWAM [16] follows the direct-policy design of Fast-WAM [3] to avoid test-time video rollout, yet the remaining multi-step denoising is still impractical for onboard driving. In contrast, our proposed OSWAM inherits large-scale video generative priors without test-time video rollout, and reduces action generation to a single denoising step.

World Action Models for Embodied Control

Unlike autonomous driving, where compact BEV representations naturally encode road topology and traffic interactions, embodied control needs to model fine-grained object interactions and appearance changes in close-range scenes. As these phenomena are naturally captured in video, embodied WAMs are predominantly built on pretrained video generators for image-domain future modeling in two ways: explicit future imagination and direct-policy inference. Explicit-imagination methods [13, 12] generate future visual representations at test time to guide action prediction, incurring substantial rollout latency. Direct-policy WAMs [3] retain video prediction as a training signal but generate actions directly from current observations, avoiding test-time visual rollout; however, its iterative action denoising remains a bottleneck for high-frequency closed-loop control. In contrast, OSWAM preserves large-scale video-generative priors without test-time visual rollout and reduces action generation to a single denoising step.

Effective World Action Models

The main deployment bottleneck for WAMs lies in the high computational cost of diffusion-based video and action generators, whose test-time latency is prohibitive for onboard driving and closed-loop manipulation. Drawing one-step distillation from video generation, Flash-WAM [11] compresses joint video-action diffusion to a single denoising step per modality, but still synthesizes future video at inference. Fast-WAM [3] removes test-time video rollout by predicting actions from current observations only, yet retains iterative action denoising. Light-WAM [17] further replaces the generative action expert with a lightweight fusion head for single-pass action prediction, but underperforms on challenging benchmarks such as RoboTwin. SimWAM [16] follows Fast-WAM to drop test-time video rollout, but its multi-step action denoising keeps inference latency well above the budget required for onboard deployment. Building on this video-rollout-free direct-policy formulation, OSWAM addresses the remaining action-generation bottleneck through one-step action denoising.

Method

We present OSWAM, a direct-policy world action model that enables high-quality one-step action generation for autonomous driving and embodied control. The preliminary concepts associated with our method are provided in § 3.1. In the following sections, we first present the overall architecture of OSWAM, then describe in detail the fundamental components of OSWAM, including imagination tokens, Role-aware 3D RoPE, and privileged-information distillation.

Preliminary

A world action model predicts an action chunk by introducing future visual observations as an intermediate variable. Let $o$ be the current observation, $l$ an optional command, $a_{1:H}$ an action chunk, and $v_{1:T}$ the future visual states. Inverse-dynamics WAMs [13] first model the distribution of future visual states conditioned on the current observation and command, and then infer actions from the generated future:

Joint-WAMs [12], in contrast, model future video and actions within a shared generative process, $p(a_{1:H},v_{1:T}\mid o,l)$. Both paradigms tightly couple control prediction with iterative future-video generation, which introduces substantial computational overhead.

Direct-policy WAMs [3, 16] remove explicit test-time video generation by retaining video prediction only as a training objective. At inference, the video DiT produces a latent representation $z(o)$ in a single forward pass, and actions are predicted directly as:

However, action generation is iterative due to flow matching. Given $a_t=(1-t)a+t\epsilon$ with velocity target $\epsilon-a$, inference integrates the learned velocity field from noise ($t_K=1$) to data ($t_0=0$) as:

High-quality actions typically require multiple velocity evaluations, making $K$ a remaining inference bottleneck.

Our proposed OSWAM aims to preserve this direct-policy interface while reducing action generation to a single velocity evaluation, $K=1$. A one-step update from an arbitrary time $t\in(0,1]$ is:

When the start is pure noise, $t=1$ and $\hat{a}=a_1-v_\theta(a_1,1)$. OSWAM instead starts from a structured intermediate state with $t<1$, which makes this single update sufficient for applications that require latency-critical control, such as autonomous driving.

Overall Architecture

An overview of OSWAM is shown in Figure 3. As in previous works [3, 16], we consider predicting an action chunk $\mathbf{a}_{1:H}$ from the current frame $f_0$ and an optional task instruction $l$, with ego-state information additionally provided for driving tasks. We follow Fast-WAM in the world model backbone configuration and direct-policy setting. Specifically, the Video DiT uses the Wan2.2-5B backbone [32], accompanied by its pretrained T5 text encoder and video VAE. The Action DiT is initialized by interpolating and downsampling the pretrained Video DiT weights. For autonomous driving, we add an ego encoder to encode the ego state and the driving command. The task languages and ego states are both provided to all tokens through cross-attention. The full model is organized as a Mixture-of-Transformer (MoT) architecture with shared attention between the video and action blocks. We also follow the direct-policy protocol: video modeling is used only as a training signal, and future video is not rolled out at inference.

Unlike prior direct-policy methods, which use the Video DiT at inference solely to encode the first-frame visual feature for Action DiT through shared MoT attention, OSWAM additionally introduces two groups of imagination tokens into the Video DiT: Action Imagination Tokens (AIT) and Scene Imagination Tokens (SIT). As shown in Figure 3, both AIT and SIT are initialized from noise in the Video DiT and interact with the first-frame visual features. The resulting AIT representation then replaces the unstructured Gaussian initialization of the Action DiT, providing a compact, structured start point for action denoising. SIT provides additional latent slots that aggregate scene context for AIT; they are updated indirectly through privileged distillation and the auxiliary action loss on AIT. Together, AIT and SIT comprise far fewer tokens than the first-frame visual input and therefore introduce only a small computational overhead during inference. To strengthen future-aware control under this first-frame-only interface, we introduce privileged-information distillation. The privileged teacher comprises a Video DiT with AIT that attends to visual tokens over the full clip, while the AIT in OSWAM attends only first-frame visual context. By aligning OSWAM’s intermediate imagination token features with those of the teacher, we transfer this privileged information to improve action generation, while OSWAM remains first-frame-only at inference.

Imagination Tokens

Direct-policy WAMs eliminate future video generation at inference but still require multi-step denoising. To enable one-step action generation, we introduce two groups of imagination tokens.

Action and Scene Imagination Tokens. As mentioned in the last subsection, we augment the Video DiT input with two groups of noise-initialized latent tokens: AIT and SIT. We adapt the Video DiT using LoRA [34] to support these additional tokens. Their combined token count is much smaller than that of the first-frame visual input tokens, introducing only a small computational overhead during inference.

During training, video tokens follow the standard diffusion process with randomly sampled noise levels, whereas AIT are constructed by encoding an initial action noise sample $\boldsymbol{\epsilon}\sim\mathcal{N}(0,I)$ and assigned the fixed timestep $t_{\max}$. Within the Video DiT, AIT attends to first-frame visual tokens, SIT, and other AIT, while receiving task conditioning through cross-attention. These interactions integrate visual evidence and scene context into AIT. An auxiliary decoder maps the resulting AIT representations to an action velocity $\hat{\mathbf{v}}^{A}$, yielding a structured action prior $\bar{\mathbf{a}}=\boldsymbol{\epsilon}-\hat{\mathbf{v}}^{A}$. We then detach this prior and renoise it with the same initial noise that was encoded into AIT:

where $\sigma_r$ controls the noise level and $\operatorname{sg}$ prevents gradients from propagating through the structured action prior. Conditioned on first-frame visual features and task inputs, the Action DiT predicts a refinement velocity $\hat{\mathbf{v}}^{\mathrm{act}}$ and performs the single denoising step:

Its training target is $(\mathbf{a}^{\mathrm{init}}_{\tau_r}-\mathbf{a}_{\mathrm{gt}})/\sigma_r$, enabling the Action DiT to refine the structured initialization into the target action in a single denoising step.

To complement the action prior encoded by AIT, SIT acts as latent context slots that aggregate scene context for AIT to attend to. They are initialized from Gaussian noise independent of the AIT, projected into the Video DiT input space through a separate encoder, and assigned the timestep $t_{\max}$. As shown in Figure 4 (a), under the structured attention mask, SIT attends to first-frame visual tokens and other SIT, while AIT additionally attends to SIT to access the aggregated scene context. SIT is learned indirectly through the auxiliary loss and feature distillation on AIT, with gradients flowing through the attention connections from AIT to SIT.

Role-aware 3D RoPE. As shown in Figure 4 (b), we introduce Role-aware 3D RoPE to reflect the distinct functions of video tokens, AIT, and SIT within a shared positional coordinate system. Video tokens represent local visual content and retain their original spatiotemporal coordinates. AIT encodes action structure and is assigned spatial coordinates outside the video grid to distinguish them from visual tokens. SIT serves as scene context aggregators and is anchored at the grid center, providing a common spatial reference for gathering information across the scene:

where $H_v$ and $W_v$ are the video token grid dimensions. For each action position $n$, the corresponding AIT and SIT share the temporal coordinate $t_n=n s_a/(s_v d_t)$. Here, $s_a$ and $s_v$ are the action and video sampling strides, and $d_t$ is the VAE temporal downsampling factor. This converts action positions to the video latent timeline. Thus, AIT and SIT share temporal coordinates but use distinct spatial anchors, preserving temporal alignment while providing role-specific positional cues for action initialization and scene context aggregation, respectively.

Privileged-Information Distillation

Although imagination tokens provide a structured action prior, long-horizon context is only implicitly encoded in the first-frame representation. We therefore introduce privileged-information distillation to transfer representations learned from privileged future observations into OSWAM’s AIT, strengthening their predictive capability without requiring future observations at inference.

We separately train a Video DiT teacher with auxiliary AIT, following the same training settings as OSWAM except for the token groups and attention masks. As shown in Figure 4 (a), the teacher contains video tokens and AIT, with AIT attending to the full video clip. OSWAM additionally includes SIT, and its AIT attends to first-frame visual tokens, SIT, and other AIT, while SIT attends only to first-frame visual tokens and other SIT.

During distillation, we freeze the teacher and provide it with ground-truth video latents over the full future horizon. The first-frame latents remain clean, while all future-frame latents are corrupted with noise at a fixed diffusion timestep $\tau_v^{T}$ before the teacher forward pass. The teacher’s AIT remains initialized from pure noise at $t_{\max}$ and acquires privileged information by attending to these future video tokens. OSWAM retains its standard training procedure, using a randomly sampled video timestep and the fixed AIT timestep $t_{\max}$. The teacher and OSWAM share the same video noise and action noise samples, respectively. We then align OSWAM’s AIT features with the corresponding teacher features at a set of Video DiT blocks $\mathcal{J}$, providing multi-depth distillation supervision. SIT is supervised indirectly through their attention connections to AIT. At inference, the teacher is discarded, preserving OSWAM’s first-frame-only interface without additional computation.

Training Objective

OSWAM jointly learns visual dynamics, structured action DiT start points, and one-step action denoising. The Video DiT is supervised by video flow matching, while its AIT representations receive complementary action and privileged feature supervision to learn the structured initialization consumed by the Action DiT:

where $\mathbf{z}^{+}$ denotes clean future video latents and $\boldsymbol{\epsilon}^{V}$ is the video noise. To learn an observation-conditioned action prior, the AIT decoder is supervised with the action velocity target at $t_{\max}$:

Privileged feature distillation further guides the prior by aligning the OSWAM’s AIT representations with teacher across the selected blocks $\mathcal{J}$:

where $\mathbf{h}^{\mathrm{stu}}_{j}$ and $\mathbf{h}^{\mathrm{tea}}_{j}$ denote corresponding AIT features after block $j$. Together, $\mathcal{L}_{\mathrm{AIT}}$ and $\mathcal{L}_{\mathrm{dist}}$ support learning an observation-conditioned structured initialization for the Action DiT. Given $\mathbf{a}^{\mathrm{init}}_{\tau_r}$ constructed from the AIT prediction, the Action DiT learns to predict the demonstrated action chunks in one denoising step:

Scaling the action residual by $\sigma_r^{-1}$ ensures that one Euler update with step size $\sigma_r$ directly moves $\mathbf{a}^{\mathrm{init}}_{\tau_r}$ to the ground-truth action $\mathbf{a}_{\mathrm{gt}}$. The overall objective is:

Experiments

Datasets and Evaluation Metrics

Autonomous Driving. For autonomous driving, we evaluate OSWAM on the NAVSIM [35], which provides real-world driving data sampled at 2Hz, with eight cameras covering a 360-degree field of view and LiDAR observations. The dataset comprises 1192 scenarios for training and 136 scenarios for testing. NAVSIM introduces practical and reliable metrics that effectively bridge open-loop and closed-loop evaluation. Specifically, NAVSIM evaluates model performance using the Predictive Driver Model Score (PDMS), which is calculated based on five factors: No At-Fault Collision (NC), Drivable Area Compliance (DAC), Time-to-Collision (TTC), Comfort (Comf.), and Ego Progress (EP). We report PDMS, which is calculated as $\mathrm{PDMS}=\mathrm{NC}\times\mathrm{DAC}\times({5\times \mathrm{EP} + 5\times \mathrm{TTC} + 2\times \mathrm{C}})/{12}.$

Embodied Control. For embodied control, we evaluate OSWAM on RoboTwin 2.0 [36], which is a challenging bimanual manipulation benchmark featuring more than 50 tasks that require coordinated dual-arm control. We adopt the multi-task setting of Fast-WAM, using data from over 50 tasks. The training data consist of 2,500 demonstrations from clean scenes and 25,000 demonstrations collected with heavy scene randomization. Performance is evaluated by the average success rate over 100 trials per task under both clean and randomized conditions.

Implementation Details

Following Fast-WAM, we initialize OSWAM from Wan2.2-5B and adopt the Action DiT with a reduced hidden dimension of 1024. The imagination-token LoRA adapters use rank 8 and $\alpha=4$. For NAVSIM, we follow [16] to use a single front camera at $384\times672$ resolution and predict eight waypoints over four seconds at 2Hz. We use the current frame and eight future frames in training. For RoboTwin, we concatenate the head and two wrist views into a $384\times320$ image and use an action horizon of 32, with nine video frames sampled at a temporal stride of 4. For temporal RoPE, we set $s_a=1$ and $d_t=4$ on both benchmarks, with $s_v=1$ for NAVSIM and $s_v=4$ for RoboTwin. AIT and SIT therefore use $t_n=n/4$, ($n=1,\ldots,8$) on NAVSIM and $t_n=n/16$, ($n=1,\ldots,32$) on RoboTwin; video latents use temporal coordinates $\{0,1,2\}$ in training.

We train with AdamW using a learning rate of $10^{-4}$, weight decay of $0.01$, and cosine decay with 5% linear warmup. NAVSIM training runs for 40,000 steps with a global batch size of 256, while RoboTwin training runs for 29,335 steps with a global batch size of 1024. For distillation, the frozen teacher uses a video timestep of 200, and feature alignment is applied after Video DiT blocks $[8, 16, 24, 30]$. We set $\lambda_{\mathrm{dist}}=10$. Distillation starts after 500 steps, with its weight increasing linearly over the next 500 steps. Across both benchmarks, we set the AIT diffusion timestep to $t_{\max}=1000$ and the re-noising timestep for Action DiT initialization to $\tau_r=500$ ($\sigma_r=0.5$). During inference, we use a single Action DiT denoising step. Unless otherwise specified, all latency measurements are conducted on a single NVIDIA H20 GPU and Intel Xeon Max 9468 CPU.

Main Results

NAVSIM. Table tab:navsim compares OSWAM with state-of-the-art planners on NAVSIM. OSWAM achieves a competitive PDMS of $90.3$ with a latency of $136$ ms. Compared with SimWAM, OSWAM matches its PDMS while reducing action denoising from ten steps to one, yielding a $3.01\times$ speedup. OSWAM also outperforms Metis by $1.2$ PDMS with a $3.24\times$ speedup. Compared with DriveVA, OSWAM achieves a $4.36\times$ speedup with only a $0.2$ lower PDMS. These results demonstrate the effectiveness of structured initialization in enabling efficient one-step action generation while maintaining competitive planning performance.

RoboTwin 2.0. Table tab:main_results reports success rates across 50 tasks under Clean and Randomized settings. OSWAM achieves an average success rate of $90.00\%$ with a latency of $142$ms. Compared with Fast-WAM, OSWAM reduces action denoising from ten steps to one and achieves a $2.80\times$ speedup, with only $1.83\%$ decrease in average success rate. Compared with the one-step Flash-WAM variant, OSWAM improves average success by $8.59\%$ and achieves a $4.29\times$ speedup. While Light-WAM achieves lower latency, OSWAM substantially improves the average success at the cost of only $35$ms additional latency. Overall, OSWAM is competitive with multi-step direct-policy WAMs while substantially accelerating inference on embodied control tasks.

Inference Efficiency. Table tab:inference_efficiency compares inference efficiency with cached language context under the same environment. On RoboTwin 2.0, OSWAM reduces action GPU time from $67.61$ to $6.78$ms compared with Fast-WAM, while imagination tokens increase visual GPU time by $5.48$ms. This yields a $3.39\times$ overall speedup, reducing latency from $323.67$ to $95.35$ms with nearly unchanged peak memory. On NAVSIM, OSWAM similarly achieves a $3.32\times$ speedup over SimWAM, reducing overall latency to $99.70$ms. With compiled inference using torch.compile and CUDA Graph optimization, OSWAM achieves latencies of $40.81$ms on RoboTwin 2.0 and $52.64$ms on NAVSIM, supporting inference rates above $10$Hz and facilitating deployment in latency-critical applications such as autonomous driving. These results show that the reduction in action denoising cost outweighs the additional computation introduced by imagination tokens.

Method	NAVSIM	RoboTwin
	Latency$\downarrow$	PDMS$\uparrow$	ADE$\downarrow$	Latency$\downarrow$	Clean$\uparrow$	Rand.$\uparrow$
Baseline	108.35	71.04	1.049	112.84	57.38	55.80
+ AIT	135.92	89.81	0.594	139.82	87.42	88.16
+ Distillation	135.92	89.96	0.589	139.82	89.42	89.62
+ SIT	136.09	90.28	0.587	142.67	90.04	89.96

Ablation Studies

Effectiveness of Each Component. To more clearly illustrate the effectiveness of the different components in OSWAM, we conduct experiments by adding each of them to a baseline, which is set to one-step Fast-WAM. As shown in Table 1, AIT substantially improves performance across both benchmarks, raising NAVSIM PDMS from $71.04$ to $89.81$, and increasing RoboTwin success rates from $57.38\%$ to $87.42\%$ under Clean settings and from $55.80\%$ to $88.16\%$ under Randomized settings. The gains demonstrate the importance of structured start points for one-step action generation. Besides, privileged-information distillation further improves PDMS to $89.96$ and RoboTwin success rates to $89.42\%$ and $89.62\%$, without increasing inference latency. Finally, incorporating SIT yields the best performance on both benchmarks, demonstrating the complementary benefit of aggregating scene context for learning effective action priors.

Distillation Weight. Table tab:dist_weight studies the effect of the distillation weight $\lambda_{\mathrm{dist}}$ on NAVSIM. While a small weight of $0.01$ slightly underperforms the baseline, increasing the distillation strength generally improves performance. The best result is obtained with $\lambda_{\mathrm{dist}}=10$, reaching a $90.28$ PDMS and a $0.587$ ADE. This suggests that sufficiently strong distillation supervision is important for transferring future context into AIT. We therefore use $\lambda_{\mathrm{dist}}=10$ as the default setting.

$\tau_r$	PDMS$\uparrow$	ADE$\downarrow$
Rectified-flow refinement	89.97	0.609
400	90.17	0.587
500	90.28	0.587
600	90.08	0.598
Uniform (400-600)	90.02	0.591

Action Initialization. To evaluate different strategies for utilizing the structured action prior, we compare rectified-flow refinement [29] with re-noising at different noise levels. As shown in Table 2, all re-noising variants outperform rectified-flow refinement, achieving higher PDMS and lower ADE. Increasing $\tau_r$ from 400 to 500 improves PDMS without affecting ADE, whereas a further increase to 600 slightly degrades both metrics. Uniformly sampling $\tau_r$ from 400 to 600 also underperforms the fixed setting of 500. These results suggest that a fixed intermediate noise level provides a favorable trade-off between preserving the structured prior and allowing sufficient action refinement; we therefore set $\tau_r=500$ by default.

Teacher Video Timestep. To investigate how the teacher noise level affects AIT distillation, we vary the teacher video timestep $\tau_v^T$ from 0 to 800. As shown in Table tab:teacher_video_timestep, performance changes only slightly across the evaluated settings, with PDMS ranging from $90.16$ to $90.28$. Even at $\tau_v^T=800$, where future video inputs are pure noise, OSWAM achieves a PDMS of $90.16$, close to the clean-future setting ($\tau_v^T=0$). This suggests that the teacher’s learned future-modeling representations can provide useful supervision for AIT even without access to clean ground-truth future frames. We use $\tau_v^T=200$ by default, as it achieves the highest PDMS.

Real-World Evaluation

We further evaluate OSWAM on a PiPER robotic arm using two real-world tabletop manipulation tasks: placing a packaged bread into a bowl and placing a board eraser into an open drawer. Both tasks require coordinated grasping, transport, and accurate placement, while the drawer task additionally imposes clearance constraints from the drawer opening and surrounding cabinet. We use 50 training demonstrations for bread and eraser placement, respectively. Each method is evaluated over 20 trials per task. A trial is successful if the object is released into the target container and remains inside at the end of the execution.

As shown in Table tab:real_world, OSWAM achieves success rates of 85% on bread placement and 80% on eraser placement, compared with 80% and 75% for Fast-WAM, respectively. These results indicate that OSWAM’s structured action initialization followed by a single Action DiT denoising step supports effective physical manipulation, with higher observed success rates than Fast-WAM on both evaluated tasks. Among successful episodes, OSWAM reduces the mean video-estimated task completion time by 34.2% for bread placement and 21.8% for eraser placement compared with Fast-WAM. Figure 5 presents representative successful rollouts, showing the progression from grasping and transport to final placement. OSWAM places the bread inside the bowl and releases the board eraser into the target drawer. Together, the quantitative results and execution examples complement our simulation evaluation and demonstrate the applicability of OSWAM to real-world manipulation. Additional execution videos are available at https://gwenzhang.github.io/OSWAM.

Conclusion

In this paper, we presented OSWAM, an efficient world-action model for one-step action generation. To address the reliance of direct-policy WAMs on iterative action denoising, we introduced Action Imagination Tokens to learn a structured initialization, supported by Scene Imagination Tokens that aggregate complementary scene context. We further developed privileged-information distillation to transfer full future context into the Action Imagination Tokens. Together, these designs enable action generation with a single Video DiT pass and one Action DiT denoising step. Experiments on NAVSIM and RoboTwin 2.0 demonstrated the competitive planning and manipulation performance of OSWAM and its substantial latency reductions over direct-policy baselines. OSWAM offers a practical approach to efficient world-action modeling for latency-critical control.

Limitations. Our study is focused on learning action DiT structured initialization. Although SIT and privileged-information distillation support this objective, their potential to directly enhance the Action DiT remains unexplored. Extending scene-context conditioning and privileged supervision to the action expert is a promising direction for future work.

Appendix

We provide the following materials in this appendix:

• § 6: More experiment results.

• § 7: Qualitative visualizations.

More experiment results

More Implementation Details

nuScene. We also evaluate OSWAM on the nuScenes [25] dataset for autonomous driving. nuScenes comprises 1,000 real-world driving scenes, with keyframes annotated at 2 Hz. The scene images are captured by 6 cameras covering 360 degree FOV horizontally. The dataset is split into 700 scenes for training, 150 for validation, and 150 for testing. Following previous works [42, 46], $L_2$ Displacement Error (DE) and Collision Rate (CR) are adopted to comprehensively evaluate the planning performance.

For nuScenes, we use a single front camera at $384\times672$ resolution. Each training sample contains the current frame and eight future frames sampled at 2Hz. We use eight action steps to match the four-second video horizon, with supervision and evaluation limited to the first three seconds. Each action is represented as $(x,y,\theta)$ in the current LiDAR coordinate frame, with padded and unsupervised positions excluded from the action loss. For Role-aware 3D RoPE, we set $s_a=s_v=1$ and $d_t=4$, giving AIT and SIT temporal coordinates $t_n=n/4$ for $(n=1,\ldots,8)$. Video latents use temporal coordinates $\{0,1,2\}$ during training.

We train for 10,000 steps with a global batch size of 128 using AdamW, a learning rate of $10^{-4}$, weight decay of $0.01$, and cosine decay with 5% linear warmup. For distillation, the frozen teacher uses a video timestep of 200, and feature alignment is applied after Video DiT blocks $[8,16,24,30]$. We set $\lambda_{\mathrm{dist}}=10$. Distillation starts after 500 steps, with its weight increased linearly over the next 500 steps. We set the AIT diffusion timestep to $t_{\max}=1000$ and the re-noising timestep for Action DiT initialization to $\tau_r=500$ ($\sigma_r=0.5$).

Method	Auxiliary Supervision	L2 (m)$\downarrow$	Collision Rate (%)$\downarrow$
		1s	2s	3s	gray!30Avg.	1s	2s	3s	gray!30Avg.
ST-P3	Map & Box & Depth	1.33	2.11	2.90	gray!302.11	0.23	0.62	1.27	gray!300.71
UniAD	Map & Box & Motion & Tracklets & Occ	0.48	0.96	1.65	gray!301.03	0.05	0.17	0.71	gray!300.31
OccNet	3D-Occ & Map & Box	1.29	2.13	2.99	gray!302.14	0.21	0.59	1.37	gray!300.72
OccWorld	3D-Occ	0.52	1.27	2.41	gray!301.40	0.12	0.40	2.08	gray!300.87
VAD-Tiny	Map & Box & Motion	0.60	1.23	2.06	gray!301.30	0.31	0.53	1.33	gray!300.72
VAD-Base	Map & Box & Motion	0.54	1.15	1.98	gray!301.22	0.04	0.39	1.17	gray!300.53
GenAD	Map & Box & Motion	0.36	0.83	1.55	gray!300.91	0.06	0.23	1.00	gray!300.43
Doe-1	QA	0.50	1.18	2.11	gray!301.26	0.04	0.37	1.19	gray!300.53
Epona	None	0.61	1.17	1.98	gray!301.25	0.01	0.22	0.85	gray!300.36
DriveLaW	None	0.44	1.10	1.91	gray!301.15	0.15	0.10	0.48	gray!300.24
PWM	Text	0.20	0.38	0.65	gray!300.41	0.01	0.02	0.09	gray!300.04
OSWAM (Ours)	None	0.19	0.36	0.66	gray!300.41	0.00	0.01	0.07	gray!300.03

More Results

Navsim v2. Compared with Navsim v1, Navsim v2 introduces the Extended Predictive Driver Model Score (EPDMS). EPDMS incorporates Driving Direction Compliance (DDC), Traffic Light Compliance (TLC), Lane Keeping (LK), and Extended Comfort (EC), while extending the original comfort metric to History Comfort (HC). As shown in Table tab:navsim_v2, OSWAM achieves the highest EPDMS of $90.3$ among the compared planners with only one action denoising step. It slightly surpasses SimWAM $90.2$ and outperforms Metis $89.5$, both of which require ten action denoising steps. These results further demonstrate the effectiveness of OSWAM under more comprehensive evaluation criteria encompassing safety, traffic-rule compliance, and driving comfort.

nuScene. Table 3 reports open-loop planning performance on nuScenes. Without auxiliary supervision, OSWAM achieves an average $0.41$ m L2 error, matching PWM, while reducing the average collision rate from $0.04\%$ to $0.03\%$. It also achieves the lowest L2 errors at 1 and 2 seconds and the lowest collision rates across all evaluated horizons. These results demonstrate that OSWAM not only accelerates action generation but also delivers strong open-loop planning performance, matching the best average L2 error while achieving the lowest average collision rate.

RoboTwin Per-task Success Rates. Table 5 reports per-task success rates across 50 manipulation tasks on RoboTwin 2.0. OSWAM matches or exceeds Light-WAM on $49$ tasks under the clean setting and $46$ tasks under the randomized setting. OSWAM also matches or exceeds Fast-WAM on $22$ tasks under the clean setting and $15$ tasks under the randomized setting. Notably, OSWAM achieves substantial gains over Fast-WAM on Open Microwave under both settings, as well as on Move Stapler Pad, Hanging Mug, and Move Can Pot under the randomized setting. Together with the $2.80\times$ inference speedup over Fast-WAM, these findings show that OSWAM maintains competitive task-level performance while substantially reducing inference latency.

One-Step Action Generation. Table 4 and tab:one_step_robotwin compare OSWAM with direct policy baselines with a one-step denoising setting. On NAVSIM, OSWAM achieves a PDMS of $90.3$, outperforming the reported one-step results of SimWAM, Metis, and DriveVA by $21.4$, $3.3$, and $77.1$ points, respectively. On RoboTwin 2.0, OSWAM improves the average success rate from $56.59\%$ to $90.00\%$ over our one-step Fast-WAM baseline, with consistent gains under both Clean and Randomized settings. These results demonstrate that OSWAM enables effective one-step action generation across both autonomous driving and embodied control.

Method	PDMS$\uparrow$	NC$\uparrow$	DAC$\uparrow$
SimWAM (1-step)	68.9	97.4	91.3
Metis (1-step)	87.0	--	--
DriveVA (1-step)	13.2	61.8	36.9
OSWAM (Ours)	90.3	98.6	98.0

Qualitative visualizations

Figure 6 compares SimWAM and OSWAM under the one-step action generation setting. SimWAM struggles to produce smooth trajectories, particularly when turning or navigating complex interactions with surrounding agents, leading to drivable-area violations and at-fault collisions. In contrast, OSWAM handles these challenging examples with a single action denoising step, producing smoother trajectories while satisfying both no-at-fault collision and drivable-area compliance criteria. Together with the real-time inference speed reported in Table tab:inference_efficiency, these results highlight the potential of OSWAM for practical deployment in autonomous driving.

Task	OSWAM	Fast-WAM	Light-WAM	LingBot-VA	Motus	$\pi_{0.5}$
	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean	Rand.
Adjust Bottle	100	99	100	100	100	100	90	94	89	93	100	99
Beat Block Hammer	99	92	99	97	83	80	96	98	95	88	96	93
Blocks Ranking RGB	99	97	100	100	96	91	99	98	99	97	92	85
Blocks Ranking Size	84	81	94	98	57	54	94	96	75	63	49	26
Click Alarmclock	100	99	100	100	100	100	99	100	100	100	98	89
Click Bell	97	99	100	100	100	100	100	100	100	100	99	66
Dump Bin Bigbin	92	95	97	96	81	75	89	96	95	91	92	97
Grab Roller	100	100	100	100	100	98	100	100	100	100	100	100
Handover Block	84	79	95	81	71	59	99	78	86	73	66	57
Handover Mic	96	100	99	100	90	94	94	96	78	63	98	97
Hanging Mug	47	68	58	62	25	17	40	28	38	38	18	17
Lift Pot	100	98	100	100	93	93	100	99	96	99	96	85
Move Can Pot	74	94	90	88	57	74	94	97	34	74	51	55
Move Pillbottle Pad	97	96	100	99	69	74	99	99	93	96	84	61
Move Playingcard Away	99	99	100	100	93	92	100	99	100	96	96	84
Move Stapler Pad	69	72	77	64	26	34	91	79	83	85	56	42
Open Laptop	99	98	98	100	91	97	92	94	95	91	90	96
Open Microwave	82	76	62	45	76	59	82	86	95	91	34	77
Pick Diverse Bottles	80	80	80	85	61	57	89	82	90	91	81	71
Pick Dual Bottles	97	92	100	96	90	63	100	99	96	90	93	63
Place A2B Left	89	95	95	93	84	83	97	93	88	79	87	82
Place A2B Right	96	94	93	99	89	85	97	95	91	87	87	84
Place Bread Basket	96	96	91	93	81	80	97	95	91	94	77	64
Place Bread Skillet	92	85	90	93	92	82	95	90	86	83	85	66
Place Burger Fries	97	99	96	99	95	98	97	95	98	98	94	87
Place Can Basket	79	65	71	69	57	57	81	84	81	76	62	62
Place Cans Plasticbox	97	94	99	96	37	68	100	99	98	94	94	84
Place Container Plate	99	99	96	100	99	94	99	97	98	99	99	95
Place Dual Shoes	83	87	94	88	53	51	94	89	93	87	75	75
Place Empty Cup	98	100	100	100	91	95	100	100	99	98	100	99
Place Fan	95	89	96	96	77	74	99	93	91	87	87	85
Place Mouse Pad	83	83	83	89	58	62	93	96	66	68	60	39
Place Object Basket	86	78	89	88	81	72	91	88	81	87	80	76
Place Object Scale	88	90	90	97	72	74	96	95	88	85	86	80
Place Object Stand	91	94	90	94	78	85	99	96	98	97	91	85
Place Phone Stand	99	97	97	99	85	88	97	97	87	86	81	81
Place Shoe	93	96	96	99	84	87	98	98	99	97	92	93
Press Stapler	80	86	90	97	65	76	85	82	93	98	87	83
Put Bottles Dustbin	90	89	95	90	65	65	87	91	81	79	84	79
Put Object Cabinet	84	85	94	89	80	68	85	87	88	71	80	79
Rotate QRcode	83	84	93	89	72	84	96	91	89	73	89	87
Scan Object	85	81	89	92	60	52	96	91	67	66	72	65
Shake Bottle Horizontally	100	100	100	100	100	98	100	99	100	98	99	99
Shake Bottle	100	100	100	100	100	99	100	97	100	97	99	97
Stack Blocks Three	94	95	95	97	65	67	99	98	91	95	91	76
Stack Blocks Two	99	100	100	100	94	91	100	98	100	98	97	100
Stack Bowls Three	84	75	80	81	65	72	86	83	79	87	77	71
Stack Bowls Two	95	96	92	98	91	97	94	98	98	98	95	96
Stamp Seal	92	87	90	94	60	63	96	97	93	92	79	55
Turn Switch	60	65	61	59	33	39	44	45	84	78	62	54
Average	90.0	90.0	91.9	91.8	76.4	76.3	92.9	91.5	88.7	87.0	82.7	76.8