Summary (Overview)

• SimWAM is a novel World-Action Model (WAM) for end-to-end autonomous driving that uses video generation purely as a training signal, eliminating the need for costly future-frame synthesis at inference time.
• The method co-trains a pretrained video Diffusion Transformer (Wan2.2-5B) and a lightweight action DiT (1.02B) via joint flow matching, with an isolated attention mask that keeps action prediction independent of future frames.
• SimWAM achieves 91.5 PDMS on the NAVSIM benchmark, surpassing state-of-the-art WAM-based planners (DriveWAM: 90.1, DriveLaW: 89.1) with substantially lower inference latency, and transfers zero-shot to nuScenes with a 0.04% average collision rate.
• The decoupled architecture enables video backbone flexibility (Wan, Cosmos, LTX-Video) and independent scaling of the action expert without modifying the learning objective or inference pipeline.
• Reinforcement learning via a marginal-preserving SDE with GRPO further optimizes a compositional driving reward, improving PDMS by 1.2 points over imitation-only training.

---

Introduction and Theoretical Foundation

Background and Motivation

End-to-end autonomous driving maps raw sensor observations directly to planned trajectories using a unified network, eliminating hand-crafted interfaces and reducing error propagation in traditional perception-prediction-planning pipelines. However, existing end-to-end planners remain primarily imitation policies that reproduce logged behavior while capturing traffic semantics, user intent, and scene dynamics only implicitly.

The World-Action Model Paradigm

World-Action Models (WAMs) address this limitation by providing an explicit prior over how the environment evolves under motion. Recent driving WAMs such as DriveLaW and DriveWAM follow an imagine-then-act pipeline where the planner conditions its output on generated future frames. This design places costly video synthesis inside the real-time planning loop, substantially increasing inference latency.

Key Theoretical Insight

"Explicit future synthesis is unnecessary for effective world-action learning. Fast-WAM shows that video co-training benefits action prediction primarily through training-time representation learning rather than test-time future imagination."

SimWAM builds on this insight: the planner directly predicts trajectories at inference while video generation serves only as a training signal to transfer traffic-dynamics priors into the action expert's observation representation.

Flow Matching Foundation

SimWAM models both trajectories and future frames with rectified flow. Given a clean target $x$ and Gaussian noise $\epsilon \sim \mathcal{N}(0, I)$, the linear interpolation $x_\tau = (1-\tau)x + \tau\epsilon$ ($\tau \in [0,1]$) has constant velocity $\epsilon - x$, which a network $v_\theta$ learns to predict under conditioning $c$:

$$\mathcal{L}_{\mathrm{FM}} = \mathbb{E}_{x, \epsilon, \tau} \left[ \| v_\theta(x_\tau, \tau, c) - (\epsilon - x) \|_2^2 \right].\tag{1}$$

---

Methodology

Problem Formulation

SimWAM formulates trajectory planning from a front-camera observation $o_t$, ego state $s_t$ (velocity, acceleration, yaw rate), and navigation command $l$. The planner predicts an ego trajectory $a_{t+1:t+H}$ in the ego-vehicle coordinate frame, where each waypoint $a_i = (x_i, y_i, \theta_i)$ specifies position and heading.

Key architectural difference from existing WAMs:

• Existing WAMs: $p_\theta(a_{t+1:t+H} \mid o_t, s_t, l) = \int p_\theta(z_{t+1:t+N} \mid o_t, s_t, l) p_\theta(a_{t+1:t+H} \mid o_t, s_t, l, z_{t+1:t+N}) \, dz_{t+1:t+N}$ (Eq. 3)
• SimWAM: $p_\theta(a_{t+1:t+H} \mid o_t, s_t, l) = p_\theta(a_{t+1:t+H} \mid z(o_t), s_t, l)$ (Eq. 4)

Model Architecture

Video Expert: A video Diffusion Transformer initialized from Wan2.2-5B with its VAE and T5 text encoder. The VAE maps driving frames into latent tokens; navigation commands enter through T5 cross-attention. The current frame serves as a clean condition, and N future frames are noised and reconstructed with flow matching.

Action Expert: A lightweight Diffusion Transformer with hidden size $d_a = 1024$. Conditioned on $c = \{z(o_t), s_t, l\}$, it predicts the trajectory velocity field $v_{\theta_a}(a_{t+1:t+H}^\tau, \tau, c)$ via flow matching.

Co-training Objective:

$$\mathcal{L} = \mathcal{L}_{\mathrm{FM}}^{\mathrm{act}} + \lambda \mathcal{L}_{\mathrm{FM}}^{\mathrm{vid}},\tag{5}$$

Isolated Attention Mask

The shared attention stream contains current observation latents $z(o_t)$, future frame latents $z_{t+1:t+N}$, and action tokens. Both future frame tokens and action tokens attend to $z(o_t)$ while remaining mutually invisible. This is the only structural modification required to separate the two experts, allowing the video branch to be discarded at inference.

From ODE to SDE for Reinforcement Learning

