Towards Generalizable Reinforcement Learning via Causality-Guided Self-Adaptive Representations

## Abstract

Abstract General intelligence requires quick adaption across tasks. While existing reinforcement learning (RL) methods have made progress in generalization, they typically assume only distribution changes between source and target domains. In this paper, we explore a wider range of scenarios where not only the distribution but also the environment spaces may change. For example, in the CoinRun environment, we train agents from easy levels and generalize them to difficulty levels where there could be new enemies that have never occurred before. To address this challenging setting, we introduce a causality-guided self-adaptive representation -based approach, called CSR, that equips the agent to generalize effectively across tasks with evolving dynamics. Specifically, we employ causal representation learning to characterize the latent causal variables within the RL system. Such compact causal representations uncover the structural relationships among variables, enabling the agent to autonomously determine whether changes in the environment stem from distribution shifts or variations in space, and to precisely locate these changes. We then devise a three-step strategy to fine-tune the causal model under different scenarios accordingly. Empirical experiments show that CSR efficiently adapts to the target domains with only a few samples and outperforms state-of-the-art baselines on a wide range of scenarios, including our simulated environments, CartPole, CoinRun and Atari games.

## 1 Introduction

In recent years, deep reinforcement learning (DRL, (Arulkumaran et al., )) has made incredible progress in various domains (Silver et al., ; Mirowski et al., ). Most of these works involve learning policies separately for fixed tasks. However, many practical scenarios often have a sequence of tasks with evolving dynamics. Instead of learning each task from scratch, humans possess the ability to discover the similarity between tasks and quickly generalize learned skills to new environments (Pearl & Mackenzie, ; Legg & Hutter, ). Therefore, it is essential to build a system where agents can also perform reliable and interpretable generalizations to advance towards general artificial intelligence (Kirk et al., ).

A straightforward solution is policy adaptation, i.e., leveraging the strategies developed in source tasks and adapting them to the target task as effective as possible (Zhu et al., ). Approaches along this line include but are not limited to finetuning (Mesnil et al., ), reward shaping (Harutyunyan et al., ), importance reweighting (Tirinzoni et al., ), learning robust policies (Taylor et al., ; Zhang et al., ), sim2real (Peng et al., ), adaptive RL (Huang et al., ), and subspace building (Gaya et al., ). However, these algorithms often rely on an assumption that all the source and target domains have the same state and action space while ignoring the out-of-distribution scenarios which are more common in practice (Taylor & Stone, ; Zhou et al., ).

Figure: Figure 1: Different levels in CoinRun: (a) Level 1: Traverse obstacles to collect coins; (b) Level 2: The number and shapes of obstacles change; (c) Level 3: Enemies emerge and attack the agents.
Refer to caption: x1.png

In this paper, we expand the application of RL beyond its traditional confines by exploring its adaptability in broader contexts. Specifically, our investigation focuses on policy adaptation in two distinct scenarios:

- 1.
Distribution shifts: the source and target data originate from the same environment space but exhibit differences in their distributions, e.g. changes in transition, observation or reward functions;
- 2.
State/Action space expansions: the source and target data are collected from different environment spaces, e.g. they differ in the latent state or action spaces.

These scenarios frequently occur in practical settings. To illustrate, we reference the popular CoinRun environment (Cobbe et al., ). As shown in Fig. 1, the goal of CoinRun is to overcome various obstacles and collect the coin located at the end of the level. The game environment is highly dynamic, with variations in elements such as background colors and the number and shape of obstacles (see Fig. 1 and Fig. 1) — this exemplifies distribution shifts. Additionally, CoinRun offers multiple difficulty levels. In lower difficulty settings, only a few stationary obstacles are present, while at higher levels, a variety of enemies emerge and attack the agents (see Fig. 1). To prevail, agents must learn to adapt to these new enemies — this scenario illustrates state/action space expansions.

We propose a Causality-guided Self-adaptive Representation-based approach, termed CSR, to address this problem for partially observable Markov decision processes (POMDPs). Considering that the raw observations are often reflections of the underlying state variables, we employ causal representation learning (Schölkopf et al., ; Huang et al., ; Wang et al., ) to identify the latent causal variables in the RL system, as well as the structural relationships among them. By leveraging such representations, we can automatically determine what and where the changes are. To be specific, we first augment the world models (Ha & Schmidhuber, ; Hafner et al., ) by including a task-specific change factor $\boldsymbol{\theta}$ to capture distribution shifts, e.g., $\boldsymbol{\theta}$ can characterize the changes in observations due to varying background colors in CoinRun. If the introduction of $\boldsymbol{\theta}$ can well explain the current observation, it is enough to keep previously learned causal variables and merely update a few parameters in the causal model. Otherwise, it implies that the current task differs from previously seen ones in the environment spaces, we then expand the causal graph by adding new causal variables and re-estimate the causal model. Finally, we remove some irrelevant causal variables that are redundant for policy learning according to the identified causal structures. This three-step strategy enables us to capture the changes in the environments for both scenarios in a self-adaptive manner and make the most of learned causal knowledge for low-cost policy transfer. Our key contributions are summarized below:

- •
We investigate a broader scenario towards generalizable reinforcement learning, where changes occur not only in the distributions but also in the environment spaces of latent variables, and propose a causality-guided self-adaptive representation-based approach to tackle this challenge.
- •
To characterize both the causal representations and environmental changes, we construct a world model that explicitly uncovers the structural relationships among latent variables in the RL system.
- •
By leveraging the compact causal representations, we devise a three-step strategy that can identify where the changes of the environment take place and add new causal variables autonomously if necessary. With this self-adaptive strategy, we achieve low-cost policy transfer by updating only a few parameters in the causal model.

## 2 World Model with Causality-Guided Self-adaptive Representations

We consider generalizable RL that aims to effectively transfer knowledge across tasks, allowing the model to leverage patterns learned from a set of source tasks while adapting to the dynamics of a target task. Each task $\mathcal{M}_{i}$ is characterized by $\langle\mathcal{S}_{i},\mathcal{A}_{i},\mathcal{O}_{i},R_{i},T_{i},\phi_{i},
\gamma_{i}\rangle$, where $\mathcal{S}_{i}$ represents the latent state space, $\mathcal{A}_{i}$ is the action space, $\mathcal{O}_{i}$ is the observation space, $R_{i}\colon\mathcal{S}_{i}\times\mathcal{A}_{i}\rightarrow\mathbb{R}$ is the reward function, $T_{i}\colon\mathcal{S}_{i}\times\mathcal{A}_{i}\rightarrow P(\mathcal{S}_{i})$ is the transition function, $\phi_{i}\colon\mathcal{S}_{i}\times\mathcal{A}_{i}\rightarrow P(\mathcal{O}_{i})$ is the observation function, and $\gamma_{i}$ is the discount factor. By leveraging experiences from previously encountered tasks $\{\mathcal{M}_{j}\}_{j=1}^{i-1}$, the objective is to adapt the optimal policy $\pi^{\star}$ that maximizes cumulative rewards to the target task $\mathcal{M}_{i}$. Here, we consider tasks arriving incrementally in the sequence $\langle\mathcal{M}_{1},\ldots,\mathcal{M}_{N}\rangle$ over time periods $\langle\mathcal{T}_{1},\ldots,\mathcal{T}_{N}\rangle$. In each period $\mathcal{T}_{i}$, only a replay buffer containing sequences $\{\langle o_{t},a_{t},r_{t}\rangle\}_{t=1}^{\mathcal{T}_{i}}$ from the current task $\mathcal{M}_{i}$ is available, representing an online setting. While tasks can also be presented offline with predefined source and target tasks, the online framework more closely mirrors human learning, making it a crucial step towards general intelligence.

In this section, we first construct a world model that explicitly embeds the structural relationships among variables in the RL system, and then we show how to encode the changes in the environment by introducing a domain-specific embedding into the model and leveraging it for policy adaptation.

Figure: Figure 2: The overall CSR framework. For each target task, we first use the prediction error, $\mathcal{L}_{\text{pred}}$, to determine whether it involves distribution shifts or space shifts. We then adjust the model accordingly by updating the task-specific change factor $\boldsymbol{\theta}_{i}$, or by adding new variables. Finally, we conduct causal graph pruning that removes variables unnecessary for the current task. Based on such compact causal representations, we can efficiently implement policy adaptation in a self-adaptive manner.
Refer to caption: extracted/5894886/figures/framework.png

### 2.1 Augmenting World Models with Structural Relationships

In POMDPs, extracting latent state representations from high-dimensional observations is crucial for enhancing the efficiency of the decision-making process. World models address this challenge by learning a generative model, which enables agents to predict future states through imagination. These methods typically consider all extracted representations of state variables equally important for policy learning, thereby utilizing all available information regardless of its relevance to the current task. However, real-world tasks often require a focus on specific information. For instance, in the Olympics, swimming speed is crucial in competitive swimming events, but it is less important in synchronized swimming, where grace and precision are prioritized. Hence, it is essential for agents to understand and focus on task-specific aspects to facilitate effective knowledge transfer by selectively using minimal sufficient information.

To this end, we adopt a causal state representation learning approach that not only enables us to extract state representations, but also to discover structural relationships over the variables. Suppose we observe sequences $\{\langle o_{t},a_{t},r_{t}\rangle\}_{t\in\mathcal{T}_{i}}$ for task $\mathcal{M}_{i}$, and denote the underlying causal latent states by $s_{t}\in\mathcal{S}_{i}$, we formulate the world model into:

