🤖 machine learning

AdaJEPA: An Adaptive Latent World Model

AdaJEPA is an adaptive latent world model that enhances planning robustness under distribution shifts by performing test-time self-supervised adaptation within the model predictive control loop, using observed state transitions to continuously recalibrate the model without requiring additional expert demonstrations.

Original authors: Ying Wang, Oumayma Bounou, Yann LeCun, Mengye Ren

Published 2026-07-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ying Wang, Oumayma Bounou, Yann LeCun, Mengye Ren

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are learning to drive a car in a simulator. You spend weeks practicing, and the simulator teaches you exactly how the car responds to the steering wheel and gas pedal. You become a pro at navigating the virtual city.

Now, imagine you take that same car out onto a real road. But here's the catch: the real road is slightly different. Maybe the tires are a bit more slippery, the wind is stronger, or the road is covered in a layer of dust that wasn't in the simulator.

The Problem with "Frozen" Models
Most AI systems today work like a driver who memorized the simulator perfectly but refuses to learn anything new once they step into the real car. In technical terms, these systems use a "world model" (a mental map of how the world works) that is frozen after training.

If the real world shifts even a little bit—like the car handling differently due to rain or a new obstacle—the AI's mental map becomes inaccurate. It tries to drive based on its old, perfect simulator knowledge, but the car doesn't respond as expected. The AI keeps making the same "perfect" moves, but because the reality has changed, it crashes or fails to reach its destination.

The Solution: AdaJEPA (The Adaptable Driver)
The paper introduces AdaJEPA, a new way of driving that acts more like a human. Instead of sticking rigidly to its old training, AdaJEPA has a special "Plan–Act–Adapt–Replan" loop.

Here is how it works, using a simple analogy:

Plan: The AI looks at the road ahead and plans a route based on its current mental map.
Act: It takes the first step (like turning the wheel slightly).
Observe: It immediately looks at what actually happened. Did the car slide more than expected? Did it turn sharper?
Adapt (The Magic Step): Before it even plans the next move, the AI uses that single observation to update its mental map. It says, "Okay, my map said the car would turn left, but it actually turned right. I need to tweak my understanding of how this car handles right now."
Replan: With this slightly corrected map, it plans the next move.

This happens in a split second, over and over again, as the AI moves. It's like a driver who constantly adjusts their steering based on the immediate feedback of the road, rather than relying on a static map from a week ago.

Why This Matters
The researchers tested this in two main ways:

Visual Changes: Imagine the lighting in the room changes, or the object you are pushing changes color. A frozen AI might get confused because it was trained on a gray object in bright light. AdaJEPA quickly realizes, "Oh, this object is red now, and the light is dim," and adjusts its plan instantly.
Physics Changes: Imagine the floor is suddenly much more slippery (low friction) or the object is much heavier. A frozen AI would push with the same force it used on the dry floor and fail. AdaJEPA pushes, sees the object slide too far, and immediately learns, "I need to push less," for the very next step.

The Results
The paper shows that AdaJEPA is incredibly efficient. It doesn't need to relearn everything from scratch. It often only needs one tiny adjustment (a single "gradient step") after every action to fix its mental map.

In the "Safe" Zone: Even when the world is exactly as the AI expected, AdaJEPA doesn't hurt performance; it just stays steady.
In the "Changed" Zone: When the world changes (new shapes, different colors, slippery floors), AdaJEPA dramatically outperforms the frozen models. It recovers from mistakes much faster and reaches its goals more often.

The Bottom Line
The paper argues that AI shouldn't be a static statue that stops learning once it leaves the training lab. Instead, it should be a living system that continuously calibrates its understanding of the world while it acts. AdaJEPA proves that by making these tiny, real-time adjustments, an AI can handle a changing, unpredictable world much better than one that refuses to adapt.

Technical Summary: AdaJEPA: An Adaptive Latent World Model

Problem Statement

Latent world models, particularly Joint-Embedding Predictive Architectures (JEPAs), enable efficient planning by predicting future states in a compact latent space. These models are typically trained on reward-free offline trajectories and paired with Model Predictive Control (MPC) for goal-conditioned tasks. However, a critical limitation exists: the world model is usually frozen after training.

When deployed, if the model's predictions become inaccurate due to test-time distribution shifts (e.g., visual changes like noise or lighting, or dynamic changes like friction and mass), the MPC planner optimizes an incorrect objective. Actions that appear effective in "latent imagination" may fail in the real environment. Small one-step errors compound over the planning horizon, leading to substantial performance degradation. Current approaches often require additional expert demonstrations, target-domain fine-tuning data, or complex outer loops to address this, hindering real-world application.

Methodology: AdaJEPA

The authors propose AdaJEPA, an adaptive latent world model that performs test-time adaptation within the closed loop of MPC. Instead of treating the model as static, AdaJEPA continuously recalibrates the world model using self-supervised signals derived from the agent's own interactions.

