Before We Adapt: The Zero-Shot Generalization Bottleneck in Continual World Models
Authors: Amir Mesbah, Di Wu
Organizations: Independent Researcher · Department of Mechanical, Industrial and Aerospace Engineering, Concordia University · Department of Electrical and Computer Engineering, McGill University
Continual reinforcement learning (CRL) requires agents to operate under environmental changes while reusing prior knowledge whenever possible. We argue that model-based reinforcement learning provides a promising framework for this setting because learned world models (WMs) can, in principle, support change detection, knowledge reuse, adaptation, and acquisition of new knowledge. However, such benefits depend heavily on whether the learned model can generalize beyond the contexts in which it was trained. In this work, we propose to formulate this problem with a contextual Markov decision process in which context is decomposed into behavioral and physical components. This view organizes continual WM learning into reuse, adaptation, and acquisition of new knowledge. As a first step, we study direct reuse when behavioral rules remain unchanged. Using Dyna-Q and TD-MPC in controlled FourRooms tasks, we find that zero-shot reuse can degrade substantially under physical-context changes and that this behavior depends strongly on the observation representation used by the WM. Additional ablation studies and latent-space diagnostics suggest that learned representations do not automatically provide the invariances needed for reuse. These results identify representation generalization as an important bottleneck for reliable world-model reuse in CRL.
Figures & tables
Figure 1: Zero-shot one-step dynamics generalization across the four FourRooms contexts for three observation representations. Results are shown for Dyna-Q and TD-MPC and averaged over actions, agent directions, and six seeds. Colors indicate the log generalization gap relative to each model’s training-context error. Position + direction contains the agent’s global position and orientation, distance contains local distances from obstacles in the agent’s relative directions, and combined contains both types of information.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 2: Contextual FourRooms example. Tasks may differ in physical context cP while sharing behavioral rules cB . When the relevant behavioral structure remains applicable, this corresponds to the reuse regime in which previously learned knowledge should remain useful rather than being learned again from scratch.
Hyperparameter
Value
Dyna
Total timesteps
100,000
Q-network learning rate
3×10−4
Dynamics-model learning rate
1×10−3
Learning starts
5,000
Training frequency
4
Appendix
Table 1: Training and model hyperparameters used for Dyna and TD-MPC.
Figure 3: Controlled physical-context shifts for Dyna-Q across the three observation representations. Translation conditions preserve the training layout while changing its global spatial realization; the layout condition changes local geometry at the training location; and layout + translation changes both. Colors show the one-step dynamics generalization gap.
Figure 4: Controlled physical-context shifts for TD-MPC across the three observation representations. The experiments separate translation, layout change, and their combination while keeping the underlying action and reward rules fixed. Colors show the one-step dynamics generalization gap.
Diagnostic
Probe input X
Target y
Metric
Context decoding
latent z
room identity
accuracy
Direction transfer
latent z
agent direction
accuracy
Position transfer
latent z
normalized local (x,y)
R2
Latent alignment
matched latent sets (Z0,Zc)
–
Linear CKA
Appendix
Table 2: Summary of the TD-MPC latent-space diagnostics.
Figure 5: Diagnostics of TD-MPC’s learned latent representations. Direction and position probes are trained on the training context and transferred to unseen contexts without updating their parameters. Context decoding instead measures whether room identity is linearly recoverable from the learned representation. Points show individual seeds and larger markers show the mean with 95% confidence intervals. The results show that representation generalization depends strongly on the observation mode, with the largest failures under spatial translation occurring when the representation is built from absolute position information.
A desirable property of a world model is the ability to learn continually across tasks, adapting to new environments without forgetting what the agent has already learnt. In particular, the ability to retain and reuse knowledge obtained from prior experiences underpins an agent's ability to efficiently adapt to novel environments, as the dynamics of the physical world can often be described in recurring mechanisms. However, the world model's measure of adaptation entangles two abilities: the speed and capacity to learn unseen tasks, and the reuse of knowledge already acquired, since incoming tasks carry novel content alongside what recurs. In order to isolate knowledge reuse from prior experiences, we propose a compositional continual learning benchmark for world models in robot manipulation. Specifically, we design each task curriculum with compositional tasks that combine aspects of the tasks seen in the sequence. We further factorise this composition along the axes of action and perception to better understand how different input modalities bottleneck knowledge reuse. We evaluate state-of-the-art world models under canonical continual learning methods, alongside a modular world model whose dynamics backbone contains explicitly reusable components. Results show that modularity balances reuse against forgetting better than conventional methods, but none solve the problem fully, leaving clear room for continual world models built to reuse without forgetting. More details are available on our project website: https://object814.github.io/Compositional-Continual-Learning/.
Haoyu Zhou, Joe Watson, Anson Lei +1
Applied Artificial Intelligence (A2I) Lab, Department of Engineering Science, University of Oxford
Continual learning treats degradation on previously seen data as evidence of failure, a convention inherited from settings with a stationary prediction target, where a correct label remains correct indefinitely. World models do not satisfy this condition. Their prediction target is the environment, which changes, so knowledge that was accurate when acquired may later become false, and discarding it is required behavior rather than a defect. Non-stationary ground truth is well studied in the concept drift literature and in the temporal factuality of language models, but has not been formulated for world models, which are distinctive in that they also encode knowledge that must never be revised. We argue that continual world models require retention stratified by invariance timescale, separating invariants such as physics and object permanence, which must never be revised, from instance-level facts that should be revised as soon as the environment changes. Standard forgetting metrics cannot distinguish a world model that has correctly revised outdated knowledge from one that has suffered catastrophic forgetting, and consequently rank a frozen model highest, while existing physical-reasoning benchmarks evaluate only frozen checkpoints. We propose differential retention, which reports invariant regression testing across the adaptation stream jointly with revision latency, without aggregation.
World models let agents plan by predicting the consequences of their actions, but changes in the environment can make them inaccurate. We study the problem of adapting a world model to an unknown test-time environment, drawn from a known environment family, using only a few episodes of interaction. Existing approaches trade off computational cost against expressivity, i.e., the range of models a method can produce. For example, in-context learning is computationally cheap but limited in expressivity, and gradient-based adaptation is expressive but computationally expensive. We present CLAW (Context-conditioned Low-rank Adaptation of World models), which addresses this tradeoff by using a hypernetwork to generate low-rank (LoRA) adapters at test time. During pretraining, we simulate adaptation to a variety of environments and jointly train the hypernetwork and base world model. At test time, we freeze the base model and use a forward pass of the hypernetwork to generate adapters from a small batch of test-time transitions. We evaluate CLAW in locomotion and manipulation environment families that vary in dynamics, embodiment, and reward. We show that, using only seconds of test-time data, CLAW outperforms gradient-based adaptation and in-context learning during online adaptation. We also show that CLAW avoids overfitting in data-scarce regimes, that its advantage comes from the expressive adapters rather than context conditioning, and that pretraining the hypernetwork jointly with the base model outperforms training it post hoc.