适应之前：持续世界模型中的零样本泛化瓶颈

# Before We Adapt: The Zero-Shot Generalization Bottleneck in Continual World Models
Source: [https://arxiv.org/html/2610.41041](https://arxiv.org/html/2610.41041)
\\workshoptitle

Continual World Models

Di WuAffiliation:Department of MechanicalIndustrial and Aerospace EngineeringConcordia UniversityAffiliation:Department of Electrical and Computer EngineeringMcGill UniversityEmail:[di\.wu@concordia\.ca](mailto:)

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

## 1Introduction

Continual reinforcement learning \(CRL\) considers a setting in which an agent continuously interacts with its environment and encounters changing tasks or conditions over time\([Khetarpal et al\., 2022](https://arxiv.org/html/2610.41041#bib.bib2);[Abel et al\., 2023](https://arxiv.org/html/2610.41041#bib.bib3)\)\. Effective continual agents should not simply relearn after every such change\. Instead, they should preserve previously acquired knowledge and determine whether it can be reused directly, adapted with limited additional learning, or whether genuinely new knowledge must be acquired\. Achieving this requires addressing several challenges, including catastrophic forgetting, the stability\-plasticity dilemma, limited interaction budgets, and detecting when the environment has changed\([Ahn et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib29);[Tang et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib27);[Hu et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib30);[Zuffer et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib4)\)\.

Model\-based reinforcement learning \(MBRL\) provides a promising foundation for enabling these capabilities through learned world models\. By learning an internal world model \(WM\) of the environment, an agent can reason about the consequences of its actions and generate imagined experience, reducing its dependence on costly interaction\([Ha and Schmidhuber, 2018](https://arxiv.org/html/2610.41041#bib.bib6);[M\. Moerland et al\., 2023](https://arxiv.org/html/2610.41041#bib.bib5)\)\. More importantly, a learned WM can also provide signals of environmental change when new observations become inconsistent with its predictions\([Zollicoffer et al\., 2026](https://arxiv.org/html/2610.41041#bib.bib7)\)\. Furthermore, previously learned dynamics provide a form of knowledge that should not need to be relearned whenever the environment changes\. Recent work has begun to explore world models for CRL, for example by using online world models\([Liu et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib28)\), or considering a generative model as a replay buffer\([Fu et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib21)\)to reduce catastrophic forgetting across sequential tasks\. However, model\-based continual reinforcement learning remains comparatively less studied than its model\-free counterpart\([Zuffer et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib4)\)\.

To study these capabilities, we need a formulation that captures the relationships among tasks encountered over time\. We view continual learning as interaction with a*contextual world*, rather than as a sequence of completely independent Markov decision processes \(MDPs\)\. Different tasks correspond to different contexts and may share different amounts of structure\. A contextual MDP provides a natural way to describe these relationships\([Hallak et al\., 2015](https://arxiv.org/html/2610.41041#bib.bib8)\)\. This view allows us to distinguish changes in the underlying behavioral rules, such as action effects, motion rules, or reward mechanisms, from changes in their physical realization, such as layout, geometry, or spatial location\. How much of this structure remains shared with previous experience can then guide whether existing knowledge should be reused, adapted, or replaced by newly learned knowledge\.

This perspective makes generalization a prerequisite for knowledge reuse\. If a new task preserves behavioral rules that have already been learned, a world model should remain useful even when their physical realization changes\. However, this requires its predictions and representations to generalize beyond the exact states and contexts encountered during training\. If they instead encode context\-specific information that is not relevant to the shared behavioral rules, even a task that should allow direct reuse may require substantial relearning\. Generalization is therefore central to whether previously learned WM knowledge can actually be reused\.

In this work, we develop this perspective using a contextual MDP formulation that separates behavioral and physical sources of task variation and organizes continual WM learning into three regimes:*reuse*,*adaptation*, and*acquisition of new knowledge*\. We focus experimentally on the first and most favorable regime, where behavioral rules remain shared and previously learned dynamics should be reusable without further learning\. Using Dyna\-Q\([Sutton, 1991](https://arxiv.org/html/2610.41041#bib.bib10)\)and TD\-MPC\([Hansen et al\., 2022](https://arxiv.org/html/2610.41041#bib.bib1)\)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 or learned representation\. These results identify the ability to learn context\-robust representations as an important requirement for reliable WM reuse\.

## 2Contextual formulation of continual WMs in reinforcement learning

We model continual reinforcement learning as interaction with a contextual family of environments rather than as a sequence of unrelated MDPs\. Following the contextual MDP \(CMDP\) perspective of[Kirk et al\. \(2023\)](https://arxiv.org/html/2610.41041#bib.bib17), building on[Hallak et al\. \(2015\)](https://arxiv.org/html/2610.41041#bib.bib8), we considerℳ=\(𝒮′,𝒜′,𝒪′,𝒞,𝒫,ℛ,ϕ,p𝒞\)\\mathcal\{M\}=\(\\mathcal\{S\}^\{\\prime\},\\mathcal\{A\}^\{\\prime\},\\mathcal\{O\}^\{\\prime\},\\mathcal\{C\},\\mathcal\{P\},\\mathcal\{R\},\\phi,p\_\{\\mathcal\{C\}\}\), where𝒮′\\mathcal\{S\}^\{\\prime\},𝒜′\\mathcal\{A\}^\{\\prime\}, and𝒪′\\mathcal\{O\}^\{\\prime\}are global state, action, and observation spaces,𝒞\\mathcal\{C\}is the context space,𝒫\\mathcal\{P\},ℛ\\mathcal\{R\}, andϕ\\phidefine the transition, reward, and observation mechanisms across contexts, andp𝒞p\_\{\\mathcal\{C\}\}is a distribution over contexts\. A contextci∈𝒞c\_\{i\}\\in\\mathcal\{C\}induces a taskτi=ℳ⁡\(ci\)\\tau\_\{i\}=\\mathcal\{M\}\(c\_\{i\}\), whose task\-specific mechanisms are obtained by fixing the context, i\.e\.,𝒫​i=𝒫\|ci\\mathcal\{P\}i=\\mathcal\{P\}\|\{c\_\{i\}\},ℛ​i=ℛ\|ci\\mathcal\{R\}i=\\mathcal\{R\}\|\{c\_\{i\}\}, andϕi=ϕ\|ci\\phi\_\{i\}=\\phi\|\_\{c\_\{i\}\}\. A continual learning problem then presents a context sequence𝐜=\(c1,c2,…\)\\mathbf\{c\}=\(c\_\{1\},c\_\{2\},\\ldots\)and the corresponding task sequence𝝉=\(τ1,τ2,…\)\\boldsymbol\{\\tau\}=\(\\tau\_\{1\},\\tau\_\{2\},\\ldots\)\. We make no independence or identical\-distribution assumption on this sequence\. The complete formulation is given in Appendix[A](https://arxiv.org/html/2610.41041#A1)\.

We decompose context asc=\(cB,cP\)∈𝒞B×𝒞Pc=\(c\_\{B\},c\_\{P\}\)\\in\\mathcal\{C\}\_\{B\}\\times\\mathcal\{C\}\_\{P\}, wherecBc\_\{B\}denotes*behavioral context*andcPc\_\{P\}denotes*physical context*\. Behavioral context represents the rules governing how the environment responds to the agent, whereas physical context represents how those rules are instantiated in a particular task, including its geometry, layout, and spatial configuration\. This decomposition describes the type of context change, while allowing either context component to affect the transition, reward, or observation mechanisms\.

This view suggests three learning regimes for continual world models\. When relevant behavioral structure remains valid across a context change, previously learned knowledge may be directly*reused*\. When only part of the relevant structure remains applicable, prior knowledge should support efficient*adaptation*\. When a context introduces previously unseen behavioral structure, the model must*acquire*new knowledge while preserving previously learned structure\. These regimes depend on the structure shared between the current context and prior experience\.

In our preliminary study, we examine the reuse regime in a contextual FourRooms environment by changing physical context while keeping behavioral context fixed\. A detailed example of this construction is provided in Appendix[B](https://arxiv.org/html/2610.41041#A2)\. We then ask whether learned world models generalize sufficiently across these physical\-context changes to support knowledge reuse\.

## 3Preliminary study of world models’ zero\-shot generalization

Zero\-shot generalization of world models remains relatively underexplored in continual model\-based RL\. Closest to our setting,[Prasanna et al\. \(2024\)](https://arxiv.org/html/2610.41041#bib.bib13);[Röder et al\. \(2026\)](https://arxiv.org/html/2610.41041#bib.bib20)study DreamerV3\([Hafner et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib9)\)and propose context\-aware extensions in which context is either provided or inferred\. Our goal is different\. Rather than focusing on a single architecture or proposing a remedy, we investigate which factors contribute to poor reuse through ablations over physical context and observation representation\. Our decomposition of physical and behavioral context also lets us study cases where behavioral rules remain shared while physical context changes\.

![Refer to caption](https://arxiv.org/html/2610.41041v1/fourroom_average_actions_all_observations.png)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\.We evaluate Dyna\-Q\([Sutton, 1991](https://arxiv.org/html/2610.41041#bib.bib10)\)and TD\-MPC\([Hansen et al\., 2022](https://arxiv.org/html/2610.41041#bib.bib1)\), training each model in one FourRooms context and evaluating it zero\-shot in the other rooms without parameter updates\. We report the log generalization gapGc=log10⁡\(Ec/E¯train\)G\_\{c\}=\\log\_\{10\}\(E\_\{c\}/\\overline\{E\}\_\{\\mathrm\{train\}\}\), where prediction error is normalized by the corresponding mean training\-context error; positive values indicate degradation relative to the training baseline\. We compare three observation representations to examine how representation affects reuse\. Full environment, training, and evaluation details are provided in Appendix[C](https://arxiv.org/html/2610.41041#A3)\.

Figure[1](https://arxiv.org/html/2610.41041#S3.F1)shows that zero\-shot generalization depends strongly on how the environment is represented\. With*position \+ direction*, both methods show increased prediction error in unseen rooms, especially as observations move further outside the training range\. Part of this degradation may arise from the representation itself, since spatial translation changes the global coordinates even when the local dynamics remain unchanged\. In contrast, distance\-based observations remain substantially more stable, while the combined representation shows intermediate behavior\. This suggests that sensitivity to physical\-context variation is an important contributor to the reuse bottleneck\. The distance representation should, however, be viewed as an optimistic reference because it is designed to be less sensitive to the spatial translations considered here\. Real observations are more likely to mix context\-stable and context\-dependent information, which is approximated by the combined representation\.

The FourRooms comparison changes both spatial location and layout, so it does not fully separate coordinate extrapolation from changes in local geometry\. Appendix[D\.1](https://arxiv.org/html/2610.41041#A4.SS1)therefore includes controlled translation\-only and layout\-only changes\. The translation conditions provide a cleaner test of reuse by keeping the layout and behavioral rules fixed\. Some error under layout changes is expected because new obstacle locations are not explicitly available from position\-based observations\. From a continual learning perspective, this is still the most favorable regime, since the relevant behavioral knowledge has already been learned\. Our results therefore suggest that successful reuse depends strongly on the representation\. This issue also appears in learned representations\. TD\-MPC diagnostics in Appendix[D\.2](https://arxiv.org/html/2610.41041#A4.SS2)show that linear probes trained on the training context do not always transfer to new contexts, and that representations of matched states can become less aligned after a context shift\. Reliable world model reuse therefore requires representations that preserve shared behavioral structure while being less sensitive to irrelevant physical variation\.

## 4Discussion and conclusion

We view continual world model learning through three regimes,*reuse*,*adaptation*, and*acquisition of new knowledge*\. Our contextual MDP view provides a common way to describe these cases\. A task change can be treated as a change in context, and the amount of useful structure shared with previous experience can guide how prior knowledge should be used\. In this work, we focused on the most favorable case, where behavioral rules remain shared, and previously learned dynamics should be reusable\. Our results show that even this case can be difficult when the representation is sensitive to changes in physical context\. More invariant observations generalize considerably better, while the representations learned by TD\-MPC do not automatically remove this dependence\. This suggests that reliable reuse requires representations that preserve shared behavioral structure while being less sensitive to physical changes that are not relevant to that structure\.

This observation also changes how we think about the role of representation learning in continual world models\. Learning a latent state does not by itself guarantee that information useful across contexts will be represented in a stable way\. A useful continual representation should preserve the parts of the environment that remain valid across tasks while allowing context\-specific information to change\. Our results therefore motivate further study of context\-invariant or factorized representations, although how such structure should be learned without removing information needed for prediction remains an open question\. World model prediction errors may also provide useful signals for detecting when the current context has changed and when previously learned knowledge is no longer sufficient\.

Several important parts of the framework remain to be studied\. We plan to extend the evaluation to DreamerV3 and other modern world model approaches and to include changes in behavioral context, where zero\-shot reuse may no longer be sufficient\. We also plan to study how an agent can detect context changes, estimate how much previous knowledge is still useful, and adapt when only part of the learned structure remains valid\. When genuinely new knowledge must be learned, plasticity becomes important, but it must be balanced with preserving previously acquired knowledge and avoiding negative transfer\. Another direction is to identify which parts of current world model agents, including encoders, latent dynamics models, recurrent components, and planners, are most responsible for these limitations\. Together, these questions move toward world models that can make better use of previous experience as their context changes over time\.

## Acknowledgments and Disclosure of Funding

We thank Benoit Boulet for his valuable support during this work\. The authors received no external funding for this work\.

## References

- Abelet al\.\(2023\)D\. Abel, A\. Barreto, B\. Van Roy, D\. Precup, H\. van Hasselt, and S\. SinghA definition of continual reinforcement learning\.InAdvances in Neural Information Processing Systems,A\. Oh, T\. Naumann, A\. Globerson, K\. Saenko, M\. Hardt, and S\. Levine \(Eds\.\),Vol\.36,pp\. 50377–50407\.External Links:[Document](https://dx.doi.org/10.52202/075280-2192)Cited by:[§1](https://arxiv.org/html/2610.41041#S1.p1.1)\.
- Ahnet al\.\(2025\)H\. Ahn, J\. Hyeon, Y\. Oh, B\. Hwang, and T\. MoonPrevalence of negative transfer in continual reinforcement learning: analyses and a simple baseline\.InThe Thirteenth International Conference on Learning Representations,Cited by:[§1](https://arxiv.org/html/2610.41041#S1.p1.1)\.
- Alyahyaet al\.\(2026\)A\. Alyahya, A\. A\. Siyabi, M\. R\. Ernst, L\. Yang, L\. Kuhlmann, and G\. KowadloARROW: augmented replay for RObust world models\.Transactions on Machine Learning Research\.External Links:ISSN 2835\-8856Cited by:[§E\.2](https://arxiv.org/html/2610.41041#A5.SS2.p1.1)\.
- Beukmanet al\.\(2023\)M\. Beukman, D\. Jarvis, R\. Klein, S\. James, and B\. RosmanDynamics generalisation in reinforcement learning via adaptive context\-aware policies\.External Links:2310\.16686Cited by:[§E\.3](https://arxiv.org/html/2610.41041#A5.SS3.p1.1)\.
- Fuet al\.\(2025\)H\. Fu, Y\. Sun, M\. Littman, and G\. KonidarisKnowledge retention in continual model\-based reinforcement learning\.InProceedings of the 42nd International Conference on Machine Learning,A\. Singh, M\. Fazel, D\. Hsu, S\. Lacoste\-Julien, F\. Berkenkamp, T\. Maharaj, K\. Wagstaff, and J\. Zhu \(Eds\.\),Proceedings of Machine Learning Research, Vol\.267,pp\. 17832–17851\.Cited by:[Appendix C](https://arxiv.org/html/2610.41041#A3.SS0.SSS0.Px1.p1.1),[§E\.2](https://arxiv.org/html/2610.41041#A5.SS2.p1.1),[§1](https://arxiv.org/html/2610.41041#S1.p2.1)\.
- Fuet al\.\(2022\)H\. Fu, S\. Yu, M\. Littman, and G\. KonidarisModel\-based lifelong reinforcement learning with bayesian exploration\.Advances in Neural Information Processing Systems35,pp\. 32369–32382\.Cited by:[§E\.2](https://arxiv.org/html/2610.41041#A5.SS2.p1.1)\.
- Ha and Schmidhuber \(2018\)D\. Ha and J\. SchmidhuberWorld models\.arXiv preprint arXiv:1803\.101222\(3\),pp\. 440\.Cited by:[§1](https://arxiv.org/html/2610.41041#S1.p2.1)\.
- Hafneret al\.\(2025\)D\. Hafner, J\. Pasukonis, J\. Ba, and T\. LillicrapMastering diverse control tasks through world models\.Nature,pp\. 1–7\.Cited by:[§E\.1](https://arxiv.org/html/2610.41041#A5.SS1.p1.1),[§3](https://arxiv.org/html/2610.41041#S3.p1.1)\.
- Hallaket al\.\(2015\)A\. Hallak, D\. Di Castro, and S\. MannorContextual markov decision processes\.arXiv preprint arXiv:1502\.02259\.Cited by:[§E\.3](https://arxiv.org/html/2610.41041#A5.SS3.p1.1),[§1](https://arxiv.org/html/2610.41041#S1.p3.1),[§2](https://arxiv.org/html/2610.41041#S2.p1.1)\.
- Hansenet al\.\(2022\)N\. Hansen, X\. Wang, and H\. SuTemporal difference learning for model predictive control\.International Conference on Machine Learning\.Cited by:[§E\.1](https://arxiv.org/html/2610.41041#A5.SS1.p1.1),[§1](https://arxiv.org/html/2610.41041#S1.p5.1),[§3](https://arxiv.org/html/2610.41041#S3.p2.1)\.
- Huet al\.\(2025\)J\. Hu, Z\. Lian, Z\. Wen, C\. Li, G\. Chen, X\. Wen, B\. Xiao, and M\. TanContinual knowledge adaptation for reinforcement learning\.InAdvances in Neural Information Processing Systems,D\. Belgrave, C\. Zhang, H\. Lin, R\. Pascanu, P\. Koniusz, M\. Ghassemi, and N\. Chen \(Eds\.\),Vol\.38, Main Conference,pp\. 36283–36313\.External Links:[Document](https://dx.doi.org/10.52202/085713-1219)Cited by:[§1](https://arxiv.org/html/2610.41041#S1.p1.1)\.
- Huanget al\.\(2021\)Y\. Huang, K\. Xie, H\. Bharadhwaj, and F\. ShkurtiContinual model\-based reinforcement learning with hypernetworks\.InIEEE International Conference on Robotics and Automation \(ICRA\),Cited by:[§E\.2](https://arxiv.org/html/2610.41041#A5.SS2.p1.1)\.
- Janneret al\.\(2021\)M\. Janner, J\. Fu, M\. Zhang, and S\. LevineWhen to trust your model: model\-based policy optimization\.External Links:1906\.08253Cited by:[§E\.1](https://arxiv.org/html/2610.41041#A5.SS1.p1.1)\.
- Kessleret al\.\(2023\)S\. Kessler, M\. Ostaszewski, M\. P\. Bortkiewicz, M\. Żarski, M\. Wolczyk, J\. Parker\-Holder, S\. J\. Roberts, and P\. MiłośThe effectiveness of world models for continual reinforcement learning\.InProceedings of The 2nd Conference on Lifelong Learning Agents,S\. Chandar, R\. Pascanu, H\. Sedghi, and D\. Precup \(Eds\.\),Proceedings of Machine Learning Research, Vol\.232,pp\. 184–204\.Cited by:[§E\.2](https://arxiv.org/html/2610.41041#A5.SS2.p1.1)\.
- Khetarpalet al\.\(2022\)K\. Khetarpal, M\. Riemer, I\. Rish, and D\. PrecupTowards continual reinforcement learning: a review and perspectives\.Journal of Artificial Intelligence Research75,pp\. 1401–1476\.Cited by:[§1](https://arxiv.org/html/2610.41041#S1.p1.1)\.
- Kirket al\.\(2023\)R\. Kirk, A\. Zhang, E\. Grefenstette, and T\. RocktäschelA survey of zero\-shot generalisation in deep reinforcement learning\.Journal of Artificial Intelligence Research76,pp\. 201–264\.External Links:ISSN 1076\-9757,[Document](https://dx.doi.org/10.1613/jair.1.14174)Cited by:[§E\.3](https://arxiv.org/html/2610.41041#A5.SS3.p1.1),[§2](https://arxiv.org/html/2610.41041#S2.p1.1)\.
- Kirkpatricket al\.\(2017\)J\. Kirkpatrick, R\. Pascanu, N\. Rabinowitz, J\. Veness, G\. Desjardins, A\. A\. Rusu, K\. Milan, J\. Quan, T\. Ramalho, A\. Grabska\-Barwinska,et al\.Overcoming catastrophic forgetting in neural networks\.Proceedings of the national academy of sciences114\(13\),pp\. 3521–3526\.Cited by:[§E\.2](https://arxiv.org/html/2610.41041#A5.SS2.p2.1)\.
- Leeet al\.\(2020\)K\. Lee, Y\. Seo, S\. Lee, H\. Lee, and J\. ShinContext\-aware dynamics model for generalization in model\-based reinforcement learning\.External Links:2005\.06800Cited by:[§E\.3](https://arxiv.org/html/2610.41041#A5.SS3.p1.1)\.
- Liuet al\.\(2025\)Z\. Liu, G\. Fu, C\. Du, W\. S\. Lee, and M\. LinContinual reinforcement learning by planning with online world models\.InProceedings of the 42nd International Conference on Machine Learning,A\. Singh, M\. Fazel, D\. Hsu, S\. Lacoste\-Julien, F\. Berkenkamp, T\. Maharaj, K\. Wagstaff, and J\. Zhu \(Eds\.\),Proceedings of Machine Learning Research, Vol\.267,pp\. 38397–38423\.Cited by:[§E\.2](https://arxiv.org/html/2610.41041#A5.SS2.p1.1),[§1](https://arxiv.org/html/2610.41041#S1.p2.1)\.
- M\. Moerlandet al\.\(2023\)T\. M\. Moerland, J\. Broekens, A\. Plaat, and C\. M\. JonkerModel\-based reinforcement learning: a survey\.Foundations and Trends in Machine Learning16\(1\),pp\. 1–118\.Cited by:[§E\.1](https://arxiv.org/html/2610.41041#A5.SS1.p1.1),[§1](https://arxiv.org/html/2610.41041#S1.p2.1)\.
- Maeset al\.\(2026\)L\. Maes, Q\. L\. Lidec, D\. Scieur, Y\. LeCun, and R\. BalestrieroLeWorldModel: stable end\-to\-end joint\-embedding predictive architecture from pixels\.External Links:2603\.19312,[Link](https://arxiv.org/abs/2603.19312)Cited by:[§E\.1](https://arxiv.org/html/2610.41041#A5.SS1.p1.1)\.
- Prasannaet al\.\(2024\)S\. Prasanna, K\. Farid, R\. Rajan, and A\. BiedenkappDreaming of many worlds: learning contextual world models aids zero\-shot generalization\.External Links:2403\.10967Cited by:[§E\.3](https://arxiv.org/html/2610.41041#A5.SS3.p1.1),[§3](https://arxiv.org/html/2610.41041#S3.p1.1)\.
- Rahimi\-Kalahroudiet al\.\(2023\)A\. Rahimi\-Kalahroudi, J\. Rajendran, I\. Momennejad, H\. van Seijen, and S\. ChandarReplay buffer with local forgetting for adapting to local environment changes in deep model\-based reinforcement learning\.InConference on Lifelong Learning Agents,pp\. 21–42\.Cited by:[§E\.2](https://arxiv.org/html/2610.41041#A5.SS2.p1.1)\.
- Röderet al\.\(2026\)F\. Röder, J\. Benad, M\. Eppe, and P\. BanerjeeDynamics\-aligned latent imagination in contextual world models for zero\-shot generalization\.Advances in Neural Information Processing Systems38,pp\. 72736–72771\.Cited by:[§E\.3](https://arxiv.org/html/2610.41041#A5.SS3.p1.1),[§3](https://arxiv.org/html/2610.41041#S3.p1.1)\.
- Schwarzet al\.\(2018\)J\. Schwarz, W\. Czarnecki, J\. Luketina, A\. Grabska\-Barwinska, Y\. W\. Teh, R\. Pascanu, and R\. HadsellProgress & compress: a scalable framework for continual learning\.InInternational conference on machine learning,pp\. 4528–4537\.Cited by:[§E\.2](https://arxiv.org/html/2610.41041#A5.SS2.p2.1)\.
- Shimizu and Tomizuka \(2024\)Y\. Shimizu and M\. TomizukaBisimulation metric for model predictive control\.External Links:2410\.04553Cited by:[§E\.3](https://arxiv.org/html/2610.41041#A5.SS3.p3.1)\.
- Sutton \(1991\)R\. S\. SuttonDyna, an integrated architecture for learning, planning, and reacting\.SIGART Bull\.2\(4\),pp\. 160–163\.External Links:ISSN 0163\-5719,[Document](https://dx.doi.org/10.1145/122344.122377)Cited by:[§E\.1](https://arxiv.org/html/2610.41041#A5.SS1.p1.1),[§1](https://arxiv.org/html/2610.41041#S1.p5.1),[§3](https://arxiv.org/html/2610.41041#S3.p2.1)\.
- Tanget al\.\(2025\)H\. Tang, J\. Obando\-Ceron, P\. S\. Castro, A\. Courville, and G\. BersethMitigating plasticity loss in continual reinforcement learning by reducing churn\.InProceedings of the 42nd International Conference on Machine Learning,A\. Singh, M\. Fazel, D\. Hsu, S\. Lacoste\-Julien, F\. Berkenkamp, T\. Maharaj, K\. Wagstaff, and J\. Zhu \(Eds\.\),Proceedings of Machine Learning Research, Vol\.267,pp\. 58883–58904\.Cited by:[§1](https://arxiv.org/html/2610.41041#S1.p1.1)\.
- Zollicofferet al\.\(2026\)G\. Zollicoffer, T\. Chopra, M\. Yan, X\. Ma, K\. Eaton, and M\. RiedlWorld model robustness via surprise recognition\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 3146–3155\.Cited by:[§1](https://arxiv.org/html/2610.41041#S1.p2.1)\.
- Zufferet al\.\(2025\)A\. Zuffer, M\. Burke, and M\. HarandiAdvancements and challenges in continual reinforcement learning: a comprehensive review\.arXiv preprint arXiv:2506\.21899\.Cited by:[§1](https://arxiv.org/html/2610.41041#S1.p1.1),[§1](https://arxiv.org/html/2610.41041#S1.p2.1)\.

## Appendix ADetailed definition of contextual continual world model formulation

We consider a contextual family of environmentsℳ=⟨𝒮′,𝒜′,𝒪′,𝒞,𝒫,ℛ,ϕ,p𝒞⟩\\mathcal\{M\}=\\langle\\mathcal\{S\}^\{\\prime\},\\mathcal\{A\}^\{\\prime\},\\mathcal\{O\}^\{\\prime\},\\mathcal\{C\},\\mathcal\{P\},\\mathcal\{R\},\\phi,p\_\{\\mathcal\{C\}\}\\rangle, where𝒮′\\mathcal\{S\}^\{\\prime\},𝒜′\\mathcal\{A\}^\{\\prime\}, and𝒪′\\mathcal\{O\}^\{\\prime\}are the global state, action, and observation spaces,𝒞\\mathcal\{C\}is the context space,𝒫\\mathcal\{P\},ℛ\\mathcal\{R\}, andϕ\\phidefine the transition, reward, and observation mechanisms across contexts, andp𝒞p\_\{\\mathcal\{C\}\}is a distribution over contexts\.

We decompose the context asc=\(cB,cP\)∈𝒞B×𝒞Pc=\(c\_\{B\},c\_\{P\}\)\\in\\mathcal\{C\}\_\{B\}\\times\\mathcal\{C\}\_\{P\}, wherecBc\_\{B\}is the*behavioral context*andcPc\_\{P\}is the*physical context*\. Behavioral context describes the rules governing how the environment responds to the agent, while physical context describes how these rules are instantiated in a particular environment\. Either type of context change may affect the transition, reward, or observation mechanisms\.

The state, action, and observation spaces of a task can depend on its context\. We represent them using the set\-valued mappings𝒮:𝒞→2𝒮′\\mathcal\{S\}:\\mathcal\{C\}\\rightarrow 2^\{\\mathcal\{S\}^\{\\prime\}\},𝒜:𝒞→2𝒜′\\mathcal\{A\}:\\mathcal\{C\}\\rightarrow 2^\{\\mathcal\{A\}^\{\\prime\}\}, and𝒪:𝒞→2𝒪′\\mathcal\{O\}:\\mathcal\{C\}\\rightarrow 2^\{\\mathcal\{O\}^\{\\prime\}\}\. For a contextcic\_\{i\}, the corresponding task\-specific spaces are therefore𝒮i=𝒮⁡\(ci\)⊆𝒮′\\mathcal\{S\}\_\{i\}=\\mathcal\{S\}\(c\_\{i\}\)\\subseteq\\mathcal\{S\}^\{\\prime\},𝒜i=𝒜⁡\(ci\)⊆𝒜′\\mathcal\{A\}\_\{i\}=\\mathcal\{A\}\(c\_\{i\}\)\\subseteq\\mathcal\{A\}^\{\\prime\}, and𝒪i=𝒪⁡\(ci\)⊆𝒪′\\mathcal\{O\}\_\{i\}=\\mathcal\{O\}\(c\_\{i\}\)\\subseteq\\mathcal\{O\}^\{\\prime\}\. These spaces do not need to be disjoint, so different tasks may share states, actions, or observations\. A contextcic\_\{i\}induces a taskτi=ℳ⁡\(ci\)\\tau\_\{i\}=\\mathcal\{M\}\(c\_\{i\}\), which we write as

τi=⟨𝒮i,𝒜i,𝒪i,𝒫i,ℛi,ϕi⟩\.\\tau\_\{i\}=\\left\\langle\\mathcal\{S\}\_\{i\},\\mathcal\{A\}\_\{i\},\\mathcal\{O\}\_\{i\},\\mathcal\{P\}\_\{i\},\\mathcal\{R\}\_\{i\},\\phi\_\{i\}\\right\\rangle\.The task\-specific mechanisms are obtained by fixing the context:𝒫i\(⋅∣s,a\)=𝒫\(⋅∣s,a,ci\)\\mathcal\{P\}\_\{i\}\(\\cdot\\mid s,a\)=\\mathcal\{P\}\(\\cdot\\mid s,a,c\_\{i\}\),ℛi​\(s,a,s′\)=ℛ⁡\(s,a,s′,ci\)\\mathcal\{R\}\_\{i\}\(s,a,s^\{\\prime\}\)=\\mathcal\{R\}\(s,a,s^\{\\prime\},c\_\{i\}\), andϕi​\(s\)=ϕ⁡\(s,ci\)\\phi\_\{i\}\(s\)=\\phi\(s,c\_\{i\}\)\. Thus, the context determines both the spaces that are valid for a task and its transition, reward, and observation mechanisms\.

A continual learning problem consists of a context sequence𝐜=\(c1,c2,…,cN\)\\mathbf\{c\}=\(c\_\{1\},c\_\{2\},\\ldots,c\_\{N\}\), withci∈𝒞c\_\{i\}\\in\\mathcal\{C\}andN∈ℕ∪\{∞\}N\\in\\mathbb\{N\}\\cup\\\{\\infty\\\}, which induces the task sequence𝝉=\(τ1,τ2,…,τN\)\\boldsymbol\{\\tau\}=\(\\tau\_\{1\},\\tau\_\{2\},\\ldots,\\tau\_\{N\}\), whereτi=ℳ⁡\(ci\)\\tau\_\{i\}=\\mathcal\{M\}\(c\_\{i\}\)\. We do not assume that the contexts are independent or identically distributed\. More generally, the context at taskiimay depend on previous contexts throughp\(ci∣c1:i−1\)p\(c\_\{i\}\\mid c\_\{1:i\-1\}\)\. For a finite task sequence, this can be written as

pseq\(c1:N\)=p\(c1\)∏i=2Np\(ci∣c1:i−1\)\.p\_\{\\mathrm\{seq\}\}\(c\_\{1:N\}\)=p\(c\_\{1\}\)\\prod\_\{i=2\}^\{N\}p\(c\_\{i\}\\mid c\_\{1:i\-1\}\)\.Independent sampling from the CMDP context distribution is a special case in whichp\(ci∣c1:i−1\)=p𝒞\(ci\)p\(c\_\{i\}\\mid c\_\{1:i\-1\}\)=p\_\{\\mathcal\{C\}\}\(c\_\{i\}\)\. Sinceci=\(cB,i,cP,i\)c\_\{i\}=\(c\_\{B,i\},c\_\{P,i\}\), a change between task sequences can occur in three ways:

cB,i\+1≠cB,i,cP,i\+1=cP,i\\displaystyle c\_\{B,i\+1\}\\neq c\_\{B,i\},\\quad c\_\{P,i\+1\}=c\_\{P,i\}\(behavioral change\),\\displaystyle\\text\{\(behavioral change\)\},cB,i\+1=cB,i,cP,i\+1≠cP,i\\displaystyle c\_\{B,i\+1\}=c\_\{B,i\},\\quad c\_\{P,i\+1\}\\neq c\_\{P,i\}\(physical change\),\\displaystyle\\text\{\(physical change\)\},cB,i\+1≠cB,i,cP,i\+1≠cP,i\\displaystyle c\_\{B,i\+1\}\\neq c\_\{B,i\},\\quad c\_\{P,i\+1\}\\neq c\_\{P,i\}\(joint change\)\.\\displaystyle\\text\{\(joint change\)\}\.The transition, reward, and observation functions may in general depend on both context components, i\.e\.,𝒫\(⋅∣s,a,cB,cP\)\\mathcal\{P\}\(\\cdot\\mid s,a,c\_\{B\},c\_\{P\}\),ℛ⁡\(s,a,s′,cB,cP\)\\mathcal\{R\}\(s,a,s^\{\\prime\},c\_\{B\},c\_\{P\}\), andϕ⁡\(s,cB,cP\)\\phi\(s,c\_\{B\},c\_\{P\}\)\. However, a change in context does not require all of these functions to change\. For example, a function can remain unchanged with respect to physical context if its value is the same for differentcPc\_\{P\}whilecBc\_\{B\}is fixed, and similarly for behavioral context\. Therefore, each task change may affect only some of the task\-specific spaces or functions, while the remaining parts can be shared across tasks\.

## Appendix BFourRooms contextual\-world example

Figure[2](https://arxiv.org/html/2610.41041#A2.F2)illustrates how the contextual\-world formulation is instantiated in FourRooms\. Each room can be viewed as a task with its own physical contextcPc\_\{P\}, while the underlying behavioral contextcBc\_\{B\}can remain shared\. Physical context may vary through spatial location, layout, or obstacle configuration, whereas the rules governing actions and rewards remain unchanged\. For example, the action*forward*follows the same rule in every context\. The agent moves one cell when the target cell is free and remains in place when it is blocked\.

Figure 2:Contextual FourRooms example\. Tasks may differ in physical contextcPc\_\{P\}while sharing behavioral rulescBc\_\{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\.
## Appendix CExperimental details

#### Environment\.

We use a contextual version of the FourRooms environment\. The environment and the*position \+ direction*observation follow the continual model\-based RL setup of[Fu et al\. \[2025\]](https://arxiv.org/html/2610.41041#bib.bib21), while we modify the task contexts and evaluation procedure to study zero\-shot generalization\. The transition dynamics are deterministic\. The agent can turn left, turn right, or move forward by one grid cell; forward movement succeeds when the target cell is free and otherwise leaves the agent in place\. We use the same dense reward mechanism across contexts,rt=1−∥pt−pg∥2/dgoalr\_\{t\}=1\-\\lVert p\_\{t\}\-p\_\{g\}\\rVert\_\{2\}/d\_\{\\mathrm\{goal\}\}, wheredgoal=12d\_\{\\mathrm\{goal\}\}=12\. Thus, action semantics, motion rules, and the reward mechanism remain shared while physical properties such as spatial location, layout, and obstacle configuration may vary\.

We evaluate three observation representations\. The*position \+ direction*observation isotpos=\(x~t,y~t,cos⁡θt,sin⁡θt\),o\_\{t\}^\{\\mathrm\{pos\}\}=\(\\tilde\{x\}\_\{t\},\\tilde\{y\}\_\{t\},\\cos\\theta\_\{t\},\\sin\\theta\_\{t\}\),where\(xt,yt\)\(x\_\{t\},y\_\{t\}\)is the agent’s global position andθt\\theta\_\{t\}its orientation\. The global coordinates are min–max normalized to\[−1,1\]\[\-1,1\]using fixed global bounds\[−26,130\]\[\-26,130\]and clipped to this range\. The*distance*observation isotdist=\(cos⁡θt,sin⁡θt,dtfront,dtback,dtleft,dtright\),o\_\{t\}^\{\\mathrm\{dist\}\}=\(\\cos\\theta\_\{t\},\\sin\\theta\_\{t\},d\_\{t\}^\{\\mathrm\{front\}\},d\_\{t\}^\{\\mathrm\{back\}\},d\_\{t\}^\{\\mathrm\{left\}\},d\_\{t\}^\{\\mathrm\{right\}\}\),where the four distances are measured to the nearest obstacle or room boundary in each relative direction and normalized by2​r2r, with room sizer=13r=13\. The*combined*observation isotcomb=\(x~t,y~t,cos⁡θt,sin⁡θt,dtfront,dtback,dtleft,dtright\),o\_\{t\}^\{\\mathrm\{comb\}\}=\(\\tilde\{x\}\_\{t\},\\tilde\{y\}\_\{t\},\\cos\\theta\_\{t\},\\sin\\theta\_\{t\},d\_\{t\}^\{\\mathrm\{front\}\},d\_\{t\}^\{\\mathrm\{back\}\},d\_\{t\}^\{\\mathrm\{left\}\},d\_\{t\}^\{\\mathrm\{right\}\}\),and therefore contains both global\-position and local\-distance information\. The distance representation is less sensitive to global spatial translation, while the combined representation mixes both global and local information\. The position \+ direction representation is not invariant to spatial translation\. Under this observation, translating the environment changes\(xt,yt\)\(x\_\{t\},y\_\{t\}\)and may move observations outside the range seen during training, even when the underlying behavioral rules remain unchanged\. Thus, some degradation under translation may arise from the representation choice itself\. Under position \+ direction observations, the locations of newly introduced obstacles are also not explicitly provided to the model\.

#### Training protocol\.

Dyna\-Q and TD\-MPC are trained only in theroom00context\. For each algorithm and observation representation, we train six models using random seeds\{1,2,43,100,256,314\}\\\{1,2,43,100,256,314\\\}for100,000100\{,\}000environment steps\. During evaluation, all model parameters are frozen, and no adaptation or fine\-tuning is performed in the test contexts\. The hyperparameters used for both methods are summarized in Table[1](https://arxiv.org/html/2610.41041#A3.T1)\. The planning horizon remains fixed at33throughout training\.

Table 1:Training and model hyperparameters used for Dyna and TD\-MPC\.
#### Linear schedules\.

We use the notationlinear⁡\(a,b,T\)\\mathrm\{linear\}\(a,b,T\)for a schedule that linearly interpolates fromaatobboverTTenvironment steps and remains atbbafterward\. The standard\-deviation schedule islinear⁡\(1\.0,0\.1,250000\)\\mathrm\{linear\}\(1\.0,0\.1,250000\)\. Since training lasts100000100000environment steps, only the first part of this schedule is used\.

#### Generalization regimes\.

We organize the four base rooms according to how their spatial observation support relates to the training room\.room00is the*in\-distribution reference*, since it is the context used during training\.room10androom01form an*intermediate*regime where one spatial coordinate lies in a range related to the training context while the other moves outside it\. Finally,room11provides the*full extrapolation*case, in which both spatial dimensions lie outside the training range\. The layouts also differ between these rooms; controlled translation\-only and layout\-only experiments in Appendix[D\.1](https://arxiv.org/html/2610.41041#A4.SS1)are used for more investigation\.

#### One\-step dynamics evaluation\.

For every valid state in an evaluation context, we test the three actions\{left,right,forward\}\\\{\\text\{left\},\\text\{right\},\\text\{forward\}\\\}for all four agent directions\. For Dyna\-Q, next\-state prediction error is measured in observation space\. TD\-MPC predicts in its learned latent space, so its predicted latent next state is compared with the encoding of the true next observation\. We use mean squared error \(MSE\) as the one\-step dynamics error\. Because the two methods operate in different representation spaces, we do not compare their raw MSE values directly\.

#### Generalization gap\.

For each trained run and action, we first compute the mean one\-step dynamics error over the complete training context,E¯train\\overline\{E\}\_\{\\mathrm\{train\}\}\. For an evaluated state with errorEcE\_\{c\}, we reportGc=log10⁡\(Ec/E¯train\)G\_\{c\}=\\log\_\{10\}\(E\_\{c\}/\\overline\{E\}\_\{\\mathrm\{train\}\}\)\. Positive values indicate higher error than the training\-context average and negative values indicate lower error\. Importantly, individual cells in the training\-context heatmap are not expected to equal zero because each cell is normalized by the*mean*error over the whole training context rather than by its own error\. For the reported heatmaps, values are averaged over agent directions and seeds and then equally across the three actions\.

#### Evaluation contexts\.

In addition to the four base FourRooms tasks shown in the main paper, we use controlled physical\-context changes to separate different sources of generalization error\. Positions are expressed in grid\-cell coordinates, and each room has sizer=13r=13\. Context offsets are converted to global\-coordinate offsets by multiplying them byrr\. Thus, adjacent base rooms are separated by1313grid cells along the corresponding axis\. The two translation\-only conditions preserve the training layout, obstacles, and behavioral rules while shifting the global coordinates by2​r=262r=26and3​r=393r=39grid cells along thexx\-axis, respectively\. The*layout*condition changes the obstacle layout while keeping the original spatial location\. Finally, the*combined*\(layout \+ translation\) condition changes both\. The translation\-only conditions therefore isolate sensitivity to global spatial translation while preserving the local environment structure\.

## Appendix DComplementary generalization results

### D\.1Controlled physical\-context shifts

The FourRooms results in the main text show that generalization depends strongly on the observation representation, but the unseen rooms change both spatial location and layout\. We therefore use controlled physical\-context shifts to separate these effects\. Figures[3](https://arxiv.org/html/2610.41041#A4.F3)and[4](https://arxiv.org/html/2610.41041#A4.F4)show the results for Dyna\-Q and TD\-MPC, respectively\.

![Refer to caption](https://arxiv.org/html/2610.41041v1/spatial_average_actions_all_observations_dyna.png)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\.The translation conditions provide the cleanest comparison because the layout, obstacles, action rules, and reward mechanism are unchanged\. Only the global spatial realization differs from training\. Position \+ direction observations become increasingly sensitive to these shifts, while the distance\-based representation remains considerably more stable\. This supports the main\-text observation that generalization depends on whether the representation contains information tied to a particular physical context\.

The layout condition tests a different source of change\. Here the global spatial region remains the same, but obstacle locations differ from training\. Errors around changed obstacles are expected because their new locations are not explicitly represented by position \+ direction observations\. This case therefore measures both reuse of the shared motion rules and the difficulty of predicting transitions under previously unseen geometry\.

![Refer to caption](https://arxiv.org/html/2610.41041v1/spatial_average_actions_all_observations_tdmpc.png)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\.The combined \(layout \+ translation\) condition combines both sources of difficulty\. The model must operate in a new spatial region while also encountering a different local geometry\. We therefore interpret this condition as a harder physical context change rather than as a pure test of coordinate extrapolation\.

### D\.2TD\-MPC latent representation diagnostics

Because TD\-MPC predicts dynamics in a learned latent space, we further examine whether the learned representation preserves useful information across contexts\. Table[2](https://arxiv.org/html/2610.41041#A4.T2)summarizes the training targets and evaluation settings for the representation diagnostics\. Direction and local\-position probes are trained only on the training context and then evaluated unchanged in unseen contexts\. Context decodability is instead a post\-hoc diagnostic trained across contexts to test whether context identity can be recovered from the latent representation\.

For all probe plots, individual points correspond to independently trained seeds, while the larger marker and error bar show the mean and 95% confidence interval across seeds\. Below, we describe each probing experiment in more details\.

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\.#### Context decodability\.

We first ask whether the latent representation contains enough information to identify the room in which a state was observed\. We pool latent vectors from the four base rooms \(room00,room10,room01, androom11\) and use the latent vector as input,X=zX=z, and the room identity as the classification target\. We use an 80/20 stratified train/test split with random seed 42\. The probe is a standardization layer followed by logistic regression with a maximum of 3000 iterations\. Since there are four rooms, chance accuracy is0\.250\.25\. This probe measures whether context identity is linearly decodable from the latent representation; it is not a cross\-context transfer test\.

#### Direction transfer\.

We next ask whether the agent’s direction is represented in a way that remains stable across contexts\. A linear classifier is trained only on latent vectors from the training context,room00, withX=zX=zand the four\-valued agent direction as the target\. We again use an 80/20 split, stratified by direction, with random seed 42\. The classifier consists of standardization followed by logistic regression with a maximum of 3000 iterations\. We first report accuracy on the held\-out part ofroom00\. We then refit the probe using allroom00states and evaluate it on the unseen contexts without updating the probe\. Chance accuracy is0\.250\.25\.

#### Local\-position transfer\.

We similarly test whether local position is encoded consistently across contexts\. The probe input is again the latent vector,X=zX=z, while the regression target is the normalized local position,y=\(xlocal/S,ylocal/S\)y=\(x\_\{\\mathrm\{local\}\}/S,y\_\{\\mathrm\{local\}\}/S\), whereSSis the room size\. The same 80/20 split used for the direction probe is used here\. We fit a standardization layer followed by ridge regression with regularizationα=1\\alpha=1\. After measuring performance on the held\-out training states, the probe is refitted using all training\-room states and transferred unchanged to the unseen contexts\. We report the coefficient of determination,R2R^\{2\}, in Fig\.[5](https://arxiv.org/html/2610.41041#A4.F5);R2=0R^\{2\}=0corresponds to the constant\-mean baseline, while negative values mean that the transferred probe performs worse than this baseline\.

#### Latent alignment\.

The final diagnostic does not train a probe\. Instead, we directly compare the geometry of latent representations between the training context and each unseen context\. We match states that have the same local position and agent direction in both contexts and compare their latent matrices\. We use linear centered kernel alignment \(CKA\), a similarity measure between two representations\. A value near11indicates that the matched latent representations preserve very similar geometry, while lower values indicate that the representation changes across contexts\. Only states that are valid in both contexts are included in this comparison\.

Overall, the diagnostics show that learning a latent representation does not automatically make it invariant to physical context\. The amount and stability of semantic information in TD\-MPC’s latent space depend on the observation representation and the evaluation context\. At the same time, the stronger results with distance\-based observations show that reliable reuse is possible when the input already removes much of the irrelevant physical variation\. These results motivate representations that preserve shared behavioral structure while reducing sensitivity to physical changes that are not needed for predicting the reusable dynamics\.

## Appendix EAdditional related work

This section reviews broader work in continual reinforcement learning, its intersection with model\-based reinforcement learning, and zero\-shot generalization\.

### E\.1Model\-based RL

Model\-based reinforcement learning \(MBRL\) has recently shown strong performance across a wide range of single\-task problems, making it a promising direction for building more general agents\[[M\. Moerland et al\., 2023](https://arxiv.org/html/2610.41041#bib.bib5)\]\. Dyna\[[Sutton, 1991](https://arxiv.org/html/2610.41041#bib.bib10)\]as a classic algorithm, and Dreamer\[[Hafner et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib9)\], TD\-MPC\[[Hansen et al\., 2022](https://arxiv.org/html/2610.41041#bib.bib1)\], and methods based on joint\-embedding predictive architectures \(JEPA\)\[[Maes et al\., 2026](https://arxiv.org/html/2610.41041#bib.bib11)\]are examples of this growing family\. These methods learn an internal model of the environment and use it for prediction, planning, or imagined experience\. They mainly differ in how the world model is represented and how it is used for decision making\. However, model errors can affect policy improvement and planning\[[Janner et al\., 2021](https://arxiv.org/html/2610.41041#bib.bib12)\]\. Most work in this direction focuses on strong performance within a single task, while reuse and generalization across a sequence of changing tasks have received less attention\.

### E\.2Continual model\-based RL

Work at the intersection of continual reinforcement learning and model\-based reinforcement learning has explored several ways of preserving and transferring knowledge across tasks\.*Replay\-based approaches*modify how previous experience is stored and sampled from the replay buffer, with the goal of reducing catastrophic forgetting while preserving plasticity for new tasks\[[Rahimi\-Kalahroudi et al\., 2023](https://arxiv.org/html/2610.41041#bib.bib23),[Kessler et al\., 2023](https://arxiv.org/html/2610.41041#bib.bib24),[Fu et al\., 2025](https://arxiv.org/html/2610.41041#bib.bib21),[Alyahya et al\., 2026](https://arxiv.org/html/2610.41041#bib.bib22)\]\.*Dual\-architecture approaches*maintain two different models, often with task\-specific and task\-agnostic roles, and transfer information between them when a new task is encountered\[[Huang et al\., 2021](https://arxiv.org/html/2610.41041#bib.bib18),[Fu et al\., 2022](https://arxiv.org/html/2610.41041#bib.bib19)\]\. More recently,[Liu et al\. \[2025\]](https://arxiv.org/html/2610.41041#bib.bib28)study continual RL directly through online world model learning and model predictive control, with the goal of learning new tasks while avoiding catastrophic forgetting\.

More broadly, continual reinforcement learning has also been studied through*regularization\-based methods*, which constrain parameter updates to preserve previously learned knowledge\[[Kirkpatrick et al\., 2017](https://arxiv.org/html/2610.41041#bib.bib26)\], and*network\-expansion methods*, which increase model capacity as new tasks arrive and dedicate additional parameters to learning them\[[Schwarz et al\., 2018](https://arxiv.org/html/2610.41041#bib.bib25)\]\.

Many of these approaches assume that task boundaries are known\. In contrast, a learned world model may provide useful signals for detecting changes when new observations become inconsistent with its predictions\. Dual\-architecture methods are particularly related to our perspective, since a task\-agnostic component may support transfer and retention across tasks\. However, our view also asks whether a new learning phase is needed at all\. When previously learned behavioral knowledge is still applicable, the world model should ideally be reused directly rather than relearned or adapted unnecessarily\.

### E\.3Zero\-shot generalization in \(deep\) RL

Prior studies of the zero\-shot generalization problem in RL agents mainly focus on model\-free approaches[Kirk et al\. \[2023\]](https://arxiv.org/html/2610.41041#bib.bib17)\. Most of the work in this direction tries to incorporate a definition of ‘context’ \(long\-term properties of the environment, which is fixed during a learning episode\) in their methods and model their problem in a Contextual Markov Decision Process \(CMDP\)[Hallak et al\. \[2015\]](https://arxiv.org/html/2610.41041#bib.bib8)\. Agents learn a task under a specific defined context\. To measure an agent’s generalization, it is then tested in turn under an unseen new context\. Train and test contexts in most cases share structure and similar properties, and the agent is supposed to leverage this shared structure from the learned task to unseen ones\. Methods in this direction mostly differ in how they try to use the context and also their access level to context in the environment \(fully observable or they need to be inferred\)\. Recently, model\-based approaches have also considered this setting\.[Prasanna et al\. \[2024\]](https://arxiv.org/html/2610.41041#bib.bib13)assume access to context information and, after observing weak zero\-shot generalization of Dreamer, introduce a contextual version of the model\.[Röder et al\. \[2026\]](https://arxiv.org/html/2610.41041#bib.bib20)instead infer context through a self\-supervised module\. Related work has also used context inference with policy\-learning methods\[[Lee et al\., 2020](https://arxiv.org/html/2610.41041#bib.bib14)\]or hypernetworks that condition model parameters on observed context\[[Beukman et al\., 2023](https://arxiv.org/html/2610.41041#bib.bib16)\]\.

Our focus is complementary to these approaches\. Rather than proposing a context\-aware mechanism, we study why previously learned world\-model knowledge may fail to transfer\. We separate physical and behavioral context and use controlled changes in physical context and observation representation to study which factors limit zero\-shot reuse in continual RL\.

Another line of work studies invariant representation learning\. These methods aim to preserve task\-relevant information while reducing sensitivity to variations that should not affect the agent’s behavior\. This can improve generalization when different tasks share common structure\.[Shimizu and Tomizuka \[2024\]](https://arxiv.org/html/2610.41041#bib.bib15), for example, introduce bisimulation\-based representation learning for TD\-MPC to encourage more robust latent representations\. However, the zero\-shot generalization of this method across contextual changes has not been studied directly\. This direction is closely related to our observation that reliable reuse depends strongly on whether the representation preserves shared behavioral structure across physical\-context changes\.