JEPA-TTT: Persistent Test-Time Training of Latent World Models for Planning under Dynamics Shifts

JEPA-TTT. Persistent Test-Time Training of Latent World Models for Planning under Dynamics Shifts. By Zheyuan Zhang and seven others, from Honda Research Institute USA and Johns Hopkins University. Published on the thirtieth of September. No code repository is linked to it.

World models let an agent plan by predicting future states of its environment. Those predictions can become unreliable when the dynamics at test time differ from the dynamics seen during training. The agent is then planning against a world that no longer behaves the way its model expects, and a plan is only as sound as the predictions it rests on. The question is whether a pretrained world model can be brought back in line with the new dynamics while the agent is still out there acting.

JEPA-TTT keeps adapting the model while it is in use. It takes a pretrained, action-conditioned Joint-Embedding Predictive Architecture world model, and adapts its latent dynamics predictor throughout test time. The updates are self-supervised, and they accumulate across episodes rather than starting over each time. Meanwhile the visual encoder and the reward head stay fixed, which preserves the pretrained representation and the task objective. Planning needs neither a goal image nor online reward from the environment.

The updates are fed by dense replay. Prediction windows are formed at every temporal offset, retained in a buffer that keeps growing, and minibatches are sampled from that buffer for each predictor update. Because the buffer persists, what was seen in earlier episodes keeps informing the updates made in later ones.

The test covers eight dynamics shifts across four continuous-control environments, and JEPA-TTT improves planning on every one of them. After five hundred test-time episodes, it cuts autoregressive latent prediction error by eighty-three percent on average. Over the frozen JEPA world model, it improves planning performance by one hundred and fifty-three percent. The two numbers move together: the predictor becomes more accurate about the shifted world, and plans built on it get better.

The division of labour is the design. The encoder and reward head hold what the model learned and what the task asks for, while only the predictor moves. Persistent, self-supervised training at test time, confined to that one part, is what the authors show can adapt a pretrained latent world model when the dynamics change underneath it.