PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Akshay Mete, Dileep Kalathil, P. R. Kumar

arXiv:2610.35636v1cs.LGcs.AIcs.RO

TL;DR

Frozen latent world models can plan less effectively under distribution shift, and prediction-only test-time adaptation can fit observed outcomes without learning which actions reach the goal. OpTTA combines observed-transition prediction loss with a planning loss on imagined rollouts; across shifted settings, it improves final success over both the frozen model and prediction-only adaptation.

  • Problem

    Under distribution shift, frozen world models can plan less effectively, while prediction-only adaptation may fit outcomes of its own actions without identifying actions that reach the goal.

    Source: §1 · pp. 1–2

  • Method

    OpTTA adapts model parameters by balancing prediction loss on observed transitions with the planner’s goal-reaching cost on imagined rollouts.

    Source: CONTRIBUTIONS · p. 2, §3 · p. 4, §4 · p. 4

  • Results

    OpTTA improves final success over prediction-only adaptation in most settings; gains over the frozen model average 13.5 points versus 6.9 for AdaJEPA with only the predictor adapted, and 13.5 versus 7.5 with the encoder also adapted.

    Source: CONTRIBUTIONS · p. 3, §7 · p. 9

  • Takeaways & Limitations

    The results suggest that test-time world-model adaptation for planning should account for the downstream planning objective, not prediction alone.

    Source: §7 · p. 9

    On TwoRoom, neither method improves over the pretrained model; the best parameters to update vary across tasks, and OpTTA adds computation per replan.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Test-time distribution shifts can degrade planning with frozen latent world models, while prediction-only adaptation may fit observed outcomes without identifying actions that reach the goal. OpTTA adds a task-aware planning cost to prediction loss, and improves final planning success over prediction-only adaptation across the reported shifted settings.

1 INTRODUCTION

  • OpTTA gains 13.5 points over the frozen model versus 6.9 for AdaJEPA with only the predictor adapted, and 13.5 versus 7.5 when the encoder is also adapted.These averages span three JEPA world models and 15 shifted settings; Figure 1 summarizes gains over visual corruptions or environment shifts.Source: Figure 1 · p. 1, CONTRIBUTIONS · p. 3
  • Prediction-only adaptation can fit outcomes of the planner’s own actions while remaining wrong about actions that would reach the goal, creating a closed-loop identifiability problem.The resulting mismatch is between prediction loss and the planning objective.Source: §1 · p. 2
  • Inspired by reward-biased maximum likelihood, OpTTA favors models with lower imagined goal-reaching cost while retaining a prediction-loss anchor to observed transitions.The planning loss is differentiated with respect to world-model parameters; setting α_k to zero recovers AdaJEPA’s prediction-only objective.Source: §1 · p. 2, CONTRIBUTIONS · p. 2
  • OpTTA balances prediction accuracy on observed transitions against the planner’s goal-reaching cost on imagined rollouts.At each replanning step, adaptation uses observed transitions and the task-aware loss evaluated on candidate rollouts already generated by the planner.Source: CONTRIBUTIONS · p. 2, Figure 2 · p. 2

2 RELATED WORK

Related work connects test-time world-model adaptation to distribution-shift adaptation, objective-mismatch remedies, and JEPA planning models. OpTTA brings reward-biased optimism to pretrained world-model adaptation for goal-conditioned planning.

2 RELATED WORK

  • AdaJEPA adapts latent world models for planning by minimizing prediction error on observed transitions, while online dynamics updates are also standard in adaptive control and model-based reinforcement learning.The cited works include AdaJEPA (Wang et al., 2026a) and adaptive-control and model-based-RL approaches (Borkar & Varaiya, 1979; Kumar & Becker, 1982; Hafner et al., 2020; 2024; Hansen et al., 2024).Source: §2 · p. 3
  • Value-aware model-learning methods address objective mismatch by adding reward or value information to prediction error, but require a learned reward or value function unavailable for this goal-conditioned test-time adaptation setting.The passage cites Lambert et al. (2020) on objective mismatch and Farahmand et al. (2017); Farahmand (2018); Hansen et al. (2024) on value-aware methods.Source: §2 · p. 3
  • RBMLE favors higher-reward models while anchoring estimates to observed transitions, without requiring calibrated uncertainty estimates that are difficult under shift and limited deployment data.RBMLE is contrasted with UCB methods and uncertainty-based optimistic approaches; OpTTA applies RBMLE-style optimism to pretrained world models for goal-conditioned planning, unlike reward-based online RL (Mete et al., 2026).Source: §2 · p. 3
  • JEPA world models predict target embeddings rather than reconstructing observations, and apply action-conditioned prediction to goal-conditioned planning.Examples listed are PLDM, DINO-WM, temporal straightening, and LeWM (Maes et al., 2026).Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries describe a latent world model that rolls forward imagined states and a goal-conditioned planner that scores candidate action sequences by goal distance. Test-time adaptation updates the model from recent observed transitions within a plan–execute–adapt–replan loop.