The Plan–Execute–Adapt–Replan Loop

The core mechanism operates as follows at each MPC step:

Plan: The agent plans a sequence of actions using the current (potentially outdated) world model to minimize a latent goal-reaching cost.
Execute: The agent executes the first action chunk ( $a_t$ ) in the environment.
Observe: The agent receives the next observation ( $o_{t+1}$ ).
Adapt: The transition $\{o_t, a_t, o_{t+1}\}$ is added to a replay buffer. The model parameters are updated via a small number of gradient steps (often just one) to minimize the latent prediction error on this new transition.
Replan: The updated model is immediately used for the next planning iteration.

Technical Components

Architecture: The underlying world model consists of a sensory encoder ( $\mathcal{E}_s$ ), an action encoder ( $\mathcal{E}_a$ ), and a predictor ( $f_\theta$ ). It is trained offline to predict future latent representations ( $z_{t+1}$ ) rather than reconstructing pixels.
Adaptation Loss: The adaptation uses the same self-supervised latent prediction objective as pretraining:
$\mathcal{L}_{ada}(\mathcal{B}) = \frac{1}{|\mathcal{B}|} \sum_{(o_i, a_i, o_{i+1}) \in \mathcal{B}} \ell(f_\theta(z_i, \mathcal{E}_a(a_i)), \text{sg}(z_{i+1}))$
where $\mathcal{B}$ is a buffer of recent transitions, $\ell$ is a latent prediction loss (e.g., MSE), and $\text{sg}(\cdot)$ is a stop-gradient operator to prevent collapse.
Parameter Updates: To ensure efficiency, adaptation is lightweight. The authors update only a small subset of parameters (e.g., the final layers of the encoder and predictor) rather than the entire network.
Buffer Strategies: The replay buffer $\mathcal{B}$ can be managed via "recent-N" (keeping the most recent $N$ transitions) or "hard-N" (keeping transitions with the largest prediction errors).

Key Contributions

First JEPA Adaptation during Planning: To the authors' knowledge, this is the first work to adapt a JEPA world model during the planning phase, rather than relying solely on pre-training or offline fine-tuning.
Closed-Loop Self-Supervision: The method utilizes the immediate consequences of executed actions as a self-supervised signal, eliminating the need for reward labels, expert demonstrations, or separate data collection phases.
Sample and Compute Efficiency: The framework achieves significant performance gains with minimal computational overhead. In many experiments, a single gradient step per MPC replanning step is sufficient.
Robustness to Distribution Shifts: The approach demonstrates robustness across visual shifts (noise, blur, color changes), dynamic shifts (mass, damping), and structural shifts (unseen maze layouts).

Experimental Results

The authors evaluated AdaJEPA on PushT (block pushing) and PointMaze (navigation) benchmarks under various train-test distribution shifts.

In-Distribution Performance: Even on tasks similar to training data, AdaJEPA improves success rates (e.g., >20% gain on specific PushObj shapes) or maintains baseline performance without degradation, proving it is safe to apply in-distribution.
Out-of-Distribution (OOD) Performance:
- Shape Shifts: On unseen block shapes, AdaJEPA nearly doubled the planning success rate compared to the frozen model.
- Visual Shifts: The method showed clear gains under Gaussian blur, salt-and-pepper noise, and dark lighting.
- Dynamics & Layout Shifts: In PointMaze, AdaJEPA recovered from failures caused by low mass or high damping and significantly improved success rates on unseen maze layouts (e.g., +25.3% gain on specific layout shifts).
Latency: The adaptation step introduces negligible latency (approx. 0.01–0.03s per replan) because it updates only a small subset of parameters. Furthermore, by recovering from prediction errors faster, the agent often reaches the goal in fewer total replans, reducing overall execution time.
Generalization to Data Scarcity: In low-data regimes, test-time adaptation compensated for limited training coverage. For instance, with limited training data, AdaJEPA surpassed a frozen model trained with 16x more trajectories.

Significance and Claims

The paper argues that latent world models should not remain fixed after training but should continue to improve from experience encountered during deployment.

Biological Inspiration: The approach is grounded in biological sensorimotor adaptation, where systems (like the cerebellum) adjust behavior based on changing inputs and dynamics.
Complement to Data Scaling: While larger and more diverse training data improves generalization, AdaJEPA provides a sample-efficient route to refine models during deployment, acting as a complement to data scaling.
Future Direction: The authors suggest that the next step is combining this lightweight test-time adaptation with continual and active learning to expand the world model's coverage over time, enabling more resilient perception and planning in changing environments.

The authors emphasize that their method is agnostic to the specific world model implementation, showing consistent improvements across different architectures (ResNet, Transformer, DINO-based) and planners (Gradient Descent, CEM).