DA-WAM: Decision-Aligned Future Latents for Driving World Models

Ruiguo Zhong; Xiaolong Chen; Benshan Ma; Lang Zhang; Mingyue Feng; Yaonong Wang; Pei Liu; Jun Ma

cs.ROAug 19, 2026

DA-WAM: Decision-Aligned Future Latents for Driving World Models

Authors: Ruiguo Zhong, Xiaolong Chen, Benshan Ma, Lang Zhang, Mingyue Feng, Yaonong Wang, Pei Liu, Jun Ma

Organizations: The Hong Kong University of Science and Technology (Guangzhou) · Leapmotor · The Hong Kong University of Science and Technology

Abstract

World models for autonomous driving should anticipate the consequences that matter for driving decisions rather than reconstruct every detail of future scenes. However, predictive representations alone are not guaranteed to distinguish candidate actions that lead to different planning outcomes.We propose DA-WAM, a decision-aligned world action model that couples a Joint-Embedding Predictive Architecture (JEPA) with candidate-level trajectory scoring. During joint training, the JEPA objective grounds future representations in observed scene dynamics, while the planning objectives encourage these representations to capture distinctions relevant to safety and utility. An action-conditioned predictor generates a distinct future latent for each candidate trajectory, and a multi-head trajectory scorer uses the corresponding latent to estimate planning factors and overall utility. In addition, safety-critical post-training with hard negatives further improves discrimination between geometrically similar trajectories with different safety outcomes. DA-WAM achieves state-of-the-art performance among learned camera-based planners on NAVSIM-v1, NAVSIM-v2, and Bench2Drive. Ablation studies demonstrate the effectiveness of each component in driving decision-making.

Figures & tables

Figure 1: Candidate–representation coupling in driving decisions. Schematic comparison of (a) trajectory prediction without an explicit future-representation input to the scorer, (b) joint trajectory and latent prediction without explicit candidate-wise pairing, (c) multiple candidates evaluated using a shared future representation, and (d) DA-WAM, which computes a candidate-conditioned representation for each trajectory and scores the corresponding trajectory–representation pair. The scorer receives the trajectory together with its associated representation and scene context. The panels illustrate conditioning structures rather than an exhaustive classification of prior work.

Figure 2: Overview of DA-WAM. The online visual encoder produces shared scene tokens. A shared action-conditioned predictor constructs compact tokens for each candidate trajectory, and the scorer evaluates each candidate together with its associated tokens and scene context. A training-only dense decoder and a stop-gradient EMA target provide future-feature supervision for the expert-matched candidate.

Figure 3: Safety-focused supervision with retrieved trajectory alternatives. A planning evaluator assigns utility and safety labels to candidate trajectories. Compared with generated-candidate supervision alone (a), the augmented training set (b) includes expert-proximate alternatives with different collision and time-to-collision outcomes. Generated and retrieved trajectories use the same candidate-conditioned predictor and scorer. Retrieved alternatives are supervised by planning labels.

Method	Venue	NC	DAC	TTC	Comfort	EP	PDMS
PDM-Closed Dauner et al. (2023)	CoRL’23	94.6	99.8	89.9	86.9	99.9	89.1
Human driver Dauner et al. (2024)	NeurIPS’24	100.0	100.0	100.0	99.9	87.5	94.8
Ego-stat. MLP Dauner et al. (2024)	NeurIPS’24	93.0	77.3	83.6	100.0	62.8	65.6
UniVLA Wang et al. (2026d)	ICLR’26	96.9	91.1	91.7	96.7	76.8	81.7
VAD-v2 Jiang et al. (2026)	ICLR’26	98.1	94.8	94.3	100.0	80.6	86.2
PRIX Wozniak et al. (2026)	RA-L’26	98.1	96.3	94.1	100.0	82.3	87.8

Table 1: NAVSIM-v1 benchmark comparison.

