Technical Summary: Discrete-WAM

Problem Statement

Autonomous driving fundamentally requires reasoning about how ego actions shape the evolution of the surrounding world, rather than merely reacting to instantaneous observations. Current end-to-end (E2E) systems typically rely on direct state-to-action mappings via behavior cloning, capturing statistical correlations without explicitly modeling action-conditioned dynamics. While continuous-latent world models offer strong continuity, they often lack the compositional structure necessary for causal reasoning across counterfactual futures. Furthermore, existing approaches often treat world modeling, policy learning, and decision-making as separate tasks or auxiliary objectives, failing to establish consistent causal correspondences between observations, actions, and induced future states within a unified generative process.

Methodology: Discrete-WAM

The authors propose Discrete-WAM, a unified latent vision-action world policy that represents future visual states and ego actions as aligned discrete tokens. This framework enables compositional causal reasoning across alternative futures through a shared discrete diffusion architecture.

1. Unified Discrete Representation

Discrete-WAM converts continuous observations and actions into a shared discrete token space:

• Vision Tokenization: Continuous camera observations are encoded into discrete semantic tokens using a pre-trained VQ-VAE tokenizer.
• Action Tokenization: Continuous future motion is converted into a discrete-token-compatible representation over an acceleration vocabulary. The authors employ a soft-label interpolation strategy over a 2D acceleration grid to avoid deterministic hard-quantization errors, allowing continuous control values to be preserved through weighted sums of neighboring prototypes.
• Decision Tokens: High-level decision conditions (e.g., maneuver intent, target lane) are represented as sparse, low-frequency decision tokens ($D_t$).

2. Hierarchical Policy Modeling

Instead of generating dense future actions directly from the scene context, Discrete-WAM decomposes policy generation into a hierarchical process:

• High-Level: The model predicts a sparse decision skeleton ($D_t$) capturing coarse maneuver intent and multi-modal driving choices.
• Low-Level: A discrete diffusion-based action-token planner generates smooth, temporally consistent trajectories conditioned on the predicted decision.
  This separation allows the model to select multi-modal decisions at the high level while realizing smooth trajectories at the low level, preventing the collapse of alternative driving modes.

3. Unified Generative Framework

Discrete-WAM formulates autonomous driving as sequence modeling over a shared discrete token space, supporting three complementary training modes under a single architecture:

• World Modeling: Predicts future visual tokens ($V_{t+1:t+H}$) conditioned on scene context ($C_t$) and future action tokens ($A_{t+1:t+H}$).
• World-Policy Modeling: Jointly trains vision and action prediction within an interleaved token sequence, encouraging the coupling of policy generation with world evolution.
• Policy Modeling: Predicts the high-level decision ($D_t$) followed by the action sequence ($A_{t+1:t+H}$) conditioned on the decision.

4. Training and Post-Training

• Unified Pretraining: The model undergoes multi-stage pretraining using a token-editing objective. It learns to recover corrupted tokens (visual or action) based on historical context and conditioning inputs. The loss function combines cross-entropy for token classification with continuous motion reconstruction losses (acceleration, trajectory, and position) to ensure physical consistency.
• Post-Training (RL): To address rare or safety-critical behaviors, the authors apply Grouped Relative Policy Optimization (GRPO). The model samples candidate trajectories, evaluates them with an online reward function (incorporating safety and rule-critical metrics), and updates the policy to refine decision-trajectory combinations.

5. Inference and Editing

During inference, Discrete-WAM utilizes iterative planning generation and editing over future states. A scheduling strategy (e.g., replace_confidence) selectively updates tokens based on prediction confidence and entropy, allowing for non-causal, globally consistent reasoning over action-conditioned futures without committing to a single forward rollout.

Key Contributions

1. Unified Vision-Action World Policy: The introduction of a framework that represents observations and ego actions in a shared discrete latent space, enabling aligned semantic modeling and structured reasoning over action-conditioned futures.
2. Discrete Diffusion Framework: A unified pretraining approach for world and action modeling that jointly captures action-conditioned future dynamics within a single generative process, moving beyond separate objectives.
3. Hierarchical Decision-Conditioned Planning: A demonstration that decomposing policy into high-level decision skeletons and low-level action refinement improves planning performance while enabling controllable future generation, counterfactual reasoning, and safety-aware evaluation.

Experimental Results

Experiments were conducted on large-scale autonomous driving benchmarks, specifically NAVSIM-v1 and NAVSIM-v2.

• Planning Performance: Discrete-WAM achieved competitive or superior planning performance compared to strong end-to-end baselines. On NAVSIM-v2, it achieved an EPDMS score of 90.4, outperforming methods like WAM-Flow (87.0) and SparseDriveV2 (90.1). On NAVSIM-v1, it achieved a PDMS of 92.2, surpassing state-of-the-art planners.
• World Generation: In terms of generative quality (FID and FVD), Discrete-WAM outperformed existing generative driving world models (e.g., Vista, DrivingWorld), achieving an FID of 6.6 and FVD of 80.0, despite being optimized for short-horizon planning.
• Ablation Studies:
  • Unified Pretraining: Models initialized with unified pretraining and fine-tuned via LoRA (LoRA-SFT) outperformed those trained from scratch or with full fine-tuning, indicating the value of preserving pretrained world-policy knowledge.
  • Decision Modeling: The inclusion of decision tokens ($D_t$) significantly improved planning metrics, confirming that decision-level guidance provides an effective behavioral prior.
  • Post-Training: RL post-training further improved EPDMS from 90.0 to 90.4, demonstrating the benefit of exploring decision-trajectory combinations.
  • Counterfactual Reasoning: The model demonstrated a strong negative correlation between "surprise" (KL divergence between factual and counterfactual rollouts) and planning safety metrics (PDMS), indicating its ability to capture safety-critical scene dynamics.

Significance and Claims

The paper claims that Discrete-WAM offers a principled path toward more reliable decision-making in autonomous driving by moving beyond reactive imitation toward structured world-policy reasoning. By unifying world modeling, policy learning, and decision-making within a discrete token space, the framework supports:

• Controllable Generation: The ability to edit specific aspects of future observations and trajectories.
• Counterfactual Reasoning: The capacity to reason about alternative futures and assess safety risks before execution.
• Safety-Aware Prediction: The ability to identify high-surprise scenarios that correspond to collisions or rule violations.

The authors posit that this approach addresses fundamental bottlenecks in representation formulation and generative modeling, providing a scalable foundation for embodied intelligence that can generalize across diverse driving scenarios.