3 PRELIMINARIES

  • A predictor rolls forward latent states under an action sequence, producing a trajectory differentiable with respect to actions and predictor parameters.The model uses latent states and action embeddings, and the trajectory is denoted τθ.Source: §3 · p. 4
  • Under MPC, candidate action sequences are scored by the weighted goal-distance cost of their imagined latent-state trajectories.The step weights set relative importance, and distance is squared Euclidean distance.Source: §3 · p. 4
  • AdaJEPA adapts at test time by minimizing one-step latent prediction loss on a buffer of up to N recent observed transitions.The buffer contains observation, action, and next-observation tuples.Source: §3 · p. 4
  • A selected set of encoder and/or predictor parameters is updated for U gradient steps, and the adapted model is used at the next replanning step.After each task, the model resets to pretrained parameters and the buffer is cleared.Source: §3 · p. 4

4 OPTIMISTIC TEST-TIME ADAPTATION FOR LATENT WORLD MODELS

OpTTA adapts a latent world model by combining prediction loss on observed transitions with goal-reaching cost on candidate imagined rollouts. It retains the AdaJEPA protocol while specifying which parameters and rollouts contribute to adaptation and how the optimism weight is set.

4 OPTIMISTIC TEST-TIME ADAPTATION FOR LATENT WORLD MODELS

  • OpTTA minimizes prediction error on observed transitions together with the goal-reaching cost of candidate rollouts.The planning component is the average goal-reaching cost of the candidate rollouts.Source: §4 · p. 4
  • The procedure executes a planned action, adds its observed transition to the buffer, then computes both losses for adaptation.The paper identifies parameter selection, rollout selection, and optimism-weight choice as configurable design choices.Source: §4 · p. 5
  • With α_k > 0, the planning term favors model estimates with lower candidate-rollout goal costs, while prediction loss anchors them to observed transitions.The objective is instantiated within AdaJEPA by replacing its prediction-only objective while retaining its planner, buffer, adapted parameter set, and update schedule.Source: §4 · p. 5
  • OpTTA updates selected encoder and/or predictor parameters while keeping action-embedding parameters fixed.Prediction loss is differentiated through the selected set; planning loss is differentiated only through a selected subset of predictor parameters.Source: §4 · p. 5
  • The rollout population uses CEM’s final-iteration candidates, whose action sequences stay fixed while their predicted trajectories are differentiated through the model.The planning gradient is with respect to model parameters, not the candidate actions.Source: §4 · p. 5
  • The optimism weight is scaled to target a scheduled planning-gradient magnitude relative to the prediction gradient, increasing optimism over replanning steps.The schedule limits optimism early; a stabilizer prevents division by zero, and the weight is treated as constant during differentiation.Source: §4 · p. 5

5 EXPERIMENTS

The experiments compare optimistic and prediction-only test-time adaptation across three latent world models, planning tasks, distribution shifts, adaptation scopes, and planners. OpTTA improves shifted-setting planning overall, while ablations examine its design choices, diagnostic behavior, and computational cost.

5.1 SETUP

  • The evaluation spans LeWM, temporal straightening, and DINO-WM across manipulation, reaching, and navigation, with visual and environment shifts and predictor-only or joint adaptation.The comparison uses the same initialization, planner, transition buffer, adapted parameters, and update schedule for AdaJEPA and OpTTA; DINO-WM uses predictor-only adaptation.Source: §5.1 · p. 6