Method	Img. Backbone	NC	DAC	DDC	TLC	EP	TTC	LK	HC	EC	EPDMS
Ego Status MLP	ResNet-34	93.1	77.9	92.7	99.6	86.0	91.5	89.4	98.3	85.4	64.0
TransFuser Chitta et al. (2022)	ResNet-34	96.9	89.9	97.8	99.7	87.1	95.4	92.7	98.3	87.2	76.7
Hydra-MDP++ Li et al. (2025a)	ResNet-34	97.2	97.5	99.4	99.6	83.1	96.5	94.4	98.2	70.9	81.4
DriveSuprim Yao et al. (2026)	ResNet-34	97.5	96.5	99.4	99.6	88.4	96.6	95.5	98.3	77.0	83.1
ARTEMIS Feng et al. (2026)	ResNet-34	98.3	95.1	98.6	99.8	81.5	97.4	96.5	98.3	98.3	83.1
DiffusionDriveV2 Zou et al. (2025)	ResNet-34	97.7	96.6	99.2	99.8	88.9	97.2	96.0	97.8	91.0	87.5

Table 2: NAVSIM-v2 benchmark comparison.

Method	Avg. L2 (m) $\downarrow$	Efficiency $\uparrow$	Comfortness $\uparrow$	Success Rate (%) $\uparrow$	Driving Score $\uparrow$
TCP-traj w/o distillation Wu et al. (2022)	1.96	78.78	22.96	20.45	49.30
AD-MLP Zhai et al. (2023)	3.64	48.45	22.63	0.00	18.05
UniAD-Tiny Hu et al. (2023)	0.80	123.92	47.04	13.18	40.73
UniAD-Base Hu et al. (2023)	0.73	129.21	43.58	16.36	45.81
VAD Jiang et al. (2023)	0.91	157.94	46.01	15.00	42.35
GenAD Zheng et al. (2024)	–	–	–	15.90	44.81

Table 3: Bench2Drive closed-loop comparison.

Configuration	Safety sup.	PDMS	NC	DAC	EP	TTC	Comfort
No Future Prediction	✗	93.31	98.45	98.27	91.36	95.48	99.99
Shared Global Future	✗	92.81	99.02	98.46	88.68	96.54	99.99
Current-Latent Conditioning	✗	93.25	98.44	98.19	91.38	95.49	99.94
Action-Conditioned Future	✗	93.46	98.88	98.58	90.47	96.33	99.69
Action-Conditioned Future	✓	93.68	99.11	98.88	89.97	96.81	99.77

Table 4: Ablation on action-conditioned futures and safety-focused refinement.

Figure 4: Safety–progress trade-offs during refinement. Epoch 0 denotes Stage 1, and epochs 1–3 denote completed refinement epochs; circles mark the retained checkpoint.

Table 9

Figure 5: Future-feature evolution in a truck scene. Predictions from the 1.5-second visualization checkpoint. Each column shows an independent prediction.

Figure 6: Safety-critical attention and safer decision making.

Figure 7: Direction Following and Safety-Critical Interaction. Camera views, bird’s-eye-view trajectories, and per-scene metrics compare DiffusionDrive, DrivoR, and DA-WAM in a large left turn, tight traffic, and a yielding conflict.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Spatial reorganization under LoRA adaptation. Identical input frames are processed with the learned adapters disabled and enabled in the same trained model. Attention maps show a redistribution of spatial response, including clearer local responses near several annotated vehicles and signals; the color scale is relative to uniform attention. Shared PC1 projections show complementary changes in feature organization while preserving broader scene context.

Figure 9: Future-feature prediction during a crossing interaction. Each column is an independent 1.5-second-ahead prediction. During the transition, some predicted structures remain closer to the current map than to the target, whereas the $t=5$ prediction captures part of the target’s broader shift toward lower PC1 values. The dense branch therefore follows some coarse future changes but incompletely recovers local and transitional structure.

