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

Akshay Mete; Dileep Kalathil; P. R. Kumar

cs.LGOct 26, 2026

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

Authors: Akshay Mete, Dileep Kalathil, P. R. Kumar

Organizations: ECE Department, Texas A&M University, College Station, Texas, USA.

Abstract

Latent world models are typically pre-trained offline and kept frozen at test time, so planning with them can degrade under distribution shift. Test-time adaptation (TTA) can mitigate this degradation by updating the model during deployment. Unlike pre-training on fixed data, TTA for planning is closed-loop: the planner uses model rollouts to select actions, and the resulting transitions drive the next update. This raises the question: which objective should guide TTA for planning? 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. Inspired by reward-biased maximum likelihood from adaptive control, we introduce Optimistic Test-Time Adaptation (OpTTA). Its objective has two terms, the prediction loss on observed transitions and the planner's goal-reaching cost on its candidate rollouts, both minimized over the model parameters. The planning term biases adaptation toward models that are optimistic about reaching the goal, without any additional environment interaction. Across visual and environment shifts on manipulation, reaching and navigation tasks, OpTTA raises final planning success over the frozen model by 13.5 percentage points on average, compared with 6.9 for the prediction-only objective (AdaJEPA), when only the predictor is adapted. When the encoder is adapted as well, the gains are 13.5 and 7.5 points. These experiments span three JEPA world models with different encoders and training objectives and, for temporal straightening, both CEM and gradient-based planners. Our results suggest that test-time adaptation of world models for planning benefits from an optimistic, task-aware term in addition to the prediction loss.

Figures & tables

Figure 1: Gain in final planning success over the frozen model for prediction-only adaptation (AdaJEPA) and optimistic adaptation (OpTTA), averaged over the visual corruptions or environment shifts of each setting (Table 1 ). Error bars show

\pm 1

s.e. over seeds.

Figure 2: Optimistic Test-Time Adaptation. At every replanning step the world model is adapted on the observed transitions, using the prediction loss plus the task-aware optimistic loss evaluated on the candidate rollouts the planner has already generated.

Table 1: Planning success (%, mean

\pm

s.e. over seeds). Final is success by replanning step

30

, and Mean over replans is the average cumulative success rate across replanning steps

1

–

30

. Parentheses report the adaptation gain over Frozen, paired by seed and computed before rounding. Panel (a) adapts only the predictor; panel (b) adapts both the predictor and encoder. Shifted rows average over the number of shifts in their labels, each visual corruption over four severity levels (DINO-WM: one level and three seeds); per-shift results are in Appendix A .

Panel of Figure 3 — Figure 3: Ablations of OpTTA’s design choices, varied one at a time (blue: AdaJEPA; orange: OpTTA; ∗ : default). Rows show temporal-straightening PushT under darkening $0.7$ and LeWM Reacher under darkening $0.85$ , with the predictor adapted; in (c), top- $n$ keeps the $n$ lowest-cost candidates. Cumulative success is averaged over three seeds of $50$ episodes; shading shows $\pm 1$ s.e.

	Frozen	AdaJEPA	OpTTA	OpTTA 30	AdaJEPA + enc.	OpTTA + enc.
LeWM PushT	118	141 (+23)	163 (+45)	165 (+47)	149 (+31)	169 (+51)
TS PushT	379	389 (+10)	410 (+31)	410 (+31)	394 (+15)	415 (+36)
DINO-WM PushT	2955	2971 (+16)	5233 (+2278)	3304 (+349)	NA	NA

Table 2: Time per replan (ms) on one RTX PRO 6000 GPU: planning, which does not depend on the method and is averaged over each row, plus the adaptation update (parentheses). OpTTA 30 restricts

\mathcal{P}_{\theta}

to the

30

CEM elites (Figure 3(c) ). LeWM times are wall time and include the environment step; its adaptation time is the difference to Frozen.

Figure 4: Adaptation diagnostics for LeWM Reacher (top, gamma

1.5

) and temporal-straightening PushT (bottom, darkening

0.85

), predictor adapted; mean

\pm

s.e. over three seeds of 50 episodes. Own-trajectory loss is measured on the newest transition before adapting on it; the goal-reaching trajectory is the dataset’s start-to-goal trajectory.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Planning (CEM; all methods)
Hyperparameter	LeWM	TS	DINO-WM
Candidates per iteration $M$	$200$	$200$	$200$
Elites	$30$	$30$	$30$
CEM iterations	$10$	$10$	$10$
Planning horizon	$5$ latent steps ( $25$ environment steps)
Executed per replan	first latent step ( $5$ environment steps)

Appendix

