WALT: Learning World-Model-Aligned Latent Trajectories for Autonomous Driving

Mingkai Jia, Jiaxin Guo, Zhijian Shu, Jiawei Xu, Mingxiao Li, Jintao Cheng, Ping Tan, Wei Yin · N/A · 2026

Summary

Abstract Summary

Driving world models learn rich predictive representations of the surrounding environment from visual observations, yet accurate visual prediction does not necessarily translate into effective trajectory planning. We argue that a key bottleneck lies in the mismatch between visual world states and raw geometric trajectories, which may limit the planner's ability to exploit action-relevant semantics encoded by the world model. To address this issue, we propose World-Model Alignment for Latent Trajectories (WALT), which learns a compact generative trajectory latent space by transferring information from a frozen pretrained driving world model without modifying the world model itself. Rather than directly generating raw waypoints, WALT maps them into compact representations through a dual-branch trajectory autoencoder and transfers semantic knowledge from the frozen visual world model into this trajectory space, encouraging the learned action representation to capture scene-level cues relevant to future motion and planning. Beyond our proposed formulation, we systematically study latent learning based on Joint-Embedding Predictive Architectures (JEPA) and feature alignment following Representation Alignment (REPA) to investigate how trajectory-only representation learning affects downstream planning. We evaluate WALT on the NAVSIM benchmarks. Relative to the raw-waypoint baseline, WALT improves PDMS from 89.4 to 89.8 on NAVSIMv1 and EPDMS from 87.3 to 87.9 on NAVSIMv2 while reducing trajectory planner FLOPs by 30.5%. These results suggest that preserving world representations while extracting action-relevant information provides an effective interface for world-model-based trajectory planning.

Key Points

Driving world models learn rich predictive representations of the surrounding environment from visual observations, yet accurate visual prediction does not necessarily translate into effective trajectory planning.
We argue that a key bottleneck lies in the mismatch between visual world states and raw geometric trajectories, which may limit the planner's ability to exploit action-relevant semantics encoded by the world model.
To address this issue, we propose World-Model Alignment for Latent Trajectories (WALT), which learns a compact generative trajectory latent space by transferring information from a frozen pretrained driving world model without modifying the world model itself.
Rather than directly generating raw waypoints, WALT maps them into compact representations through a dual-branch trajectory autoencoder and transfers semantic knowledge from the frozen visual world model into this trajectory space, encouraging the learned action representation to capture scene-level cues relevant to future motion and planning.
Beyond our proposed formulation, we systematically study latent learning based on Joint-Embedding Predictive Architectures (JEPA) and feature alignment following Representation Alignment (REPA) to investigate how trajectory-only representation learning affects downstream planning.

Method

The WALT pipeline is illustrated in Fig. 2 . In Stage 1, we train a dual-branch trajectory autoencoder to compress raw waypoints into a compact latent representation (Sec. III-B ). The training objective includes a reconstruction loss and a world-model alignment loss that transfers information from the frozen world model into the trajectory latent space (Sec. III-C ).

In Stage 2, we employ the aligned trajectory latent as the generation space for the action head of the frozen world model (Sec. III-D ).

Method — WALT: Learning World-Model-Aligned Latent Trajectories for Autonomous Driving — Fig. 3 : Qualitative trajectories and world-model alignment. Columns show cases for braking, straight driving, turn, and curve. From top to bottom: front-view images, world–trajectory correspondence maps, and predicted trajectories (orange) with ground truth (black). The maps visualize cosine similarity between projected semantic tokens from planner-generated latents and frozen visual world-model tokens, with warmer colors indicating higher similarity.

Experiments

Benchmarks and metrics. We evaluate on NAVSIMv1 navtest using the official planning protocol [ 23 ] . We report PDMS together with no-at-fault collision (NC), drivable-area compliance (DAC), ego progress (EP), time-to-collision (TTC), and comfort (C).

We additionally evaluate on NAVSIMv2 [ 24 ] , reporting EPDMS and its associated component metrics, including traffic light compliance (TLC) and lane keeping (LK), while expanding comfort (C) into history comfort (HC) and extended comfort (EC).

Experiments — WALT: Learning World-Model-Aligned Latent Trajectories for Autonomous Driving — Fig. 4 : Trajectory representation analysis on NAVSIMv1. Columns compare raw waypoints, JEPA-Traj semantic latents, and WALT semantic latents on identical samples. From top to bottom: t-SNE embeddings, within-class and between-class cosine-similarity distributions, and class-pair mean-similarity matrices. Learned representations use only their semantic components.

Sources and Demonstrations

Figures are reproduced from Jia et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.