Adapt for Planning, Not Just Prediction Optimistic Test-Time World Model Adaptation

Optimism improves planning under distribution shift

Both methods share the checkpoint, planner, episodes, buffer, adapted parameters and update schedule; only the adaptation objective differs. OpTTA reaches higher final success in 14 of 15 shifted settings and higher mean success over replans in 13 of 15. TwoRoom is the exception, where both objectives fall below the frozen model.

Which objective should guide test-time adaptation for planning?

Latent world models are pre-trained offline and kept frozen at test time, so planning with them degrades under distribution shift. Test-time adaptation updates the model during deployment, in a closed loop: the planner acts on the model's rollouts, and the resulting transitions drive the next update. Under a prediction-only objective the model is fit only to the outcomes of its own actions, so it can remain wrong about the actions that would reach the goal, and low prediction error need not translate into better planning.

Optimistic test-time adaptation

OpTTA formulates test-time adaptation as task-aware, optimistic model estimation, inspired by reward-biased maximum likelihood from adaptive control. Optimism enters through the estimation objective itself, without confidence sets or explicit uncertainty estimates, and the objective assumes no particular world-model architecture or planner.

How optimism changes adaptation

The advantage does not come from fitting the agent's own experience better: OpTTA's prediction loss on the transitions it executes is not consistently lower. Its model is more accurate on the goal-reaching trajectory on Reacher and more optimistic about its chosen plan on both tasks. Planning with an optimistic model either reaches the goal or produces transitions that the prediction loss then uses to correct the model.

Conclusion

Across three JEPA world models, both adaptation scopes, visual and environment shifts, and CEM and gradient-based planners, OpTTA improves final success over prediction-only adaptation in most settings. World-model adaptation should account for the downstream planning objective rather than prediction alone, and optimism offers a simple way to do so. Limitations: on TwoRoom neither objective improves over the pretrained model, and OpTTA adds computation per replan.

BibTeX

Copy@article{mete2026optta,
  title   = {Adapt for Planning, Not Just Prediction: Optimistic Test-Time World Model Adaptation},
  author  = {Mete, Akshay and Kalathil, Dileep and Kumar, P. R.},
  journal = {arXiv preprint arXiv:2610.35636},
  year    = {2026}
}