AdaJEPA: An Adaptive Latent World Model
Ying Wang1,2, Oumayma Bounou1, Yann LeCun1,2*, and Mengye Ren1*
1New York University, 2AMI Labs
*Equal advising
CoRR
TL;DR: AdaJEPA adapts a latent world model inside closed-loop MPC, using each observed transition as a self-supervised signal before the next replan.
Abstract
Latent world models enable planning from high-dimensional observations by predicting future states in a compact latent space. However, these models are typically kept frozen at test time: when their predictions become inaccurate, planning can fail, especially under test-time distribution shift. To address this, we propose AdaJEPA, an adaptive latent world model that performs test-time adaptation within the closed loop of model predictive control (MPC). After training, AdaJEPA plans and executes the first action chunk, uses the observed next-state transition as a self-supervised adaptation signal, and replans with the updated model. This closed-loop update continuously recalibrates the world model without additional expert demonstrations. Across a range of goal-reaching tasks, AdaJEPA substantially improves planning success with as few as one gradient step per MPC replanning step.
Overview
Latent world models make planning from high-dimensional observations tractable by predicting future states in a compact representation space. However, world models are usually kept frozen after training. Inaccurate predictions, especially severe under test distribution shift, can make MPC optimize actions for the wrong imagined future, hindering planning.
AdaJEPA addresses this by adapting the world model during deployment. Each action executed by MPC produces a new transition $(o_t, a_t, o_{t+1})$, which becomes a self-supervised training signal before the next replan. This couples learning and planning in a simple loop: plan, act, adapt, and replan.
Method
AdaJEPA starts from a pretrained JEPA world model with a sensory encoder $\mathcal{E}^s_\phi$, an action encoder $\mathcal{E}^a_\psi$, and a latent predictor $f_\theta$. Given a goal observation $o_g$, MPC plans in latent space by rolling out the predictor and minimizing the distance to the goal representation $z_g = \mathcal{E}^s_\phi(o_g)$:
$$ a^*_{t:t+H-1} = \arg\min_{a_{t:t+H-1}} \sum_{k=1}^{H} \alpha_k\, d(\hat z_{t+k}, z_g). $$After executing the first action, AdaJEPA stores the observed transition in a small online buffer and adapts the model by minimizing the latent prediction objective:
$$ \mathcal{L}_{\rm ada}(\mathcal{B}) = \frac{1}{|\mathcal{B}|} \sum_{(o_i,a_i,o_{i+1})\in\mathcal{B}} \ell\!\left( f_\theta\!\left(z_i,\mathcal{E}^a_\psi(a_i)\right), \operatorname{sg}(z_{i+1}) \right). $$Our proposed adaptation is very lightweight: by default, we use one gradient step per MPC replan, a replay buffer of five recent transitions, and updates restricted to the final layers of the visual encoder and predictor. See the ablations in the paper for alternative adaptation targets, learning rates, update steps, and buffer choices.
Results
AdaJEPA improves planning in both in-distribution and out-of-distribution settings. For in-distribution environments, AdaJEPA improves performance when the frozen model is suboptimal and preserves strong baselines when the frozen model is already near-optimal. Under distribution shift, AdaJEPA gives consistent gains because each observed transition helps recalibrate the model before the next replan.
Across PushObj shape shifts and PushT visual shifts, adaptation improves planning by recalibrating the latent world model to the object or observation stream encountered at test time.
On PointMaze dynamics shifts, the frozen model is already strong, likely because the three-frame history gives it some in-context learning to the current dynamics. AdaJEPA still improves beyond this baseline by updating the world model from the observed transition. On unseen layouts, adaptation improves success and makes trajectories closer to shortest paths.
Data scaling improves both frozen and adaptive models, but test-time adaptation is especially valuable when offline data is limited. On low-data seen shapes, AdaJEPA can more than double frozen-model success and even outperform frozen models trained with much more data.
Visualization
These examples illustrate how adaptation leads to better prediction and planning. Blue denotes the frozen model and red denotes AdaJEPA. A star (★) marks an unseen test shape or configuration, where decoded rollouts can still reconstruct reasonably while retaining training-domain structure. Although the decoder is frozen after training on default data, it still produces meaningful rollouts after lightweight test-time adaptation. This suggests that AdaJEPA improves planning by exploiting shared latent structure and recalibrating predictions, while remaining close to the learned latent manifold.
Takeaways
AdaJEPA suggests that world models should continue learning during deployment rather than remain frozen after training. By improving predictions from the transitions encountered while planning and acting, adaptive world models can support more resilient perception and planning in a changing world.
BibTeX
@misc{wang2026adajepaadaptivelatentworld,
title={AdaJEPA: An Adaptive Latent World Model},
author={Ying Wang and Oumayma Bounou and Yann LeCun and Mengye Ren},
year={2026},
eprint={2606.32026},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2606.32026},
}