World-model test-time adaptation stops fitting prediction alone: OpTTA adds an optimistic planning term and beats AdaJEPA on final success in 14 of 15 shifted settings
Synopsis
Pretrained latent world models are usually frozen at test time, and distribution shift can degrade planning with them; updating only on prediction loss fits the outcomes of the agent's own actions and can stay wrong about actions that would reach the goal. OpTTA adds the planner's goal-reaching cost on candidate imagined rollouts to the prediction loss, biasing adaptation toward model parameters that are optimistic about reaching the goal. Across three JEPA world models and 15 shifted settings, adapting only the predictor raises final success over the frozen model by 13.5 percentage points on average, versus 6.9 for prediction-only AdaJEPA; with the encoder also adapted, the gains are 13.5 versus 7.5 points.
Interpretation
The online update objective for a deployed world model expands from a single prediction loss to prediction loss plus planning cost, so adaptation targets whether the model can reach the goal rather than only fitting observed transitions. AdaJEPA previously minimized only the one-step latent prediction loss on recently observed transitions at test time, so the model was fit to the outcomes of its own actions and could fit them well while remaining wrong about actions that would reach the goal. OpTTA's final success exceeds AdaJEPA's in 14 of 15 shifted settings; across the eight predictor-only settings it improves over the frozen model by 13.5 points on average versus 6.9 for AdaJEPA, and across the seven joint predictor-encoder settings the gains are 13.5 versus 7.5 points.
The adaptation objective has two terms, the prediction loss on observed transitions and the planner's goal-reaching cost on candidate imagined rollouts, both minimized over world-model parameters. The second term borrows reward-biased maximum likelihood from adaptive control, needs no confidence sets or explicit uncertainty estimates, and adds no extra environment interaction; the optimism weight is scheduled by the relative magnitude of planning and prediction gradients and grows over replanning steps. Experiments span three JEPA world models, LeWM, DINO-WM and temporal straightening, on PushT and PushObj manipulation, Reacher reaching, and TwoRoom and maze navigation, under seven visual corruptions plus color, geometry, mass and damping changes, with both CEM and gradient-based planners.
Perspective
The demonstrated scope is goal-conditioned planning with three JEPA latent world models on manipulation, reaching and navigation tasks, under visual corruptions and environment-parameter shifts, with CEM and gradient-based planners; adaptation is task-local, resetting the model to pretrained parameters and clearing the buffer after each task.
On TwoRoom under visual corruptions, both AdaJEPA and OpTTA reduce final success relative to the frozen model under both adaptation scopes, and the best parameters to update vary across tasks, a task dependence left open; OpTTA adds computation per replan, its update on DINO-WM costs nearly as much as planning, the default configuration need not be optimal, and the design of the optimism weight remains an open question.