Figure 10: Future-feature reorganization during a turn. Each column is an independent 1.5-second-ahead prediction. At $t=2$ , a dominant higher-valued region shifts from the right in the current map to the left in both prediction and target.

Backbone	E2E	Frozen
ViT-B ( vitb )	86.9958	85.1258
ViT-L ( vitl )	88.1473	85.5403
ViT-g ( vitg )	87.1119	84.6653
ViT-G ( vitG )	87.9113	81.4653

Appendix

Table 9: Backbone comparison in the preliminary planning baseline. E2E and Frozen denote end-to-end adaptation and a fixed encoder. ViT-g and ViT-G are distinct checkpoint variants. PDMS is reported on a 0–100 scale.

Method	Candidates	Forward time (ms) $\downarrow$
DiffusionDrive Liao et al. (2025)	–	24.68
DrivoR Kirby et al. (2026)	–	295.56
	1	38.11
	4	36.89
	8	36.52
	16	39.00

Appendix

Table 10: Inference latency and candidate-count scaling. Reported forward-pass time in milliseconds; lower is better.

Symbol	Meaning
$X_{t}$ , $X_{t+\Delta}$	Current and observed future visual inputs.
$u_{t}$	Ego-state input used by the trajectory generator.
$\mathcal{T}=\{\tau_{i}\}_{i=1}^{N}$	Set of $N$ candidate ego trajectories.
$\tau^{\mathrm{exp}}$ , $\tau_{j}^{-}$ , $\tau^{\star}$	Expert, hard-negative, and finally selected trajectories.
$Q_{i}$	Native pose-token sequence produced with candidate $\tau_{i}$ .
$E_{\theta}$ , $E_{\bar{\theta}}$	Online encoder and EMA target encoder.

Appendix

Table 11: Notation. Symbols used in the main formulation and implementation details. Compact candidate-conditioned tokens

\widehat{z}_{i}

are distinguished from dense future-feature predictions

\widetilde{Z}_{i^{\mathrm{exp}}}

Parameter	Value
Visual encoder	V-JEPA 2.1
LoRA rank	32
LoRA target projections	Query and value
Generated candidates	32
External counterfactual candidates	128
Future prediction horizon	$0.5\,\mathrm{s}$

Appendix

Table 12: Default architecture and Stage 1 training configuration. Unless stated otherwise, the model uses 32 generated candidates, rank-32 LoRA, and a 0.5-second predictive target. External candidates are used only for training supervision.

Loss	Stage 1 weight	Stage 2 weight	Supervision target	Updated modules
Future-feature regression	$0.1$	$0$	Expert-matched generated candidate	Stage 1 dense decoder, compact predictor, encoder adapters, and proposal path through native pose tokens
Factor prediction	$1.0$ generated; $\lambda_{\mathrm{ext}}(e)$ external	$1.0$ generated; $0.5$ external	Five planning factors for all generated and sampled external candidates	Generated Stage 1 path: compact branch, scorer, encoder adapters, and proposal generator through $Q_{i}$ ; external Stage 1 and all Stage 2 paths: compact branch and scorer
Utility prediction	$1.0$ generated; $\lambda_{\mathrm{ext}}(e)$ external	$1.0$ generated; $0.5$ external; unsafe BCE $\times 4$	Learned-utility target for all generated and sampled external candidates	Generated Stage 1 path: compact branch, scorer, encoder adapters, and proposal generator through $Q_{i}$ ; external Stage 1 and all Stage 2 paths: compact branch and scorer
Safe–unsafe ranking	$0$	$0.1$ generated; $0.05$ external	Highest-target safe candidate versus unsafe candidates among predicted Top- $8$	Stage 2 compact branch and scorer
Inherited proposal loss	$1.0$	$0$	Generated trajectories against the expert trajectory	Proposal generator and its trainable upstream path
Inherited auxiliary losses	$1.0$ outer weight	$0$	Agent and area targets from the base planner	Corresponding base-planner modules

