Before We Adapt: The Zero-Shot Generalization Bottleneck in Continual World Models

Amir Mesbah; Di Wu

cs.LGOct 29, 2026

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

Abstract

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

c_{P}

while sharing behavioral rules

c_{B}

. 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.

Dyna
Hyperparameter	Value
Total timesteps	$100{,}000$
Q-network learning rate	$3\times 10^{-4}$
Dynamics-model learning rate	$1\times 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)$	$R^{2}$
Latent alignment	matched latent sets $(Z_{0},Z_{c})$	–	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.

Continual RL

Continual reinforcement learning studies how agents adapt to changing tasks while retaining skills learned earlier.

See all 54 papers

Benchmarking World Models for Continual Learning on Compositional Tasks
Sep 18, 2026Haoyu Zhou, Joe Watson, Anson Lei +1World Models World Model Learning
Applied Artificial Intelligence (A2I) Lab, Department of Engineering Science, University of Oxford
What Should World Models Forget? Stratified Retention for Continual Adaptation
Oct 2, 2026Nishit Anand, Ramani Duraiswami, Dinesh ManochaWorld Model Learning Continual Learning
University of Maryland, College Park, USA
Amortized Low-Rank Adaptation for Model-Based Reinforcement Learning
Sep 10, 2026Fernando Palafox, David Fridovich-KeilTest-Time Adaptation World Model Learning
University of Texas at Austin Austin, TX, USA

Sep 18, 2026J/K move · Enter open · S save

Benchmarking World Models for Continual Learning on Compositional Tasks

Haoyu Zhou, Joe Watson, Anson Lei, Ingmar Posner

Applied Artificial Intelligence (A2I) Lab, Department of Engineering Science, University of Oxford

World Models World Model Learning World Models for Robotics World Model Evaluation

Open paper AskPDF arXiv

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/.

Abstract

Continual RL

Explore similar work

Benchmarking World Models for Continual Learning on Compositional Tasks

What Should World Models Forget? Stratified Retention for Continual Adaptation

Amortized Low-Rank Adaptation for Model-Based Reinforcement Learning