ALRA: Action-Latent Residual Adaptation for World Models

Latent world models degrade when deployed under unseen dynamics shifts. Recent test-time adaptation methods such as AdaJEPA address this by fine-tuning selected sub-networks of the predictor and visual encoder online, updating on the order of $10^7$ parameters within every episode. We show that such updates are unnecessary when the shift acts on the agent's own dynamics: a 210-parameter affine residual applied to the action embedding, with all encoders and the predictor left entirely frozen, recovers as much planning performance as sub-network fine-tuning. On PointMaze-Medium, this residual matches the strongest AdaJEPA baseline under all three shifts: a $50\times$ damping increase, a $0.2\times$ mass reduction, and a $10\times$ mass increase, improving over the frozen baseline in each while leaving in-distribution performance unchanged, and updating four orders of magnitude fewer parameters. Our results suggest that dynamics shifts which rescale how commanded actions translate into motion can be corrected entirely in action space, leaving the learned visual and proprioceptive representation untouched.\newline \url{https://alra-neurips.github.io/}