Skip to content
Bud's Wiki

AdaJEPA: An Adaptive Latent World Model

TL;DR

Section titled “TL;DR”

AdaJEPA is a latent world model (JEPA-family) that performs self-supervised test-time adaptation inside the MPC loop. After standard training, the planner executes the first action chunk, uses the resulting next-state transition as a self-supervised signal to update the world model on the fly, and then replans. This closes the loop between prediction error and model recalibration without requiring expert demonstrations, and reportedly improves planning success across goal-reaching tasks with as little as one gradient step per replan. It’s from the LeCun / Mengye Ren line, extending the JEPA world-model program into the test-time-adaptation regime.

Key claims

Section titled “Key claims”
  • Frozen latent world models degrade under test-time distribution shift, and this failure is diagnosable through prediction error observed after each MPC action chunk [Abstract].
  • Test-time adaptation using the observed next-state transition as a self-supervised target recalibrates the world model without expert data or reward relabeling [Abstract].
  • AdaJEPA closes the MPC loop: plan → execute first chunk → adapt on observed transition → replan with updated model [Abstract].
  • Substantial gains in planning success on goal-reaching tasks are reported with as few as one gradient step per MPC replanning step [Abstract].

Method

Section titled “Method”

The architecture follows the JEPA recipe — encoder + latent-space predictor trained to match the embedding of the next observation from the current embedding and action. Planning is standard sampling-based MPC in the latent space: propose action chunks, roll out the learned predictor, score against a goal latent, execute the first chunk of the winner. The novelty is what happens between chunks. Once the first chunk executes, the actually-observed next state is encoded and compared to what the predictor had produced; the discrepancy is used as a self-supervised loss to update the predictor (and possibly the encoder) with a small number of gradient steps. The next MPC replan then uses the freshly-adapted model. Because the adaptation signal is just the next-state transition, no additional labeling or demonstration is needed — the closed loop supplies its own supervision.

Results

Section titled “Results”

The abstract commits to one headline: “substantially improves planning success” on goal-reaching tasks with as few as one gradient step per replan. Specific environments, tables, and comparisons are not in the fetched abstract; the full body is behind the PDF. The claim to track once numbers surface is the size of the adaptation gain relative to (a) the frozen JEPA baseline and (b) offline fine-tuning on trajectories from the shifted distribution — the latter being the natural upper bound.

Why it’s interesting

Section titled “Why it’s interesting”

AdaJEPA is a direct answer to one of the open sub-problems named on the World Foundation Models page: how to make an action-conditional world model useful under distribution shift, without paired action–observation pretraining data covering the shift. Rather than solving the chicken-and-egg by scaling paired data (Waymo’s approach) or by burning the observed context as in-context tokens (In-Context World Modeling for Robotic Control‘s ICWM), AdaJEPA writes the adaptation signal into the parameters through gradient steps at test time. The obvious parallel is End-to-End Test-Time Training for Long Context — long-context language modeling as continual learning over a sliding window — same recipe transferred from token prediction to latent dynamics prediction.

Complements Diffusion Model Predictive Control (D-MPC) which uses factorized diffusion dynamics + sample-score-rank inside MPC but keeps the model frozen at deployment. AdaJEPA proposes the missing axis: the dynamics model itself should update. Also complements Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning: QGF pushes a value gradient into the flow-policy sampler at test time; AdaJEPA pushes a prediction-error gradient into the world-model weights at test time. Both are in the LeCun/Levine tradition of “put the signal at inference rather than in training,” applied respectively to the policy and to the dynamics prior. The LeCun co-authorship also situates this squarely as a follow-on to LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels (LeWorldModel) — the end-to-end-stable pixel-JEPA — and to V-JEPA 2.1: Unlocking Dense Features in Video Self-Supervised Learning on the representation side.

See also

Section titled “See also”