Structured Residual Adaptation for Multi-Task World Models

Abstract

Multi-task model-based reinforcement learning (MBRL) typically relies on a single shared world model, but heterogeneity in task dynamics and physical properties often leads to brittle training and negative transfer. We propose Structured Residual Adaptation (SRA), a parameter-efficient approach that adapts a frozen world-model backbone by injecting a lightweight, task-conditioned residual into the latent dynamics transition prior to normalization. To control adaptation capacity, SRA constrains the residual to a low-capacity structured function class, instantiated via either physics-inspired basis primitives or Random Fourier Features (RFF), and combined through a lightweight task-routing module. On the MT30 benchmark, SRA-RFF introduces approximately 2M trainable parameters on top of a frozen 48M-parameter TD-MPC2 backbone, achieving a mean score of 76.3 and outperforming a substantially larger 317M-parameter monolithic baseline under the same protocol. In a targeted out-of-distribution (OOD) physics-shift case study on two representative tasks, SRA suggests improved stability under mass and friction perturbations. Finally, gradient-alignment diagnostics show patterns consistent with reduced off-diagonal interference in the trainable parameter space.