AdaJEPA: An Adaptive Latent World Model

Latent world models are typically frozen after training, which makes them vulnerable to inaccuracies when faced with test-time distribution shifts. This paper asks whether a world model can be made more robust by continuously adapting its parameters during deployment, using the agent's own real-time experiences as a self-supervised learning signal.

The authors propose AdaJEPA, a framework that integrates test-time adaptation into the closed-loop cycle of Model Predictive Control (MPC). In this loop, the agent plans using its current world model, executes an action, and observes the resulting state transition. This new transition is then used to perform a small number of gradient updates on the world model's encoder and predictor before the next replanning step. This process allows the model to recalibrate its internal dynamics to the current environment without needing external labels or a separate training phase.

AdaJEPA demonstrates consistent and significant performance gains across a variety of challenging scenarios, including unseen object shapes, visual corruptions (such as noise and lighting changes), and altered physical dynamics. The method is highly efficient, often requiring only a single gradient step per MPC replanning cycle, which adds negligible latency. The authors show that this adaptive approach is particularly powerful in low-data regimes, where it can outperform frozen models trained on significantly larger datasets. Furthermore, AdaJEPA is shown to be agnostic to the specific architecture of the underlying world model, providing consistent improvements across different JEPA implementations.

This work challenges the standard paradigm of keeping world models static after training. By demonstrating that lightweight, continuous adaptation is both feasible and effective, the authors provide a scalable path toward building autonomous agents that can remain reliable in dynamic, real-world environments where training-time assumptions may not hold.