The deterministic flow ODE lacks stochasticity for exploration. Following Flow-GRPO, SimWAM transforms the ODE into a marginal-preserving SDE:

$$\mathrm{d}x_\tau = \left[ v_\theta(x_\tau, \tau) + \frac{\sigma_\tau^2}{2\tau}\big(x_\tau + (1-\tau)v_\theta(x_\tau, \tau)\big) \right] \mathrm{d}\tau + \sigma_\tau \, \mathrm{d}w, \qquad \sigma_\tau = a\sqrt{\frac{\tau}{1-\tau}},\tag{2}$$

Each Euler-Maruyama step yields an isotropic Gaussian transition $\pi_\theta(x_{\tau-\Delta\tau} \mid x_\tau) = \mathcal{N}\big(\mu_\theta(x_\tau, \tau), \sigma_\tau^2 \Delta\tau I\big)$ with tractable log-likelihoods for importance sampling. G=8 candidate trajectories are sampled per scenario, evaluated with the NAVSIM PDM reward, and optimized via clipped policy updates on rank-32 LoRA adapters.

---

Empirical Validation / Results

Main Results (NAVSIM navtest)

Table 1: Comparison with state-of-the-art planners on the NAVSIM navtest benchmark.

SimWAM surpasses the strongest VLM-based planner (SGDrive) by 0.4 points and outperforms DriveLaW and DriveWAM by 2.4 and 1.4 points respectively, using only a single front camera.

Component Analysis

Table 2: Component analysis.

Video co-training and RL contribute complementary gains, improving PDMS by 4.9 points while preserving efficient inference.

Attention Mask Analysis

Table 3: Attention mask analysis.

The isolated mask achieves the best PDMS while enabling efficient inference without future generation.

Video Backbone Flexibility

Table 4: Video backbone flexibility.

Cosmos-Predict2.5, pretrained on driving videos, achieves the best PDMS (90.4), demonstrating SimWAM's ability to absorb stronger domain-relevant priors.

Zero-Shot Generalization (nuScenes)

Table 6: Zero-shot generalization on the nuScenes open-loop planning benchmark.

SimWAM achieves the lowest average collision rate (0.04%) without nuScenes supervision or auxiliary annotations.

Ablation Studies

• Exploration sampler: The marginal-preserving SDE (91.5 PDMS) outperforms random noise perturbation (91.3) by maintaining better safety (NC: 98.4 vs 97.7).
• RL training dynamics: Training on the hard subset (PDMS < 90) consistently outperforms full-set training, peaking at 91.5 PDMS at 15k steps.
• Prediction horizon: Broad temporal coverage (4s at 2Hz) is more important than dense frame sampling (4s at 1Hz: 90.2 vs 2s at 2Hz: 89.9).
• Input resolution: 384×672 provides the best balance (90.3 PDMS, 518ms); 768×1344 adds only 0.3 points with 55ms more latency.
• Sampling steps: 10 steps achieves the highest PDMS (90.3); 20 steps provides no improvement while nearly doubling latency.

---

Theoretical and Practical Implications

Theoretical Contributions

1. Training-time world modeling is sufficient: SimWAM provides strong evidence that future-scene generation need not be part of the inference loop for world-action learning to be effective. The motion prior is internalized during training through joint flow matching.

2. Decoupling via attention masking: The isolated attention mask demonstrates a minimal structural modification that cleanly separates action prediction from future-frame dependence while preserving the benefits of joint learning.

3. ODE-to-SDE reformulation for RL: Converting the deterministic flow ODE to a marginal-preserving SDE enables policy-gradient optimization with tractable transition likelihoods, bridging generative modeling and reinforcement learning.

Practical Implications

• Efficient deployment: SimWAM retains only a 1.02B action DiT at inference, avoiding the multi-billion-parameter video generator entirely, resulting in substantially lower latency than imagine-then-act planners.
• Modular scalability: The video backbone can be upgraded to more powerful generators without retraining the action expert, and the action expert can be resized independently to meet latency budgets.
• Cross-domain robustness: Zero-shot transfer to nuScenes with the lowest collision rate suggests the learned dynamics prior generalizes beyond the training benchmark.

---

Conclusion

SimWAM demonstrates that training-time world modeling can deliver strong real-time planning without costly test-time imagination. The key takeaways are:

1. Video generation serves purely as a training signal via joint flow matching, transferring traffic-dynamics priors to a lightweight action expert.
2. The isolated attention mask decouples action prediction from future frames, enabling efficient direct trajectory inference.
3. The decoupled architecture supports video backbone flexibility and independent scaling of both experts.
4. Reinforcement learning via a marginal-preserving SDE further aligns the planner with driving quality beyond imitation.

Future directions include exploring more advanced video generators as they become available, extending the framework to multi-modal sensor inputs, and investigating closed-loop evaluation settings. SimWAM is positioned as a simple yet solid baseline that can readily benefit from rapid advances in video generation for efficient autonomous driving.

"These results show that training-time world modeling could support strong real-time planning without costly test-time imagination."