Warm take: Your world model should never stop learning

Introducing AdaJEPA, an adaptive WM that plans, acts, and adapts in a closed loop. Every action leads to a new observation, and every transition refines the latent representation and prediction.

📝: https://t.co/85WWwn5QDS https://t.co/kC8E0TO6aW

AdaJEPA: An Adaptive Latent World Model | Agentic Learning AI Lab

Source: https://agenticlearning.ai/adajepa/

Ying Wang1,2, Oumayma Bounou1,Yann LeCun1,2*, andMengye 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.

Contents:OverviewMethodResultsVisualizationTakeaways## Overview

AdaJEPA performs a closed-loop plan-act-adapt-replan cycle. At each MPC step, the agent plans with the current world model, executes the first action, observes the next transition, updates the model using latent prediction error, and replans with the adapted model. 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.

AdaJEPA closed-loop plan-and-adapt algorithm.

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.

Shape shifts change the shapes in PushObj; stars mark held-out object shapes. Visual shifts corrupt PushT observations with blur, salt-and-pepper noise, dark lighting, and color changes. 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.

Dynamics shifts change PointMaze physics; layout shifts test held-out maze layouts. 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.

Training data scale varies PushObj shape diversity K and trajectories per shape N. 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.

Ex. 1: PushObj +, frozen.Ex. 1: PushObj +, AdaJEPA.Ex. 2: PushObj I★, frozen.Ex. 2: PushObj I★, AdaJEPA.Ex. 3: PushT blur★, frozen.Ex. 3: PushT blur★, AdaJEPA.Ex. 4: PushT red★, frozen.Ex. 4: PushT red★, AdaJEPA.

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}, 
}