Table 3: Hyperparameters of the planning and adaptation protocol. Adaptation follows AdaJEPA ( Wang et al., 2026a ) and is identical for AdaJEPA and OpTTA; OpTTA adds only the optimism hyperparameters, which are the same for every world model, task, and shift. A latent step covers five environment steps (frameskip).

	Final			Mean over replans
Setting	Frozen	AdaJEPA	OpTTA	Frozen	AdaJEPA	OpTTA
LeWM PushT
Blur (2, 4, 6, 8)	77.6 $\pm$ 1.3	78.7 $\pm$ 1.8 (+1.1)	85.5 $\pm$ 1.1 (+7.9)	58.7 $\pm$ 0.7	60.3 $\pm$ 1.1 (+1.6)	67.4 $\pm$ 0.7 (+8.7)
Salt & pepper (0.025, 0.05, 0.075, 0.1)	53.5 $\pm$ 1.4	53.9 $\pm$ 0.9 (+0.4)	60.5 $\pm$ 1.3 (+7.0)	40.5 $\pm$ 1.0	41.5 $\pm$ 0.7 (+1.1)	46.5 $\pm$ 0.8 (+6.0)
Darkening (0.95, 0.9, 0.85, 0.8)	54.8 $\pm$ 1.2	58.2 $\pm$ 1.7 (+3.4)	61.8 $\pm$ 0.9 (+7.0)	41.8 $\pm$ 0.8	44.8 $\pm$ 1.2 (+2.9)	48.2 $\pm$ 0.7 (+6.4)
Gamma (0.5, 0.75, 1.25, 1.5)	76.7 $\pm$ 1.7	82.6 $\pm$ 0.5 (+5.9)	87.8 $\pm$ 1.0 (+11.1)	57.5 $\pm$ 1.2	62.4 $\pm$ 0.7 (+4.8)	68.4 $\pm$ 0.8 (+10.9)

Appendix

Table 4: Predictor adapted. Planning success (%, mean

\pm

s.e. over seeds, five seeds unless noted) per shift, in the layout of Table 1 . Each visual corruption averages the severity levels listed in its label. Parentheses give the gain over Frozen, paired by seed; bold marks the best method at the printed precision.

	Final			Mean over replans
Setting	Frozen	AdaJEPA	OpTTA	Frozen	AdaJEPA	OpTTA
LeWM PushT
Blur (2, 4, 6, 8)	77.6 $\pm$ 1.3	80.0 $\pm$ 1.6 (+2.4)	86.5 $\pm$ 1.0 (+8.9)	58.7 $\pm$ 0.7	60.4 $\pm$ 1.3 (+1.7)	67.6 $\pm$ 0.7 (+8.9)
Salt & pepper (0.025, 0.05, 0.075, 0.1)	53.5 $\pm$ 1.4	53.5 $\pm$ 1.0 (+0.0)	58.3 $\pm$ 2.7 (+4.8)	40.5 $\pm$ 1.0	41.3 $\pm$ 1.1 (+0.8)	45.2 $\pm$ 1.9 (+4.7)
Darkening (0.95, 0.9, 0.85, 0.8)	54.8 $\pm$ 1.2	59.5 $\pm$ 2.4 (+4.7)	63.9 $\pm$ 1.8 (+9.1)	41.8 $\pm$ 0.8	45.4 $\pm$ 1.8 (+3.6)	50.1 $\pm$ 1.6 (+8.3)
Gamma (0.5, 0.75, 1.25, 1.5)	76.7 $\pm$ 1.7	81.6 $\pm$ 1.4 (+4.9)	87.7 $\pm$ 0.9 (+11.0)	57.5 $\pm$ 1.2	61.7 $\pm$ 1.3 (+4.1)	68.5 $\pm$ 0.9 (+10.9)

Appendix