$$ $\left\{\begin{array}[]{lll}{\text{observation model:}}&{p_{\phi}(o_{t}\mid D^{ \boldsymbol{s}\to o}\odot\boldsymbol{s}_{t})}\\ {\text{reward model:}}&{p_{\phi}(r_{t}\mid D^{\boldsymbol{s}\to r}\odot \boldsymbol{s}_{t})}\\ {\text{transition model:}}&{p_{\beta}(s_{k,t}\mid D_{k}^{\boldsymbol{s}\to \boldsymbol{s}}\odot\boldsymbol{s}_{t-1},D_{k}^{a\to\boldsymbol{s}}\odot a_{t- 1}),\text{ for }k=1,\ldots,d}\\ {\text{representation model:}}&{q_{\alpha}(\boldsymbol{s}_{t}\mid\boldsymbol{s }_{t-1},a_{t-1},o_{t}),}\end{array}\right.$ (1) $$

where $\boldsymbol{s}_{t}=(s_{1,t},\cdots,s_{d,t})$, $\odot$ is the element-wise product, and $D^{\cdot\to\cdot}$ denote binary masks indicating structural relationships over variables. For instance, if the $j$-th element of $D^{\boldsymbol{s}\to o}\in\{0,1\}^{d\times 1}$ in Eq. (1) is $1$, it indicates a causal edge from the state variable $\boldsymbol{s}_{j,t}$ to the current observation signal $o_{t}$, i.e., $\boldsymbol{s}_{j,t}$ is one of the parents of $o_{t}$. Consequently, we are supposed to retain $\boldsymbol{s}_{j,t}$ for the observation model. Otherwise, if $D^{\boldsymbol{s}\to o}_{j}=0$, then $\boldsymbol{s}_{j,t}$ should be removed from the causal model. Section 3.3 further discusses the estimation procedures for the structural matrices $D$, as well as the corresponding pruning process of the causal model. By learning such causal representations, we can explicitly characterize the decisive factors within each task. However, given that the underlying dynamics often vary across tasks, merely identifying which variables are useful is insufficient. We must also determine how these variables change with the environment for better generalization.

### 2.2 Characterization of Environmental Changes in a Compact Way

To address the above need, we now shift our focus to demonstrating how the world model can be modified to ensure robust generalization across the two challenging scenarios, respectively.

Characterization of Distribution Shifts. It is widely recognized that changes in the environmental distribution are often caused by modifications in a few specific factors within the data generation process (Ghassami et al., ; Schölkopf et al., ). In the CoinRun example, such shifts might be due to alterations in background colors ($\varrho$), while other elements remain constant. Therefore, to better characterize these shifts, we introduce a domain-specific change factor, $\boldsymbol{\theta}_{i}$, that captures the variations across different domains. Concurrently, we leverage $\boldsymbol{s}_{t}$ to identify the domain-shared latent variables of the environments. This leads us to reformulate Eq. (1) as follows:

$$ $\left\{\begin{array}[]{lll}{\text{observation model:}}&{p_{\phi}(o_{t}\mid D^{ \boldsymbol{s}\to o}\odot\boldsymbol{s}_{t},D^{\boldsymbol{\theta}_{i}\to o} \odot\theta_{i}^{o})}\\ {\text{reward model:}}&{p_{\phi}(r_{t}\mid D^{\boldsymbol{s}\to r}\odot \boldsymbol{s}_{t},D^{\boldsymbol{\theta}_{i}\to r}\odot\theta_{i}^{r})}\\ {\text{transition model:}}&{p_{\beta}(s_{k,t}\mid D_{k}^{\boldsymbol{s}\to \boldsymbol{s}}\odot\boldsymbol{s}_{t-1},D_{k}^{\boldsymbol{\theta}_{i}\to \boldsymbol{s}}\odot\theta_{i}^{\boldsymbol{s}},D_{k}^{a\to\boldsymbol{s}} \odot a_{t-1}),\text{ for }k=1,\ldots,d}\\ {\text{representation model:}}&{q_{\alpha}(\boldsymbol{s}_{t}\mid\boldsymbol{s }_{t-1},\boldsymbol{\theta}_{i},a_{t-1},o_{t}),}\end{array}\right.$ (2) $$

where $\boldsymbol{\theta}_{i}=\{\theta_{i}^{o},\theta_{i}^{r},\theta_{i}^{
\boldsymbol{s}}\}$ captures essential changes in the observation model, reward model, and transition model, respectively. In CoinRun, this enables us to make quick adaptations by re-estimating $\boldsymbol{\theta}_{i}^{o}=\varrho$ in the target task. We assume that the value of $\boldsymbol{\theta}_{i}$, as well as the structural matrices $D$, remains constant within the same task, but may differ across tasks.

Characterization of State/Action Space Expansions. In scenarios where the state or action space expands, we are supposed to add new variables to the existing causal model. The key challenge here is to determine whether the changes stem from distribution shifts or space variations. This dilemma can be addressed using $\boldsymbol{\theta}_{i}$: If the introduction of $\boldsymbol{\theta}_{i}$ can well capture the changes in the current observations, it implies that previous tasks $\{\mathcal{M}_{j}\}_{j=1}^{i-1}$ and $\mathcal{M}_{i}$ share the same causal variables but exhibit sparse changes in some certain parameters (i.e., distribution shifts). So we only need to store the specific part $\boldsymbol{\theta}_{i}$ of the causal model for $\mathcal{M}_{i}$. If it is not the case, then the causal graph must be expanded by adding new causal variables to explain the features unique to $\mathcal{M}_{i}$.

Benefits of Explicit Causal Structure. Upon detecting changes in the environment, we can further leverage the structural constraints $D$ to prune the causal graph. Essentially, we temporarily disregard variables that are irrelevant to the current task. However, for subsequent tasks, we reassess the structural relationships among the variables, enabling potential reuse. This approach allows us to not only preserve previously acquired information but also maintain the flexibility needed to customize the minimal sufficient state variables for each task. Details of this strategy are given in Section 3.

### 2.3 Identifiablity of Underlying World Models

In this section, we provide the identifiability theory under different scenarios in this paper: (1) For source task $\mathcal{M}_{1}$, Theorem 1 establishes the conditions under which the latent variable $\boldsymbol{s}_{t}$ and the structural matrices $D$ can be identified; (2) For the target task with distribution shifts, Theorem 2 outlines the identifiability of the domain-specific factor $\boldsymbol{\theta}_{i}$ in linear cases; (3) For the target task with state space shifts, Theorem 3 specifies the identifiability of the newly added state variables $\boldsymbol{s}^{\text{add}}_{t}$; (4) For the target task that includes both distribution shifts and state space shifts, Corollary 1 demonstrates the identifiability of both $\boldsymbol{\theta}_{i}$ and $\boldsymbol{s}^{\text{add}}_{t}$. The proofs are presented in Appendix A. We also discuss the possibility and challenges of establishing the identifiability of $\theta_{i}^{\boldsymbol{s}}$ in nonlinear cases in Appendix A.4, followed by empirical results where the learned $\hat{\theta}_{i}^{\boldsymbol{s}}$ demonstrates a monotonic correlation with the true values. Below we first introduce the definition of component-wise identifiability, related to Yao et al. (), and then we present the theoretical results.

###### Definition 1 .

(Component-wise identifiability). Let $\hat{\boldsymbol{s}}_{t}$ be the estimator of the latent variable $\boldsymbol{s}_{t}$. Suppose there exists a mapping $h$ such that $\boldsymbol{s}_{t}=h(\hat{\boldsymbol{s}}_{t})$. We say $\boldsymbol{s}_{t}$ is component-wise identifiable if $h$ is an invertible, component-wise function.

###### Theorem 1 .

(Identifiablity of world model in Eq. (1)). Assume the data generation process in Eq. (3). If the following conditions are satisfied, then $\boldsymbol{s}_{t}$ is component-wise identifiable: (1) for any $k_{1},k_{2}\in\{1,\ldots,d\}$ and $k_{1}\neq k_{2}$, $\hat{\boldsymbol{s}}_{k_{1},t}$ and $\hat{\boldsymbol{s}}_{k_{2},t}$ are conditionally independent given $\hat{\boldsymbol{s}}_{t-1}$; (2) for every possible value of $\boldsymbol{s}_{t}$, the vector functions defined in Eq. (8) are linearly independent. Furthermore, if the Markov condition and faithfulness assumption hold, then the structural matrices $D$ are also identifiable:

$$ $\left\{\begin{array}[]{cll}\left[o_{t},r_{t+1}\right]&{=}&{g(\boldsymbol{s}_{t },\epsilon_{t})}\\ {\boldsymbol{s}_{t}}&{=}&{g^{\boldsymbol{s}}(\boldsymbol{s}_{t-1},a_{t-1}, \epsilon_{t}^{\boldsymbol{s}}),}\end{array}\right.$ (3) $$

where

$$ $\left\{\begin{array}[]{lll}{o_{t}}&{=}&{g^{o}(\boldsymbol{s}_{t},\epsilon^{o}_ {t})}\\ {r_{t+1}}&{=}&{g^{r}(\boldsymbol{s}_{t},\epsilon_{t+1}^{r})}.\end{array}\right.$ (4) $$

The $\epsilon_{t},\epsilon_{t}^{\boldsymbol{s}},\epsilon_{t}^{o},\epsilon_{t+1}^{r}$ terms are corresponding independent and identically distributed (i.i.d.) random noises. Following Kong et al. (), here we only assume that the global mapping $g$ is invertible.

###### Theorem 2 .

(Identifiability of $\boldsymbol{\theta}_{i}$ in Eq. (2)). Assume the data generation process in Eq. (5), where the state transitions are linear and additive. If the process encounters distribution shifts and $\boldsymbol{s}_{t}$ has been identified according to Theorem 1, then $\boldsymbol{\theta}_{i}$ are component-wise identifiable:

$$ $\left\{\begin{array}[]{cll}\left[o_{t},r_{t+1}\right]&{=}&{g(\boldsymbol{s}_{t },\theta_{i}^{o},\theta_{i}^{r},\epsilon_{t})}\\ {\boldsymbol{s}_{t}}&{=}&{\mathbf{A}\boldsymbol{s}_{t-1}+\mathbf{B}a_{t-1}+ \mathbf{C}\theta_{i}^{\boldsymbol{s}}+\epsilon_{t}^{\boldsymbol{s}},}\end{ array}\right.$ (5) $$

where

$$ $\left\{\begin{array}[]{lll}{o_{t}}&{=}&{g^{o}(\boldsymbol{s}_{t},\theta_{i}^{o },\epsilon^{o}_{t})}\\ {r_{t+1}}&{=}&{g^{r}(\boldsymbol{s}_{t},\theta_{i}^{r},\epsilon_{t+1}^{r})}. \end{array}\right.$ (6) $$

Following Yao et al. (), here we assume that $\mathbf{A}$ is full rank, and $\mathbf{C}$ is full column rank. We further assume that $\boldsymbol{s}_{0}=\hat{\boldsymbol{s}}_{0}$. Moreover, if the Markov condition and faithfulness assumption hold, the structural matrices $D^{\boldsymbol{\theta}_{i}\to\cdot}$ are also identifiable.

###### Theorem 3 .

(Identifiability of Expanded State Space). Assume the data generation process in Eq. (3). Consider the expansion of the state space $\mathcal{S}$ by incorporating additional dimensions. Suppose $\boldsymbol{s}_{t}$ has already been identified according to Theorem 1, then the component-wise identifiability of the newly added variables $\boldsymbol{s}^{\text{add}}_{t}$ and the additional structural matrices, i.e., $D^{\boldsymbol{s}^{\text{add}}\to\cdot}$ and $D^{\cdot\to\boldsymbol{s}^{\text{add}}}$, can be established if $\boldsymbol{s}^{\text{add}}_{t}$ (1) represents a differentiable function of $[o_{t},r_{t+1}]$, i.e., $\boldsymbol{s}^{\text{add}}_{t}=f(o_{t},r_{t+1})$, and (2) fulfills conditions (1) and (2) specified in Theorem 1.

###### Corollary 1 .

(Identifiability under Multiple Shifts). Assume the data generation process in Eq. (5) involves both distribution shifts and state space shifts that comply with Theorem 2 and Theorem 3, respectively. In this case, both the domain-specific factor $\boldsymbol{\theta}_{i}$ and the newly added state variable $\boldsymbol{s}^{\text{add}}_{t}$ are component-wise identifiable.

###### Definition 1 .

###### Theorem 1 .

###### Theorem 2 .

###### Theorem 3 .

###### Corollary 1 .

## 3 A Three-Step Self-Adaptive Strategy for Model Adaptation

In this section, we provide a detailed description of CSR, a strategy aimed at addressing the environmental changes between source and target tasks, thereby ensuring that models can effectively respond to evolving dynamics. Specifically, we proceed with a three-step strategy: (1) Distribution Shifts Detection and Characterization, (2) State/Action Space Expansions, and (3) Causal Graph Pruning. The overall process of our three-step strategy are described in Fig. 2 and Algorithm 1.

Before this, we first give the estimation procedures for the world models defined in Eq. (2), which follows the state-of-the-art work Dreamer (Hafner et al., ; ). Given the observations in period $\mathcal{T}_{i}$, we maximize the object function $\mathcal{J}$, defined as $\mathcal{J}=\mathcal{J}_{\text{rec}}-\mathcal{J}_{\text{KL}}+\mathcal{J}_{
\text{reg}}$, for model optimization. The reconstruction part $\mathcal{J}_{\text{rec}}$ is commonly used to minimize the reconstruction error for the perceived observation $o_{t}$ and the reward $r_{t}$, which is defined as

$$ $\mathcal{J}_{\text{rec}}=\mathbb{E}_{q_{\alpha}}\left(\sum\limits_{t\in \mathcal{T}_{i}}\{\log p_{\phi}(o_{t}\mid D^{\boldsymbol{s}\to o}\odot \boldsymbol{s}_{t},D^{\boldsymbol{\theta}_{i}\to o}\odot\theta_{i}^{o})+p_{ \phi}(r_{t}\mid D^{\boldsymbol{s}\to r}\odot\boldsymbol{s}_{t},D^{\boldsymbol{ \theta}_{i}\to r}\odot\theta_{i}^{r})\}\right).$ $$

We also consider the KL-divergence constraints $\mathcal{J}_{\text{KL}}$ that helps to ensure that the latent representations attain optimal compression of the high-dimensional observations, which is formulated as:

$$ $\mathcal{J}_{\text{KL}}=\mathbb{E}_{q_{\alpha}}\left(\sum\limits_{t\in\mathcal {T}_{i}}\{\lambda_{\text{KL}}\cdot\text{KL}(q_{\alpha}(s_{k,t}\mid\boldsymbol{ s}_{t-1},\boldsymbol{\theta}_{i},a_{t-1},o_{t})\|p_{\beta}(s_{k,t}\mid D_{k}^{ \boldsymbol{s}\to\boldsymbol{s}}\odot\boldsymbol{s}_{t-1},D_{k}^{\boldsymbol{ \theta}_{i}\to\boldsymbol{s}}\odot\theta_{i}^{\boldsymbol{s}},D_{k}^{a\to \boldsymbol{s}}\odot a_{t-1})\}\right),$ $$

where $\lambda_{\text{KL}}$ is the regularization term. Moreover, as explained below in the Section 3.3, we further use $\mathcal{J}_{\text{reg}}$ as sparsity constraints that help to identify the binary masks $D$ better. Upon implementation, these three components are jointly optimized for model estimation. During the first task $\mathcal{M}_{1}$, we focus on developing the world models from scratch to capture the compact causal representations effectively. Then, for any subsequent target task $\mathcal{M}_{i}$ (where $i\geq 2$), our objective shifts to continuously refining the world model to accommodate new tasks according to the following steps.

### 3.1 Distribution Shifts Detection and Characterization

For each task $\mathcal{M}_{i}$, our first goal is to determine if it exhibits any distribution shifts. Therefore, in this step, we exclusively updates the domain-specific part $\boldsymbol{\theta}_{i}$, while keeping all other parameters unchanged from the previous task $\mathcal{M}_{i-1}$, to detect whether the distributions have changed. Recall that the effect of each edge in the structural matrices $D$ can differ from one task to another. By adjusting the values of $\boldsymbol{\theta}_{i}$, we can also easily characterize the task-specific influence of these connections. Particularly, when $\boldsymbol{\theta}_{i}$ is set to zero, we temporarily switch the related edges off in task $\mathcal{M}_{i}$.

Here we adopt forward prediction error (Guo et al., ) as the criteria to determine whether the re-estimated model well explains the observations in current task, defined as

$$ $\mathcal{L}_{\text{pred}}=\mathbb{E}_{\hat{o}_{t+1}\sim p_{\phi}}\|\hat{o}_{t+ 1}-o_{t+1}\|_{2}^{2}.$ (7) $$

A corresponding threshold, $\tau^{\star}$, is established. Upon implementation, we use the final prediction loss of the model on the source task $\mathcal{M}_{1}$ as the threshold value, thereby avoiding the need for manual setting. If the model’s performance $\tau$ is below this expected threshold, it means that the current task $\mathcal{M}_{i}$ shares the same causal variables with previous tasks, requiring only sparse changes of some parameters in the world model, and then we only need to re-estimate the specific part $\boldsymbol{\theta}_{i}$ to effectively manages these distribution changes. Otherwise, we proceed to the next step.

### 3.2 State/Action Space Expansions

When the involved domain-specific features $\boldsymbol{\theta}_{i}$ fail to accurately represent the target task $\mathcal{M}_{i}$, it becomes essential to incorporate additional causal variables into the existing causal graph to account for the features encountered in the new task. Given that the action variables are observable, we can directly obtain the relevant information when the action space expands. Thus, in this step, we focus on developing strategies to effectively manage state expansions. Let $d^{\prime}$ denote the number of causal variables to be added, we first determine the value of $d^{\prime}$ and introduce new causal features. Following this decision stage, we extend the causal representations from $\boldsymbol{s}_{t}$ to $\boldsymbol{s}^{\prime}_{t}=(\boldsymbol{s}_{t},\boldsymbol{s}^{\text{add}}_{t})$, where $\boldsymbol{s}^{\text{add}}_{t}=(s_{d+1,t},\ldots,s_{d+d^{\prime},t})$, by incorporating the additional $d^{\prime}$ causal variables. This is implemented by increasing the dimensions of input/output layers of the world models. For instance, the state input of the transition model will increase from $d$ to $d+d^{\prime}$. Accordingly, only the newly incorporated components are trained from scratch, while the existing parts of the model are fine-tuned. This approach allows us to leverage prior knowledge effectively and achieve low-cost knowledge transfer. Specifically, we propose the following three implementations for state space expansion:

- 1.
Random (Rnd): $d^{\prime}$ is randomly sampled from a uniform distribution.
- 2.
Deterministic (Det): It sets a constant value for $d^{\prime}$. However, this approach may overlook task-specific differences, potentially leading to either insufficient or redundant expansions. To address this, we employ group sparsity regularization (Yoon et al., ) on the added parameters after deterministic expansion, which allows for expansion while retaining the capability of shrinking.
- 3.
Self-Adaptive (SA): It searches for the value of $d^{\prime}$ that best fits the current task. To achieve this, we transform expansion into a decision-making process by considering the number of causal variables added to the graph as actions. Inspired by Xu & Zhu (), we define the state variable to reflect the current causal graph and derive the reward based on changes in predictive accuracy, which is calculated as the differences of the model’s prediction errors before and after expansion. Details are given in Appendix B.

It is noteworthy that our method allows for flexibility in the choice of expansion strategies. Intuitively, in our case, the Self-Adaptive approach is most likely to outperform others, and the experimental results in Section 4 further verify this point.

### 3.3 Causal Graph Pruning

As discussed in Section 2.1, not all variables contribute significantly to policy learning, and the necessary subset also differs between tasks. Therefore, during the generalization process, it is essential to identify the minimal sufficient state variables for each task. Fortunately, with the estimated causal model that explicitly encodes the structural relationships between variables, we can categorize these state variables $\boldsymbol{s}_{t}$ into the following two classes:

- 1.
Compact state representation $\boldsymbol{s}_{k,t}^{c}$: A variable that either affects the observation $o_{t}$, or the reward $r_{t+1}$, or influences other state variables $\boldsymbol{s}_{j,t+1}$ ($k\neq j$) at the next time step (i.e., $D^{\boldsymbol{s}\to o}=1$, or $D^{\boldsymbol{s}\to r}=1$, or $D_{j,k}^{\boldsymbol{s}\to\boldsymbol{s}}=1$, e.g., $s_{1,t}$ in Fig. 2).
- 2.
Non-compact state representation $\boldsymbol{s}_{k,t}^{\bar{c}}$: A variable that does not meet the criteria for a compact state representation (e.g., $s_{2,t}$ in Fig. 2).

Similarly, the change factors $\boldsymbol{\theta}_{i}$ can be classified in the same manner. These definitions allow us to selectively remove non-compact ones, thereby pruning the causal graph. That is, a variable is retained only when its corresponding structural constraints $D$ are non-zero. Hence, to better characterize the binary masks $D$ and the sparsity of $\boldsymbol{\theta}_{i}$, we define a regularization term $\mathcal{J}_{\text{reg}}$ by leveraging the edge-minimality property (Zhang & Spirtes, ), formulated as

$$ $\mathcal{J}_{\text{reg}}=-\lambda_{\text{reg}}\left[\|D^{\boldsymbol{s}\to o} \|_{1}+\|D^{\boldsymbol{s}\to r}\|_{1}+\|D^{\boldsymbol{s}\to\boldsymbol{s}}\| _{1}+\|D^{a\to\boldsymbol{s}}\|_{1}+\|D^{\boldsymbol{\theta}_{i}\to\boldsymbol {s}}\|_{1}+\|\boldsymbol{\theta}_{i}\|_{1}\right],$ $$

where $\lambda_{\text{reg}}$ represents the regularization term. Incorporating the regularization term $\mathcal{J}_{\text{reg}}$ directly into the objective function $\mathcal{J}$ confers a notable advantage: it enables concurrent pruning and model training. Specifically, the presence of $\mathcal{J}_{\text{reg}}$ induces certain entries of $D$ to transition from $1$ to 0 0 during model estimation, thereby promoting sparsity naturally without additional training phases.

### 3.4 Low-Cost Policy Generalization under Different Scenarios

After identifying what and where the changes occur, we are now prepared to perform policy generalization to the target task $\mathcal{M}_{i}$. Given that the number of state variables varies between the distribution shifts and state space expansion scenarios, the strategy for policy transfer also differs.

According to above definitions, in all these tasks, we incorporate both the domain-shared state representation $\boldsymbol{s}^{c}_{t}$ and the domain-specific change factor $\boldsymbol{\theta}_{i}^{c}$ as inputs to the policy $\pi^{\star}$, represented as: $a_{t}=\pi^{\star}(\boldsymbol{s}_{t}^{c},\boldsymbol{\theta}_{i}^{c})$. This approach enables the agent to accommodate potentially variable aspects of the environment during policy learning. Consequently, given the re-estimated value of the compact state representation $\boldsymbol{s}_{t}^{c}$ and the compact change factor $\boldsymbol{\theta}_{i}^{c}$ for task $\mathcal{M}_{i}$, if the model’s prediction error in the Distribution Shifts Detection and Characterization step meets expectations, we can directly transfer the learned policy $\mathcal{M}_{i}$ by applying $a_{t}=\pi^{\star}(\boldsymbol{s}_{t}^{c},\boldsymbol{\theta}_{i}^{c})$.

Otherwise, if the state variables expand from $\boldsymbol{s}_{t}^{c}$ to ${\boldsymbol{s}_{t}^{c}}^{\prime}$, along with the updated change factor ${\boldsymbol{\theta}_{i}^{c}}^{\prime}$, we then relearn the policy ${\pi^{\star}}^{\prime}$ basing on $\pi^{\star}$. Similarly, we train the newly added structures in the policy network while finetuning the original parameters, thereby updating the policy to $a_{t}={\pi^{\star}}^{\prime}({\boldsymbol{s}_{t}^{c}}^{\prime},{\boldsymbol{
\theta}_{i}^{c}}^{\prime})$.

Figure: Figure 3: (a) Simulation results; (b) CoinRun results; (c) Averaged training curves of CSR with/without structural embeddings $D$ in Atari games; (d) Normalized average training episodic return of CSR with different expansion strategies in our experiments.
Refer to caption: x3.png

## 4 Experiments

We evaluate the generalization capability of CSR on a number of simulated and well-established datasets, including the CartPole, CoinRun and Atari environments. For simulated experiments, we generate synthetic datasets that satisfy the two scenarios with different types of environmental changes. For CartPole, we consider distribution shifts in the task domains with different gravity or cart mass, and space variations by adding cart friction as the new state variable and additional force values for action expansion. For CoinRun, we train agents from easy levels and generalize them to difficulty levels where there could be new enemies that have never occurred before. For Atari games, we design the experiments by generating tasks with different game mode and difficulty levels. Such mode and difficulty switches lead to different consequences that changes the latent game dynamics or introduces new actions into the environment (Machado et al., ). For all these benchmarks, we evaluate the POMDP case, where the inputs are high-dimensional observations. We compare our approach against several baselines: Dreamer (Hafner et al., ), which handles fixed tasks without integrating causal knowledge; AdaRL (Huang et al., ), which employs simple scenario-based policy adaptation without space expansion considerations; and the traditional model-free DQN (Mnih et al., ) and SPR (Schwarzer et al., ). Additionally, for the Atari games, we benchmark against the state-of-the-art method, EfficientZero (Ye et al., ). All results are averaged over $5$ runs, more implementation details can be found in Appendix D.

Simulated Experiments. In our simulated experiments, we conducted a sequence of four tasks following the procedures outlined in Appendix D.2. To evaluate the performance of different methods across varying scenarios, Task 2 focuses exclusively on distribution shifts, Task 3 addresses changes solely within the environment space, and Task 4 combines both distribution and space changes. As shown in Fig. 3, CSR consistently outperforms the baselines in adapting to new tasks, particularly excelling in scenarios with space variations. This demonstrates CSR’s ability to accurately detect and adjust to the changes in the environment. Moreover, CSR tends to converge faster toward higher rewards, underscoring its efficiency in data utilization during generalization.

CartPole Experiments. CartPole is a classic control task where players move a cart left or right to balance a pendulum attached on it. In our experiments, we consider four consecutive tasks. Task 2 focuses exclusively on distribution shifts by randomly selecting the cart mass and the gravity from $\{0.5,1,2.5,3.5,4.5\}$ and $\{5,9.8,20,30,40\}$, respectively. In Tasks 1 and 2, we disregard the influence of the friction force between the cart and the track. In Task 3, however, we introduce this friction into the environment, and vary it over time, which simulates a game scenario where the cart moves on different surfaces, such as ice or grass (see Fig. 5 and Fig. 6). We reflect these changes in the observations by visualizing the track with different colored segments. To explore the generalization capabilities of the proposed method in scenarios with action expansion, we further designed Task 4, in which we expand the action space by incorporating additional possible force values that can be applied to the cart. The evaluation outcomes for these models are summarized in Table 1. We find that CSR consistently achieves the highest scores across all tasks, demonstrating its capability to promptly detect and adapt to environmental changes. In contrast, other baseline methods struggle to adjust to the introduction of the new friction variable and actions.

**Table 1: Evaluation results of different approaches in CartPole, with a maximum episode length of 500. The term ’Minimum Adaptation Steps’ refers to the minimal amount of data required for models to generalize to new tasks, as illustrated in Fig. 7. ✗denotes that a method fails to adapt under limited training steps due to non-convergence or suboptimal performance.**
| Model | Scores | Minimum Adaptation Steps |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Task 1 | Task 2 | Task 3 | Task 4 | Task 1 | Task 2 | Task 3 | Task 4 |  |
| DQN | 102.8 ($\pm$ 7.8) | 65.6 ($\pm$ 15.6) | 107.4 ($\pm$ 17.0) | 104.0 ($\pm$ 12.0) | ✗ | ✗ | ✗ | ✗ |
| Dreamer | 500.0 ($\pm$ 0.0) | 397.6 ($\pm$ 7.2) | 311.6 ($\pm$ 23.3) | 356.3 ($\pm$ 59.9) | 50k | ✗ | ✗ | ✗ |
| AdaRL | 500.0 ($\pm$ 0.0) | 468.0 ($\pm$ 2.3) | 410.0 ($\pm$ 3.4) | 407.5 ($\pm$ 34.5) | 50k | 4k | ✗ | ✗ |
| CSR (ours) | 500.0 ($\pm$ 0.0) | 500.0 ($\pm$ 0.0) | 500.0 ($\pm$ 0.0) | 500.0 ($\pm$ 0.0) | 50k | 2k | 4k | 10k |

**Table 2: Average final scores on the listed Atari games.**
| Task | Random | SPR | Dreamer | AdaRL | EfficientZero | CSR (ours) |
| --- | --- | --- | --- | --- | --- | --- |
| Alien | 291.9 | 970.3 | 1010.1 | 1147.5 | 557.4 | 1586.9 |
| ($\pm$ 83.5) | ($\pm$ 311.1) | ($\pm$ 339.0) | ($\pm$ 125.7) | ($\pm$ 185.9) | ($\pm$ 127.0) |  |
| Bank Heist | 18.4 | 110.1 | 1313.7 | 1285.4 | 181.0 | 1454.1 |
| ($\pm$ 2.5) | ($\pm$ 128.2) | ($\pm$ 341.5) | ($\pm$ 131.8) | ($\pm$ 90.6) | ($\pm$ 178.8) |  |
| Crazy Climber | 9668.0 | 32723.5 | 68026.5 | 62565.3 | 56408.3 | 88306.5 |
| ($\pm$ 2286.0) | ($\pm$ 12125.2) | ($\pm$ 15628.6) | ($\pm$ 15162.2) | ($\pm$ 13388.0) | ($\pm$ 18029.6) |  |
| Gopher | 235.6 | 294.0 | 5607.3 | 5359.6 | 1083.2 | 6718.6 |
| ($\pm$ 42.5) | ($\pm$ 312.1) | ($\pm$ 1982.9) | ($\pm$ 1736.2) | ($\pm$ 784.8) | ($\pm$ 1703.1) |  |
| Pong | -20.2 | -6.8 | 18.0 | 17.6 | 6.8 | 19.6 |
| ($\pm$ 0.1) | ($\pm$ 14.3) | ($\pm$ 3.1) | ($\pm$ 2.7) | ($\pm$ 7.3) | ($\pm$ 1.1) |  |

CoinRun Experiments. The learning curves in Fig. 3 depict our method’s consistent superiority over the baselines during knowledge transferring from low to high difficulty levels in CoinRun. We also observe that model-based methods tend to generalize more quickly than model-free ones in our experiments. This finding suggests that the forward-planning capabilities of world models confer significant advantages in adaptation. Visualizations of the reconstructed observations from various methods are presented in Appendix D.4.

Atari Experiments. We also conduct a series of interesting experiments on the Atari 100K games, which includes 26 games with a budget of 400K environment steps (Kaiser et al., ). Specifically, we select five representative games for evaluation: Alien, Bank Heist, Carzy Climber, Gopher, and Pong. The modes and difficulties available in each game are summarized in Table 6. For each of these games, we perform experiments among a sequence of four tasks, where each task randomly assigns a (mode, difficulty) pair. We then train these models on the source task and generalize them to downstream target tasks. Table 2 summarizes the average final scores across these tasks. We see that CSR achieves the highest mean scores in all the five games. Moreover, Fig. 14 illustrates the average generalization performance of various methods on downstream target tasks, while Fig. 15 to Fig. 19 present the training curves for each game, respectively. The reconstructions, as well as the estimated structural matrices, are provided in Appendix D.5.

Ablation Study. We also investigate the impact of integrating causal knowledge by explicitly embedding structural matrices $D$ into the world model on the generalization ability. Figure 3 illustrates the average performance of CSR with and without $D$ in Atari games. We observe a significantly faster and higher increase in the cumulative reward when taking structural relationships into consideration. This demonstrates the efficiency enhancement in policy learning through the removal of redundant causal variables, which accelerates the extraction and utilization of knowledge during the generalization process. Furthermore, to explore the differences among various expansion strategies, we conduct comparative experiments using the three methods described in Section 3.2 for all the environments and average them into Fig. 3. The results demonstrate that seeking for the optimal structure significantly improves expansion performance, leading us to apply the Self-Adaptive approach. However, we also observe that each search step requires extensive training time for models with different expansion scales, making the search process highly time-consuming. Therefore, it is crucial to consider this trade-off in practical applications.

## 5 Related Work

Recently, extensive research efforts have been invested in learning abstract representations in RL, employing methodologies such as image reconstruction (Watter et al., ), contrastive learning (Sermanet et al., ; Mazoure et al., ), and the development of world models (Sekar et al., ). A prominent research avenue within this domain is causal representation learning, which aims to identify high-level causal variables from low-level observations, thereby enhancing the accuracy of information available for decision-making processes (Schölkopf et al., ). Approaches such as ASRs (Huang et al., ) and IFactor (Liu et al., ) leverage causal factorization and structural constraints within causal variables to develop more accurate world models. Moreover, CDL (Wang et al., ), GRADER (Ding et al., ) and Causal Exploration (Yang et al., ) seek to boost exploration efficiency by learning causal models. Despite these advancements, many of these studies are tailored to specific tasks and struggle to achieve the level of generalization across tasks where human performance is notably superior (Taylor & Stone, ; Zhou et al., ).

To overcome these limitations, Harutyunyan et al. () develop a reward-shaping function that captures the target task’s information to guide policy learning. Taylor et al. () and Zhang et al. () aim to map tasks to invariant state variables, thereby learning policies robust to environmental changes. AdaRL (Huang et al., ) is dedicated to learning domain-shared and domain-specific representations to facilitate policy transfer. Distinct from these works, CSP (Gaya et al., ) approaches from the perspective of policy learning directly, by incrementally constructing a subspace of policies to train agents. However, most of these works assume a constant environment space, which is often not the case in practical applications. Therefore, in this paper, we investigate the feasibility of knowledge transfer when the state space can also change. Furthermore, the approach we propose is also related to the area of dynamic neural networks, where various methods have been developed to address sequences of tasks that require dynamical modifications to the network architecture, such as DEN (Yoon et al., ), PackNet (Mallya & Lazebnik, ), APD (Yoon et al., ), CPG (Hung et al., ), and Learn-to-Grow (Li et al., ).

## 6 Conclusions and Future Work

In this paper, we explore a broader range of scenarios for generalizable reinforcement learning, where changes across domains arise not only from distribution shifts but also space expansions. To investigate the adaptability of RL methods in these challenging scenarios, we introduce CSR, an approach that uses a three-step strategy to enable agents to detect environmental changes and autonomously adjust as needed. Empirical results from various complex environments, such as CartPole, CoinRun and Atari games, demonstrate the effectiveness of CSR in generalizing across evolving tasks. The primary limitation of this work is that it only considers generalization across domains and does not account for nonstationary changes over time. Therefore, a future research direction is to develop methods to automatically detect and characterize nonstationary changes both over time and across tasks.

## Appendix A Proofs of Identifiablity Theory

Before presenting the proofs, we first introduce the relevant notations and assumptions.

#### Notations.

We denote the underlying state variable by $\boldsymbol{s}_{t}=\{s_{1,t},\ldots,s_{d,t}\}$ and denote $o_{t}$ as the observation. Also, we denote the mapping from state estimator $\hat{\boldsymbol{s}}_{t}$ to state $\boldsymbol{s}_{t}$ by $\boldsymbol{s}_{t}=h(\hat{\boldsymbol{s}}_{t})$, and denote the Jacobian matrix of $h$ as $\mathbf{J}^{h}_{t}$. Let $\zeta_{k,t}\triangleq\log p(s_{k,t}\mid\boldsymbol{s}_{t-1})$, we further denote

$$ $\boldsymbol{\omega}_{k,t}\triangleq\left(\frac{\partial^{2}\zeta_{k,t}}{ \partial s_{k,t}\partial s_{1,t-1}},\frac{\partial^{2}\zeta_{k,t}}{\partial s_ {k,t}\partial s_{2,t-1}},\ldots,\frac{\partial^{2}\zeta_{k,t}}{\partial s_{k,t }\partial s_{d,t-1}}\right)^{\top},\stackrel{{\scriptstyle\circ}}{{\boldsymbol {\omega}}}_{k,t}\triangleq\left(\frac{\partial^{3}\zeta_{k,t}}{\partial s_{k,t }^{2}\partial s_{1,t-1}},\frac{\partial^{3}\zeta_{k,t}}{\partial s_{k,t}^{2} \partial s_{2,t-1}},\ldots,\frac{\partial^{3}\zeta_{k,t}}{\partial s_{k,t}^{2} \partial s_{d,t-1}}\right)^{\top}.$ (8) $$

###### Assumption 1 .

$\zeta_{k,t}$ is twice differentiable with respect to $s_{k,t}$ and differentiable with respect to $s_{l,t-1}$, for all $l\in\{1,\ldots,d\}$.

###### Assumption 2 (Faithfulness assumption) .

For a causal graph $\mathcal{G}$ and the associated probability distribution $P$, every true conditional independence relation in $P$ is entailed by the Causal Markov Condition applied to $\mathcal{G}$ (Spirtes et al., ).

###### Assumption 1 .

###### Assumption 2 (Faithfulness assumption) .

### A.1 Proof of Theorem 1

Based on aforementioned assumptions and definitions, Theorem 1 establishes the conditions for the component-wise identifiablity of the state variable $\boldsymbol{s}_{t}$ and the structural matrices $D$ in Eq. (1).

###### Theorem 1 .

(Identifiablity of world model in Eq. (1)). Assume the data generation process in Eq. (9). If the following conditions are satisfied, then $\boldsymbol{s}_{t}$ is component-wise identifiable: (1) for any $k_{1},k_{2}\in\{1,\ldots,d\}$ and $k_{1}\neq k_{2}$, $\hat{\boldsymbol{s}}_{k_{1},t}$ and $\hat{\boldsymbol{s}}_{k_{2},t}$ are conditionally independent given $\hat{\boldsymbol{s}}_{t-1}$; (2) for every possible value of $\boldsymbol{s}_{t}$, the vector functions defined in Eq. (8) are linearly independent. Furthermore, if the Markov condition and faithfulness assumption hold, then the structural matrices $D$ are also identifiable:

$$ $\left\{\begin{array}[]{cll}\left[o_{t},r_{t+1}\right]&{=}&{g(\boldsymbol{s}_{t },\epsilon_{t})}\\ {\boldsymbol{s}_{t}}&{=}&{g^{\boldsymbol{s}}(\boldsymbol{s}_{t-1},a_{t-1}, \epsilon_{t}^{\boldsymbol{s}}),}\end{array}\right.$ (9) $$

where

$$ $\left\{\begin{array}[]{lll}{o_{t}}&{=}&{g^{o}(\boldsymbol{s}_{t},\epsilon^{o}_ {t})}\\ {r_{t+1}}&{=}&{g^{r}(\boldsymbol{s}_{t},\epsilon_{t+1}^{r})}.\end{array}\right.$ (10) $$

The $\epsilon_{t},\epsilon_{t}^{\boldsymbol{s}},\epsilon_{t}^{o},\epsilon_{t+1}^{r}$ terms are corresponding independent and identically distributed (i.i.d.) random noises. Following Kong et al. (), here we only assume that the global mapping $g$ is invertible.

###### Proof.

The proof proceeds in two steps. First, we demonstrate that the data generation process in Eq. (9) is equivalent to the noiseless data distribution. Second, we summarize the proof steps of the identifiablity of the state variables $\boldsymbol{s}_{t}$ under the noiseless distribution, which has already been provided in Yao et al. ().

Step 1: transform into noise-free distributions.

Let $y_{t}=\left[o_{t},r_{t+1}\right]$. We denote $p(y_{t})=\int p_{g}(y_{t}|\boldsymbol{s}_{t})p_{\kappa}(\boldsymbol{s}_{t})d
\boldsymbol{s}_{t}$ where $g,\kappa$ are the parameters of the probability functions. Suppose $p_{g,\kappa}(y_{t})=p_{\hat{g},\hat{\kappa}}(y_{t})$ holds for all $y_{t}$, where $(g,\kappa)$ and $(\hat{g},\hat{\kappa})$ are two sets of parameters. We complete the proof primarily by following Khemakhem et al. ().

By applying the law of total probability, we have

$$ $\int_{\mathcal{S}}p_{\kappa}(\boldsymbol{s}_{t})\cdot p_{g}(y_{t}|\boldsymbol{ s}_{t})d\boldsymbol{s}_{t}=\int_{\mathcal{S}}p_{\hat{\kappa}}(\boldsymbol{s}_{ t})\cdot p_{\hat{g}}(y_{t}|\boldsymbol{s}_{t})d\boldsymbol{s}_{t}.$ (11) $$

Further define $p_{g}(y_{t}|\boldsymbol{s}_{t})=p_{\epsilon_{t}}(y_{t}-g(\boldsymbol{s}_{t}))$, we get

$$ $\int_{\mathcal{S}}p_{\kappa}(\boldsymbol{s}_{t})\cdot p_{\epsilon_{t}}(y_{t}-g (\boldsymbol{s}_{t}))d\boldsymbol{s}_{t}=\int_{\mathcal{S}}p_{\hat{\kappa}}( \boldsymbol{s}_{t})\cdot p_{\epsilon_{t}}(y_{t}-\hat{g}(\boldsymbol{s}_{t}))d \boldsymbol{s}_{t}.$ (12) $$

Replacing $\overline{y}_{t}=g(\boldsymbol{s}_{t})$ on the left hand side, and similarly on the right hand side, we obtain

$$ $\int_{\mathcal{O}}p_{\kappa}(g^{-1}(\overline{y}_{t}))~{}\text{vol}~{}\mathbf{ J}^{g^{-1}}_{t}(\overline{y}_{t})\cdot p_{\epsilon_{t}}(y_{t}-\overline{y}_{t} )d\overline{y}_{t}=\int_{\mathcal{O}}p_{\hat{\kappa}}(\hat{g}^{-1}(\overline{y }_{t}))~{}\text{vol}~{}\mathbf{J}^{\hat{g}^{-1}}_{t}(\overline{y}_{t})\cdot p_ {\epsilon_{t}}(y_{t}-\overline{y}_{t})d\overline{y}_{t},$ (13) $$

where $\text{vol}~{}\mathbf{J}=\sqrt{\det\mathbf{J}^{\top}\mathbf{J}}$.

By introducing $\tilde{p}_{g,\kappa}(\overline{y}_{t})=p_{\kappa}(g^{-1}(\overline{y}_{t}))~{}
\text{vol}~{}\mathbf{J}^{g^{-1}}_{t}(\overline{y}_{t})\mathds{1}(\overline{y}_
{t})$ on both sides, we can rewrite Eq. (13) as

$$ $\int_{\mathbb{R}^{\upsilon}}\tilde{p}_{g,\kappa}(\overline{y}_{t})\cdot p_{ \epsilon_{t}}(y_{t}-\overline{y}_{t})d\overline{y}_{t}=\int_{\mathbb{R}^{ \upsilon}}\tilde{p}_{\hat{g},\hat{\kappa}}(\overline{y}_{t})\cdot p_{\epsilon_ {t}}(y_{t}-\overline{y}_{t})d\overline{y}_{t},$ (14) $$

where $\upsilon=\dim\mathcal{O}+\dim R$. By the definition of convolution, Eq. (14) is equivalent to

$$ $(\tilde{p}_{g,\kappa}*p_{\epsilon})(y_{t})=(\tilde{p}_{\hat{g},\hat{\kappa}}*p _{\epsilon})(y_{t}),$ (15) $$

where $*$ denote the convolution operator. Denote $F[.]$ the Fourier transform and $\varphi_{\epsilon}=F[p_{\epsilon}]$, we have

$$ $F[\tilde{p}_{g,\kappa}](\Omega)\varphi_{\epsilon}(\Omega)=F[\tilde{p}_{\hat{g} ,\hat{\kappa}}](\Omega)\varphi_{\epsilon}(\Omega).$ (16) $$

Assume set $\{\boldsymbol{x}\in\mathcal{X}|\varphi_{\epsilon}(\boldsymbol{x})=0\}$ has measure zero, we can drop $\varphi_{\epsilon}({\Omega})$ from both sides, which obtains

$$ $F[\tilde{p}_{g,\kappa}](\Omega)=F[\tilde{p}_{\hat{g},\hat{\kappa}}](\Omega).$ (17) $$

Therefore, for all $y_{t}\in\mathcal{O}\times R$, we have

$$ $\tilde{p}_{g,\kappa}(y_{t})=\tilde{p}_{\hat{g},\hat{\kappa}}(y_{t}).$ (18) $$

This indicates that the noise-free distributions must coincide for the overall distributions to remain identical after adding noise, effectively reducing the noisy case in Eq. (9) into the noiseless case.

Step 2: establish identifiability of state variables.

After transforming the problem into the noise-free case, we proceed by summarizing the key proof steps of the identifiability of the state variables $\boldsymbol{s}_{t}$, following Yao et al. (), which are:

- •
First, by making use of the conditional independence of the components of $\hat{\boldsymbol{s}}_{t}$ given $\hat{\boldsymbol{s}}_{t-1}$, it is shown that:
$\frac{\partial^{2}\log p(\hat{\mathbf{s}}_{t}\mid\hat{\mathbf{s}}_{t-1})}{
\partial\hat{s}_{i,t}\partial\hat{s}_{j,t}}=0.$
(19)
- •
Second, by utilizing the Jacobian matrix $\mathbf{J}^{h}_{t}$ to calculate Eq. (19), it is derived that
$\frac{\partial^{3}\log p\left(\hat{\boldsymbol{s}}_{t}\mid\hat{\boldsymbol{s}}
_{t-1}\right)}{\partial\hat{s}_{i,t}\partial\hat{s}_{j,t}\partial s_{l,t-1}}=
\sum_{k=1}^{d}\left(\frac{\partial^{3}\zeta_{k,t}}{\partial s_{k,t}^{2}
\partial s_{l,t-1}}\cdot\mathbf{J}^{h}_{k,i,t}\mathbf{J}^{h}_{k,j,t}+\frac{
\partial^{2}\zeta_{k,t}}{\partial s_{k,t}\partial s_{l,t-1}}\cdot\frac{
\partial\mathbf{J}^{h}_{k,i,t}}{\partial\hat{s}_{j,t}}\right)\equiv 0.$
(20)
- •
Finally, it is established that $\boldsymbol{s}_{t}$ is identifiable, up to an invertible, component-wise nonlinear transformation of a permuted version of $\hat{\boldsymbol{s}}_{t}$, if the linear independence of vector funtions defined in Eq. (8) holds and the the Jacobian matrix $\mathbf{J}^{h}_{t}$ satisfies Eq. (20).

Moreover, the proofs for the identifiablity of the structural matrices $D$ are presented in Huang et al. (). Based on these steps, we next provide the proofs of Theorem 2-3, and Corollary 1.
∎

###### Theorem 1 .

###### Proof.

### A.2 Proof of Theorem 2

Different from existing methods, to capture the changing dynamics in the environment, we have introduced a task-specific change factor, $\boldsymbol{\theta}_{i}$, into the world model as defined in Eq. (2). Accordingly, Theorem 2 presents the identifiability of $\boldsymbol{\theta}_{i}$ and the corresponding structural matrices $D^{\boldsymbol{\theta}_{i}\to\cdot}$ for scenarios involving only distribution shifts in linear cases.

###### Theorem 2 .

(Identifiability of $\boldsymbol{\theta}_{i}$ in Eq. (2)). Assume the data generation process in Eq. (21), where the state transitions are linear and additive. If the process encounters distribution shifts and $\boldsymbol{s}_{t}$ has been identified according to Theorem 1, then $\boldsymbol{\theta}_{i}$ are component-wise identifiable:

$$ $\left\{\begin{array}[]{cll}\left[o_{t},r_{t+1}\right]&{=}&{g(\boldsymbol{s}_{t },\theta_{i}^{o},\theta_{i}^{r},\epsilon_{t})}\\ {\boldsymbol{s}_{t}}&{=}&{\mathbf{A}\boldsymbol{s}_{t-1}+\mathbf{B}a_{t-1}+ \mathbf{C}\theta_{i}^{\boldsymbol{s}}+\epsilon_{t}^{\boldsymbol{s}},}\end{ array}\right.$ (21) $$

where

$$ $\left\{\begin{array}[]{lll}{o_{t}}&{=}&{g^{o}(\boldsymbol{s}_{t},\theta_{i}^{o },\epsilon^{o}_{t})}\\ {r_{t+1}}&{=}&{g^{r}(\boldsymbol{s}_{t},\theta_{i}^{r},\epsilon_{t+1}^{r})}. \end{array}\right.$ (22) $$

Following Yao et al. (), here we assume that $\mathbf{A}$ is full rank, and $\mathbf{C}$ is full column rank. We further assume that $\boldsymbol{s}_{0}=\hat{\boldsymbol{s}}_{0}$. Moreover, if the Markov condition and faithfulness assumption hold, the structural matrices $D^{\boldsymbol{\theta}_{i}\to\cdot}$ are also identifiable.

###### Proof.

The proofs for the identifiablity of $D^{\boldsymbol{\theta}_{i}\to\cdot}$ are given in Huang et al. (). Here, we provide the proof for the identifiablity of $\boldsymbol{\theta}_{i}$, which is done in the following four steps:

- •
In step 1, we prove that $\theta^{o}_{i}$ can be identified up to component-wise transformation when only the observation function exhibits distribution shifts.
- •
In step 2, we demonstrate that $\theta^{r}_{i}$ is component-wise identifiable when only the reward function experiences distribution shifts.
- •
In step 3, we show that $\theta^{\boldsymbol{s}}_{i}$ can be identified component-wisely when only the transition function undergoes distribution shifts.
- •
In step 4, we establish that in the general case, where the observation, reward, and transition functions may undergo distribution shifts simultaneously, $\boldsymbol{\theta}_{i}=\{\theta_{i}^{o},\theta_{i}^{r},\theta_{i}^{
\boldsymbol{s}}\}$ is identifiable.

Step 1: prove the identifiability of $\theta^{o}_{i}$.

According to Eq. (22), we have

$$ $y_{t}=g(\boldsymbol{s}_{t},\theta_{i}^{o},\epsilon_{t}).$ (23) $$

Denote $\boldsymbol{x}_{t}=(\boldsymbol{s}_{t},\theta_{i}^{o})$. There exists

$$ $\boldsymbol{x}_{t}=h^{\prime}(\hat{\boldsymbol{x}}_{t}),$ (24) $$

where $h^{\prime}=g^{-1}\circ\hat{g}$, and $\hat{\boldsymbol{x}}_{t}$ is the estimator of $\boldsymbol{x}_{t}$. Since both $g$ and $\hat{g}$ are invertible, $h^{\prime}$ is invertible. Therefore, we have

$$ $\mathbf{J}^{h^{\prime}}_{t}=\begin{bmatrix}\frac{\partial\boldsymbol{s}_{t}}{ \partial\hat{\boldsymbol{s}}_{t}}&\frac{\partial\boldsymbol{s}_{t}}{\partial \hat{\theta}_{i}^{o}}\\ \frac{\partial\theta_{i}^{o}}{\partial\hat{\boldsymbol{s}}_{t}}&\frac{\partial \theta_{i}^{o}}{\partial\hat{\theta}_{i}^{o}}\end{bmatrix},$ (25) $$

where $\mathbf{J}^{h^{\prime}}_{t}$ is full rank. Note that $\frac{\partial\boldsymbol{s}_{t}}{\partial\hat{\theta}_{i}^{o}}=0$ and $\frac{\partial\theta_{i}^{o}}{\partial\hat{\boldsymbol{s}}_{t}}=0$. Further recall that we assume the identifiability of $\boldsymbol{s}_{t}$, which means that $\mathbf{J}^{h}_{t}=\frac{\partial\boldsymbol{s}_{t}}{\partial\hat{\boldsymbol{
s}}_{t}}$ is full rank. We can derive that $\frac{\partial\theta_{i}^{o}}{\partial\hat{\theta}_{i}^{o}}$ must be full rank. That is, $\theta_{i}^{o}$ is component-wise identifiable.

Step 2: prove the identifiability of $\theta^{r}_{i}$.

If only the reward function $g^{r}$ exhibits distribution shifts, we have

$$ $y_{t}=g(\boldsymbol{s}_{t},\theta_{i}^{r},\epsilon_{t}).$ (26) $$

It is straightforward to see that $\theta_{i}^{r}$ is blockwise identifiable using the same technique in Step 1.

Step 3: prove the identifiability of $\theta^{\boldsymbol{s}}_{i}$.

Recall that we have

$$ $\boldsymbol{s}_{t}=\mathbf{A}\boldsymbol{s}_{t-1}+\mathbf{B}a_{t-1}+\mathbf{C} \theta_{i}^{\boldsymbol{s}}+\epsilon_{t}^{\boldsymbol{s}}.$ (27) $$

By leveraging the recursive property of the state transition process, we can derive that

$$ $\boldsymbol{s}_{t}=\mathbf{A}^{t}\boldsymbol{s}_{0}+\left(\sum_{k=0}^{t-1} \mathbf{A}^{k}\mathbf{B}\right)a_{t-1-k}+\left(\sum_{k=0}^{t-1}\mathbf{A}^{k} \right)\mathbf{C}\theta_{i}^{\boldsymbol{s}}+\left(\sum_{k=0}^{t-1}\mathbf{A}^ {k}\right)\epsilon_{t-k}.$ (28) $$

Similarly for $\hat{\boldsymbol{s}}_{t}$, we have

$$ $\hat{\boldsymbol{s}}_{t}=\hat{\mathbf{A}}^{t}\hat{\boldsymbol{s}}_{0}+\left( \sum_{k=0}^{t-1}\hat{\mathbf{A}}^{k}\hat{\mathbf{B}}\right)a_{t-1-k}+\left( \sum_{k=0}^{t-1}\hat{\mathbf{A}}^{k}\right)\hat{\mathbf{C}}\theta_{i}^{ \boldsymbol{s}}+\left(\sum_{k=0}^{t-1}\hat{\mathbf{A}}^{k}\right)\epsilon_{t-k}.$ (29) $$

Note that $\boldsymbol{s}_{0}=\hat{\boldsymbol{s}}_{0}$. Therefore, combining Eq. (28) and Eq. (29) gives

$$ $\left(\sum_{k=0}^{t-1}\mathbf{A}^{k}\right)\mathbf{C}\theta_{i}^{\boldsymbol{s }}=\boldsymbol{s}_{t}-\mathbf{A}^{t}\hat{\mathbf{A}}^{-t}\left[\hat{ \boldsymbol{s}}_{t}-\left(\sum_{k=0}^{t-1}\hat{\mathbf{A}}^{k}\right)\hat{ \mathbf{C}}\theta_{i}^{\boldsymbol{s}}\right]+\Theta,$ (30) $$

where $\Theta$ is a constant term. Taking the derivative w.r.t $\hat{\theta}_{i}^{\boldsymbol{s}}$ on both sides, we obtain

$$ $\mathbf{M}\frac{\partial\theta_{i}^{\boldsymbol{s}}}{\partial\hat{\theta}_{i}^ {\boldsymbol{s}}}=\frac{\partial\boldsymbol{s}_{t}}{\partial\hat{\boldsymbol{s }}_{t}}\frac{\partial\hat{\boldsymbol{s}}_{t}}{\partial\hat{\theta}_{i}^{ \boldsymbol{s}}}-\mathbf{A}^{t}\hat{\mathbf{A}}^{-t}\left[\frac{\partial\hat{ \boldsymbol{s}}_{t}}{\partial\hat{\theta}_{i}^{\boldsymbol{s}}}-\hat{\mathbf{M }}\right],$ (31) $$

where $\mathbf{M}=\left(\sum_{k=0}^{t-1}\mathbf{A}^{k}\right)\mathbf{C}$ and $\hat{\mathbf{M}}=\left(\sum_{k=0}^{t-1}\hat{\mathbf{A}}^{k}\right)\hat{\mathbf
{C}}$. Note that we further have $\frac{\partial\hat{\boldsymbol{s}}_{t}}{\partial\hat{\theta}_{i}^{\boldsymbol{
s}}}=\hat{\mathbf{M}}$ according to Eq. (29). That is,

$$ $\mathbf{M}\frac{\partial\theta_{i}^{\boldsymbol{s}}}{\partial\hat{\theta}_{i}^ {\boldsymbol{s}}}=\frac{\partial\boldsymbol{s}_{t}}{\partial\hat{\boldsymbol{s }}_{t}}\hat{\mathbf{M}}.$ (32) $$

Recall that $\frac{\partial\boldsymbol{s}_{t}}{\partial\hat{\boldsymbol{s}}_{t}}=\mathbf{J}
^{h}_{t}$ is full rank. Moreover, the full column rank of $\mathbf{M}$ and $\hat{\mathbf{M}}$ is guaranteed by the full rank of $\mathbf{A}$ and the full column rank of $\mathbf{C}$ (see Proposition 1). Therefore, we can derive

$$ $\text{rank}(\mathbf{M}\frac{\partial\theta_{i}^{\boldsymbol{s}}}{\partial\hat{ \theta}_{i}^{\boldsymbol{s}}})=\text{rank}(\frac{\partial\boldsymbol{s}_{t}}{ \partial\hat{\boldsymbol{s}}_{t}}\hat{\mathbf{M}})=\dim\theta_{i}^{\boldsymbol {s}}.$ (33) $$

Due to the rank inequality property of matrix products, we have

$$ $\text{rank}(\mathbf{M}\frac{\partial\theta_{i}^{\boldsymbol{s}}}{\partial\hat{ \theta}_{i}^{\boldsymbol{s}}})\leq\min\left(\mathbf{M},\text{rank}(\frac{ \partial\theta_{i}^{\boldsymbol{s}}}{\partial\hat{\theta}_{i}^{\boldsymbol{s}} })\right)=\min\left(\dim\theta_{i}^{\boldsymbol{s}},\text{rank}(\frac{\partial \theta_{i}^{\boldsymbol{s}}}{\partial\hat{\theta}_{i}^{\boldsymbol{s}}})\right )\leq\dim\theta_{i}^{\boldsymbol{s}}.$ (34) $$

Eq. (33) and Eq. (34) show that

$$ $\dim\theta_{i}^{\boldsymbol{s}}\leq\min\left(\dim\theta_{i}^{\boldsymbol{s}}, \text{rank}(\frac{\partial\theta_{i}^{\boldsymbol{s}}}{\partial\hat{\theta}_{i }^{\boldsymbol{s}}})\right)\leq\dim\theta_{i}^{\boldsymbol{s}}.$ (35) $$

To make the above equation hold true, it must have $\text{rank}(\frac{\partial\theta_{i}^{\boldsymbol{s}}}{\partial\hat{\theta}_{i
}^{\boldsymbol{s}}})=\dim\theta_{i}^{\boldsymbol{s}}$. That is, $\frac{\partial\theta_{i}^{\boldsymbol{s}}}{\partial\hat{\theta}_{i}^{
\boldsymbol{s}}}$ must be full rank, then $\theta_{i}^{\boldsymbol{s}}$ must be component-wise identifiable.

Step 4: prove the identifiability of $\boldsymbol{\theta}_{i}$ in the general case.

This step can be easily demonstrated by directly combining the proofs above.
∎

###### Proposition 1 .

Suppose $\mathbf{A}$ is a matrix with full rank, and $\mathbf{C}$ is a matrix with full column rank. Define $\mathbf{M}=\left(\sum_{k=0}^{t-1}\mathbf{A}^{k}\right)\mathbf{C}$. Then $\mathbf{M}$ is full column rank.

###### Proof.

To establish that $\mathbf{M}$ is full column rank, it suffices to show that $\mathbf{N}=\sum_{k=0}^{t-1}\mathbf{A}^{k}$ is of full rank. Assume the eigenvalues of $\mathbf{A}$ are denoted by $\Lambda$. Given that $\mathbf{A}$ is of full rank, the eigenvalues of $\mathbf{A}^{k}$ are $\Lambda^{k}$. Let $\Lambda^{\mathbf{N}}$ represent the eigenvalues of $\mathbf{N}$, then we have

$$ $\Lambda^{\mathbf{N}}=\sum_{k=0}^{t-1}\Lambda^{k}.$ (36) $$

It is evident that there exists at least one $t$ such that for any non-zero $\Lambda$, Eq. (36) is non-zero, thereby confirming that $\mathbf{N}$ has no zero eigenvalues. This verification ensures that $\mathbf{N}$ is full rank. Hence, $\mathbf{M}$ maintains full column rank.
∎

###### Theorem 2 .

###### Proof.

###### Proposition 1 .

###### Proof.

### A.3 Proof of Theorem 3

In Theorem 3, we foucus on the identifiablity of the newly added variables $\boldsymbol{s}^{\text{add}}_{t}$ and their corresponding structural matrices, when the state space is expanded by incorporating additional dimensions.

###### Theorem 3 .

(Identifiability of Expanded State Space). Assume the data generation process in Eq. (9). Consider the expansion of the state space $\mathcal{S}$ by incorporating additional dimensions. Suppose $\boldsymbol{s}_{t}$ has already been identified according to Theorem 1, then the component-wise identifiability of the newly added variables $\boldsymbol{s}^{\text{add}}_{t}$ and the additional structural matrices, i.e., $D^{\boldsymbol{s}^{\text{add}}\to\cdot}$ and $D^{\cdot\to\boldsymbol{s}^{\text{add}}}$, can be established if $\boldsymbol{s}^{\text{add}}_{t}$ (1) represents a differentiable function of $[o_{t},r_{t+1}]$, i.e., $\boldsymbol{s}^{\text{add}}_{t}=f(o_{t},r_{t+1})$, and (2) fulfills conditions (1) and (2) specified in Theorem 1.

###### Proof.

For the newly added variables, since they also satisfy the conditional independence condition, we can derive the same properties as described in Eq. (19) and Eq. (20) using the same technique in the proof steps of Theorem 1. Additionally, since the vector functions corresponding to $\boldsymbol{s}^{\text{add}}_{t}$ also satisfies linear independence condition, it is straightforward that $\boldsymbol{s}^{\text{add}}_{t}$ can also be identified component-wisely. As for the additional structural matrices introduced, the Markov condition and faithfulness assumptions required for their identifiability have already been demanded in the identifiability properties of the existing structural matrices, thus no additional proof is needed.
∎

###### Corollary 1 .

(Identifiability under Multiple Shifts). Assume the data generation process in Eq. (21) involves both distribution shifts and state space shifts that comply with Theorem 2 and Theorem 3, respectively. In this case, both the domain-specific factor $\boldsymbol{\theta}_{i}$ and the newly added state variable $\boldsymbol{s}^{\text{add}}_{t}$ are component-wise identifiable.

###### Proof.

This corollary can be directly derived by leveraging the conclusions from Theorems 1-3.
∎

###### Theorem 3 .

###### Proof.

###### Corollary 1 .

###### Proof.

### A.4 Extension to Nonlinear Cases: Challenges and Empirical Validation

The main challenge in extending the identifiability of $\theta_{i}^{\boldsymbol{s}}$ to nonlinear scenarios lies in the fact that $g^{\boldsymbol{s}}$, in this context, represents a general nonparametric transition dynamic. This makes it difficult to disentangle $\theta_{i}^{\boldsymbol{s}}$ from $(\boldsymbol{s}_{t-1},\theta_{i}^{\boldsymbol{s}})$, as we do in the proofs of Theorem 2. Although recent works have made significant progress in establishing the identifiability of causal processes in nonparametric settings (Yao et al., ; ; Kong et al., ), they typically rely on the assumption of invertibility. However, as noted in Liu et al. (), while assuming the invertibility of the mixing function $g$ is reasonable, we cannot make the same assumption for $g^{\boldsymbol{s}}$, as this often does not hold in practice. But this does not imply that $\theta_{i}^{\boldsymbol{s}}$ is unidentifiable. On the contrary, the empirical results in Fig. 4 demonstrate that even in nonlinear settings, the learned $\hat{\theta}_{i}^{\boldsymbol{s}}$ remains a monotonic function of the actual change factors $\theta_{i}^{\boldsymbol{s}}$, corroborated by findings from Huang et al. (). This motivates us to extend Theorem 2 to broader nonlinear scenarios in our future research, a task that is challenging yet promising.

Figure: Figure 4: Correlation between the ground-truth $\boldsymbol{\theta}_{i}^{\boldsymbol{s}}$ and the estimated values in our simulations.
Refer to caption: x7.png

## Appendix B Self-Adaptive Expansion Strategy

We design three different approaches for state space expansion: Random, Deterministic, and Self-Adaptive. In Random, the number of causal variables expanded is randomly chosen. For Deterministic, we follow the approach used in DEN (Yoon et al., ), first adding a predefined number of variables to the causal graph and then applying group sparsity regularization to the network parameters corresponding to the newly added variables. Table 3 provides the final expansion results of various methods across these tasks, as well as the scope of expansion. Next, we introduce the Self-Adaptive method.

Different from prior methods, Self-Adaptive integrates state space expansion into the reinforcement learning framework. To achieve this, the first thing that needs to be done is to transform the expansion concept into a decision-making process. Since our goal is to determine how many causal variables should be incorporated into the causal graph, the action $u_{t}$ can be intuitively represented as the number of variables to add. Regarding the state variable $v_{t}$, it is designed to reflect the current state of the system. Given that the model’s expansion is inseparable from its original structure and adaptability to the target task $\mathcal{M}_{i}$, we formulate the state variable as a reflection of both the original network size and its predictive capability on $\mathcal{M}_{i}$, denoted as $v_{t}=(x_{t},\Delta_{\tau})$. To be specific, $x_{t}=(x^{o}_{t},x^{r}_{t},x^{\boldsymbol{s}}_{t})$, where $x^{o}_{t}$, $x^{r}_{t}$, and $x^{\boldsymbol{s}}_{t}$ represent the combination of the number of nodes for each layer in the models defined in Eq. (2), respectively. If the transition model is an $m$-layer network, then $x^{\boldsymbol{s}}_{t}$ is an $m$-dimensional vector, with the $l$-th element representing the number of nodes in the $l$-th layer. Moreover, $\Delta_{\tau}=\tau-\tau^{\star}$ represents the difference between the model’s predictive performance $\tau$ and the threshold $\tau^{\star}$.

Whenever the controller takes an action $u_{t}$, we correspondingly extend the model by augmenting it with additional components, and train the newly added parts from scratch with a few amount of data. For instance, if $u_{t}=d^{\prime}$, it implies that $d^{\prime}$ causal state variables will be incorporated. Then for the observation model, we only need to focus on learning the mapping from $\boldsymbol{s}^{\text{add}}_{t}$ to $o_{t}$, together with the structural constraints $D^{\boldsymbol{s}^{\text{add}}_{t}\to o}$. A similar principle applies to the reward and transition model. Finally, we re-estimate the performance of the expanded model, denoted as $\tau^{\prime}_{t}$, and derive the reward as:

$$ $r_{t}=(\tau-\tau^{\prime}_{t})-\lambda_{r}u_{t},$ (37) $$

where $\tau-\tau^{\prime}_{t}$ reflects the change in the model’s representational capacity before and after expansion, the term $-\lambda_{r}u_{t}$ acts as a regularization penalty that imposes a cost on model expansion, and $\lambda_{r}$ is the corresponding scaling factor, which is set to
$0.01$ in our experiments.

Building upon the above foundation, it becomes feasible to develop and train a policy aimed at dynamically enhancing the causal model in adaptation to current task $\mathcal{M}_{i}$ through strategic expansion.

**Table 3: The expansion results of different methods in our experiments.**
| Experiments | Random | Deterministic | Self-Adaptive | Expansion Scope |
| --- | --- | --- | --- | --- |
| Simulated | 4.0 | 5.0 | 4.2 | (0, 8] |
| CartPole | 6.2 | 6.0 | 4.0 | (0, 8] |
| CoinRun | 9.4 | 8.0 | 6.8 | (0, 10] |
| Atari | 7.8 | 8.0 | 6.4 | (0, 10] |

## Appendix C The CSR Algorithm

Figure: Algorithm 1 Towards Generalizable RL through Causality-Guided Self-Adaptive Representations

## Appendix D Complete Experimental Details

All experiments are conducted using an Nvidia A100 GPU. Training from scratch on the simulated and CartPole environments take less than 4 hours, training on Atari required approximately one day, and training on CoinRun takes about 4 days. Below, we provide detailed implementation specifics for the experiments, including a thorough description of the environment, the selection of hyperparameters, and additional experimental results.

### D.1 Hyperparameters

**Table 4: Architecture and hyperparameters for the simulated environment.**
| Simulated Environment | Architecture | Hyper Parameters |
| --- | --- | --- |
| Change factor $\theta^{\boldsymbol{s}}$ | - | Uniform, [-1, 1] |
| Random noise $\epsilon_{t}^{\boldsymbol{s}}$ | - | Gaussian, $\mathcal{N}(0,0.2I)$ |
| Reward function $g^{r}$ | Dense | $128$, he uniform, relu |
| Dense | $64$,   he uniform, relu |  |
| Dense | $1$,  glorot uniform |  |
| Transition function $g^{\boldsymbol{s}}$ | Dense | $4$,  glorot uniform, tanh |
| Observation function $g^{o}$ | Dense | $128$, glorot uniform |

**Table 5: Hyperparameters of CSR for CartPole, CoinRun and Atari games.**
| Hyper Parameters | Values in CartPole | Values in CoinRun and Atari |
| --- | --- | --- |
| Action repeat | 1 | 4 |
| Batch size | 20 | 16 |
| Imagination horizon | 8 | 15 |
| Sequence length | 30 | 64 |
| Size of $\theta$ | 2 | 20 |
| Size of $h_{t}$ | 30 | 512 |
| Size of $z_{t}$ | 4 | 32 |
| Size of hidden nodes | 100 | 512 |
| Size of hidden layers | 2 | 2 |
| Regularization terms $\lambda_{\text{KL}},\lambda_{\text{reg}}$ | 0.02 | 0.1 |

**Table 6: Available modes and difficulties in each game of our Atari experiments.**
| Game | Modes | Difficulties |
| --- | --- | --- |
| Alien | [0, 1, 2, 3] | [0, 1, 2, 3] |
| Bank Heist | [0, 4, 8, 12, 16, 20, 24, 28] | [0, 1, 2, 3] |
| Crazy Climber | [0, 1, 2, 3] | [0, 1] |
| Gopher | [0, 2] | [0, 1] |
| Pong | [0, 1] | [0, 1] |

### D.2 Simulated Environment

We construct the simulated environment based on the following POMDP framework:

$$ $\displaystyle\boldsymbol{s}_{1}$ $\displaystyle\sim\mathcal{N}(0,I_{0}),$ (38) $\displaystyle o_{t}$ $\displaystyle=g^{o}(\boldsymbol{s}_{t-1}),$ $\displaystyle s_{t}$ $\displaystyle=g^{\boldsymbol{s}}(\theta^{\boldsymbol{s}},\boldsymbol{s}_{t-1}, a_{t-1})+\epsilon_{t}^{\boldsymbol{s}},\quad\epsilon_{t}^{\boldsymbol{s}}\sim \mathcal{N}(0,I_{\epsilon})$ $\displaystyle r_{t}$ $\displaystyle=g^{r}(\boldsymbol{s}_{t-1}),$ $$

where $\boldsymbol{s}_{1}$ and $\epsilon_{t}^{\boldsymbol{s}}$ are sampled from Gaussian distributions, and functions $g^{o}$, $g^{\boldsymbol{s}}$, and $g^{r}$ are implemented using MLPs. To simulate scenarios of distribution shifts, we generate random values for $\theta^{\boldsymbol{s}}$ in different tasks. To model changes in the state space $\mathcal{S}$, we randomly augment it with $n$ dimensions, where $n$ is uniformly sampled from the range [3, 7]. Moreover, to introduce structural constraints into the data generation process, we initialize the network parameters for $g^{o}$, $g^{\boldsymbol{s}}$, and $g^{r}$, by randomly dropping them out with a probability of $0.5$. The network weights then remain constant throughout the learning process. For each task, agents are allowed to collect data over 100 episodes, each consisting of $256$ time steps. Table 4 provides the corresponding network architecture and hyperparameters.

### D.3 CartPole Environment

Based on the conclusions in Florian (), we modify the CartPole game to introduce changes in the distribution and state space. Specifically, for Task 1 and Task 2, the transition processes adhere to the following formulas:

$$ $\begin{gathered}\ddot{\psi}=\frac{g\sin\psi+\cos\psi\left(\frac{-F-m_{p}l\dot{ \psi}^{2}\sin\psi}{m_{c}+m_{p}}\right)-\frac{\mu_{p}\dot{\psi}}{m_{p}l}}{l \left(\frac{4}{3}-\frac{m_{p}\cos^{2}\psi}{m_{c}+m_{p}}\right)}\\ \ddot{x}=\frac{F+m_{p}l\left(\dot{\psi}^{2}\sin\psi-\ddot{\psi}\cos\psi\right) }{m_{c}+m_{p}},\end{gathered}$ (39) $$

where the parameters used are the same as those defined in Section 2 of Florian (), except that $\psi$ is used in place of $\theta$. By altering the values of $m_{c}$ and $g$, we can simulate distribution shifts. For Task 3, we introduce the friction between tha cart and the track into the game, denoted as $\mu_{c}$, thus altering Eq. (39) to Eq. (21) and Eq. (22) in Florian (), which is:

$$ $\begin{gathered}\ddot{\psi}=\frac{g\sin\psi+\cos\psi\left\{\frac{-F-m_{p}l\dot {\psi}^{2}\left[\sin\psi+\mu_{c}\operatorname{sgn}\left(N_{c}\dot{x}\right) \cos\psi\right]}{m_{c}+m_{p}}+\mu_{c}g\operatorname{sgn}\left(N_{c}\dot{x} \right)\right\}-\frac{\mu_{p}\dot{\psi}}{m_{p}l}}{l\left\{\frac{4}{3}-\frac{m_ {p}\cos\psi}{m_{c}+m_{p}}\left[\cos\psi-\mu_{c}\operatorname{sgn}\left(N_{c} \dot{x}\right)\right]\right\}}\\ \ddot{x}=\frac{F+m_{p}l\left(\dot{\psi}^{2}\sin\psi-\ddot{\psi}\cos\psi\right) -\mu_{c}N_{c}\operatorname{sgn}\left(N_{c}\dot{x}\right)}{m_{c}+m_{p}}.\end{gathered}$ (40) $$

Note that $\mu_{c}$ varies cyclically every $5$ steps among {3e-4, 5e-4, 7e-4}, so that the agent must continually monitor it throughout the process to achieve higher and stable rewards, This helps us assess whether the agent has detected the newly introduced variable. Additionally, we also visualize these changes in the image inputs; Fig. 6 presents examples under different friction coefficients. In the typical CartPole setup, the action values represent the direction of the force $F$. Specifically, 0 0 denotes a leftward force, while $1$ indicates a rightward force, with a default magnitude of $F_{\text{mag}}=10$. In Task 4, we have expanded the possible values of $F$ to include $\{0.5\times F_{\text{mag}},F_{\text{mag}},1.5\times F_{\text{mag}}\}$, thereby extending the action dimension to $6$. Our implementation is built upon Dreamer Hafner et al. (), Table 5 lists the hyperparameters that are specifically set in our experiments. Fig. 7 illustrates the corresponding training results. Moreover, Fig. 8 shows a comparison of the reconstruction effects of different world models across the three tasks. Note that the transition model is an RSSM (Hafner et al., ) in our implementation. Consequently, we divide the state $\boldsymbol{s}_{t}$ into a deterministic state $\boldsymbol{h}_{t}$ and a stochastic state $\boldsymbol{z}_{t}$. With this setup, the identified structural matrices in our experiment is shown in Fig. 9.

Figure: Figure 5: An illustration of the CartPole environment.
Refer to caption: extracted/5894886/figures/env_cartpole.png

Figure: Figure 6: An illustration of the CartPole game under different friction coefficients.
Refer to caption: extracted/5894886/figures/cartpole_fric_illu.png

Figure: Figure 7: Training results of our CartPole experiments.
Refer to caption: x8.png

Figure: Figure 8: The reconstructed observations of different world models in CartPole.
Refer to caption: extracted/5894886/figures/cartpole_recon_illu.png

Figure: Figure 9: Estimated causal structural matrices in the experiments: (a) CartPole Task 1; (b) High difficulty CoinRun games; (c) An illustration of the state space in Atari and how $D_{k}^{z\to h}$ evolves across tasks. Here, $d$ represents the size of the state space, with $d^{\prime}=0$ indicating distribution shifts, and non-zero $d^{\prime}$ signifying state space expansion.
Refer to caption: x9.png

### D.4 CoinRun Environment

CoinRun serves as an apt benchmark for studying generalization, owing to its simplicity and sufficient level diversity. Each level features a difficulty coefficient ranging from 1 to 3. Following Cobbe et al. (), we utilize a set of 500 levels as source tasks and generalize the agents to target tasks with higher difficulty levels outside these 500 levels. We maintain all environmental parameters consistent with those reported in Cobbe et al. (). For the world models of CoinRun and Atari games, we employ the same hyperparameters, which are listed in Table 5. Fig. 12 visualizes the reconstructions generated by various methods when generalizing to high-difficulty CoinRun games, where the first row displays the ground truth observations, the second row illustrates the model-generated reconstructions, and the third row highlights the differences between the ground truth and the reconstructions. We find that our proposed CSR method effectively captures newly emerged enemies, which the baseline methods fail to do. Moreover, Fig. 9 presents the estimated structural matrices.

Figure: Figure 10: Different modes of the game Crazy Climber.
Refer to caption: extracted/5894886/figures/env_atari.png

Figure: Figure 11: Visualization of reconstructions when generalizing to high-difficulty CoinRun games using various methods: (a) Dreamer; (b) AdaRL; (c) CSR (ours).
Refer to caption: x12.png

### D.5 Atari Environment

Atari serves as a classic benchmark in reinforcement learning, with most studies using it to evaluate the performance of proposed methods on fixed tasks. However, as mentioned in Machado et al. (), many Atari tasks are quite similar, also allowing for the assessment of a reinforcement learning method’s generalization capabilities. Specifically, within the same game, we can adjust its modes and difficulty levels to alter the game dynamics. Although the goals of the game remain unchanged, increasing the complexity of modes and difficulty necessitates consideration of more variables, thus posing challenges for knowledge generalization.

According to Table 10 in Machado et al. (), we select five games that feature different modes and levels of difficulty, and set the task sequence as four in our experiments. The corresponding modes and difficulties available in these five games are given in Table 6. Fig. 10 gives an example in Crazy Climber. For Task 1, the agent is trained from scratch. For Tasks 2 to 4, different methods are employed to maximize the generalization of acquired knowledge to new tasks. Fig. 15 to Fig. 19 illustrates the training returns in these five games with different methods, Fig. 14 are the average generalization performances, and Fig. 13 are the corresponding reconstructions. Besides, we also illustrate how the structural matrices evolve during model adaptation in Fig. 9.

Note that changes in the latent state space in Atari games are not as straightforward as previous tasks, because variations in mode and difficulty typically influence latent state transitions rather than introducing new entities which are directly observable in the game environment. Hence, to further explore what the newly added variables represent, we first deactivate them and deduce their representations within the model, and then generate the reconstructions. Fig. 12 displays the reconstructed observations using the full state representations and after removing the newly added variables. By comparing the observations before and after removal, it is evident that although the model can still reconstruct most of the buildings, it loses the precise information about the climber (the colorful person in the lower left corner of the image). Such disappearance in the reconstructions demonstrates the success of introducing the newly added variables in capturing the changing aspects in the latent state transitions. This further illustrates that CSR is also capable of handling general generalization tasks, even when the target domains do not exhibit significant space variations.

Furthermore, an interesting direction for future research is the exploration of the ability of RL methods to generalize across different games, such as Space Invaders and Demon Attack. These games feature distinct visuals but share similar gameplay and rules. While humans can easily transfer knowledge between these tasks, this remains a challenging feat for artificial intelligence. Additionally, investigating the relationship between generalization performance and task similarity presents a valuable research topic.

Figure: (a) Dreamer
Refer to caption: x17.png

Figure: Figure 14: Average generalization performance of different methods in Atari games.
Refer to caption: x29.png

Figure: Figure 15: Training results of various methods in game Alien.
Refer to caption: x30.png

Figure: Figure 16: Training results of various methods in game Bank Heist.
Refer to caption: x31.png

Figure: Figure 17: Training results of various methods in game Crazy Climber.
Refer to caption: x32.png

Figure: Figure 18: Training results of various methods in game Gopher.
Refer to caption: x33.png

Figure: Figure 19: Training results of various methods in game Pong.
Refer to caption: x34.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - Given that we can describe any two tasks as a single MDP with an SCM that contains a task variable that takes values in ${0,1}$ and which can be used to activate/deactivate other state variables, couldn't the method drop the expansion mechanism? Shouldn't learning $D$ and $\theta$ be sufficient to identify the underlying changes in the MDP?
- Is the Dreamer baseline configured to account for the extra learned parameters of CSR (eg. does it have the same state size as CSR eventually reaches)?
soundness: 2
strengths: - the problem tackled here is difficult and highly relevant. Generalization under variations of the underlying dynamics is an area that has shown little progress
- paper is well structured, fairly easy to follow and the description of the method appears to be complete
- I take the results on the toy synthetic task in fig 3a and the modified CartPole experiment to be the most clear evidence in supporting C1 and C2.
- the ablation on the expansion mechanism suggests that indeed the performance gain is not merely correlated with the increased capacity of the estimators but with the quality of the approximation of the underlying causal graph.
summary: This work augments model-based methods such Dreamer with a mechanism for facilitating policy adaption. The method is evaluated on a synthetic task, a modified CartPole with changing dynamics, CoinRun (Procgen) and 5 Atari games with modes and difficulty levels enabled.

**Claims**

C1. The method distinguishes between and addresses two situations: i) distributions shifts ("changes in transition, observation or reward functions") and ii) state/action space expansions ("differences in latent or action spaces")

C2. Causal representation learning (as employed here) uncovers the structural relationships among variables, enabling the agent to autonomously determine in which of the two situations they are in.

**Assumptions**

Identifiability of the model parameters assumes conditional independence between estimated states at different time-steps (Markov condition) and the (stronger?) faithfulness assumption.

Identifiability of the "expanded state space" further assumes linearity, but provides empirical evidence of monotonic correlation between estimated parameters and the ground-truth for simple cases.
weaknesses: - The paper would benefit from more qualitative analysis in support of C1 and C2. For example graphs showing the evolution of $\mathcal{L}_{\text{pred}}$ and the ratio between the two modes of the algorithm (distribution shift detection and space expansion). What I am going for is whether it would be possible to visualize a correlation between a task and which mechanism is selected during adaption to that task.
- I take the empirical results on Atari to be somewhat weak, with the baseline (Dreamer) often performing statistically insignificantly from the proposed method. Furthermore I would encourage the authors to make a more principled selection of environment when training on a small subset of ALE, for example Atari-5 or Atari-10 games that support tasks [1].
- I have some reservations regarding the way the distinction between the two scenarios is presented. Given that for any number of tasks we could write down the full MDP using a single transition (and reward) matrix, I believe it could be better exemplified not through CoinRun (which is fairly opaque) but maybe through a simple Factored MDP or a structural causal model.

[1] https://arxiv.org/abs/2210.02019

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 3
questions: Refer to the questions above.
soundness: 2
strengths: - This paper is well-written. The structure of CSR is clearly explained, and the authors provide some technical details. 

- The three-step CSR framework (detection, expansion, and pruning) is well-structured. It provides a systematic way to handle different environment changes, contributing to effective and efficient policy adaptation.

- Including ablation studies and comparative evaluations for various expansion strategies (random, deterministic, self-adaptive) provides valuable insights into the effectiveness of different design choices.
summary: This paper proposes a novel model, CSR to enhance RL generalization across diverse and dynamic environments. CSR tackles distributional changes and environmental space variations, such as introducing new elements or challenges. CSR uses causal representation learning to identify latent variables and structural relationships, enabling agents to recognize and respond to environment changes more efficiently. The approach is structured around a three-step process: detecting shifts, expanding the model when necessary, and pruning irrelevant variables for task-specific learning. Experimental results show that CSR achieves faster adaptation and better performance than existing methods on tasks like CartPole, CoinRun, and Atari games.
weaknesses: - CSR's effectiveness depends on accurately identifying latent causal variables and relationships. In noisy or complex environments, causal discovery can be error-prone, impairing the model's adaptability and leading to suboptimal or even incorrect adaptations. Therefore, it is necessary to see whether CSR can work in a complex environment or task.

- Techniques like causal representation learning for domain adaptation, adaptive policy updates, and structure discovery in RL are not entirely new. Many recent papers explore causal inference within RL and adaptive mechanisms for generalization, such as AdaRL. I cannot tell if this paper's novelty differs from AdaRL.

- Based on my understanding, $\theta$ is trained on the data collected from the target domain. Therefore, how does CSR adapt to an unseen task?

- The self-adaptive strategy, which searches for the optimal structure of causal variables, can be computationally intensive. Can authors clarify the training cost for each step of CSR?

- CSR concentrates on improving the performance of an RL agent with causal representation in dynamic change environments. I am inquisitive about how CSR compares to other curriculum RL methods regarding sample efficiency and final performance, such as ALPGMM.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: * how does CSR compare with the other methods in terms of added parameters and training/inference time ?
soundness: 4
strengths: **Interesting and novel motivation**. The idea to use causality to distinguish different types of changes is interesting, and might be one of the most promising directions to allow RL agents to generalize past their training environments.

**Soundness of the method**. The method is overall sound.

**Overall clarity**. The overall structure of the paper makes it quite easy to follow and understand the method. The clarity ~can be~ has been improved (as detail led below).

**Broad experimental evaluation**. The experimental evaluation is broad (Procgen+Atari), and demonstrate the superiority of CSR in terms of pure performances to the baselines. The Atari modifications of the environments have been introduced in [1], it would be nice to cite them.

[1] Farebrother, J., Machado, M. C., & Bowling, M. Generalization and regularization in DQN. arXiv (2018).
summary: The authors introduce *causality-guided self-adaptive representation* (CSR), for RL agents. It introduces a causal world model into RL agents, that allow them to distinguish two types of environmental changes:
1. distribution shifts
2. State/Action space expansion. 
The causal world model generates the next observation and can be adjusted if the predicted observation differs too much from the observed one, allowing to detect that the environment has changed.

Thus, the authors evaluate their baselines on different sets of environments such as the famous Atari and Procgen.
weaknesses: Several points are unclear: 
* **Implicit world models** CSR still encodes an implicit world model. The learned world models cannot be interpreted. The use of causality had me first think of the contrary. It could be nice to insist on this point. Explicit RL reasoning agents have also been developed ([1, 2, 5]). 
* **The difference between the two types of change is not clear**. Beside the apparent effort of the authors to explain this point, I'm still unsure of the difference, and it's a core contribution. I would advise using the MDP in the introduction, explaining how the MDP is modified for this type of change.
* **What are the answered scientific questions?** The experimental evaluations can be reorganized in terms of scientific questions (Q1 learning ability, Q2 adaptivity, Q3 core components for performances, i.e. ablation)... 

A potentially striking point for the paper would be:
Could a causal world model help mitigate misgeneralisation/misalignment problems ([2, 3]) ? The misalignment detected in [2] on the *LazyEnemy* variation of Pong [4] (where they show that RL agents tend to follow the vertical position of the enemy instead of the ball) is quite striking. A causal world model, that could then update itself to detect that the Enemy is following the ball instead of the contrary could there help, couldn't it ?
A discussion on this issue or its use in introduction/future work could demonstrate the potential benefit of causal WM in RL.

------------
* [1] Luo et al. "INSIGHT: End-to-End Neuro-Symbolic Visual Reinforcement Learning with Language Explanations." ICML (2024).
* [2] Delfosse et al. "Interpretable concept bottlenecks to align reinforcement learning agents." Advances in Neural Information Processing Systems 37 (2024).
* [3] Di Langosco et al. "Goal misgeneralization in deep reinforcement learning." International Conference on Machine Learning. PMLR (2022).
* [4] Delfosse et al. "HackAtari: Atari Learning Environments for Robust and Continual Reinforcement Learning." arXiv (2024).
* [5] Marton, et al. "SYMPOL: Symbolic Tree-Based On-Policy Reinforcement Learning." arXiv (2024).

I am willing to augment my scores if these are addressed.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: (140) - "In each period, only a replay buffer containing sequences from the current task is available, representing an online setting." I'm not sure I understand the reasoning here. If one maintains a replay buffer for each task, and tasks are encountered sequentially, why can't you access a replay buffer for a previous task during the current task? 

(276) - In the expression for $J_{\rm rec}$, is there meant to be a log around the second probability term? 

(276-283) - Can you clarify what exactly you mean by the expectation under $q_\alpha$ in the definition of $J_{\rm KL}$ and $J_{\rm rec}$? Importantly, are the $s_{t-1}$ and $s_t$ sampled afresh according to $q_\alpha$, or recalled from the memory replay buffer? Are gradients passed through these samples, or are they held fixed? Do you pass gradients through the samples using analytic expressions for the KL-divergence, via the reparametrisation trick, or by a REINFORCE/score-function style estimator? 

(276-283) - I would also appreciate clarity on what the objectives are a function of. I'm assuming they're functions of $\phi, \beta, \alpha$, but they could also be functions of $\theta$. This just isn't clear from the main text. 

(293) - "Therefore, in this step, we exclusively updates [sic.] the domain-specific part $\theta_i$, while keeping all other parameters unchanged". Can you be more specific about how these updates are done? Are they done using the objective $J$? How many updates are performed? Do you train to convergence? With early stopping? Using a train/validation split? 

(303) - "Upon implementation, we use the final prediction loss of the model on the source task as the threshold value". This threshold seems to stringent in my mind. If there is some noise in $L_{\rm pred}$ (which would occur simply because the loss is sample based), then if $M_2 = M_1$, you would expect that on half of occasions the threshold would be violated, simply because of the noise. Am I missing something here? If not, can you confirm that this doesn't happen when $M_2 = M_1$.
soundness: 3
strengths: Overall, I liked this paper a lot, and thought it contained many good ideas. I was impressed by the experimental results. 

Originality: 
- The core of the method proposed by the paper (uncovering structural relationships between latent variable to allow for quick adaptation to novel environments) is, to my knowledge, a new and creative contribution. 
- Moreover, the algorithm is non-trivial, and involves multiple steps. The authors have solved numerous problems in order to create a pipeline which works.  

Quality: 
- The algorithm proposed by the paper, CSR, works well, and can outperform policy transfer with existing algorithms. 
- The authors perform numerous experiments across different domains to demonstrate the performance of CSR. 
- Solutions to various problems are not ad-hoc, and seem well-motivated. 
- Finally, the authors provide multiple theoretical results which justify their approach.  

Clarity: 
- The writing was in general good, and the paper had a nice overall flow.

Significance: 
- Transfer learning is an increasingly important topic in ML, and the authors do a good job of tackling the problem head-on.
summary: The authors introduce a new algorithm, Causality-guided Self-adaptive Representation-based approach, or CSR. CSR aims to allow an RL agent to quickly adapt to novel tasks which are similar to previous tasks, but involve changes in various latent variables. These latent variables are divided into a state and a "task-specific change factor"; changes in these correspond to distributional shifts and shifts in the environmental structure, respectively. 

CSR uses a world model, which consists of probabilistic mappings that approximate observations, rewards, transitions, and next-state inference functions. Within these models are various causal masks, which screen-off some conditioning variables, inducing a sparse causal structure. 

Three theorems are given which show that the world models and their sub-components are identifiable under certain assumptions. 

Details are then given for how CSR operates:
* How the world model is trained
* How distributional shifts are detected
* How the state-action space is modified when distributional shifts are insufficient to capture changes in observations

Lastly, experiments are done showing that, compared to various existing methods, CSR performs favorably following changes in the environment
weaknesses: The main weakness of the paper was that, at various points, the paper under-specified various implementation details for the algorithm. This significantly hurts the reproducibility of the work. Leaving various details under-specified in the main text would have been acceptable if the algorithm in the appendix (page 19, line 1013) filled in these details, and the main text included pointers to that appendix. Currently, the algorithm in the appendix is even less helpful for clarifying implementation details than the main text is! If the text and appendix are modified so that another research could *actually* reproduce the algorithm from the descriptions given, I think it likely I would increase my rating from borderline reject (5) to accept, good paper (8). A reader should not have to guess what you are doing - it ought to be clear from the text; otherwise it makes it very challenging for others to build upon your work. The method is very complicated - and while this is not a bad thing in and of itself, it does mean that care is needed to be clear about how the many aspects of the method work, and how they fit together. At present, that care has not been shown. If you are low on space in the main text then this is exactly the purpose of an appendix. My questions have illuminated at least some, but not all, of the points where implementation details are unclear for the paper. 

(315) - "Let $d'$ denote the number of causal variables to be added, we first determine the value of $d'$ and introduce new causal features." This is not good grammar. Replace the comma with a period. 

(319) - You say that the "existing parts of the world model are fine-tuned". Does your method suffer from catastrophic forgetting of previous tasks, or does it retain performance on task $M_1$ after going through tasks $M_2, \dots, M_9$ (for example)? Additionally, can you clarify why fine-tuning takes place, given that the previous model is intended to capture the existing relationships between variables? 

(361) - You say that $D$ can "[transition] from $1$ to $0$ during model estimation". But earlier, in line (168), you say $D$ is a binary mask, i.e., each element is in $\{ 0, 1 \}$. It's unclear to me whether $D$ is a continuous or discrete variable, or whether you threshold $D$ at some point, and if so, at which point and how. Please clarify this, and (if necessary), remove the claim that $D$ is binary. Also, is $D$ constrained in $J_{\rm reg}$? Can it be negative?

(371) - The policy is conditioned on the state - how is the state generated? Presumably Thompson sampling through your world model. Is that correct? This detail should be in the paper. 

(377) - "we train the newly added structure in the policy network while finetuning the original parameters". Train how? You haven't given sufficient details as to how you learn a policy. It looks deterministic, so is trained through DDPG? Or some other method? This is a very important detail! 

A final note - it is not exactly clear to me where to position your method within the literature, or what problem it seeks to pose. At first, it seems as though it is a continual learning algorithm. But you don't test on a continual learning problem - you test on sequential learning. So is the imagined use-case for the algorithm simply adaptation to tasks in sequence?  

Additionally, it appears that CSR requires notification whenever the task has changed. But this is surely a weakness compared to policy adaptation methods, which can just continue training on the incoming data without being informed of a distributional shift.

---------------------------------------------------------------------

As rebuttal to this review, the authors have made an effort to clarify the implementation details of the algorithm. I now feel comfortable saying that the paper could be reproduced by other researchers, and therefore view it as a solid contribution to the field.

Meta-review: This paper aims to improve RL method’s ability to handle domain shifts and state/action space changes. The core aspect of the proposed method is creating a transition model in a latent space with an expanded state space that maintains some parameters to model the current model dynamics. The approach also handles state and action space changes by shrinking and expanding the latent space and model parameters. A casual modeling approach is used to identify the relationships between relevant parameters. The approach is applied to a modified version of the Dreamer architecture. Experiments demonstrate that the model can adapt to domain shifts and state and action space changes. Theoretical results are also presented that show that the latent space parameters can be identified in the case of linear models. 

The modeling of latent space parameters is reminiscent of Bayes adaptive MDPs [1] and other model-based Bayesian RL methods. The authors should review this literature and add references to this line of work where relevant. 

[1] Duff, M.O., & Barto, A.G. (2002). Optimal learning: computational procedures for Bayes-adaptive Markov decision processes.

The reviewers all agree this paper is worthy of acceptance. The paper expands on recent adaptive RL methods, most notably regarding the adaptive latent space dimensions. As this successfully addresses the limitations of previous methods, it warrants publication.

Decision: Accept (Poster)