5.2 PLANNING PERFORMANCE

  • OpTTA beats AdaJEPA in 14 of 15 shifted settings and gains 13.5 versus 6.9 points over Frozen predictor-only, and 13.5 versus 7.5 jointly.Final success is measured at replan 30.Source: §5.2 · p. 6, Table 1 · p. 7
  • On PushT visual corruptions, OpTTA delivers about 2.5× AdaJEPA’s gain for LeWM and temporal straightening; on PushT recolorings, its gains are 17.1 and 16.8 points.On DINO-WM PushT visual corruptions, the respective AdaJEPA and OpTTA gains are 11.8 and 13.4 points.Source: §5.2 · p. 6
  • OpTTA has higher Mean over replans in 13 of 15 shifted settings; without shift, both methods improve PushT, while Reacher is near ceiling and TwoRoom ends slightly below Frozen.With a gradient-based planner, OpTTA’s gains over Frozen on temporal-straightening PushT are similar to those with CEM.Source: §5.2 · pp. 6–7

5.3 ABLATIONS AND SENSITIVITY ANALYSIS

  • OpTTA’s final success changes little between the tested predictor blocks, while AdaJEPA is more block-sensitive on PushT; OpTTA also remains robust across tested optimism weights.For AdaJEPA on PushT, adapting the first rather than last block raises final success from 44.7% to 68.7%; reversing the planning-term sign lowers success below AdaJEPA on both tasks.Source: §5.3 · pp. 7–8
  • Using only the top-1 candidate yields OpTTA’s lowest final success on both tasks, while all 200 candidates give higher Mean over replans than the 30 elites.On Reacher, top-1 reaches 65.3% final success, below AdaJEPA’s 70.7%; all candidates yield Mean over replans of 47.8% on PushT and 46.2% on Reacher.Source: §5.3 · p. 8, Figure 3 · p. 8
  • On LeWM and temporal straightening, OpTTA adds about 20 ms per update but has lower compute time per solved task than AdaJEPA because it often solves tasks earlier.For DINO-WM, restricting the planning term to 30 elites cuts the per-replan cost from 2.3 to 0.35 s.Source: §5.3 · p. 8, Table 2 · p. 9
  • OpTTA’s advantage does not consistently come from lower prediction loss on executed transitions; on Reacher, it predicts an ungenerated goal-reaching trajectory more accurately than AdaJEPA.On both Reacher and PushT, its model becomes more optimistic about the chosen plan closing the remaining distance to the goal.Source: §5.3 · pp. 8–9, Figure 4 · p. 9

6 LIMITATIONS

OpTTA does not improve over the pretrained model on TwoRoom, and the best parameters to update vary across tasks. Its default configuration is not necessarily optimal, and its added computation raises a performance–cost trade-off.

6 LIMITATIONS

  • On TwoRoom, neither AdaJEPA nor OpTTA improves over the pretrained model, and the best parameters to update vary across tasks.The default OpTTA configuration need not be optimal, and the method adds computation per replan; the performance–cost trade-off warrants further study.Source: §6 · p. 9

7 CONCLUSION

The paper introduces OpTTA, which adds an optimistic planning loss to prediction-based test-time adaptation. Across diverse models, shifts, and planners, it improves final success over prediction-only adaptation in most settings, while the optimism-weight design remains open.

7 CONCLUSION

  • OpTTA augments prediction-based test-time adaptation with an optimistic planning loss.The authors conclude that world-model adaptation for planning should account for the downstream planning objective, not prediction alone.Source: §7 · p. 9
  • 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.The practical design of optimism weight α_k is not necessarily optimal, and its optimal design remains an open question.Source: §7 · p. 9

A ADDITIONAL EXPERIMENTAL RESULTS

The appendix provides additional experimental setup and results supporting the comparisons in Section 5.

A ADDITIONAL EXPERIMENTAL RESULTS

A.1 EXPERIMENTAL DETAILS

The appendix details the shifted evaluation settings, adaptation protocol, and metrics used for the experimental comparisons. OpTTA and AdaJEPA share the adaptation pipeline and comparison conditions, with OpTTA adding optimism hyperparameters.