Table 5: Predictor and encoder adapted (DINO-WM has no encoder variant). Planning success (%, mean

\pm

Figure 5: LeWM PushT: cumulative success versus replanning step, in distribution and under each visual corruption, averaged over the severity levels in the panel title. Curves are means over seeds with

\pm 1

s.e. bands; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean of each curve is the table’s mean over replans, and its last point the table’s Final.

Figure 6: LeWM Reacher: cumulative success versus replanning step, in distribution and under each visual corruption, averaged over the severity levels in the panel title. Curves are means over seeds with

\pm 1

s.e. bands; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean of each curve is the table’s mean over replans, and its last point the table’s Final.

Figure 7: LeWM TwoRoom: cumulative success versus replanning step, in distribution and under each visual corruption, averaged over the severity levels in the panel title. Curves are means over seeds with

\pm 1

s.e. bands; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean of each curve is the table’s mean over replans, and its last point the table’s Final.

Figure 8: TS PushT: cumulative success versus replanning step, in distribution and under each visual corruption, averaged over the severity levels in the panel title. Curves are means over seeds with

\pm 1

s.e. bands; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean of each curve is the table’s mean over replans, and its last point the table’s Final.

Figure 9: DINO-WM PushT (one level, three seeds): cumulative success versus replanning step, in distribution and under each visual corruption, averaged over the severity levels in the panel title. Curves are means over seeds with

\pm 1

s.e. bands; all lines adapt only the predictor. The mean of each curve is the table’s mean over replans, and its last point the table’s Final.

Figure 10: TS environment shifts: cumulative success versus replanning step, each environment shift. Curves are means over seeds with

\pm 1

s.e. bands; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean of each curve is the table’s mean over replans, and its last point the table’s Final.

Figure 11: LeWM PushT: final success at each level of each visual corruption, with the uncorrupted condition (clean, shaded) placed at the identity value; levels are ordered by value. Points are means over five seeds with

\pm 1

s.e. bars; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean over a panel’s four corrupted levels is the corresponding entry of Tables 4 and 5 .

Figure 12: LeWM Reacher: final success at each level of each visual corruption, with the uncorrupted condition (clean, shaded) placed at the identity value; levels are ordered by value. Points are means over five seeds with

\pm 1

s.e. bars; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean over a panel’s four corrupted levels is the corresponding entry of Tables 4 and 5 .

Figure 13: LeWM TwoRoom: final success at each level of each visual corruption, with the uncorrupted condition (clean, shaded) placed at the identity value; levels are ordered by value. Points are means over five seeds with

\pm 1

s.e. bars; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean over a panel’s four corrupted levels is the corresponding entry of Tables 4 and 5 .

Figure 14: TS PushT: final success at each level of each visual corruption, with the uncorrupted condition (clean, shaded) placed at the identity value; levels are ordered by value. Points are means over five seeds with

\pm 1

s.e. bars; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean over a panel’s four corrupted levels is the corresponding entry of Tables 4 and 5 .

Table 6: Temporal straightening on PushT with a gradient-descent planner. Planning success (%, mean

\pm

s.e. over five seeds of

50

episodes). The gradient-descent runs use a budget of

20

replans, so Final is success by replanning step

20

and Mean over replans is the average cumulative success rate across replanning steps

1

–

20

; values are not directly comparable to the

30

-replan Table 1 . Parentheses report the adaptation gain over Frozen, paired by seed and computed before rounding. Each visual corruption averages its four severity levels.

Figure 15: Temporal straightening on PushT with a gradient-descent planner: cumulative success versus replanning step over the

20

-replan budget, one panel per row of Table 6 . Visual corruptions average the severity levels in the panel title. Curves are means over seeds with

\pm 1

s.e. bands; solid lines adapt the predictor, dashed lines the predictor and encoder. The mean of each curve is the table’s mean over replans, and its last point the table’s Final.

Test-Time Adaptation

Test-time adaptation updates or recalibrates pretrained models at inference to handle distribution shifts, often without target labels.

See all 318 papers

AdaJEPA: An Adaptive Latent World Model
Jun 30, 2026Ying Wang, Oumayma Bounou, Yann LeCun +1Continual Test-Time Adaptation World Model-Based Planning
New York University · AMI Labs
JEPA-TTT: Persistent Test-Time Training of Latent World Models for Planning under Dynamics Shifts
Sep 30, 2026Zheyuan Zhang, Suyu Ye, Nakul Agarwal +5Test-Time Adaptation Latent Dynamics Modeling
Honda Research Institute USA · Johns Hopkins University
Sandwich-Residuals: Parameter-Efficient Test-time Adaptation of World Models
Sep 18, 2026Krishnam Soni, Aditya Sehgal, Vedant Dave +1Test-Time Adaptation Latent World Models
Chair of Cyber-Physical-Systems, Montanuniversität Leoben, Austria

Jun 30, 2026J/K move · Enter open · S save

AdaJEPA: An Adaptive Latent World Model

Ying Wang, Oumayma Bounou, Yann LeCun, Mengye Ren

New York University · AMI Labs

Continual Test-Time Adaptation World Model-Based Planning Model-Based Planning Latent World Models Latent-Space Planning

Open paper AskPDF arXiv

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.

Abstract

Test-Time Adaptation

Explore similar work

AdaJEPA: An Adaptive Latent World Model

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

Sandwich-Residuals: Parameter-Efficient Test-time Adaptation of World Models