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

Krishnam Soni; Aditya Sehgal; Vedant Dave; Elmar Rueckert

cs.ROSep 18, 2026

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

Authors: Krishnam Soni, Aditya Sehgal, Vedant Dave, Elmar Rueckert

Organizations: Chair of Cyber-Physical-Systems, Montanuniversität Leoben, Austria

Abstract

Latent world models enable planning by predicting the effects of actions in a learned representation space, but their predictions can become unreliable when test-time conditions differ from training. Existing test-time adaptation methods address this by updating parts of the pretrained model, often modifying millions of parameters and requiring a choice of which internal components to adapt. We introduce Sandwich-Residuals, a lightweight alternative that keeps the pretrained world model frozen and learns only small residual corrections around the predictor. The residuals are optimized online using the model's self-supervised prediction error and require no rewards, labels, or source-domain data. Across 21 conditions on the AdaJEPA benchmark, our method achieves $1.3\times$ the success rate of the frozen model while retaining 95% of the performance of the strongest AdaJEPA variant and adapting 97-99% fewer parameters. Under compound shifts, this advantage increases to $1.9\times$ the success rate of the frozen model, while remaining comparable to internal block adaptation. We further demonstrate the same adaptation principle on a DINO-WM model for 3-D manipulation. These results suggest that effective test-time adaptation of world models does not necessarily require modifying their pretrained internal weights.

Figures & tables

Fig. 1: Overview and mechanism of Sandwich Residuals (SR). Left: At each replanning step, the world model predicts the latent representation of the next observation under the selected action $a_{t}$ , shown schematically as $\hat{o}_{t+1}$ , while executing that action yields the observed next observation $o_{t+1}$ . As test-time adaptation progresses, updates from these observed transitions reduce the prediction mismatch, illustrating progressive online recalibration of the model. Right: Circled $+$ symbols denote additive residual connections. The orange residual blocks are adapted during test-time adaptation, while the encoders and predictor remain frozen. See Section III for notation and variable definitions.

Condition	Frozen	pl	pf	Ours	$\Delta$
Medium Maze
No shift	80.8 $\pm$ 8.9	86.0 $\pm$ 5.8	84.0 $\pm$ 5.7	81.2 $\pm$ 6.7	+0.4
Damping 50 $\times$	46.4 $\pm$ 3.8	52.0 $\pm$ 4.0	59.2 $\pm$ 3.0	66.0 $\pm$ 3.7	+19.6
Density 0.2 $\times$	90.0 $\pm$ 4.5	86.4 $\pm$ 6.2	88.0 $\pm$ 7.7	88.4 $\pm$ 5.7	$-$ 1.6
Density 10 $\times$	38.0 $\pm$ 5.1	39.6 $\pm$ 6.1	42.8 $\pm$ 4.6	40.4 $\pm$ 3.0	+2.4
Blur $\sigma$ =2	83.2 $\pm$ 6.6	83.2 $\pm$ 6.4	84.0 $\pm$ 4.7	83.2 $\pm$ 7.7	+0.0

TABLE I: Success rate (%), mean

\pm

standard deviation over 5 seeds of 50 episodes. pl : predlast+enclast ; pf : predfirst+enclast ; ∗ : shape unseen during training;

\Delta

: gain of ours over the frozen model in percentage. Bold indicates the best method per condition.

Fig. 2: Example episodes under compound shifts, planned with each of the four models, from the same start state and goal: (a) Medium Maze, blur $\sigma=2$ + damping $50\times$ ; (b) PushT, blur $\sigma=2$ + controller $k_{v}$ $2\times$ ; (c) PushObj T, blur $\sigma=2$ + controller $k_{v}$ $2\times$ ; (d) PushObj square, blur $\sigma=2$ + controller $k_{v}$ $2\times$ . Dotted green lines trace the executed trajectory. Frozen fails in all four episodes; ours succeeds in all four, uniquely reaching the goal under $50\times$ damping in (a), while predfirst fails in (c) and predlast fails in (d). These are illustrative single episodes and not representative of the aggregate rates in Tables I and II .

