Planning-aware updates help world models adapt to shifts
Adapt for Planning, Not just Prediction: Optimistic Test-Time World Model Adaptation · arXiv:2610.35636
In plain words
A world model helps a planner by imagining what might happen after different actions. During planning, it learns from the outcomes of actions it has already chosen, which may leave it wrong about other actions that could reach the goal.
OpTTA is a method that updates the model using both observed outcomes and the planner’s goal-reaching costs for imagined actions.
In eight shifted settings with only the predictor updated, its average gain in episode success by replan 30 was 13.5 percentage points over a frozen model, compared with 6.9 for prediction-only updates.
The research
When a planner learns from the outcomes of its own chosen actions, is that enough to make its model useful for reaching a goal under a shift? The model may fit what it has already tried while remaining wrong about other actions that could succeed.
OpTTA adds the planner’s goal-reaching costs on imagined action sequences to prediction-based updates of pretrained world models during planning.
The authors compared both update rules from the same pretrained checkpoints, planner, data buffer and update schedule. In eight shifted settings where only the predictor was adapted, they measured the share of episodes solved by replan 30.
OpTTA’s average gain over the frozen model was 13.5 percentage points, versus 6.9 for prediction-only adaptation.
This result supports the value of adding goal-reaching costs to the update, though the average covers those settings rather than every task or shift.
The results cover the tested models, tasks, shifts and planning procedures. On TwoRoom visual corruptions, neither adaptation method beat the frozen model, and adapting the model adds computation per replan.
What the evidence shows
- A planning mismatch
Optimistic Test-Time Adaptation (OpTTA) is a method for adapting pretrained world models while a planner works toward a goal. The problem it addresses is that a planner chooses actions using imagined outcomes, then the model learns from the outcomes of those same actions.
The paper argues that this can leave the model inaccurate about other actions that might reach the goal, so fitting observed transitions alone need not improve planning.
- Two signals shape updates
OpTTA keeps prediction loss on recent observed transitions, but adds the planner’s goal-reaching cost on imagined candidate action sequences. The first term fits the model to what happened; the second favors model estimates that make candidate plans appear closer to the goal.
In the implementation, the candidate actions stay fixed while their predicted trajectories change as the model is updated.
- Gains across two scopes
The authors measured final success as the percentage of episodes solved by replanning step 30. Across eight shifted settings with only the predictor adapted, OpTTA’s average gain over a frozen model was 13.5 percentage points, compared with 6.9 for prediction-only adaptation.
With both predictor and encoder adapted, the corresponding average gains across seven settings were 13.5 and 7.5 points.
- Evidence and exception
A diagnostic found that OpTTA did not consistently reduce prediction loss on the agent’s executed transitions. On Reacher, it was more accurate than prediction-only adaptation on a start-to-goal trajectory the agent had not generated; on PushT, the methods were similar there.
The broad pattern also has a limit: under visual corruptions in TwoRoom, both adaptation methods finished below the frozen model.
Deep
- Original paper authors
- Akshay Mete, Dileep Kalathil, P. R. Kumar
- Submitted to arXiv
- License
- CC BY 4.0