AdaJEPA: An Adaptive Latent World Model

2026-06-30Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors created AdaJEPA, a system that helps a model predict future events more accurately while it's working, rather than just relying on what it learned before. Normally, these models get stuck if things change in ways they didn't expect, which makes their planning fail. AdaJEPA fixes this by letting the model learn from what actually happens step-by-step, improving its predictions on the fly without needing extra expert help. This approach helped the model succeed more in tasks where it has to reach certain goals.

latent world modelsmodel predictive control (MPC)test-time adaptationself-supervised learninggradient stepclosed-loop controldistribution shifthigh-dimensional observationsgoal-reaching tasks
Authors
Ying Wang, Oumayma Bounou, Yann LeCun, Mengye Ren
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.