AdaJEPA: An Adaptive Latent World Model

2026-06-30

AI Summary

Traditional latent world models remain fixed at test time, making them vulnerable to distribution shifts and prone to planning failures. This work proposes AdaJEPA, which introduces test-time adaptation into the model predictive control (MPC) loop for the first time. After executing an action, AdaJEPA leverages the observed state transition as a self-supervised signal to perform a single-step gradient update, thereby calibrating the latent world model in real time and replanning accordingly. Requiring no expert demonstrations, this approach significantly improves planning success across diverse goal-directed tasks, demonstrating the effectiveness and practicality of continuous test-time adaptation.

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.

Problem

• latent world models
• test-time distribution shift
• planning failure
• model adaptation

Innovation

• adaptive world model
• test-time adaptation
• model predictive control
• self-supervised learning
• latent dynamics