ALRA

Action-Latent Residual Adaptation for World Models

Krishnam Soni, Aditya Sehgal, Vedant Dave†, Elmar Rueckert†

CPS Lab, Montanuniversität Leoben

† Equal advising

NeurIPS 2026 Workshops · CL4FMAgents · TTCL · PTA

Overview

Latent world models lose planning performance when the dynamics change at test time. AdaJEPA recovers some of it by fine-tuning sub-networks of the predictor and visual encoder, on the order of 107 parameters, in every episode. ALRA leaves the entire world model frozen and learns only a 210-parameter affine correction to the action embedding, online, from the same self-supervised latent prediction error.

On PointMaze-Medium, this matches the strongest AdaJEPA baseline under all three dynamics shifts (50× damping, 0.2× mass and 10× mass) and leaves in-distribution performance unchanged, suggesting that shifts which rescale how actions translate into motion can be corrected in action space alone.

210
parameters adapted online, against 9,934,840 for the strongest AdaJEPA baseline
~47,000×
reduction in adapted parameters compared to the strongest AdaJEPA baseline
+6.6 pp
average success-rate gain over the frozen model across the three dynamics shifts (strongest AdaJEPA baseline: +6.5 pp)

Method

Side-by-side architecture diagram. In AdaJEPA, the visual encoder and the prediction network are adapted while the action and proprioception encoders stay frozen. In ALRA, every encoder and the prediction network stay frozen, and a new residual network takes the action and proprio embeddings and adds its output to the action embedding.
Architecture comparison. Snowflakes mark frozen modules, flames adapted ones. (a) AdaJEPA adapts the visual encoder’s final stage and one transformer block of the prediction network, while keeping the other encoders frozen. (b) ALRA keeps every existing module entirely frozen and instead adds a new 210-parameter affine residual to the action embedding.

When an agent encounters a dynamics shift, the environment’s visual structure is unchanged; what changes is how a commanded action translates into motion. A shift that rescales this action-to-outcome mapping leaves the predictor’s visual model intact and manifests instead as a miscalibration of the action channel.

ALRA therefore keeps the visual, proprioceptive and action encoders and the predictor entirely frozen, and computes a residual from the concatenated action and proprioceptive embeddings with an affine map gω:

Δe(at, pt; ω) = gω([Eα(at) ; Eψ(pt)]) Ẽα(at, pt; ω) = Eα(at) + Δe(at, pt; ω)W ∈ ℝ10×20, b ∈ ℝ10: 210 parameters

The residual is zero-initialized, so adaptation begins exactly at the frozen model, and it is applied to each frame of the history window with weights shared across frames. Within each episode, only these 210 parameters are optimized online, using the latent prediction term of the pre-training objective over a replay buffer of the five most recent transitions, with one gradient step per replanning step. The residual is reset to zero at the start of every episode.

Results

PointMaze-Medium success rates under test-time dynamics shifts, 5 seeds × 100 episodes.

Condition Frozen predlast+enclast predfirst+enclast ALRA (ours)
No shift (control) 82.4 ± 1.9 82.8 ± 3.6+0.4 ± 0.9 83.4 ± 3.0+1.0 ± 0.5 82.6 ± 1.9+0.2 ± 0.9
Damping 50× 47.6 ± 3.2 54.0 ± 5.3+6.4 ± 1.2 58.0 ± 2.6+10.4 ± 1.3 57.8 ± 5.1+10.2 ± 2.1
Density 0.2× 84.2 ± 3.1 86.8 ± 1.6+2.6 ± 1.3 88.0 ± 3.2+3.8 ± 1.6 89.0 ± 4.1+4.8 ± 1.5
Density 10× 36.4 ± 5.9 39.6 ± 6.7+3.2 ± 1.6 41.6 ± 6.9+5.2 ± 3.1 41.2 ± 5.9+4.8 ± 3.0
Adapted parameters 0 9,935,648 9,934,840 210

Values are mean success rates (%) ± one standard deviation across 5 seeds; the smaller line is the within-seed paired difference from the frozen baseline ± one standard error of that mean over seeds. A dot marks the best adapted method in each row and any within 2 percentage points of it; the frozen column is the reference and is excluded.

Adaptation across the episode

Four panels, one per condition, plotting success rate improvement over the frozen baseline in percentage points against MPC replanning steps 1 to 20, for predlast+enclast, predfirst+enclast and ALRA.
Success rate improvement over the frozen baseline across MPC replanning steps. Curves show cumulative success (an episode counts as successful at every step after it first reaches the goal), averaged over 5 seeds × 100 episodes; each adapted arm is differenced against the frozen run of the same seed. Shaded regions show ±1 standard error of the mean across seeds; the dashed line at zero is the frozen baseline. Endpoints correspond to the gains reported in the table.

Under the 50× damping shift, ALRA’s benefit emerges earlier in the episode: by replanning step 12 it has recovered +19.2 points over the frozen baseline, where predfirst+enclast has recovered +5.0, with the two converging by step 20. We attribute this to the size of the adapted parameter set, but do not isolate it from the learning rate, and the effect is specific to this condition.

Rollouts

One planning episode per condition, all from the same start and goal, so only the physics changes between rows. Columns compare the frozen world model, AdaJEPA’s predfirst+enclast (9,934,840 parameters adapted online) and ALRA (210).

No shiftAll three reach the goal
Frozen
predfirst+enclast
ALRA
Damping 50×The frozen model fails; both adapted methods reach the goal
Frozen
predfirst+enclast
ALRA
Density 0.2×All three reach the goal
Frozen
predfirst+enclast
ALRA
Density 10×All three reach the goal in this episode
Frozen
predfirst+enclast
ALRA

The episode was chosen to illustrate the four conditions under a common start and goal; it is not a random draw, and a single episode is not representative of aggregate performance. The quantitative comparison is the table above.

Scope and limitations

  1. A single environment. All results are on PointMaze-Medium; whether the action-space decomposition holds under contact-rich dynamics, where a shift may act on a manipulated object rather than the agent, remains open.
  2. Learning rate per shift family. ALRA uses η = 10−2 for the control and damping conditions and 10−4 for density, whereas the baselines use a single learning rate from the original implementation. This gives ALRA a tuning advantage the baselines were not given.
  3. Per-episode recalibration. The residual is reset at the start of each episode and never persisted, so ALRA does not accumulate a correction across deployment.
  4. Where it helps. ALRA’s effectiveness tracks how directly a shift rescales the mapping from action to motion. Under the 10× mass increase, which alters the agent’s momentum rather than the gain relating action to motion, no method, ALRA or the representational baselines, recovers more than a small share of the degradation.

BibTeX

@inproceedings{
soni2026alra,
title={{ALRA}: Action-Latent Residual Adaptation for World Models},
author={Krishnam Soni and Aditya Sehgal and Vedant Dave and Elmar Rueckert},
booktitle={PTA: From Pretrained Representations to Acting Agents},
year={2026},
url={https://openreview.net/forum?id=r5DF75OXlT}
}