Autonomous driving world model speeds planning with fewer patches
AD-E2E-JEPA: A Joint-Embedding Predictive Architecture For End-to-End Autonomous Driving
RoboticsArtificial IntelligenceComputer Vision and Pattern RecognitionMachine Learning
Summary
Autonomous cars need to predict what will happen around them to plan their actions safely. This paper looks at existing prediction models and finds a trade-off between accuracy and speed. The authors created a new model called AD-E2E-JEPA that keeps accurate predictions but runs much faster by simplifying the information it processes. This fast model can plan many possible driving paths quickly without needing to train a driving policy first.
What this means in practice
- •For autonomous vehicle engineers: Use fast and accurate world models to plan safe driving paths in real time without extra policy training.
- •For robotics software developers: Integrate computationally efficient predictive models to improve planning speed and accuracy in navigation tasks.
Authors
Haoran Zhu, Wancong Zhang, Yann LeCun, Anna Choromanska
Abstract
Autonomous driving requires \textit{world models} that can understand the physical world, reason and plan, and operate safely. In this paper, we first systematically evaluate existing action-conditioned joint-embedding predictive architecture (JEPA) world models, including LeWM, DINO-WM, and JEPA-WM for end-to-end autonomous driving (E2EAD). To isolate world-model quality from policy learning, we employ a goal-conditioned zero-shot planning setting that evaluates these models using ground-truth future observations as goals, without training any driving policy. We find that existing JEPA-based world models are either accurate for driving but computationally expensive, or computationally efficient but insufficient for planning. To address this trade-off, we propose \textbf{AD-E2E-JEPA}, which introduces a SIGReg-regularized learnable projector applied to projected patch embeddings. The projector reduces the number of planning patches by $16\times$ and the embedding dimension by $4\times$, achieving a $100\times$ inference speedup while retaining planning performance, with a 0.8-second runtime for an 8-frame rollout over 256 candidate trajectories. \textit{Without} training any driving policy, the world model itself reaches the goals located 20 meters away on average within the displacement of respectively 4.0/2.8 meters, using world-model rollouts over trajectory vocabularies of respectively 256/8,192 candidates. On the NAVSIMv2 benchmark, it achieves 67.3/72.9 EPDMS with multiplicative safety metrics and 84.1/86.5 EPDMS$^{\dagger}$ without them in goal-conditioned zero-shot planning. Experiments further show that the self-supervised pretrained projector improves downstream imitation learning performance from 80.2 to 85.4 EPDMS. The source code is available at https://github.com/HaoranZhuExplorer/AD-E2E-JEPA