Condition	Frozen	pl	pf	Ours	$\Delta$
Medium Maze: blur $\sigma=2$ + damping 50 $\times$	46.8 $\pm$ 7.8	54.0 $\pm$ 4.7	56.0 $\pm$ 4.0	65.6 $\pm$ 6.2	+18.8
Medium Maze: blur $\sigma=2$ + density 10 $\times$	38.8 $\pm$ 6.1	39.2 $\pm$ 3.6	37.2 $\pm$ 7.2	40.8 $\pm$ 3.0	+2.0
PushT: blur $\sigma=2$ + controller $k_{v}$ 2 $\times$	40.4 $\pm$ 4.3	76.8 $\pm$ 3.3	87.6 $\pm$ 2.6	87.6 $\pm$ 3.0	+47.2
PushT: red anchor + controller $k_{v}$ 2 $\times$	30.4 $\pm$ 4.3	42.4 $\pm$ 7.4	53.2 $\pm$ 4.1	51.2 $\pm$ 7.7	+20.8
PushT: blur $\sigma=2$ + red anchor + controller $k_{v}$ 2 $\times$	34.4 $\pm$ 8.6	41.2 $\pm$ 5.8	68.8 $\pm$ 6.7	59.6 $\pm$ 2.6	+25.2
PushObj T: blur $\sigma=2$ + controller $k_{v}$ 2 $\times$	26.8 $\pm$ 3.0	84.8 $\pm$ 2.7	85.6 $\pm$ 3.3	87.2 $\pm$ 5.0	+60.4

TABLE II: Success rate (%) under compound shifts, reported as mean

\pm

standard deviation over 5 seeds of 50 episodes. pl and pf are defined in Table I . Bold indicates the best method per condition.

\Delta

denotes the gain of ours over the frozen model in percentage.

Subset	$-$ Frozen	$-$ pl	$-$ pf
All 21 conditions	$+15.8\pm 1.2$	$+1.6\pm 0.7$	$-3.2\pm 0.8$
Medium Maze	$+4.2\pm 1.9$	$+2.4\pm 1.6$	$+0.2\pm 1.0$
Diverse Maze	$+17.6\pm 1.2$	$+9.6\pm 2.5$	$+0.4\pm 4.3$
PushObj	$+23.9\pm 1.7$	$-1.9\pm 0.9$	$-3.9\pm 1.5$
PushT	$+15.8\pm 1.8$	$+3.3\pm 0.9$	$-5.2\pm 1.2$
Compound (7)	$+31.5\pm 3.4$	$+8.2\pm 1.3$	$+0.8\pm 1.4$

TABLE III: Paired differences in success rate (percentage points), computed as ours minus each baseline and reported as mean

\pm

standard error over pairs. pl and pf are defined in Table I .

Fig. 3: Example OGBench-Cube episode for DINO-WM under a $10^{\circ}$ camera shift. All methods use the same start state, goal, and random seeds. PredLast executes ten replanning steps; however, the cube undergoes negligible displacement after the first two, causing the subsequent trajectory markers to overlap near its final position.

Condition	Frozen	pl	pf	Ours
No shift	62.7 $\pm$ 3.1	60.7 $\pm$ 3.1	64.0 $\pm$ 2.0	60.7 $\pm$ 5.0
Light 0.3	58.7 $\pm$ 4.2	60.0 $\pm$ 4.0	58.7 $\pm$ 4.6	61.3 $\pm$ 5.0
Camera $10^{\circ}$	54.0 $\pm$ 5.3	51.3 $\pm$ 5.0	52.0 $\pm$ 5.3	56.0 $\pm$ 2.0
Arm gain 0.5 $\times$	59.3 $\pm$ 5.8	59.3 $\pm$ 5.8	57.3 $\pm$ 3.1	62.0 $\pm$ 2.0
Average	58.7	57.8	58.0	60.0

TABLE IV: OGBench-Cube with our DINO-WM: success rate (%), reported as mean

\pm

standard deviation over 3 seeds of 50 episodes. pl : predlast ; pf : predfirst . Bold indicates the best method per condition.

Configuration	Medium Maze	PushT
Full	65.6 $\pm$ 6.2	87.6 $\pm$ 3.0
$-r^{a}$	63.6 $\pm$ 3.8 ( $-2.0$ )	77.6 $\pm$ 3.0 ( $-10.0$ )
$-r^{\mathrm{out}}$	56.4 $\pm$ 2.2 ( $-9.2$ )	72.0 $\pm$ 5.1 ( $-15.6$ )
$-r^{\mathrm{in}}$	64.4 $\pm$ 1.7 ( $-1.2$ )	71.2 $\pm$ 1.8 ( $-16.4$ )

TABLE V: Residual-pathway ablation under compound shifts. Success rates (%) are reported as mean

\pm

standard deviation over five seeds of 50 episodes. Parentheses indicate the change relative to the full configuration.

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

Adapt for Planning, Not just Prediction: Optimistic Test-Time World Model Adaptation
Oct 26, 2026Akshay Mete, Dileep Kalathil, P. R. KumarTest-Time Adaptation World Model-Based Planning
ECE Department, Texas A&M University, College Station, Texas, USA.
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

Oct 26, 2026J/K move · Enter open · S save

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

Akshay Mete, Dileep Kalathil, P. R. Kumar

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

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

Open paper AskPDF arXiv

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.

Abstract

Test-Time Adaptation

Explore similar work

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

AdaJEPA: An Adaptive Latent World Model

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