Latent world models enable robots to simulate the future before taking action and plan their movements accordingly. In 2022, Turing Award winner Yann LeCun's team proposed the Joint Embedding Predictive Architecture (JEPA). By predicting the evolution of future states in latent space, it propelled predictive modeling in representation space into a prominent paradigm in world model research.
The problem is that most world models stop updating their parameters after training is complete, making it difficult for them to promptly adapt to changing visual conditions and physical properties in real-world scenarios. Once the encoder or predictor becomes inaccurate, errors are progressively amplified in subsequent planning, ultimately leading to task failure.
Humans, however, are different; sensorimotor adaptation mechanisms are our core ability to adapt to environmental changes. We calibrate action predictions based on sensory feedback and continuously adjust our understanding of the environment based on new experiences.
Inspired by this biological principle, Yann LeCun's team proposed AdaJEPA, an adaptive latent world model capable of continuous learning during deployment. It embeds adaptation into the closed loop of Model Predictive Control (MPC): after executing each action, the model corrects itself using the actual observed state transition, and then replans using the updated model.
Paper link: https://arxiv.org/abs/2606.32026
Results show that AdaJEPA consistently improves planning success rates in both in-distribution tasks and across various distribution shifts. Even with only a single lightweight update before each replanning step, it generally outperforms world models whose parameters remain frozen after training.
This work opens a promising direction for adaptive world models: world models should continuously calibrate predictions and update representations based on real feedback during the course of action, thereby better adapting to changing environments.
AdaJEPA: A World Model That Understands "Adaptation"
AdaJEPA is an adaptive latent world model that continuously corrects its predictions while a robot performs a task. After completing each action step, the model uses new real observations to correct deviations and replans accordingly. The entire process requires no additional offline data, reward labels, or expert demonstrations. The workflow can be summarized in four steps: plan, execute, correct, and replan. The details are as follows:
Figure | AdaJEPA performs test-time adaptation in closed-loop MPC.
- Planning: The model internally simulates the state changes for the next few steps, compares multiple candidate plans, and selects the action sequence most likely to approach the goal.
- Execution: After planning, the model executes only the first action or a short segment of actions, and then observes the environment's true feedback. The state changes prior to execution are recorded as learning samples for subsequent adaptation.
- Adaptation: After the action is executed, AdaJEPA writes the state transition generated in this step to an online buffer and uses it to verify the model's prediction accuracy. If there is a deviation between the predicted next state and the true outcome, the model uses this error as an update signal to make a lightweight correction, preparing for the next round of planning.
- Replanning: After adaptation, the model replans future trajectories from the latest observation using the updated world model and generates a new action sequence. Throughout the task, the "plan-execute-adapt-replan" cycle repeats continuously, ensuring that each round of planning is based on the latest observations and the most current model.
Additionally, to avoid slowing down real-time planning, AdaJEPA only performs lightweight updates: adjusting a small number of parameters, maintaining a small online buffer, and retaining the objective function from the pre-training stage. The details are as follows:
- Only updating key layers: AdaJEPA does not update the entire world model. Instead, it only adjusts a few key layers in the encoder and predictor. This reduces computational overhead and minimizes disturbance to existing representations.
- Maintaining a small online buffer: The buffer defaults to storing the most recent 5 real state transitions. The research team compared two retention methods: recent-N, which keeps the most recent transitions, and hard-N, which keeps transitions with the highest prediction errors. Results showed little difference between the two, but recent-N was more stable.
- Retaining the pre-training objective function: The adaptation stage maintains the same predictive objective as pre-training, using the representation corresponding to the real observation as a supervision signal. To minimize disturbance to existing representations, the target representation serves only as a reference and does not participate in gradient backpropagation.
How Effective Is It?
Overall, AdaJEPA consistently improves planning success rates for both in-distribution tasks and under various distribution shifts. The research team evaluated the model on the PushT / PushObj object-pushing tasks and the PointMaze maze navigation task, covering variations in shape, vision, dynamics, and layout. Even with just a single lightweight update before each replanning step, AdaJEPA still generally outperforms frozen world models. The specific results are as follows:
1. In-Distribution Tasks
Results indicate that AdaJEPA's test-time adaptation does not sacrifice existing capabilities and can further improve task success rates. Whether using Gradient Descent (GD) to directly optimize the action sequence or using the Cross-Entropy Method (CEM) to search by sampling and filtering candidate actions, AdaJEPA's success rate is higher than the baseline without test-time adaptation. The most significant improvement is seen in the object-pushing tasks, with the highest success rate increasing by over 20%. In the maze navigation task, the original model already performed strongly, and AdaJEPA maintained a comparable level without obvious degradation.
Figure | Planning success rates for PointMaze under dynamics shift and layout shift conditions.
2. Out-of-Distribution Tasks
In tasks with more pronounced environmental changes, AdaJEPA's advantages become even more prominent. After each round of planning and execution, it uses new real feedback to update the world model, bringing subsequent planning closer to the current environment and improving the task success rate. In contrast, models that remain frozen after training cannot leverage these new observations, causing their success rates to plateau quickly.
Figure | Planning success rates under shape shift and visual shift conditions.
Specifically, in the multi-shape object-pushing task, if unseen object shapes appear during testing, AdaJEPA shows the most significant improvement, nearly doubling the success rate. Among visual perturbations, the gains from blurring, noise, and dark lighting are more pronounced. If only the anchor points or object colors change, AdaJEPA's advantage is relatively limited. In PointMaze navigation, AdaJEPA can also adapt to dynamics shifts and new maze layouts, planning trajectories closer to the shortest path in new layouts.
Figure | Planning trajectories in diverse mazes.
Figure | Planning trajectories for PointMaze-Medium under dynamics shift conditions.
3. AdaJEPA Shows Improvements Across Multiple JEPA Implementations
To verify whether AdaJEPA relies on a specific model implementation, the research team tested it on the PushT task by altering the representation form, model architecture, training objective, and planner. Results show that AdaJEPA improves planning success rates across all these settings. Even if the baseline model is fully trained and the evaluation remains in-distribution, test-time adaptation still provides stable benefits, adding only about 0.01–0.03 seconds of latency per replanning step.
Figure | AdaJEPA performance across different implementations.
4. AdaJEPA Calibrates Existing Predictions Rather Than Learning a New World from Scratch
Visualizations show that AdaJEPA's adaptation acts more like calibration than relearning. After decoding the adapted prediction trajectories, the research team found that even when encountering visual perturbations or unseen shapes, the decoded results tended to retain structural features from the training distribution. For example, a red block might be decoded into a gray block commonly seen during training, and unseen shapes would be decoded into similar, previously seen shapes.
Figure | Examples of AdaJEPA planning trajectories under visual and shape shifts.
5. Ablation Experiments and Analysis
Ablation results indicate that AdaJEPA requires neither large-scale updates nor complex hyperparameter tuning. Updating a few key layers, performing a single-step gradient update, and using a recent state transition buffer are sufficient to yield stable benefits.
First, AdaJEPA only updates partial layers of the encoder or predictor. When using LoRA for lightweight updates, overall performance remains superior to the baseline without test-time adaptation, indicating that retraining the entire model is unnecessary.
Second, different distribution shifts require updating different parts of the model. Under shape shifts, various update schemes show little difference, and mainly adjusting the predictor is sufficient. Under visual and layout shifts, updating only the predictor yields limited results, requiring the encoder to participate as well. In layout shifts, updating the first layer of the predictor works best, possibly because it is the earliest to fuse latent states with action information, making it easier to correct new local transition relationships.
Furthermore, the default hyperparameters are sufficiently robust. In terms of hyperparameter settings, AdaJEPA defaults to the learning rate from the training stage, performs only a single-step gradient update before each replanning, and retains recent state transitions as the buffer. A larger learning rate or more update steps might enhance adaptation effects but would also increase instability and computational overhead. Overall, the default settings achieve a good balance between performance, stability, and latency.
Figure | The impact of adaptation hyperparameters and replay buffers on planning success rates.
6. Impact of Training Data Scale and Shape Diversity on AdaJEPA
Experimental results show that AdaJEPA's effectiveness depends not only on the volume of training data but also on whether the training data is sufficiently diverse. For the PushObj multi-shape object-pushing task, shape diversity is more critical than simply stacking trajectories of the same shape. Meanwhile, test-time adaptation can compensate for some generalization gaps when data is insufficient.
Specifically, given the same total number of trajectories, allocating data across more object shapes is more conducive to generalizing to unseen shapes than concentrating on a single shape. For instance, with a total of 16k trajectories, AdaJEPA covering four shapes achieved a 51.9% success rate on unseen shapes, higher than the 45.8% achieved when covering only a single shape.
Additionally, AdaJEPA can improve success rates across different data scales, with benefits being especially pronounced in low-data scenarios. Even if the training phase covers fewer shapes and trajectories, the model can still correct predictions using new observations during deployment. For example, on seen shapes, training AdaJEPA with only 1 shape and 1k trajectories achieved a 60.8% success rate, which is higher than a model trained with 4 shapes and a total of 64k trajectories but without test-time updates.
Figure | Impact of training data scale on PushObj planning success rate: shape diversity (K) vs. number of trajectories per shape.
Limitations and Future Directions
Although AdaJEPA brings stable benefits across multiple planning tasks, the following limitations remain:
Because AdaJEPA only performs lightweight corrections during planning, its effectiveness is still constrained by the coverage of pre-trained representations. When key features not covered during training appear in the test environment, adaptation can improve planning results to a certain extent, but it still struggles to completely bridge this representation gap. In the future, lightweight test-time adaptation needs to be combined with continual learning and active learning, enabling world models to continuously accumulate new experiences during long-term deployment and gradually expand their coverage of environmental changes.
For more technical details, please refer to the original paper.
Author: Xia Qiansi
To submit an article or request reprint authorization, please leave a message directly in the comments section of this article.