AdaJEPA: An Adaptive Latent World Model
For robotic planning tasks, AdaJEPA addresses the problem of world model inaccuracy at test time, enabling robust performance without additional expert demonstrations.
AdaJEPA introduces an adaptive latent world model that performs test-time adaptation within model predictive control, improving planning success under distribution shift with as few as one gradient step per replanning step.
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.