A.1 EXPERIMENTAL DETAILS

  • The evaluation covers 43 task–shift combinations: 35 visual-corruption combinations and eight environment-specific combinations involving recolorings, held-out shapes, or dynamics changes.Visual combinations pair a world model and task with a corruption type, averaging over severity levels for that checkpoint.Source: §A.1 · p. 12
  • The evaluated models are LeWM, temporal straightening, and DINO-WM, with visual corruptions and environment shifts affecting color, geometry, mass, or damping.LeWM and temporal-straightening results average over four severity levels selected using the frozen model; DINO-WM uses one level per corruption in the aggregate comparison.Source: §A.1 · p. 12
  • The implementation builds on the public AdaJEPA codebase, ports its adaptation procedure to LeWM, and uses publicly available pretrained checkpoints.OpTTA adds the optimistic planning loss and gradient-balancing rule while retaining the AdaJEPA test-time adaptation pipeline.Source: §A.1 · p. 12
  • Comparisons use the same pretrained checkpoint, planner, start–goal pairs, transition buffer, adapted parameters, and update schedule, isolating the adaptation objective.Primary comparisons use ρ = 1 without tuning it for individual models, planners, tasks, shifts, or severity levels.Source: §A.1 · p. 12
  • CEM executes the first of five planned actions and replans for at most 30 steps; adaptation starts from pretrained weights with an empty task-local transition buffer.LeWM and temporal straightening use five seeds with 50 paired start–goal episodes per seed; DINO-WM uses three seeds with 50 episodes per seed, and its encoder stays frozen.Source: §A.1 · p. 12
  • At each replan, adaptation takes one Adam step on the last predictor block using the five most recent transitions, with AdaJEPA and OpTTA sharing these settings.OpTTA adds optimism hyperparameters that are held constant across world models, tasks, and shifts; a latent step covers five environment steps.Source: §A.1 · pp. 12–13, Table 3 · p. 13
  • Final is the percentage of episodes solved by replan 30, while Mean over replans averages cumulative success across replans 1–30.Reported means and standard errors are computed across seeds; visual-corruption rows average seven corruptions equally within each world-model–task setting, and gains over Frozen are paired by seed within adaptation scope.Source: §A.1 · p. 13

A.2 PER-SHIFT RESULTS

Tables 4 and 5 break down shifted-setting planning success by corruption type or environment change. Table 4 covers predictor-only adaptation across 43 task–shift combinations; Table 5 covers joint predictor–encoder adaptation across 36 combinations.

A.2 PER-SHIFT RESULTS

  • Table 4 covers all 43 task–shift combinations with the predictor adapted, while Table 5 covers 36 combinations with both predictor and encoder adapted.DINO-WM is excluded from Table 5 because its encoder remains frozen.Source: §A.2 · p. 13, Table 4 · p. 14, Table 5 · p. 15
  • Each visual-corruption entry averages the severity levels listed in its label.Both tables report planning success per shift as mean ± s.e. over seeds; parentheses show paired gains over Frozen, and bold marks the best method at printed precision.Source: Table 4 · p. 14, Table 5 · p. 15

A.3 SUCCESS ACROSS REPLANNING STEPS

The appendix presents success across replanning steps, corruption severities, and a separate gradient-based planning protocol. The figures show success curves or severity-level results, while the gradient-based table and figure use a 20-replan budget.

A.3 SUCCESS ACROSS REPLANNING STEPS

  • Figures report cumulative success across replanning steps, with visual-corruption curves averaging severity levels and environment shifts shown individually.Curves show means over seeds with ±1 s.e. bands; the final point is Final and the mean across 30 points is Mean over replans.Source: §A.3 · p. 16, Figure 5 · p. 16, Figure 6 · p. 17
  • Figure 10 shows temporal-straightening environment shifts individually, with the final point and mean over replans corresponding to the table’s Final and Mean.Its curves show means over seeds with ±1 s.e. bands.Source: Figure 10 · p. 19

A.4 SUCCESS ACROSS SEVERITY LEVELS

  • Figures 11–14 show final success at each visual-corruption severity for LeWM and temporal straightening; clean conditions are references, not part of corrupted-level averages.DINO-WM is omitted because its aggregate evaluation uses one severity level per corruption.Source: §A.4 · p. 19, Figure 11 · p. 19, Figure 12 · p. 20

A.5 GRADIENT-BASED PLANNER

  • The gradient-based PushT protocol directly optimizes action sequences and evaluates Lplan on the optimized plan, using 20 replans and constant ρ = 1 without the optimism ramp.The same start–goal pairs and applicable adaptation settings are used across Frozen, AdaJEPA, and OpTTA; this compares adaptation objectives but does not isolate planner choice relative to CEM.Source: §A.5 · p. 21, Figure 15 · p. 23
  • Table 6 reports Final at replan 20 and Mean over replans across steps 1–20; its values are not directly comparable to the 30-replan Table 1.Each visual-corruption result averages four severity levels.Source: Table 6 · p. 22

Related papers

Loading 2610.35636v1…