Source-linked AI summary
Discrete-WAM: Unified Discrete Vision-Action Token Editing for World-Policy Learning
Ziyang Yao, Haochen Liu, Yuncheng Jiang, Zeyu Zhu, Zibin Guo, Jingru Wang, Tianle Liu, Jianwei Cui, Kuiyuan Yang, Hongwei Xie, Jingwei Zhao, Guang Chen, Hangjun Ye
TL;DR
Autonomous driving needs to reason about how actions shape future world evolution, but existing imitation and world-model approaches are insufficiently aligned for policy generation. Discrete-WAM unifies visual, decision, action, and future-state tokens, jointly trains world-policy objectives, and uses hierarchical parallel action editing; experiments report strong planning and controllable, efficient generation, with deployment-speed comparisons bounded by decoding assumptions.
Problem
Existing end-to-end methods rely on direct observation-to-action imitation, while world models often remain weakly aligned with downstream policy generation.
Method
Discrete-WAM uses a shared discrete token space and jointly trains world modeling, world-policy modeling, and policy modeling with hierarchical decision-conditioned action-token editing.
Results
The framework achieves strong planning performance and supports controllable future generation, counterfactual evaluation, surprise-based analysis, and efficient parallel policy decoding.
Takeaways & Limitations
Discrete representation alignment, unified world-policy training, and hierarchical token editing form a supported design paradigm for physical AI.
Takeaways & Limitations
Latency comparisons reflect intrinsic parallel-versus-sequential decoding complexity under controlled non-accelerated settings, not complete deployment-level speed.
Abstract
from arXiv · showhide
Autonomous driving requires reasoning about how ego actions shape future world evolution, rather than merely mapping observations to actions. However, most end-to-end methods rely on direct state-to-action imitation, while existing world models often remain weakly aligned with downstream policy generation. We introduce Discrete-WAM, a unified discrete vision-action world-policy framework that represents visual observations, future states, high-level decisions, and ego actions within a shared token space. Built on this discrete alignment, Discrete-WAM jointly trains world modeling, world-policy modeling, and policy modeling through multi-task and multi-stage pretraining, allowing action-conditioned future prediction to directly support policy generation. For downstream planning, Discrete-WAM further decomposes policy generation into hierarchical decision prediction and parallel action-token editing, where the decision token provides a high-level planning skeleton and confidence-based scheduling refines dense future actions efficiently. Experiments on large-scale autonomous-driving benchmarks show that Discrete-WAM achieves strong planning performance while supporting controllable future generation, counterfactual evaluation, surprise-based world-model analysis, and efficient parallel policy decoding. These results suggest that discrete representation alignment, unified world-policy training, and hierarchical token editing provide a promising design paradigm for physical AI.
1 Introduction
Discrete-WAM addresses the weak alignment between action-conditioned world prediction and policy generation by unifying observations, actions, decisions, and futures in a shared discrete token space. It combines world-policy learning with hierarchical, confidence-guided action editing for physical-AI planning.
- Motivation: Existing autonomous-driving systems often map observations directly to actions, while world models and policy learning remain weakly aligned.Direct imitation captures statistical correlations without explicitly modeling action-conditioned dynamics; separate prediction and control objectives limit unified future reasoning.
- Motivation: Discrete representations provide shared compositional anchors for aligning visual observations, actions, and future evolution.The motivation is to reduce the weak alignment and entanglement associated with continuous latent spaces.
- Contributions: Discrete-WAM jointly models world dynamics, world-policy prediction, and policy generation within a shared discrete token space.The framework represents observations, decisions, actions, and future states as coupled variables rather than using actions only as conditioning signals.
- Contributions: Hierarchical decision-conditioned action-token editing uses high-level decisions as planning structure and confidence-guided parallel editing for dense future actions.This formulation links high-level decisions with low-level action generation and targets efficient parallel discrete decoding.
- Evaluation: Evaluations cover planning, world-model surprise analysis, attention grounding, scheduling dynamics, and inference efficiency.The paper presents these studies as comprehensive validation of the proposed physical-AI design paradigm.
2 Architecture
Discrete-WAM uses a shared Transformer interface for discrete visual, action, decision, and future-state tokens. Its architecture supports action-conditioned world prediction, joint world-policy modeling, and hierarchical policy generation.
- Token representation: A vision tokenizer converts camera observations into discrete visual tokens for joint sequence modeling with actions.The tokenizer provides a token-level interface between images and the Transformer while preserving scene semantics for world and policy modeling.
- Token representation: Future motion is represented by discrete tokens from a grid-structured vocabulary over longitudinal and lateral acceleration.Continuous trajectories are smoothed, differentiated into ego-centric accelerations, and assigned to vocabulary bins.
- Token representation: Soft-label interpolation replaces hard one-hot quantization and allows predicted action distributions to reconstruct continuous acceleration by weighted summation.Under exact recovery of the soft target, the weighted sum of neighboring prototypes recovers the continuous acceleration exactly.
- Unified modeling: The unified model conditions on scene context and represents future visual tokens, action tokens, and high-level decision tokens in one world-policy formulation.Decision tokens encode sparse structure such as maneuver intent, target lane, coarse waypoint, speed trend, or interaction priority.
- Unified modeling: Three training modes cover action-conditioned vision prediction, joint action and future-vision prediction, and decision-conditioned action generation.They share the token-editing interface while differing in conditioning inputs and prediction targets.
3 Method
The method trains a shared discrete token-editing framework across world modeling, policy modeling, and joint world-policy modeling. Hierarchical decision conditioning structures action generation, while joint supervision supports action-conditioned future reasoning and avoids mode-averaged continuous decoding.
- Unified pretraining: Unified pretraining uses world modeling, policy modeling, and joint world-policy modeling with a common visual-action token interface.Task families differ in conditioning streams and corrupted tokens supervised as editing targets.
- World modeling: World modeling predicts future visual tokens from scene context and future action tokens, learning how actions induce different future world evolutions.Future actions are supplied through teacher forcing and supervision is applied to future visual tokens.
- Hierarchical policy modeling: Policy modeling first predicts a high-level decision skeleton and then generates future action tokens conditioned on that decision.The decomposition separates multimodal driving choices from low-level trajectory generation.
- Hierarchical policy modeling: Decision conditioning reduces residual action-token dependence when the skeleton explains sufficient low-frequency group structure.The benefit requires positive redundancy gain and depends on the skeleton reducing dependence enough to offset its prediction cost.
- World-policy modeling: Joint world-policy modeling edits corrupted action and visual tokens so the model learns policy generation, world prediction, and action-conditioned counterfactual reasoning together.Future visual tokens always receive token-editing supervision, while the action stream can receive the same supervision.
- Decoding: Full-distribution expectation can average incompatible acceleration modes and produce physically implausible intermediate actions.The paper names this failure mode decoding-induced mode averaging.
4 Evaluation
Discrete-WAM shows strong planning and world-generation performance, while ablations and qualitative analyses examine how decisions, token scheduling, visual context, and decoding affect behavior. The evaluation also identifies efficiency benefits and a deployment-relevant limitation of the latency comparison.
- Planning results: 90.4 EPDMS on NAVSIM-v2, improving by +2.7 over WAM-Flow, with additional gains over world-model-based and reinforced cognitive planners.NAVSIM-v1 results also include +2.1 PDMS over WAM-Flow and +7.0 over a world-model-based planner.
- World generation results: FID 6.6 and FVD 80.0 give Discrete-WAM the best overall visual fidelity among the compared generative driving world models.The model targets short-horizon unified generation for downstream planning, while some competing methods support longer rollouts.
- Post-training: Ground-truth decision conditioning raises EPDMS from 89.1 to 90.0, and RL post-training further reaches 90.4 while improving comfort and safety.The post-training stage explores combinations of high-level decisions and low-level trajectory refinements rather than refining trajectories under one fixed decision.
- Decision modeling: Decision learning improves EPDMS from 84.7 to 87.2, while performance peaks at k = 16 as larger anchor sets increase optimization difficulty.The reported trade-off is between decision coverage and the quality and variance of sampled anchors.
- Scheduling and decoding: Confidence-guided schedulers refine unresolved action tokens while retaining stable predictions, and parallel discrete diffusion decoding is lower-latency than autoregressive decoding at moderate editing rounds.High-confidence tokens tend to be retained, whereas low-confidence tokens become more certain as scheduling proceeds.
- Evaluation caveat: The latency comparison measures intrinsic parallel-versus-sequential decoding complexity under controlled non-accelerated conditions, not complete deployment-level speed.Autoregressive decoding can use KV-cache acceleration, whereas discrete diffusion does not use the same mechanism.
- Qualitative results: The planner generates geometrically consistent trajectories and temporally coherent future visual states, attending to roads, lanes, vehicles, signs, and other scene structure.Attention also appears in sky regions, but the paper cautions that this may reflect attention-sink or implicit-register behavior rather than direct causal evidence.
5 Conclusion and Future Directions
Discrete-WAM unifies discrete representation alignment, world-policy training, and hierarchical policy construction for autonomous driving. The paper frames this design as an initial validation of a broader physical-AI paradigm and identifies several directions for extending its scope.
- 5 Conclusion and Future Directions: Discrete-WAM represents observations, future states, decisions, and ego actions in a shared discrete token space for unified world-policy learning.It jointly trains world modeling, world-policy modeling, and policy modeling, while using a decision token as a planning skeleton for parallel action-token editing.
- 5 Conclusion and Future Directions: The paper presents discrete representation alignment, unified world-policy training, and hierarchical policy construction with parallel token editing as an initial physical-AI design paradigm.
- 5 Conclusion and Future Directions: Future work will add language-based reasoning, adaptive serial-parallel generation schedules, intrinsic world-model value functions, broader datasets, multimodal reasoning, and interactive field tests.The authors also plan to release additional evaluation results as deployments mature.
6 Contributions and Acknowledgments
This section lists the paper's authors and contribution metadata. It records equal-contribution, project-lead, work-location, and corresponding-author annotations.
- 6 Contributions and Acknowledgments: The author list includes Zibin Guo, Jingwei Zhao, Guang Chen, and Hangjun Ye.
- 6 Contributions and Acknowledgments: Contribution metadata marks equal contribution, project lead, work done at Xiaomi EV, and corresponding author roles.
- 6 Contributions and Acknowledgments: The author list also includes Jingru Wang, Tianle Liu, Jianwei Cui, Kuiyuan Yang, and Hongwei Xie.
7.1 Related Work
Related work spans world models for action-conditioned future evolution, unified world-action learning, diffusion scheduling, and the open-loop objective gap in end-to-end driving. The section situates Discrete-WAM among methods for prediction, token decoding, and policy optimization.
- 7.1 Related Work: World-model approaches in driving have used occupancy forecasting, joint motion prediction, and future visual supervision to support downstream decision making.
- 7.1 Related Work: World-action modeling jointly formulates future observation and policy generation, including unified pretraining, shared representations, latent actions, and value-aware models.
- 7.1 Related Work: Discrete-diffusion scheduling research studies noise-level schedules, decoded-token counts, confidence or entropy budgets, position selection, and learned correction mechanisms.These lines include non-uniform time grids, balanced block schedules, entropy-based unmasking, dilated schedules, and remasking or correction.
- 7.1 Related Work: Other scheduling methods decouple position planning from token denoising or add correction mechanisms, while their theoretical guarantees are generally weaker than information-theoretic schedules.
- 7.1 Related Work: End-to-end driving systems trained with behavior cloning face an objective gap because open-loop imitation does not directly optimize sequential closed-loop evaluation objectives.Offline reinforcement learning and post-training strategies have been explored to mitigate this discrepancy.
7.2 Additional Implementation Details
The implementation details define discrete decision, action, visual, and position representations, token-editing objectives, model configuration, evaluation metrics, and comparison baselines. The design combines categorical token supervision with continuous motion and trajectory constraints.
- 7.2.1 Token Design: The decision vocabulary combines lateral path primitives and longitudinal speed profiles into 400 discrete decisions.Each decision specifies a coarse path-and-speed behavior prior; pretraining selects decisions by maximizing EPDMS, while later optimization uses a top-D subset.
- 7.2.1 Token Design: Acceleration actions use a 60 × 60 grid of longitudinal and lateral prototypes, with continuous accelerations represented by bilinearly interpolated soft labels.
- 7.2.1 Token Design: Auxiliary position supervision factorizes longitudinal and lateral coordinates into separate categorical predictions instead of a joint 2D vocabulary.This reduces classification cost while retaining trajectory-level spatial supervision.
- 7.2.2 Detailed Token-Editing Objective: Token editing conditions on scene context, optional decision tokens, and corrupted visual or action target sequences within the original discrete vocabularies.Visual corruption samples token subsets, while action corruption follows a causal suffix so earlier action tokens remain clean.
- 7.2.2 Detailed Token-Editing Objective: The editing loss trains recovery at corrupted positions and identity preservation at clean positions, providing an implicit stopping signal for refinement.
- 7.2.2 Detailed Token-Editing Objective: Training combines visual and action token classification with acceleration, trajectory, and auxiliary special-token losses.Different tasks activate different subsets of these objectives.
- 7.2.3 Model and Evaluation: The model uses a decoder-only Transformer with 18 layers, 16 attention heads, 8 key-value heads, hidden dimension 2048, and approximately 1B parameters.LoRA finetuning uses about 30M parameters.
- 7.2.3 Model and Evaluation: Planning is evaluated with PDMS and EPDMS, which combine safety, compliance, progress, time-to-collision, lane-keeping, history-comfort, and extended-comfort terms.
7.3 Analytical Results
The analysis explains when a latent skeleton can reduce dependence among future tokens and how this supports hierarchical policy modeling and soft-label action reconstruction. It also decomposes generation risk into prediction, initialization, model, and residual-dependence terms.
- Latent Skeleton Decomposition: A latent skeleton reduces residual token dependence only when its redundancy gain is positive, such as when it acts as an upstream common-cause variable.Downstream colliders or synergistic summaries may fail to reduce dependence.
- Positive Redundancy Gain under Residual Mixing: For trajectory planning, the skeleton captures shared low-frequency structure while fine tokens model local residual corrections.The condition may fail if the skeleton is difficult to predict or does not reduce residual dependence.
- KL Upper Bound with Model Error and Re-edit Schedule: The full KL bound combines skeleton prediction error, initial proposal mismatch, token-level model error, and schedule-dependent residual dependence.The bound is expressed as DKL(q(Y | C)∥pψ,θ,π(Y | C)) ≤ δZ + δinit + Bmodel(π) + Udep(π).
- Policy Modeling: The hierarchical policy factorization is justified only when decision-token prediction reduces residual action dependence enough to offset its own prediction error.The decision token represents upstream planning structure such as maneuver intent, target lane, or speed trend.
- Exact Reconstruction under Soft-label Interpolation: Soft-label interpolation removes deterministic hard-quantization error within an acceleration grid cell under exact recovery of the target distribution.With imperfect distribution prediction, the remaining reconstruction error is controlled by prediction error rather than nearest-bin quantization.
- Acceleration Reconstruction: Mode-aware decoding trades unbiased full-distribution expectation for physically consistent mode-conditioned acceleration in multimodal action prediction.It is not claimed to be error-free and helps prevent regression losses from forcing categorical predictions toward their global mean.
7.4 Additional Qualitative Results
Additional qualitative results show selective scheduling progressively stabilizes action predictions, while world-model generations preserve scene structure and surprise responds more sharply than pixel error to unsafe counterfactuals.
- Planning: Selective re-editing reduces action-token entropy across rounds while preserving tokens that have already become confident.Full replacement can perturb stable tokens and degrade long-horizon trajectory accuracy.
- World modeling: Discrete-WAM generates temporally coherent future observations that preserve static scene elements while modeling surrounding-vehicle motion and ego-induced viewpoint changes.Examples span intersections, urban roads, highway ramps, underpasses, and open-road environments.
- Counterfactual inference: Surprise increases monotonically with lateral counterfactual severity, while planning scores remain stable for plausible perturbations and drop sharply for unsafe interactions.Pixel-level L1 error rises more gradually than surprise, indicating sensitivity to semantically meaningful deviations.
- Attention map visualization: Supplementary attention visualizations examine camera-averaged maps, layer-wise front-view attention, and upper-region ablations across additional examples.They extend the analysis of driving-relevant semantics and upper-region activation beyond a single example.