Source-linked AI summary
WA-JEPA: Rethinking the Video JEPA Paradigm for World-Action Modeling in Autonomous Driving
Xinlin Wang, Yujiao Xiang, Yuheng Zhou, Jingqi Wang, Minqing Huang, Jiajie Huang, Dongxu Wei, Tingguang Zhou, Xiyang Wang, Gong Chen, Zhi Xu, Feiyang Tan, Hangning Zhou, Mu Yang
TL;DR
Autonomous-driving planning requires future-directed prediction tightly coupled with action, whereas V-JEPA's random-mask completion and regression objective are poorly suited to that need. WA-JEPA introduces future-masked pre-training, flow-based latent prediction, and joint future-action modeling, achieving strong open- and closed-loop benchmark results.
Problem
V-JEPA's random spatiotemporal completion and regression objectives are ill-suited to generating unseen futures and supporting interactive planning.
Method
WA-JEPA combines hybrid future masking, flow matching over future latents, and joint world-action modeling in a unified spatiotemporal latent space.
Results
91.7 EPDMS on NAVSIM-v2 navtest and 0.4462 HD-Score on 436 HUGSIM closed-loop scenarios demonstrate strong open-loop planning and zero-shot closed-loop generalization.
Takeaways & Limitations
WA-JEPA extends V-JEPA 2 from visual representation learning to future-predictive planning for autonomous driving.
Takeaways & Limitations
Existing V-JEPA action-conditioned approaches require a goal image and multi-round MPC optimization to recover actions, limiting interactive online planning.
Abstract
from arXiv · showhide
Video Joint Embedding Predictive Architecture (V-JEPA) learns powerful spatiotemporal representations from video through self-supervised latent feature prediction. However, V-JEPA is built around random-mask completion and deterministic regression, making it fundamentally ill-suited for autonomous driving planning that demands future-directed prediction tightly coupled with action. To address this, we rethink the V-JEPA paradigm and present WA-JEPA, a V-JEPA-native world-action model designed for autonomous driving planning. Instead of random spatiotemporal masking, WA-JEPA employs hybrid future-masked pre-training, where the model infers future latents from observed context. Departing from deterministic regression, we recast future prediction as conditional flow matching over latent futures, which substantially improves the model's ability to generate plausible future latents for downstream planning. Finally, a joint future-action predictor is proposed to denoise future scene tokens and ego trajectories together in a unified spatiotemporal latent space, allowing action supervision to directly shape planning-relevant world representations. Pre-trained on nuPlan videos and fine-tuned on NAVSIM, WA-JEPA reaches 91.7 EPDMS on NAVSIM-v2, surpassing the strongest end-to-end and world-action baselines by 1.6 and 1.3 EPDMS, and, without HUGSIM-specific fine-tuning, attains the best HD-Score of 0.4462 on the closed-loop HUGSIM benchmark under the same evaluation protocol. These results validate V-JEPA-native world-action modeling as a powerful and scalable paradigm for autonomous driving planning. Code is available at https://github.com/AFARI-Research/WA-JEPA.
Introduction
Autonomous-driving models need future-directed, action-relevant reasoning, but existing approaches face supervision bottlenecks or mismatches between representation learning and planning. WA-JEPA adapts V-JEPA with future masking, flow-based latent prediction, and joint world-action modeling, establishing strong benchmark performance.
- VLA methods map dense visual inputs to sparse actions, creating a supervision deficit that poorly constrains learning.
- WAMs use future-scene prediction to provide dense self-supervision and guide action decisions without costly perception annotations.
- V-JEPA learns rich spatiotemporal representations, but random-mask completion, deterministic regression, and indirect action recovery limit its suitability for planning.
- WA-JEPA uses hybrid future masking so observed past frames predict future spatiotemporal features, aligning pre-training with forward-predictive planning.
- Flow matching generates future latent representations, while joint world-action modeling couples visual and action representations in one spatiotemporal latent space.
- WA-JEPA establishes state-of-the-art results on open-loop NAVSIM benchmarks and transfers gains to zero-shot closed-loop HUGSIM evaluation.
Related Work
World-action models augment direct planning with future-scene prediction, while related approaches differ in their world representations and coupling between scene prediction and action generation. JEPA-based methods extend this landscape through predictive latent representations rather than pixel reconstruction.
- E2E models map sensor observations directly to ego motion, whereas WAMs augment planning with learned future-scene prediction.
- Coupled video-based WAMs jointly generate future visual content and actions, but inherit limitations of video-generative latent spaces.
- Predictive-representation methods model future latent scenes or features for trajectory planning instead of relying on pixel-level reconstruction.
Method
WA-JEPA adapts V-JEPA 2 for causal future prediction in multi-view driving videos, then jointly predicts future scene representations and ego actions. Its two-stage design combines hybrid masking, latent flow matching, and action-supervised joint modeling.
- Two-stage training: WA-JEPA follows two training stages: future-representation pre-training, followed by action-supervised joint prediction of future scenes and ego actions.Stage 2 initializes from the Stage 1 checkpoint.
- Stage 1: Hybrid Future-Masked Pre-training: Stage 1 combines causal full-future masking with patch-masked future completion on synchronized multi-view driving observations.Full masking predicts future representations from historical context alone, while patch masking retains some visible future patches.
- Stage 1: Hybrid Future-Masked Pre-training: The visual backbone is a pretrained V-JEPA 2 ViT-L that processes synchronized camera streams independently with a shared online encoder.Historical frames remain visible, while masking is applied only to future tokens.
- Latent future prediction: Conditional flow matching predicts clean future scene tokens from historical context, future conditions, noisy future tokens, and a temporal condition.The future flow predictor uses an MMDiT-style design with joint attention between context and future scene tokens.
- Stage 2: Joint World-Action Modeling: Stage 2 adds an action stream and jointly models future scene representations and actions within a unified predictor initialized from Stage 1.Historical actions and compact ego state provide ego-motion conditions, while future images supply supervision rather than student inputs.
- Stage 2: Joint World-Action Modeling: The joint predictor denoises normalized future actions with an MSE objective while preserving the future-scene prediction objective during action-supervised fine-tuning.Stop-gradient is applied to action tokens in the future-scene output stream.
Experiments
WA-JEPA is evaluated on NAVSIM and HUGSIM against established autonomous-driving paradigms, with ablations testing its encoder, masking strategy, future prediction, and scene–action coupling. It achieves strong open-loop and zero-shot closed-loop results, while flow matching and hybrid masking provide the strongest ablation outcomes.
- Benchmark Results: WA-JEPA reaches 91.7 EPDMS on NAVSIM-v2, exceeding SparseDriveV2 by 1.6 and Discrete-WAM by 1.3 EPDMS.Evaluation uses the NAVSIM-v2 navtest split and official EPDMS scoring.
- Benchmark Results: 0.4462 HD-Score is achieved on HUGSIM closed-loop evaluation, improving from 0.3252 and attaining the best NC, DAC, TTC, PDMS, RC, and HD-Score.The comparison covers zero-shot evaluation under a common protocol across 436 HUGSIM scenarios.
- Ablation Studies: V-JEPA 2 initialization reaches 89.5 EPDMS, exceeding MAE and DINOv3 by 5.7 EPDMS when Stage 2 architecture, data, and optimization are held constant.The reported gap is attributed to the V-JEPA 2 pre-training objective rather than differences among image-level objectives.
- Ablation Studies: 91.7 EPDMS is obtained by combining Patch-mask and Full-mask, outperforming Patch-mask alone by 0.7 and Full-mask alone by 0.4.The ablation reports 91.0 for Patch-mask and 91.3 for Full-mask, indicating complementary training signals.
- Ablation Studies: Flow-based future prediction with joint modeling reaches 91.7 EPDMS, compared with 90.7 for direct regression and 91.1 without explicit future-latent supervision.A separate flow-based predictor with cross-attention reaches 90.8, while the cascaded historical-latent baseline reaches 89.9.
- Ablation Studies: Flow matching reduces the directional-similarity collapse gap from 0.30 to 0.10 and raises change-magnitude collapse from 0.45 to 0.80 versus direct regression.The metrics evaluate preservation of target temporal dynamics; qualitative visualization also reports clearer spatial structures over the prediction horizon with flow matching.
Conclusion
WA-JEPA extends V-JEPA 2 into future-predictive autonomous-driving planning by jointly modeling future scene dynamics and actions in a shared spatiotemporal latent space. It achieves strong open-loop planning and zero-shot closed-loop generalization.
- WA-JEPA combines hybrid future-masked pre-training, flow-matching future-latent prediction, and joint future-action modeling.These components support shared modeling of future scene dynamics and ego actions.
- 91.7 EPDMS on NAVSIM-v2 navtest demonstrates WA-JEPA’s open-loop planning performance.
- 0.4462 HD-Score across 436 HUGSIM closed-loop scenarios demonstrates zero-shot closed-loop generalization.
A. HUGSIM Closed-Loop Evaluation
The HUGSIM evaluation uses a fixed 436-scenario protocol with consistent controllers, metrics, and aggregation, while reporting source-disjoint performance across datasets. WA-JEPA achieves the highest HD-Score on every source dataset.
- Evaluation protocol: HUGSIM evaluation uses 436 scenarios, shared ground-truth commands, aggregation, controller, evaluator, and a fixed code snapshot.The protocol applies a trajectory-to-heading coordinate-order correction and rescoring for consistent implementations.
- Aggregation: The overall HD-Score weights four difficulty levels by their scenario counts, with alternative dataset-uniform and scenario-uniform rules also reported.The primary weighting uses counts of 80, 157, 96, and 103 scenarios.
- Per-dataset results: WA-JEPA achieves the highest HD-Score on every HUGSIM source dataset despite none being used for either training stage.This indicates source-disjoint generalization across diverse visual domains.
B. Qualitative Results
Qualitative evaluations show WA-JEPA maintaining safe, geometrically coherent behavior in closed-loop HUGSIM rollouts and matching reference maneuver structure in NAVSIM predictions.
- HUGSIM closed-loop rollouts: WA-JEPA follows drivable corridors and maintains lateral clearance from reactive agents across four HUGSIM source datasets.The zero-shot rollouts cover turning, oncoming-vehicle encounters, and overtaking despite a visual domain gap.
- NAVSIM trajectory predictions: NAVSIM trajectories agree with references in maneuver direction and overall geometry across turning, fork, stopping, and straight-driving scenarios.Only minor local deviations are reported, including gateway navigation cases.
C. Additional Experimental Details
WA-JEPA uses asymmetric gradient flow in its joint predictor and evaluates stochastic flow sampling with fixed inference settings and multiple seeds. The action stream receives supervision that can shape scene representations.
- Gradient design: Scene-loss gradients are blocked at the action-token interface, preventing scene prediction from updating the action stream.The scene output stream still conditions on action tokens during the forward pass.
- Gradient design: The action output stream attends to differentiable historical context and future scene tokens, allowing action supervision to shape scene representations.
- Inference details: The flow predictor uses 12 sampling steps and ten fixed seeds for stochastic evaluations, with metrics averaged arithmetically before rounding.Model parameters, scenarios, and inference settings remain unchanged across seeds.
- Seed variability: The main experiment’s mean EPDMS is 91.7014, reported as 91.7 after rounding.
D. NAVSIM Evaluation Metrics
NAVSIM evaluation uses safety, compliance, comfort, and progress sub-metrics aggregated into PDMS or EPDMS under the official protocol. EPDMS additionally applies a human-reference penalty filter before aggregation.
- Evaluation protocol: Evaluation uses the official NAVSIM protocol and computes scores with the pseudo-simulator on navtest.
- Sub-metrics: NAVSIM evaluates no-atfault collision, drivable-area compliance, time-to-collision, and additional compliance and comfort criteria.NAVSIM-v2 adds driving-direction, traffic-light, lane-keeping, history-comfort, and extended-comfort measures.
- Aggregation: A single safety violation drives the PDMS score to zero, while remaining terms trade off safety margin, comfort, and progress.
- Aggregation: The EPDMS aggregation applies a human-reference penalty filter before combining the sub-metrics.The filter suppresses a metric-specific penalty when the human reference incurs the same violation.
- Evaluation protocol: EPDMS and EPDMS* are reported separately because they use corrected and pre-fix evaluators, respectively, and are compared only within the same protocol.
E. Temporal Representation Metrics
Temporal representation metrics assess whether predicted future scene tokens preserve target-token variation. They measure directional similarity collapse and change-magnitude collapse in a shared projected feature space.
- Metric setup: Temporal metrics compare predicted future scene-token features with EMA target features in the projected scene-token space.
- Dynamic-token selection: Dynamic-token selection ranks locations by target mean adjacent-step change to prevent static regions from dominating the statistics.
- Dynamic-token selection: The same target-selected locations are used for both methods and both metrics, with selection performed independently for each prediction instance.
- Directional similarity collapse gap: Directional similarity collapse gap measures excess mutual similarity among predicted future steps; lower values indicate better temporal preservation.
- Change-magnitude collapse: Change-magnitude collapse is a ratio where R_∆ = 1 preserves target average temporal change, and values below 1 indicate under-variation.
- Evaluation protocol: Flow-matching metrics use one-step x-predictions at sampled training flow times, whereas regression uses direct latent predictions.