Appendix

Training-loss mapping across the two stages. The Stage 1 external weight

\lambda_{\mathrm{ext}}(e)

follows the warm-up schedule in Eq. 9 .

Parameter	Value
Learning rate	$10^{-6}$
Post-training epochs	3
Retained checkpoint	Epoch 1 (zero-based)
Global batch size	64
Generated factor/utility loss weight	1.0
External factor/utility loss weight	0.5

Appendix

Table 13: Stage 2 safety-refinement configuration. The visual encoder, proposal generator, dense decoder, and EMA target are fixed; only the compact predictor and scorer are optimized. Three epochs are run, the model after two completed epochs is retained, and the future-feature loss has zero weight.

World Models for Autonomous Driving

World models for autonomous driving predict how scenes and traffic evolve under candidate actions to support planning, control, and simulation.

See all 142 papers

PlanWAM: Planning-Shaped Future Representations for End-to-End Autonomous Driving
Oct 8, 2026Jinchang Xu, Hongda Yu, Fengwei Dong +9End-to-End Autonomous Driving World Model-Based Planning
Southeast University · NIO · Nanyang Technological University +1
DriveFuture: Future-Aware Latent World Models for Autonomous Driving
May 10, 2026Yufeng Hong, Xiaotian Zhou, Yingyan Li +6NavSim World Models for Autonomous Driving
Beijing Institute of Technology · Institute of Automation, Chinese Academy of Sciences · Beihang University +5
Latent-WAM: Latent World Action Modeling for End-to-End Autonomous Driving
Mar 25, 2026Linbo Wang, Yupeng Zheng, Qiang Chen +13End-to-End Autonomous Driving World Models for Autonomous Driving
Institute of Automation, Chinese Academy of Sciences · Chongqing Chang’an Technology Co., Ltd · Zhongguancun Academy +2

Oct 8, 2026J/K move · Enter open · S save

PlanWAM: Planning-Shaped Future Representations for End-to-End Autonomous Driving

Jinchang Xu, Hongda Yu, Fengwei Dong, Wenhui Huang+8

Southeast University · NIO · Nanyang Technological University+1

End-to-End Autonomous Driving World Model-Based Planning World Action Models Predictive Representation Learning

Open paper AskPDF arXiv

World models in end-to-end autonomous driving predict future scene evolution to provide foresight for trajectory planning. Existing methods mainly study how to predict the future and how to use it, but less often ask which future representation is actually most useful for planning. To this end, we propose PlanWAM, a Planning-Shaped World Action Model. The key idea is to let the planning task shape the future-state representation, so that it retains the information most useful for planning. A latent world model then predicts this planning-shaped future latent representation from historical observations and uses it for planning, enabling foresighted planning. Specifically, we first use a Temporal Register Pyramid to compress multi-frame historical information in a recency-aware manner, learning a compact history representation oriented toward future reasoning and planning. We then introduce a privileged future posterior branch that observes ground-truth future frames, and shape its future latent representation with trajectory-planning objectives to obtain a planning-shaped future latent representation. Hindsight-to-Foresight Distillation trains a prior branch that depends only on history to predict this future latent representation. The predicted future latent representation serves as planning context and guides trajectory generation and selection. PlanWAM achieves 93.8 PDMS / 90.9 EPDMS on NAVSIM-v1/v2 navtest and reaches 38.7 HD-Score on closed-loop HUGSIM in a zero-shot setting, demonstrating leading planning performance across both open-loop and closed-loop evaluations. Extensive experiments further demonstrate that planning-shaped future representations provide an effective and deployable form of foresight for world-action models.

Abstract

World Models for Autonomous Driving

Explore similar work

PlanWAM: Planning-Shaped Future Representations for End-to-End Autonomous Driving

DriveFuture: Future-Aware Latent World Models for Autonomous Driving

Latent-WAM: Latent World Action Modeling for End-to-End Autonomous Driving