Towards Generalizable Reinforcement Learning via Causality-Guided Self-Adaptive Representations
Authors: Yupei Yang
Biwei Huang
Fan Feng
Xinyue Wang
Shikui Tu
Lei Xu
Venue: ICLR 2025 Poster (2025)
Keywords: Reinforcement Learning
Transfer Learning
Primary area: reinforcement learning

Abstract: General intelligence requires quick adaptation 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.

Published as a conference paper at ICLR 2025
TOWARDS GENERALIZABLE REINFORCEMENT LEARNING VIA CAUSALITY GUIDED SELF-ADAPTIVE
REPRESENTATIONS
Yupei Yang', Biwei Huang"; Fan Feng2,, Xinyue Wang', Shikui Tu'; Lei Xu'
Shanghai Jiao Tong University. "University of California San Diego
MMohamed bin Zayed University of Artificial Intelligence (yupei-yand tushikui, e edu. cn, (bih007, xiw159)eucsd. edu, ffengl 01 70 gmall. com
ABSTRACT
General intelligence requires quick adaptation across tasks. While existing re- inforcement learning (RL) methods have made progress in generalization, they typically assume only distribution changes between source and target domains. this paper, we explore wider range of scenarios where not only the distribution but also the environment spaces may change. For example, in the CoinRun envi- ronment, 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 causality guilled self-adaptive represen- tation-based approach, called CSR, that equips the agent to generalize effectively across tasks with evolving dynamics. Specifically, we employ causal represen- tation 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 few samples and out- performs state-of-the-art baselines on a wide range scenarios, including our simulated environments, CartPole, CoinRun and Atari games.
INTRODUCTION
In recent years, deep reinforcement learning (DRL, (Arulkumaran et al. 2017)) has made incred- ible progress in various domains (Silver et al.) 2016 Mirowski et al.) (016). 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) 2018 Legg & Hutter 2007). Therefore, essential to build system where agents also perform relibbe and interpretable generalizations to advance toward general artificial intelligence (Kirk et al. 2023).
A straightffrward 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.] (202)). Approaches along this line include, but not limited to, fine-tuning (Mesnil et al. 2012), reward shaping (Harutyunyan et al., 2015). importance reweighting (Tirinzoni et al. 2019), learning robust policies (Taylor et al. 2007| Zhang et al.] 2202)), sim2real (Peng al.) 2020), adaptive RL (Huang et al.) 2021). and subspace building (Gaya et al. 2022). However, these algorithms often rely on an as- sumption 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 2009 Zhou
al. 2023).
ccoreepponding author

---

Published as a conference paper at ICLR 2025
(a)
(b)
Figure I: Environmental changes may may not necessitate retraining RL agents, illustrated on different variations of CoinRun. Changes in the amount and shape of obstacles from (a) (b) do not prevent the agent from completing the task, while deadly holes and enemies introduced in (c) necessitate retraining.
In this paper, we expand the application of RL beyond its traditional confines by exploring its adapt- ability in broader contexts. Specifically, our investigation focuses policy adaptation in two distinct scenarios:
the source and target data originate from the same environment space but ex- hibit differences in their distributions, e.g. changes in transition, observation or reward functions;
Distribution shifts:
2. State/Action space expansions: the source and target data are collected from different environ- ment spaces, e.g. they differ the latent state or action spaces.
These scenarios frequently occur in practical settings. To illustrate, we reference the popular Coin- Run environment (Cobbe et al. 2019). As shown in Fig. [ the goal of CoinRun is 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. ((a) and Fig. I(b) this exemplifies distribution shifts. Additionally, Coin- Run offers multiple difficulty levels. In lower difficulty settings, only a few stationary obstacles are present, while at higher levels, variety of enemies emerge and attack the agents (see Fig. [(c]. To prevail, agents must learn to adapt to these new enemies - this scenario illustrates state/action space expansions.
We propose Causality- -uilde Self-adaptive Representation-basee based approach, termed CSR, to ad- dress 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 rep- resentation learning (Schölkopf al. 2021 Huang al.) 2022, Wang al. 2022) 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 (Hl & Schmidhuber 2018 Hafner et al. 2020) by including a task-specific change factor O to capture distribution shifts, c.g., can characterize the changes in observations due varying background colors in CoinRun. If the introduction of can well explain the current observation, it enough to keep previously learned causal variables and merely update few parameters in the causal model. Otherwise, it implies that the current task differs from previously seen ones 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 the environment spaces of latent variables, and propose a causality-guided self-adaptive representation--ase. approach to tackle this challenge.
To characterize both the causal representations and environmental changes, construct world model that explicitly uncovers the structural relationships among latent variables in the RL sys- tem.
2

---

Published as a conference paper at ICLR 2025
Tasks
World Models
Update task-specific /
Causal graph pruming
e
e
of
e
e
PLA
SL
.00
π2"
...
...
Fas-
Sexu
Gpoa
upt
F0-11
Gx
6
Add vuriables
Efficient policy adaptation through the CSR framework, For each target task, we first use the prediction error, Eged,, determine whether it involves distribution shifts or space shifts. We then adjust the model accordingly by updating the task-specific change factor θ:, 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.
Figure 2:
By leveraging the compact causal representations, devise three-step strategy that identify where the changes of the environment take place and add new causal variables autonomously if necessary. With this self aelf-ada strategy, we achieve low-cost policy transfer by updating only a few parameters in the causal model.
WORLD MODEL WITH CAUSALITY-GUIDEI SELF-ADAPTIVE REPRESENTATIONS
We consider generalizable RL that aims to effectively transfer knowledge across tasks, allowing the model to leverage patterns learned from set of source tasks while adapting to the dynamics of a target task. Each task M, is characterized by (S,,A,, R,, Tú, d4, where S represents the latent state space, A; is the action space, O is the observation space, R; S A; IK is the reward function, T: S; x A, > P(S,) is the transition function, Φ; S x A; + P(O,) is the observation function, and 7/2 is the discount factor. By leveraging experiences from previously encountered tasks {M, j=11 the objective is to adapt the optimal policy x* that maximizes cumu- lative rewards to the target task M₁. Here, we consider tasks arriving incrementally in the sequence (M1,..., .M),..),Mx))
time periods (Ti, Tw). In each period T, only replay buffer containing sequences {(oe or, rt) from the current task M; is available, representing an online setting. While tasks can also be presented offline with predefined source and target tasks, the online frame- work more closely mirrors human learning, making 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 encode the changes the environment by introducing a domain-specific embedding into the model and leveraging it for policy adaptation.
2.1 AUGMENTING WORLD MODELS WITH STRUCTURAL RELATIONSHIPS
In POMDPs, extracting latent state representations from high-dimensional observations 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 cur- rent task. However, real-world tasks often require focus specific information. For instance, in the Olympics, swimming speed is crucial in competitive swimming events, but it is less impor- tant in synchronized swimming, where grace and precision are prioritized. Hence, it is essential for
3

---

Published as a conference paper at ICLR 2025
agents to understand and focus task-specific aspects facilitate effective knowledge transfer by selectively using minimal sufficient information.
To this end, we adopt causal state representation learning approach that not only enables to extract state representations, but also to discover structural relationships over the variables, Suppose observe sequences {o,, ae, rt) JET, for task M,, and denote the underlying causal latent states by St E S;, we formulate the world model into:
observation model:
Po(or D8 @ 82)
reward model:
Po(Te De-T ① St)
transition model:
Ps Pa(Sn, D 81-1,, D o O0-1), for k -:,..,... I,...,d (a(s;
representation model:
01-1,0-1,01).
where s, = -51,,, ... " sa.)). @ is the element-wise product, and D denote binary masks indicat- ing structural relationships over variables. For instance, if the j-th element of D E {0, 1) 1)dx1 in Eq. 1 is 1, it indicates a causal edge from the state variable 8₁, the current observation signal Oz. le., $j.t is one of the parents of o,. Consequently, we are supposed to retain $j, for the observation model. Otherwise, if De-+ 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 insufficient. We must also determine how these variables change with the environment for better generalization.
then 8j,t should be removed from the causal model.
2.2 CHARACTERIZATION OF ENVIRONMENTAL CHANGES IN A COMPACT WAY
To address the above need, we now shift our foeus 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. 2018 Schölkopf et al. (021). In the CoinRun example, such shifts might be due alterations in background colors (Ø), while other elements remain constant. Therefore, to better characterize these shifts, we introduce a domain-specific change factor, 0,, that captures leverage 8, identify the domain-shared
the variations across different domains. Concurrently, latent variables of the environments. This leads us to reformulate Eq. [ as follows:
observation model:
Po(o: D ① St, D0:00 07)
reward model:
Pr(r, Da 3 8r, 0. 2
transition model:
Po( DE-+ ④ 8t-1, De
④ 00 00, Dg- 2 0a-1), for h: == -1, 1,...,d
representation model:
9a (s: 8:-1, 0,, o.),
(2) where 0, = = 0r, 0r} 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 0 e in the target task. We assume that the value of 6,, 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 action space expands, we are supposed to add new variables 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 0,: If the introduction 0, can well capture the changes in the current observations, implies that previous tasks [_M))'=1 and M, 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 0 of the causal model for M, If not the case, then the causal graph must be expanded by adding new causal variables to explain the features unique M.
Benefits 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 given in Section
4

---

Published as a conference paper at ICLR 2025
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 MI, Theorem [ establishes the conditions under which the latent variable 8, and the structural matrices D can be identified; (2) For the target task with distribution shifts, Theorem 2 outlines the identifiability
the domain-specific factor 0, in linear cases; (3) For the target task with state space shifts, Theorem 3 specifies the identifiability of the newly added state variables spall (4) For the target task that includes both distribution shifts and state space shifts, Corollary T demon- strates the identifiability of both 0; and sqdd The proofs are presented in Appendix discuss the possibility and challenges of establishing the identifiability of 0g in nonlinear cases in
We also
Appendix A.4 followed by empirical results where the learned 0e demonstrates a monotonic corre- lation with the true values. Below we first introduce the definition of component-wist identifiability, related to Yao et al. (2021), and then we present the theoretical results.
Definition 1. (Component-wise identifiability. Let 8t be the estimator of the latent variable st- Suppose there exists a mapping h such that s; = h(&)). We say 8t is component-wise identifiable if h is an invertible, component-wer wise function.
Theorem 1. (Identifuablity of world model in Eq. 7))- Assume the data generation process in Eq. 3. If the following conditions are satisfied, then 8, component-wise identifiale: (1) for any h1, lig E (1....d)) d} and 1 F kiz &kg. and 8 are conditionally independent given &;-1i (2) for every possible value of 8t, the vector functions defined in Eq. (7) are linearly independent. Furthermore, if the Markov condition and faithflleess assumption hold, then the structural matrices D are also identifiable:
[ou, re+1l g(st, EE) 8t = 1¹| (8 a1 (61-1,,(1,,6)))
(3)
where
{
O = g'(s 7
(4)
r(+1
(()
The ct, G. Git terms are corresponding independent and identically distributed (f.i..) random noises. Following Kong et al.((202.) (2023). here we only assume that the global mapping g is invertible.
Theorem 2. (Identifiability of 0, in Eq. (2. Assume the data generation process in Eq. (). where the state transitions are linear and additive. the process encounters distribution shifts and 8, has been identified according to Theorem
[
then O₂ are component-wise identifiable:
Jor: Tt+1]
= g(8,, 09,05 ,e)))) As + Bui-1 + CO, +
(5)
8t
where
O, = g'(s,, ( g'( ((4))
(6)
TC-1
Following Yao et al. (2027). here we assume that A is full rank, and C is full column rank. We further assume that So 8g- Moreover, the Markov condition and faithfulness assumption hold, the structural matrices D are also identifiable.
Theorem 3. (Identifiability of Expanded State Spuce). Assume the data generation process in Eq, Consider the expansion of the state space S by incorporating additional dimensions. Suppose Sy has already been identified according to Theorem then the component-wit wise identifiability of the newly added variables sqed and the additional structural matrices, i.e., D= and D D-th ccn be established if squd (1) represents a differentiable function of (eti), i.e., stuud ((o,)) and (2) fulfills conditions (1) and (2) specified in Theorem] !
3
Corollary 1. (Identifiability under Multiple Shifts). Assume the data generation process in Eq. 1
involves both distribution shifts and state space shifts that comply with Theorem 2 and Theorem respectively. In this case, both the domain speeifc factor O and the newly added state variable s are component-wiss identifiable.
3 A THREE-STEP SELF- STRATEGY FOR MODEL ADAPTATION
In this section, we provide a detailed description of CSR, a strategy aimed at addressing the envi- ronmental changes between source and target tasks, thereby ensuring that models can effectively
5

---

Published as a conference paper at ICLR 2025
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
Before this, we first give the estimation procedures for the world models defined in Eq. 2 which follows the state-of-the work Dreamer (Hafner et al. 2020 2023). Given the observations in period T, we maximize the objective function O defined as J Free - -JKL + TTeg. for model optimization. The reconstruction part Tres is commonly used minimize the reconstruction error for the perceived observation or and the reward rr. which is defined as
Jno Eyy
Σ (logpu(or D @ 81, Do to @ 0?) + log D c St, Do & 01) 1ET.
We also consider the di-divereence constraints JKL that helps to ensure that the latent representa- tions attain optimal compression of the high-dimensional observations, which is formulated as:
JKt. Enj.
I (IET,
(ARI. (L(g.(Eu,) st-1, O. ai (i))lp Dg @ $t-11, D. o D9.D$ DD C aa-1))
where AKL is the regularization term. Moreover, as explained below in the Section 3.3 we further use Try as sparsity constraints that help to identify the binary masks D better. Upon implementa- tion, these three components are jointly optimized for model estimation. During the first task M1, we focus on developing the world models from scratch to capture the compact causal representa- tions effectively. Then, for any subsequent target task M, (where 2 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 M,, our first goal is to determine if it exhibits any distribution shifts. Therefore, in this step, exclusively updates the domain-specifee part 0,. while keeping all other parameters un- changed from the previous task M,-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 0,, we can also easily characterize the task-specific influence of these connections. Particularly, when 0, is set to zero, we temporarily switch the related edges off in task Mi
Here we adopt forward prediction error (Guo et al. 2020) as the criteria to determine whether the re- estimated model well explains the observations in current task, defined as Fped Eo 1 0.+11.. corresponding threshold, 7*, established. Upon implementation, we use the final pre- diction loss of the model on the source task M₁ as the threshold value, thereby avoiding the need for manual setting. If the model's performance is below this expected threshold, it means that the current task M, 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 O 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 0, fail to accurately represent the target task M,, becomes essential to incorporate additional causal variables into the existing causal graph 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, focus on developing strategies to effectively manage state expansions. Let denote the number of causal variables to be added. We first determine the value of d' and introduce new causal features. Following this decision stage, we extend the causal representations from 8, to s = (8), s (s.41.....,, 3d4d'..) .). by incorporating the additional d' 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 to d + d'. Accordingly, we focus on learning the newly incorporated components with only few samples, already captured the existing relationships between variables. 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:
where stcii
the previous model has
Formally written as J(Ø, B, a. 0, D); arguments are omitted for brevity.
6

---

Published as a conference paper at ICLR 2025
1. Random (Rnd): d' is randomly sampled from a uniform distribution.
2. Deterministic (Det): It sets a constant value for d'. However, this approach may overlook task- specific differences, potentially leading either insufficient or redundant expansions. To address this, we employ group sparsity regularization (Yoon et al. 2017) on the added parameters after deterministic expansion, which allows for expansion while retaining the capability of shrinking.
Self-Adaptive (SA): It searches for the value of d' 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 (2018), we define the state variable 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 s prediction errors before and after expansion. Details are given in Appendix B
is noteworthy that our method allows for flexibility in the choice of expansion strategies. Intu- itively, in case, the Self-Adaptin approach is most likely to outperform others, and the experi- mental 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 nec- essary 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 s into the following two classes:
1. Compact state representation ; A variable that either affects the observation or, or the reward Tt+1, or influences other state variables 8j4+1 (k + .) at the next time step (i.e., D= 1, or D 1, 1, Dj 1, e.g., 81, in Fig. 2).
2. Non-compact state representation s A variable that does meet the criteria for a compact state representation (e.g., $2,t in Fig. 2
Similarly, the change factors 6; 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 charac- terize the binary masks D and the sparsity of O., we define a regularization term Jeeg by leveraging the edge-minimality property (Zhang & Spirtes 2011). formulated as
Treg = -
|| 101-> + DD*>>111 + D >011 + || +111 + 1 1|1 + -10:111]
where Ave represents the regularization term. Incorporating the regularization term Treg directly into the objective function confers a notable advantage: it enables concurrent pruning and model training. Specifically, the presence of Treg induces certain entries of D to transition from 1 to 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 general- ization to the target task M,, 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 rep- resentation of and the domain-specific change factor 0f as inputs to the policy x², represented as: a. = π (sf; 0.). 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 of and the compact change factor 0: for task M,, if the model's prediction error in the Distribution Shifts Detection and Characterization step meets expectations, we can directly transfer the learned policy M; by applying a; = श (s{, 0.)
Otherwise, if the state variables expand from 8f sf" along with the updated change factor 0,', then relearn the policy n² basing on ग* Similarly, we train the newly added structures in the policy network while finetuning the original parameters, thereby updating the policy to az x² < 0g¹).
7

---

Published as a conference paper at ICLR 2025
Task
Task 2
Task 3
Task
Cumuuattiv Reward
0.6
DQN
Dreamer
AdaRL
-1
-1.0
CSR (Ours)
0.2
10
20 Episode
30
40
0
10
20
30
40
50
0
10
20
40
50
0
10
20
30
40
50
Episode
Episode
Episode
10
CernaalieeCuume Norualized Cunulative Rewand
1.00
Normalized Cumulative Reward
1.0
Cumulalive Reward
0.8
0.75
0.50
DQN
0.4
Dreamer
Rnd
0.25
AdaRL
with
0.2
Det
CSR (Ours)
0.00
without D
0.0
0-
0
250k
500k
750k
1M
100k
200k
300k
400k
0.25
0.5 0,75
Steps
Steps
Normalized Steps
(b)
(c)
(d)
Figure 3: Experimental results that answer key questions in Section 4
(Q1) CSR demonstrates
best generalization capability compared to baseline methods in (a) Simulation and (b) CoinRun: (Q2) CSR with structural embeddings D significantly outperforms CSR without D in Atari games; (Q3) The SA expansion strategy yields the highest normalized average training episodic return in experiments.
EXPERIMENTS
We evaluate the generalization capability of CSR on number of simulated and well-established datasets including the CartPole, CoinRun and Atari environments, with detailed descriptions pro- vided in Appendix D.3 For all these benchmarks, we evaluate the POMDP case, where the inputs are dineesiinnii observations. Specifically, the evaluation focuses on answering the following key questions:
Q1: Can CSR effectively detect and adapt the two types of environmental changes?
Q2: Does the incorporation of causal knowledge enhance the generalization performance?
Q3: Is searching for the optimal expansion structure necessary?
We compare our approach against several baselines: Dreamer (Hafner et al. 2023). which handles fixed tasks without integrating causal knowledge; AdaRL. (Huang al. 2021), which employs sim- ple scenario-base model-free DQN (Mnih al. 2015) and SPR (Schwarzer games, we benchmark against the state-of-the-arr method, EfficientZero (Ye et al. 2021). All results are averaged over runs, more implementation details can be found in Appendix
policy adaptation without space expansion considerations: and the traditional
2020). Additionally, for the Atari
CSR consistently exhibits the best adaptation capability across all these environments (Q1). Simulated Experiments. In our simulated experiments, we conducted a sequence of four tasks following the procedures outlined in Appendix D.3.1] To evaluate the performance of different methods across varying scenarios, Task 2 focuses exclusively on distribution shifts, Task addresses changes solely within the environment space, and Task combines both distribution and space changes. As shown in Fig. 3(a) CSR consistently outperforms the baselines in adapting to new tasks, particularly excelling scenarios with space variations. This demonstrates CSR's '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.
Code available httt:://ggitaaiccccc///ACCAAAAA
8

---

Published as a conference paper at ICLR 2025
Scores
Minimum Adaptation Steps
Model
Task
Task 2
Task 3
Task 4
Task
Task 2
Task 3
Task 4
DQN
102.8 (+: 7.8)
63.6 (+: 15.6)
107.4 (+: 17.0)
104.0 (+: 12.0)
Dreamer
500.0 (+ 0.0)
397.6 (+ 7.2)
311.6 23.3)
356.3 (+ 59.9)
50k
X
X
X
AdaRL
500.0 (+ 0.0)
468.0 (+ 2.3)
410.0 (+ 3.4)
407.5 (+ 34.5)
50k
4k
CSR (ours)
500.0 (#: 0.0)
500.0 (= 0.0)
500.0 (#: 0.0)
500.0 (#: 0.0)
50k
2k
4k
10k
Table 1: Only CSR consistently adapts to successive environmental changes in CartPole. Evaluation
results of various approaches presented, with maximum episode length of 500. 'Minimum Adaptation
Steps' refers to the minimal amount of data required for model generalization, illustrated in Fig. denotes that method fails adapt under limited training steps due non-convergence suboptimal performance.
Task
Random
SPR
Dreamer
AdaRL
EfficientZero
CSR (ours)
291.9
970.3
1010.1
1147.5
557.4
1586.9
Alien
(4: 83.5)
(+: 311.1)
(+ 339.0)
(4 125.7)
(+: 185.9)
(+: 127.0)
18.4
110.1
1313.7
1285.4
181.0
1454.1
Bank Heist
(1: 2.5)
(#: 128.2)
(#: 341.5)
(#: 131.8)
(#: 90.6)
(# 178.8)
9668.0
32723.5
68026.5
62565.3
56408.3
88306.5
Crazy Climber
(1. 2286.0)
(#: 12125.2)
(1 15628.6)
(1L: 15162.2)
(#) 13388.0)
(1 18029.6)
235.6
294.0
5607.3
5359.6
1083.2
6718.6
Gopher
(#: 42.5)
(# 312.1)
(#: 1982.9)
(+ 1736.2)
(= 784,8)
(#: 1703.1)
-20.2
-6.8
18.0
17.6
6.8
19.6
Pong
(1: 0.1)
(#: 14.3)
(= 3.1)
(#: 2.7)
(+: 7.3)
( 1.1)
Table 2: CSR outperforms every baseline method on the selected set of Atari games.
CartPole Experiments. CartPole is classic control task where players move cart left or right to balance a pendulum attached on it. In our experiments, we consider four consecutive tasks. Task 2 focuses exclusively
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, 9, 20, 30, 40), respectively. In Tasks 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. and Fig. 7 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 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 adjust the introduction of the friction variable and actions.
CoinRun Experiments. The learning curves in Fig. depict our method's consistent superiority over the baselines during knowledge transferring from low to high difficulty levels 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 [.3.3 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.) 2019). 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 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 summarizes the average final scores across these tasks. We see that CSR achieves the highest mean scores all the five games. Moreover, Fig. [6 illustrates the average generalization performance of various methods on downstream target tasks, while Fig. 17

---

Published as a conference paper at ICLR 2025
to Fig. 21 present the training curves for each game, respectively. The reconstructions, as well as the estimated structural matrices, are provided in Appendix D.3.4
Integrating causal knowledge by explicitly embedding structural matrices D into the world model improves the generalization ability (Q2). Figure 3(c] illustrates the average performance of CSR with and without D in Atari games. We observe 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,
Searching for the optimal expansion structure brings notable performance gains but involves trade-off' (Q3). We conduct comparative experiments using the three methods described in Section 3.2 for all the environments and average them into Fig. 3(d) The results demonstrate that seeking for the optimal structure significantly improves expansion performance, leading 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 crucial to consider this trade-off practical applications.
RELATED WORK
Recently, extensive research efforts have been invested learning abstract representations in RL, employing methodologies such as image reconstruction (Watter et al.] 2015). contrastive learning (Sermanet et al., 2018 Mazoure et al., 2020), and the development of world models (Sekar et al., 2020). 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 ac- curacy of information available for decision-makinn processes (Schölkopf et al. 2021). Approaches such as ASRs (Huang et al.] 2022) and IFactor (Liu et al.] 2023) leverage causal factorization and structural constraints within causal variables develop more accurate world models. Moreover, CDL (Wang et al. 2022). GRADER (Ding et al. 2022) and Causal Exploration (Yang et al. 2024) 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 TTaylor & Stone 2009 Zhou et al. 2023).
To overcome these limitations, Harutyunyan et al. (015) develop a reward-shaping function that captures the target task's information to guide policy learning. [Taloor et al. (2007) and andZZhng et al.) (2020) aim to map tasks to invariant state variables, thereby learning policies robust to environmen- tal changes. AdaRL (Huang et al. (022)) is dedicated to learning domain-shares and domain-specifii representations to facilitate policy transfer. Distinct from these works, CSP (Gaya et al. 2022) ap- proaches from the perspective of policy learning directly, by incrementally constructing a subspace of policies to train agents. However, most of these works assume 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 also change. Furthermore, the approach we propose is also related to the area of dynamic neural networks, where various methods have been developed address sequences of tasks that require dynamical modifications the network archi- tecture, such as DEN (Yoon et al. 2017), PackNet (Mallya & Lazebnik 2018), APD (Yoon et al.) 2019. CPG (Hung et al. 2019). and Learn-to-Grow -row (Li et al. 2019).
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 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, demonsirate 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.
RL. methods in these challenging scenarios, we introduce CSR,
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
Yupei Yang. Shikui Tu and Lei Xu would like to acknowledge the support by the Shanghai Municipal Science and Technology Major Project. China (Grant No. 2021SHZDZZX0102), and by the National Natural Science Foundation China (62172273).
REFERENCES
Kai Arulkumaran, Marc Peter Deisenroth, Miles Brundage, and Anil Anthony Bharath. Deep rein- forcement learning: A brief survey. IEEE Signal Processing Magazine, 34(6):26-38, November 2017. ISSN 1053-5888.
Djórk-Amé Clevert.
Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:1511.07285 2015.
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generaliza- tion in reinforcement learning. International conference on machine learning, pp. 1282-1289.
PMLR, 2019.
Quentin Delfosse, Sebastian Sztwiertnia, Mark Rothermel, Wolfgang Stammer, and Kristian Ker- sting. Interpretable concept bottlenecks to align reinforcement learning agents. arXiv preprint arXiv:2401.05821, 2024.
Lauro Langosco Di Langosco, Jack Koch, Lee D Sharkey, Jacob Pfau, and David Krueger. Goal mis- generalization in deep reinforcement learning. In International Conference on Machine Learning, pp. 12009, PMLR, 2022.
Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Generalizing goal-conditioned reinforcement learning with variational causal reasoning. Advances in Neural Information Processing Systems, 35:26532-26548, 2022.
Jesse Farebrother, Marlos C Machado, and Michael Bowling. Generalization and regularization in dqn. arXiv preprint urXiv:1810.00/22. 2018.
Fan Feng, Biwei Huang, Kun Zhang, and Sara Magliacane. Factored adaptation for non-stationary reinforcement learning. Advances in Neural Information Processing Systems, 35:31957-31971, 2022.
Razvan V Florian. Correct equations for the dynamics of the cart-pole system. Center for Cognitive and Neural Studies (Coneural), Romania, pp. 63, 2007.
Bean-Baptiste Gaya, Thang Doan, Lucas Caccia, Laure Soulier, Ludovic Denoyer, and Roberta Raileanu. Building subspace of policies for scalable continual learning. arXiv preprint arXiv:2211.10445 2022.
AmirEmad Ghassami, Negar Kiyavash, Biwei Huang, and Kun Zhang. Multi-domain causal struc- ture learning in linear systems. Advances in neural information processing systems, 31, 2018.
Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, Jean-Bastien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap predic-pr representations for multi- task reinforcement learning. In International Conference on Machine Learning, pp. 3875-3886. PMLR, 2020.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: /0/22, 2018
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.02193, 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04 2023.
Anna Harutyunyan, Sam Devlin, Peter Vranex, and Ann Nowé. Expressing arbitrary reward func- tions as potential-bases advice. In Proceedings of the AAAI conference on artificial intelligence, volume 29, 2015.
11

---

Published as a conference paper at ICLR 2025
Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. arXiv preprint arXiv:2107.02729, 2021. 2021.
Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández-Lobato, Clark Glymour, Bernhard Schölkopf, and Kun Zhang. Action-sufficient state representation learning for control with struc- tural constraints. In International Conference on Machine Learning, pp. 9260-9279, PMLR, 2022.
Ching-Yi Hung, Cheng-Hao Tu, Cheng-En Wu, Chien-Hung Chen, Yi-Ming Chan, and Chu-Song Chen. Compacting, picking and growing for unforgetting continual learning, Advances in Neural Information Processing Systems, 32, 2019.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint /903.0
2019.
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ica: A unifying framework. In International conference artificial intelli- gence and statistics, pp. 2207-2217. PMLR, 2020.
Khimya Khetarpal, Matthew Riemer, Irina Rish, and Doina Precup. Towards continual reinforce- ment learning: A review and perspectives. Journal of Artificial Intelligence Research, 75:1401- 1476, 2022.
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint 2014.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot gener- alisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76:201-264, January 2023. ISSN 1076-9757.
Lingjing Kong, Biwei Huang, Feng Xie, Eric Xing, Yuejie Chi, and Kun Zhang. Identification of nonlinear latent hierarchical models. Advances in Neural Information Processing Systems, 36: 2010-2033 2023.
Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard Howard, Wayne Hub- bard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 144):441-551, 1989.
Shane Legg and Marcus Hutter. Universal intelligence: A definition of machine intelligence. Minds and machines, 1443,,444,, 2007.
Xilai Li. Yingbo Zhou, Tianfu Wu, Richard Socher, and Caiming Xiong. Learn to grow: A continual structure learning framework for overcoming catastrophic forgetting. In International Conference on Machine Learning, pp. 3925-3934. PMLR, 2019.
TP Lillicrap. Continuous control with deep reinforcement learning. arXiv:1509.0299 2015.
arXiv preprint
Yu-Ren Liu, Biwei Huang, Zhengmao Zhu, Honglong Tian, Mingming Gong, Yang Yu, and Kun Zhang. Learning world models with identifiable factorization. arXiv preprint arXiv:2306.06533 06561, 2023.
Marlos C. Machado, Mare G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
Arun Mallya and Svetlana Lazebnik. Packnet: Adding multiple tasks to a single network by iterative pruning. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, pp. 7765-7773, 2018.
Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. Advances in Neural Information Processing Systems, 33:3686-3698, 2020.
12

---

Published as a conference paper at ICLR 2025
Grégoire Mesnil, Yann Dauphin, Xavier Glorot, Salah Rifai, Yoshua Bengio, Ian Goodfellow, Er- ick Lavoie, Xavier Muller, Guillaume Desjardins, David Warde-Farke Fardee, et al. Unsupervised and transfer learning challenge: deep learning approach. In Proceedings ICML Workshop on Un- supervised and Transfer Learning, pp. 97-110. JMLR Workshop and Conference Proceedings, 2012.
Piotr Mirowski, Razvan Pascanu, Fabio Viola, Hubert Soyer, Andrew Ballard, Andrea Banino, Misha Denil, Ross Goroshin, Laurent Sifre. Koray Kavukcuoglu, et al. Learning to navigate in complex environments. arXiv preprint arXiv:1611.0367 2016.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 1(7754)):52955333 2015.
Ignavier Ng. Shengyu Zhu, Zhuangyan Fang, Haoyang Li, Zhitang Chen, and Jun Wang. Masked gradient-based causal structure learning. In Proceedings the 2022 SIAM International Confer- ence on Data Mining (SDM), pp. 424-432. SIAM, 2022.
J. Pearl and D. Mackenzie. The Book of Why: The New Science of Cause and Effect. Basic Books, 2018. ISBN 978-0-465-09761-6.
Xue Bin Peng, Erwin Coumans, Tingnan Zhang, Tsang- Wei Lee. Jie Tan, and Sergey Levine. Learn- ing agile robotic locomotion skills by imitating animals. arXiv preprint arXiv:2004.00753 007844 2020.
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 109(5):612-634, 2021.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self- celf-predictive representations. arXiv preprint arXiv:2007.05929, 2020.
Ramanan Sekar, Oleh Rybkin, Kostas Danilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self sul--uueereieed world models. In International Conference on Machine Learning, pp. 8583-8592. PMLR, 2020.
Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-contrastive networks: Self-supervised learning from video. In 2018 IEEE international conference robotics and automation (ICRA), pp. 1134-11 1141. IEEE, 2018.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):4844489, 2016.
Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2001.
Matthew Taylor and Peter Stone. Transfer learning for reinforcement learning domains: survey. Journal of Machine Learning Research, 10(7), 2009.
Matthew E Taylor, Peter Stone, and Yaxin Liu. Transfer learning via inter-task mappings for tem- poral difference learning. Journal of Machine Learning Research, 8(9), 2007.
Andrea Tirinzoni, Mattia Salvini, and Marcello Restelli. Transfer of samples in policy search via multiple importance sampling. In International Conference on Machine Learning, pp. 6264- 6274. PMLR, 2019.
Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for task-independent state abstraction. arXiv preprint arXiv:2206. 13452, 2022.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances in neural information processing systems, 28, 2015.
13

---

Published as a conference paper at ICLR 2025
Ronald J Williams. Simple statistical srudient-following algorithms for connectionist reinforcement learning. Machine learning, 1:229-256, 1992.
Ju Xu and Zhanxing Zhu. Reinforced continual learning. Advances in Neural Information Process- ing Systems, 31, 2018.
Yupei Yang. Biwei Huang. Shikui Tu, and Lei Xu. Boosting efficiency in task-agnostic exploration through causal knowledge. arXiv preprint arXiv:4007.20506, 2024.
Weiran Yao, Yuewen Sun, Alex Ho, Changyin Sun, and Kun Zhang. Learning temporally causal latent processes from general temporal data, arXiv preprint arXiv:2110.05428, 2021.
Weiran Yao, Guangyi Chen, and Kun Zhang. Temporally disentangled representation learning.
Advances in Neural Information Processing Systems, 1--26492-265033 2022.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in neural information processing systems, 11125476-25488 2021.
Jachong Yoon, Eunho Yang, Jeongtae Lee, and Sung Ju Hwang. Lifelong learning with dynamically expandable networks. arXiv preprint arXiv:1708.01544 2017.
Jaehong Yoon, Saehoon Kim, Eunho Yang, and Sung Ju Hwang. Scalable and order-robust continual learning with additive parameter decomposition. arXiv preprint arXiv: 1902.09922 2019.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint MxXiv:2006.10744 2020.
Jiji Zhang and Peter Spirtes. Intervention, determinism, and the causal minimality condition. Syn- these, 182:335-347, 2011.
Kaiyang Zhou, Ziwei Liu, Yu Qiao, Tao Xiang, and Chen Change Loy. Domain generalization: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(4):439644415, 2023.
Zhuangdi Zhu, Kaixiang Lin, Anil K Jain, and Jiyy Zhou. Transfer learning in deep reinforcement learning: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
14

---

Published as a conference paper at ICLR 2025
A PROOFS OF IDENTIHIBBLITY THEORY
Before presenting the proofs, we first introduce the relevant notations and assumptions.
Sa.e) } and denote as the Notations. We denote the underlying state variable by {s1,1,- ,,, observation. Also, we denote the mapping from state estimator & to state 8 by St = h()), and denote the Jacobian matrix of h as J Let SA.e - log p(5, 88-1),, we further denote
Pea Whet Oo Den,,002.--1 Dny,
0' 002,
O' Ds 05a
(7)
Assumption 1. CR.t is twice differentiable with respect to 8k,t and differentiable with respect to for all I E {1.....d}.
Assumption 2 (Faithfulness assumption). For a causal graph G and the associated probability distribution P, every true conditional independence relation in P is entailed by the Causal Markov Condition applied to 9 (Spirtes et 2001).
A.1 PROOF OF THEOREM 1
Based on aforementioned assumptions and definitions, component-wise identifiablity of the state variable 8, and the structural matrices D in Eq. (].
the conditions for the
Theorem 1. (Identifiablity of world model in Eq. (7)) Assume the data generation process in Eq. D If the following conditions are satisfied, then 8, component-wise identifiable: (1) for any k1, k2 E {1....,d} d} and k1 + k2, Sky, and 82z are conditionally independent given 8t-1i (2) for every possible value of 8s, the vector functions defined in Eq. (7) are linearly independent. Furthermore, the Markov condition D are also identifiable:
assumption hold, then the structural matrices
{
(or, Te+1]
g(81,(t)
(8)
18
where
{
O: = g((s
(9)
r++1 (" E(11)
The et, Ei, Eg, E[4 terms are corresponding independent and identically distributed (i.idd) random noises. Following Kong el al. (2023), here we only assume that the global mapping g is invertible.
Proof. The proof proceeds in steps. First, demonstrate that the data generation process in Eq. (8) is equivalent to the noiseless data distribution. Second, we summarize the proof steps of the identifiablity of the state variables 8, under the noiseless distribution, which has already been provided in Yao et al. (2022).
Step 1: transform into noise-free distributions.
Let y/ = [o,, rt+1]- We denote p(u:) = Po (s, pxx (px(( (s;ds, where K. are the parameters of the probability functions. Suppose Pa.x (1r) Pa.A (ys) holds for all ye. where (g, k) and (g, R) are two sets of parameters, We complete the proof primarily by following Khemakhem et al. 2020).
By applying the law of total probability, have
[ (x() Pyly 8 = Pa( (p(csss))ds
(10)
Further define Po (vels:) Pe, (vi g(8))), we get
| (si) Pe, (Vi - g(x:)) = / pra(s)) Pe, (vr ((s)))ds
(11)
Replacing V. = g(si) on the left hand side, and similarly on the right hand side, we obtain
/ palo (.,l))'(,))) vol J (T.) Pe, (vr D. d.y = pa(0 (T,)) vol Jd (V.) Pe, (yr D. dew
(12)
15

---

Published as a conference paper at ICLR 2025
where vol J = Vdet J
By introducing (U)) (v))) vol Ji (ve)l(()) on both sides, rewrite Eq. (2)
as
J (V) Pe, (y+ Vo v.)dy, = . Pg.x Pe, (ye ye dyye
(13)
where v dim O + dim R. By the definition of convolution, Eq. (13 is equivalent to
(f * (1a) = (Pg.a (vee,,,
(14)
where * denote the convolution operator. Denote F[-] the Fourier transform and yor = Fly.], we have
F (1) = F
(15)
Assume set for E (x) = 0} has measure zero, we can drop (S2) from both sides, which obtains
(16)
Therefore, for all yu E O R, we have
(v)) = pg.x
(17)
This indicates that the noise-free distributions must coincide for the overall distributions remain identical after adding noise, effectively reducing the noisy case in Eq. (8) into the noiseless case.
Step establish identifiability state variables.
After transforming the problem into the noise-firee case, we proceed by summarizing the key proof steps of the identifiability the state variables 8,, following Yao et al. (2022). which are:
First, by making use of the conditional independence of the components of 8, given &, it is shown that:
8 p(s; $1-1) 0.
(18)
05
Second, by utilizing the Jacobian matrix J to calculate Eq. (18 it is derived that
8 log 8L-1)
Cr. O52 k. 0811111
OJ K.s.e
J. kja + Osk 0s
= 0. (19)
08,
Øb5,0
1=
Finally, it is established that 8; identifiable, up to an invertible, component-wiss nonlinear transformation of permuted version of &t, if the linear independence of vector funtions defindd Eq. 7 holds and the the Jacobian matrix J satisfies Eq. (19).
Moreover, the proofs for the identifiablity of the structural matrices D are presented in Huang et al. 2021 Based on these steps, we next provide the proofs of Theorem and Corollary
A.2 PROOF OF THEOREM 2
Different from existing methods, to capture the changing dynamics in the environment, have introduced a task-specific change factor, 0,, into the world model as defined in Eq. (2). Accordingly, Theorem 2 presents the identifiability of 0, and the corresponding structural matrices D for scenarios involving only distribution shifts in linear cases.
Theorem 2. (Identifiability of 0; Eq. 2 Assume the data generation process Eq. 20 where
the state transitions are linear and additive. If the process encounters distribution shifts and 8, has been identified according to Theorem]
then 0, are component-wise identifiable:
[or, Tt+1]
S1 09, 0f Ex) As Ba C# +
(20)
81
where
O = g' 4)
(21)
rt-1
g8 0r (.1).
16

---

Published as a conference paper at ICLR 2025
Following Yao al. (2027). here assume that A is full rank, and C is full column rank. We further assume that Se So- Moreover, if the Markov condition and faithfulness assumption hold, the structural matrices D are also identifiable.
Proof. The proofs for the identifiablity of Do. are given in Huang et al. 20221 Here, we provide the proof for the identifiablity of 0, which is done in the following four steps:
In step I, we prove that 8 can be identified up to component wwse transformation when only the observation function exhibits distribution shifts.
In step 2, we demonstrate that O is component-wise identifiable when only the reward function experiences distribution shifts.
In step 3, we show that 0 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 simultancously, 0, = (0% 0% 05) is identifiable.
Step 1: prove the identifiability of 0.. According to Eq. we have
Y1 = g(s 04 E).
(22)
Denote 2; (8t, 0.). There exists
hzt)))
(23)
where h' = g-- g. and a; is the estimator of xg. Since both g and 9 are invertible, h' is invertible. Therefore, we have
Bag
B8g
J
00
88:
80
00
(24)
06:
where J is full rank. Note that da.
0 and
20
O. Further recall that we assume the identifia-
30
UR;
bility of which means that J is, 0f is component-we wise identifiable.
As. is full rank. We can derive that must be full rank. That ae.
ase 36
Step 2: prove the identifiability of O,.
If only the reward function g" exhibits distribution shifts, we have
yt = g(s,, 87, c)).
(25)
is straightforward to see that 0f blockwise identifiable using the same technique in Step 1. Step 3: prove the identifiability of 8,. Recall that we have
= Ast- + Bo.-1 + C0f + G
(26)
By leveraging the recursive property of the state transition process, can derive that
8, A'so +
Σ A
a
Σ A
CO
+
Σ
A
cr-k-
(27)
K=
K=U
Similarly for šr, we have
1-1 4 Σ
s, A' +
Σ
A
0: +
Σ
k
Čor +
ci-k-
(28)
Note that 80 8o. Therefore, combining Eq. 27 and Eq.
28
gives
6-1
1-1
Σ
A
CB = st - A'A
Σ
A
C8
+ +0,
(29)
0=
k==
17

---

Published as a conference paper at ICLR 2025
where θ is constant term. Taking the derivative w.r.t f both sides, obtain
M :8 Ös, as;
da; M
A'A AA-
(30)
06
0& 000
000
where M (Σ A
C' and M (Er=) A+)
C.
Note that we further have
M
861
according to Eq. 28 That is,
00.0 08: M.
M
(31)
00.
0&
de. Recall that 0s: J full rank. Moreover, the full column rank of M and M is guaranteed by the full rank A and the full column rank of (see Proposition [T. Therefore, we can derive
rank (M 003
0st
rank
M) dim 05.
(32)
a05
08:
Due to the rank inequality property of matrix products, we have
0
00 = min dim 07 rank
00,
rank M
S min
M, M, rank(
S dim 87.
(33)
00.
00.
a0
Eq. (32) and Eq.
33) show that
00.
dim 0 < min
dim 0º, rank(
< dim 8:
(34)
06.-
To make the above equation hold true, it must have rank( rank, then 0 must be component-wise identifiable.
00" 00.
dim 0f. That is, as"
must be full
56
Step 4: prove the identifiability of 0, in the general case.
This step can be easily demonstrated by directly combining the proofs above.
Proposition 1. Suppose A a matrix with full rank, and C matrix with full column rank. Define M = (Eto A A+)C. C. Then M is full column rank.
Proof. To establish that M full column rank, it suffices to show that N = Ex= A is of full rank, Assume the eigenvalues A are denoted by A. Given that A is of full rank, the eigenvalues of A are A², Let A represent the eigenvalues of N, then we have
t-1
AN Σ A
(35)
It is evident that there exists at least one f. such that for any non-zero A, Eq. (35) is non-zero, thereby confirming that N has no zero eigenvalues. This verification ensures that Nis full rank. Hence, M maintains full column rank.
A.3 PROOF OF THEOREM 3
In Theorem 3 we foucus on the identifiablity of the newly added variables sith and their correspond- ing 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. Consider the expansion of the state space S by incorporating additional dimensions. Suppose has already been identified according to Theorem | then the component-wis wise identifiability of the newly added variables sqdil and the additional structural matrices, i.e., D and D - can be established if shald (1) represents a differentiable function of [oe. Tt+1), 1), i.e., si = ((,,,eL))) and (2) fulfills conditions (1) and (2) specified in Theorem
18

---

Published as a conference paper at ICLR 2025
Proof. For the newly added variables, since they also satisfy the conditional independence dition, we can derive the same properties as described in Eq. (88) and Eq. (9 using the same technique in the proof steps of Theorem | Additionally, since the vector functions corresponding to si also satisfies linear independence condition, it is straightforward that s can also be identi- fied camponent-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 needed.
Corollary 1. (Identifiability under Multiple Shifts). Assume the data generation process in Eq. involves both distribution shifts and state space shifts that comply with Theorem 2 and Theorem respectively. In this case, both the domain-specific factor and the newly added state variable st are component-wise identifiable.
20
Proof. This corollary can be directly derived by leveraging the conclusions from Theorems
A.4 EXTENSION TO NONLINEAR CASES: CHALLENGES AND EMPIRICAL VALIDATION
The main challenge in extending the identifiability of to nonlinear scenarios lies in the fact that g", in this context, represents general nonparametric transition dynamic. This makes it difficult to disentangle 0 from (s,-1, (f). as we do in the proofs of Theorem 2 Although recent works have made significant progress in establishing the identifiability of causal processes in nonparamet- ric settings (Yao et al.) [021] 2022 Kong et al,, (023), they typically rely on the assumption of invertibility. However, as noted Liu al. (2023). while assuming the invertibility of the mixing function is reasonable, we cannot make the same assumption for g², as this often does not hold in practice. But this does not imply that 0rr is unidentifiable. On the contrary. the empirical results in Fig. demonstrate that even in nonlinear settings, the Iearned 0 remains a monotonic function of the actual change factors 0%, corroborated by findings from Huang et al. (202)). This motivates us to extend Theorem 2 to broader nonlinear scenarios in our future research, a task that challenging yet promising.
0.1
Estimated Value
0.0
0.1
-0.2
-1.0
-0.8
-0.6
-0.4
-0.2
0.0
0.6
Ground-truth Value
Figure 4: The estimated 0e is monotonic function of the ground-truth values in our simulations.
B SELF-ADAPTIVE EXPANSION STRATEGY
We design three different approaches for state space expansion: Random, Deterministic, and Self- Adaptive. In Random, the number causal variables expanded is randomly chosen. For Determin- istic, follow the approach used in DEN (Yoon al. 2017), first adding a predefined number of variables to the causal graph and then applying group sparsity regularization to the network pa- rameters corresponding to the newly added variables. Table provides the final expansion results
19

---

Published as a conference paper at ICLR 2025
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]
Table 3: Detailed expansion results of different methods in our experiments.
of various methods across these tasks, as well the scope of expansion. Next, 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 ex- pansion concept into a decision- peckiinn-making process. Since our goal is determine how many causal variables should be incorporated into the causal graph, the action can be intuitively represented as the number of variables to add. Regarding the state variable ut designed reflect the cur- rent state of the system. Given that the model's expansion is inseparable from its original structure and adaptability
the target task M,, we formulate the state variable as a reflection of both the original network size and its predictive capability on Mi, denoted V, (Ir, A.). To be spe- cific, 2² = (x oi xf), where xi, xy and x represent the combination of the number of nodes for each layer in the models defined in Eq. (2) respectively. If the transition model is m-layer net- work, then xf is an m-dimensional vector, with the I-th element representing the number of nodes in the i-th layer. Moreover, A, performance 7 and the threshold r'
r" represents the difference between the model's predictive
Whenever the controller takes action ue, correspondingly extend the model by augmenting with additional components, and train the newly added parts from scratch with a few amount of data. For instance, if uu d', implies that d' causal state variables will be incorporated. Then for the observation model, we only need to focus on learning the mapping from st to or, together with the structural constraints ~ A similar principle applies the reward and transition model. Finally, we re-estimate the performance of the expanded model, denoted as Ty, and derive the reward as:
(r 7() X,
(36)
where 7 reflects the change in the model's representational capacity before and after expansion, the term - acts as a regularization penalty that imposes a cost on model expansion, and X,, is the corresponding scaling factor, which is set 0.01 in our experiments.
Building upon the above foundation, it becomes feasible to develop and train a policy aimed at dy- namically enhancing the causal model in adaptation to current task M, through strategic expansion.
C DISTRIBUTION SHIFTS VS. SPACE EXPANSIONS IN A MDP SCENARIO
We further present a simple MDP scenario to illustrate the two types of envirommental changes that CSR addresses: distribution shifts and space expansions. Specifically, Fig. 55() provides graphical representation of the generative environment model for the source task, where 81, denotes the latent causal variable, and {0), 0r) represents task-specific change factors.
For distribution shift scenarios (Fig. 5(b).. the target task shares the same causal variable as the source task but differs in the value of 6. For example, CartPole, the gravity (0") might shift from 9.8 to 5. Notably, this implies that the causal diagram remains unchanged, an assumption commonly adopted in prior works (Huang et 2021 2022 Gaya et al. 2022).
In contrast, space expansion involves the emergence of new variables (e.g., 82,t in Fig. [(c) which inevitably leads to changes in the causal diagram. Consequently, world models must expand their state and action spaces accommodate these new variables. This necessity motivates the development of CSR. Algorithm presents the pseudocode for CSR, where the model estimation and policy learning processes implemented using the Dreamer framework. Notably, CSR is not restricted to Dreamer and can can also be implemented with a variety of policy-learning algorithms, such Q-learning (Mnih al. 2015) and DDPG (Lillicrap 2015).
|
20

---

Published as a conference paper at ICLR 2025
(a) Source task
(b) Distribution shifts
(c) Space expansions
Figure 5: A graphical illustration of the generative environment model and the two types of changes addressed by CSR. (a) Source task: (b) Distribution shift scenario where causal diagram remains un- changed but the value of 0" differs; (c) Space expansion scenario involving the emergence of new variable $2. Grey nodes denote observed variahles, white nodes represent unobserved variahles, and red nodes highlight the changing components in the target task compared to the source task.
Algorithm 1: Towards Generalizable RI through CSR
Input:
Maximum distribution shifts detection step Te- Initialize World Model W with parameters 6, B, randomly.
Initialize D as an all-ones matrix.
Record multiple rollouts from source task M₁ and estimate the model in Eq. (2) Obtain the optimal policy π² in Mr and calculate threshold T" using W. for target tasks Mj(i ... do
Collect multiple rollouts B from M,.
while generalization do // Model re- e-semmmataion
for training steps c c do
Draw data sequences {(o,, at, at) hET, from B.
Compute model states 8 ~ qa (st 8t-1, 02 ac-1, or).
Update 0, using J. with all other parameters fixed.
Calculate Pred using W.
// Distribution Shifts Detection and Characterization
if Eped < T" then
return Latest model W and policy x for task M,. else
// State/Action Space Expansions
Search to introduce new causal variables into the graph. while not converged do
for training steps c 1,..., do
// Model estimation (Causal Graph Pruning is concurrently implemented) Draw data sequences from B and compute model states using 4a. Update é, 8, 0,, D using J.
// Policy Learning
Imagine trajectories from each 8;.
Update policy x² from the imagined trajectories via REINFORCE gradients.
// Environment interaction
for time step t = 1,..., do
Select action a, with probability e: otherwise calculate a, using n² Execute action and receive reward rr+1 and observation o1+1- Store transition (o,, oz, IT= into replay buffer B.
return Latest model W and policy x for task Mg.
21

---

Published as a conference paper at ICLR 2025
D COMPLETE EXPERIMENTAL DETAILS
Below, we provide detailed implementation specifics for the experiments, including model architec- tures and training details, the selection of hyperparameters, thorough description of the environ- ments, and additional experimental results.
D.1 MODEL. ARCHITECTURES AND TRAINING DETAILS
2023). implement the world Model components. Following Dreamer (Hafner al. 2020
model as a Recurrent Space Model (RSSM, (Cobbe et al. 2019 ) the encoder and decoder in the representation model and observation model convolutional neural networks (LeCun al., 1989), and all other functions as multi-layer perceptrons with ELU activations (Clevert) 2015)
The implementation of D. We adopt the Gumbel-Softmax (Ng al., 2022) and Sigmoid methods to approximate the binary masks D in our experiments, which is a commonly used approach in causal representation learning.
Training details. During the generalization process, we use epsilon-greedy to balance the exploration-explootaionn trade-off, and take straight-through gradients through the sampled repre- sentations for model estimation. Since the actions are always discrete, we adopt the REINFORCE gradients (Williams) 1992) with Adam optimizer (Kingma 2014) for policy learning.
Steps for distribution shifts detection. Empirically, we set the maximum training steps for distri- bution shift detection as: 1 in Simulation, 2k CartPole, 100k in Atari, and 250k in CoinRun.
Training cost. 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 ap- proximately one day, and training on CoinRun takes about 4 days.
D.2 HYPERPARAMETERS
Simulated Environment
Architecture
Hyper Parameters Uniform, [-1, 1]
Change factor (
Random noise oe
Gaussian, N (O, 0.27) 128, he uniform, relu
Dense
Reward function g"
Dense
64, he uniform, relu
Dense
1. glorot uniform
Transition function g"
Dense
4. glorot uniform, tanh
Observation function go
Dense
128, glorot uniform
Table 4: Architecture and hyperparameters for the simulated environment.
Hyper Parameters
Values in CartPole
Values in CoinRun and Atari
Action repeat
4
Batch size
20
16
Imagination horizon
15
Sequence length
30
64
Size of Ø
2
20
Size of hr
30
512
Size of
4
32
Size of hidden nodes
100
512
Size of hidden layers
Regularization terms λKL. X
0.02
0.1
Table 5: Hyperparameters of CSR for CartPole, CoinRun and Atari games.
22

---

Published as a conference paper at ICLR 2025
Game
Modes
Diffcultiess
Alien
[0. 1, 2, 3]
10,
Bank Heist
10,4, 12, 16, 20, 24, 28]
[0. 1, 2,
Crazy Climber
[0, 1, 2, 3]
[0. 1]
Gopher
[0,
[0. 1
Pong
[0. I]
[0, 1]
Table 6: Available modes and difficulties in each game of our Atari experiments.
D.3 DETAILED DESCRIPTIONS OF THE ENVIRONMENTS
In this section, we provide detailed descriptions of the construction these environments and present additional experimental results. For simulated experiments, generate synthetic datasets that satisfy the two scenarios with different types of environmental changes. For CartPole, we con- sider distribution shifts in the task domains with different gravity cart mass, and space variations by adding cart friction as the new state variable and additional force values for action expansion. For Atari games, 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 al. 2018 Farebrother et al. 2018). For CoinRun, we train agents from casy levels and generalize them to difticulty levels where there could be new enemies that have never occurred before.
D.3.1 SIMULATED ENVIRONMENT
We construct the simulated environment based the following POMDP framework:
$1 ~ (O( 9((s:11);
(37)
= g" (0 8t-1, at-1) + t., ~~1000
(s:-1),
where 81 and ef are sampled from Gaussian distributions, and functions g", g", and g" are imple- mented using MLPs. To simulate scenarios of distribution shifts, generate random values for ge in different tasks. To model changes in the state space S, we randomly augment it with n di- mensions, where is uniformly sampled from the range [3, 7]. Moreover, introduce structural constraints into the data generation process, we initialize the network parameters for g", g", and g", 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 provides the corresponding network architecture and hyperparameters.
D.3.2 CARTPOLE ENVIRONMENT
Based on the conclusions in Florian 2007), we modify the CartPole game to introduce changes in the distribution and state space. Specifically, for Task and Task the transition processes adhere to the following formulas:
g ssin ub + cos -=F-m, m.d sin
Pp m...
mc+mp
cos? vi met-ip
(38)
F + myl 1 sin vi cos $
m. + mp
where the parameters used the same those defined in Section 2 of Florian (2077). except that s used place of 6. By altering the values of m. and g, we can simulate distribution shifts. For Task introduce the friction between tha cart and the track into the game, denoted Pe, thus
23

---

Published as a conference paper at ICLR 2025
Level
Level
Level State Space Shift
Distribution Shift
gravily
gravily 9.8
gravily
cart mass
cart mass
cart mass
Friction Se-4
Figure 6: An illustration of the CartPole environment.
Task 3 Episode Length 5
Task Episode Length 10
Task 3
Episode Length 15
Friction
Friction
Friction
Figure 7: An illustration of the CartPole game under different friction coefficients.
altering Eq. 38 to Eq. (21) and Eq. (22) in Florian 2007 which is:
g sin up + cos
- PP-mpll d [sin the sgn( (N, cos Mel sgn (Nct)
Ape my
me me my
{
mpscos
[cos 1 sgn (Naii)}
met
(39)
F + mpd (iissin sin v J cos w) MoNe sgn (Not)
mc
mp
Note that (e varies cyclically every steps among [3e-4, 5e-4, 7e-4), so that the agent must con- tinually monitor it throughout the process to achieve higher and stable rewards, This helps assess whether the agent has detected the newly introduced variable. Additionally, we also visualize these changes in the image inputs; Fig. presents examples under different friction coefficients. In the typical CartPole setup, the action values represent the direction of the force F. Specifically, de- notes leftward force, while 1 indicates rightward force, with default magnitude of / mag 10. In Task 4, we have expanded the possible values of F to include {0.5 x Fung: Fiagy 1.5 x F Fangg) ) thereby extending the action dimension 6. Our implementation
built upon Dreamer (Hafner al. 2020), Table lists the hyperparameters that are specifically set in our experiments. Fig. illustrates the corresponding training results. Moreover, Fig. 9 shows a comparison of the recon- struction effects of different world models across the three tasks. Note that the transition model is an RSSM in our implementation. Consequently. we divide the state s; into a deterministic state h, and stochastic state 2,. With this setup, the identified structural matrices in our experiment is shown in Fig. 10(a)
200
Teck
Tiak
Taak
Task
Cumutaive Roieed
2004
DQN
DQN
DGN Dreamer
DQN
Dreamse
Dreomeer
Dreamer
AdRRL
AdRL
200
CSR (Our
CSR (Ours)
CSR (Ouss)
CSR (Dus)
10k.
206
40%
50%
94.
100
24
46
61
24
4k
100
Stopp
Steps
Slzpe
Skeps
Figure Training results of our CartPole experiments.
24

---

Published as a conference paper at ICLR 2025
Task
Task 2
Task
Task
Dreamer
lobs
recon
obs
recon
obs
recon
obs
recon
AdaRL
obs
recon
obs
recon
obs
recon
obs
recon
CSR (ours)
obs
recon
obs
recon
obs
recon
obs
recon
Figure 9: The reconstructed observations different world models CartPole.
D.3.3 COINRUN ENVIRONMENT
CoinRun serves as an apt benchmark for studying generalization, owing to its simplicity and suffi- cient level diversity. Each level features difficulty coefficient ranging from to 3. Following Cobbe al. (2019), we utilize 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. (2019). For the world models CoinRun and Atari games, employ the same nyperparameters, which are listed in Table 5 Fig. 12 visualizes the reconstructions generated by various methods when generalizing to high-difficllt 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 emies, which the baseline methods fail to do. Moreover, Fig. (0(b))preeen) the estimated structural matrices.
D.3.4 ATARI ENVIRONMENT
Atari serves as a classic benchmark in reinforcement learning, with most studies using it to evaluate the performance of proposed methods fixed tasks. However, mentioned in Machado et al. (2018), 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 adjust its modes and difficulty levels to alter the game dynamics. Although the goals the game remain unchanged, increasing the complexity of modes and difficulty necessitates consideration of more variables. thus posing challenges for knowledge generalization.
According Table 10 in Machado et al. (2018), 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 these five games are given in Table 6 Fig. IT gives an example in Crazy Climber. For Task 1, the agent is trained from scratch. For Tasks to 4. different methods are employed to maximize the generalization of acquired knowledge to new tasks. Fig. [7 to Fig. 21 illustrates the training returns in these five games with different methods, Fig. [6] F the average generalization performances, and Fig. 15 are the corresponding reconstructions. Besides, we also illustrate how the structural matrices evolve during model adaptation in Fig. [0(c)
Note that changes 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 fur- ther explore what the newly added variables represent, first deactivate them and deduce their representations within the model, and then generate the reconstructions. Fig. [3] displays the re- constructed observations using the full state representations and after removing the newly added
25

---

Published as a conference paper at ICLR 2025
D
D
DE
1.00
0.75
Value
0.50
0.25
0.00
2
3
4
5
7
20
40
20
40
Dimension index
Dimension index
Dimension index
D
D Di
D
1.0
Value
0.5
0.0
20
30
200
400
Dimension index
Dimension index
Dimension index
(b)
Task 1,d
Task 2, 32,d
Task 3, 32.0
Task 4,6 36,6
30
Dimension index
Dimension index
Dimension index
Dimension index
Figure 10: Estimated causal structural matrices in the experiments: (a) CartPole Task I; (b) High difficulty CoinRun games; (c) An illustration of the state space in Pong and how DE-h evolves across tasks. Here, represents the size of the state space, with d' = 0 indicating distribution shifts, and non-zero d' signifying state space expansion.
26

---

Published as a conference paper at ICLR 2025
09700 001800
19500 001200
39506 003200
mode = 0, difficulty = 0
mode 1, difficulty =
mode = 3, difficulty =
Figure 11: Different modes game Crazy Climber.
(b)
(a)
(b)
(c)
Figure 13: Visualized reconstructions in
Figure 12: Visualization of reconstmuctions when gen- eralizing to high-diffculty CoinRun games using vari- methods: (a) Dreamer; (b) AdaRI: (c) CSR (ours).
Atari target games using CSR with: (a) Full state representations; (b) Original state rep- resentations before expansion.
27

---

Published as a conference paper at ICLR 2025
Simulation
CartPole
Atari
CoinRun
Dianturtoon Space Rsponden
26
1.27
2509
Prediction
0.15
0.15
0.05
4
low
high
Task
Difficulty Level
Simulation
CariPoll
CoinRun
Task
Task 2
Task 2
High dilficulyy
Selected Times
Task
Task 3
Task 3
Task
Task 4
Tusk 4
Distiuttion Sli Space
Dissibution Släfk Space Expussion
Distibutioon Slá Space Expunsians
Distribaion
Space Expandinss
(b)
Figure 14: Visualization of (a) Evolution of prediction error Liued at the distribution shift detection step of CSR and (b) The ratio between the two modes of CSR across experimental environments.
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 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.
lower left corner the image). Such disappearance in the reconstructions
D.4 ADDITIONAL EXPERIMENTAL RESULTS
To investigate the correlation between a task and the mechanism selected during adaptation, we further visualize the evolution of Lpred and the ratio between the two modes of CSR (distribution shift detection and space expansion) across all experimental environments in Fig. [4] Note that for Atari games, averaging across environments does not yield meaningful insights. Therefore, we use Pong as representative example, and set (mode, difficulty) sequentially across four tasks (0.0), (0,1), (1,1), and (1,0). As seen Task 2 of the simulated environment and CartPole, when only distribution shifts occur, Leped typically remains low. In contrast, the emergence of new variables. for example in Task of Pong, consistently causes substantial increase in prediction error-often by an order of magnitude or more. Besides, we also observe that the agent always opts to expand its causal world model to effectively address the added complexity and variability arising from the transition from simple to high-difficulty tasks in CoinRun. These observations collectively demonstrate that Lpred serves as a reliable indicator adaptation in our experiments.
E A DISCUSSION ON FUTURE DIRECTIONS
While CSR presents promising direction for extending RL to broader scenarios and has achieved meaningful progress, the pursuit of generalizable and interpretable RL still faces enduring chal- lenges. In this section, we outline several potential research directions to inspire further innovative studies toward general artificial intelligence.
Dynamic Graphs. CSR focuses on domain generalization but has yet to address the challenges posed by nonstationary changes both time and across tasks (Feng al. 2022). It remains
28

---

Published as a conference paper at ICLR 2025
important direction for future research to develop methods that can automatically detect and model such changes to improve adaptability and robustness.
Misalignment Problems. The goal misalignment problem, also known as shortcut behavior, refers to a situation where an agent's 's performance appears aligned with the target goal but is actually driven by side goal (Di Langosco al. 2022 Delfosse al. 2024). This often occurs because the target goal and the side goal share a common causal variable, defined as Forks in causal learning (Spirtes) al. 2001 Pearl & Mackenzie 2018). Consequently, learning a causal world model defined in Eq. (2) that captures causal relationships rather than correlations could help mitigate this misalignment issue.
Beyond Sequential Settings. While our current focus is on task adaptation in sequential settings, CSR presents promising applications in continual reinforcement learning (CRL, (Khetarpal al.) )022 ) where the agent needs to utilize replay buffer containing samples from both the current and previous tasks to identify the most similar task for policy transfer, or address domain- agnossic settings by integrating domain shift detection techniques.
Generalization across different games. Another interesting direction for future research is the investigation of RI methods' ability to generalize across very different games, such as Space In- vaders and Demon Attack. These games feature distinet visuals but share similar gameplay and rules. While humans easily transfer knowledge between these tasks, this remains challenging feat for artificial intelligence.
29

---

Published as a conference paper at ICLR 2025
a Dreamer
Stired
(b) AdaRL
(c) CSR (ours)
Figure 15: Visualization of reconstructions of various methods in the Atari games.
30

---

Published as a conference paper at ICLR 2025
Alien
Bank Heist
Crazy Climber
Average Cumulaatve Reward
1500
1500
100000
75000
1000
50000
500
500
25000
0
0
100k 200k 300k 400k
0
100k 200k 300k 400k
100k 200k 300k 400k
Episode
Episode
Episode
Gopher
Pong
8000
6000
10
SPR
Dreamer
AdaRL
2000
-10
EfficientZero
CSR (Ours)
-2
0 100k 200k 300k 400k
100k 200k 300k 400k
Episode
Episode
Figure 16: Average generalization performance different methods Atari games.
Task
Task 2
Task 3
Task
Cunulatire Reward
SPR
SPR
SPR
SPR
1000
Dreamer
Dreamer
Dreamer
Dreamer
AdaRL
AdaRL
AdaRL
AdaRL
500
ZericienZerro
EfficientZero
EdticientZero
EfficienZZere
(Curs)
CSR (Curs)
(Ouns)
(Ours)
0
100k
200k Episode
300k
400k
100k
200k 30tk
400k
100k
200k Episode
300k
400k
100k Episode
200k
300k
400%
Episode
Figure 17: Training results of various methods in game Alien.
Task
Task 2
Task
Tusk
Cunulative Rewurd
1500
SPR
SPR
SPR
SPR
Dreamer
AdaR1.
AdaRI.
AdaRT.
500
Etficieneeroo
EdiicinnZeero
EdiicimiZero
EfficienZeere
CSR (Ouns)
CSR (Ouns)
(Ours)
0
100k
200k
300k
400k
100k
200k 30th
400k
100k
2001
300k
400k
100k
200k Episodo
300k
400k
Episdde
Episodo
Episodo
Figure 18: Training results of various methods in game Bank Heist.
31

---

Published as a conference paper at ICLR 2025
Tisk
Task
Cumalative Reward
SPR
50000
Dreamer
Decamer
erficientZero
EfficientZeero CSR (Ours)
EfficientZero CSR (Ours)
EfflcientZero
0
CSR (Ours)
CSR (Ours)
0
100k 200k
300k
400k
0
100k
200k Episode
300k
400k
100k
200k Episode
200k
400k
0
100k
200k Episodo
300k
400k
Episode
Figure 19: Training results of various methods in game Crazy Climber.
Task
Task
Task
Task 4
10000
Cunalativ Reward
7500
Dreamer
AdaRL.
AdaRI.
AdaRL
AdaRL EfficientZero CSR (Ours)
EfficienZerro
Efticienizrro
EfficientZero
2500
CSR (Ouns)
CSR (Ours)
CSR (Ours)
0
100k
300k
100k
200k
400k
Episode
Episode
Episode
Episode
Figure 20: Training results of various methods in game Gopher.
Task
Task
Task
Tisk
Comulsire Reward
SPR
SPR
SPR
SPR
Drearner
Decarner
Decarner
Dreamer
0
AdaRL.
AdaRL.
AdaRL.
AdaRI.
EllicieniZero
EfficientZero
EllicientZero
EleieenZeer
-20
(Ouns)
(Ouns)
(Ours)
(Ours)
100k 200k
300k
400k
0
100k
200k 300k
400k
100k
200k Episode
300k
400k
100k
200k Episode
300k
400k
Episode
Episode
Figure 21: Training results of various methods in game Pong.
32


Reviews:
[{"review_id": "fqnyQcvVNc", "reviewer": "Reviewer_vB5A", "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.\n\n**Claims**\n\nC1. 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\")\n\nC2. 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.\n\n**Assumptions**\n\nIdentifiability of the model parameters assumes conditional independence between estimated states at different time-steps (Markov condition) and the (stronger?) faithfulness assumption.\n\nIdentifiability 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.", "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?\n- 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)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "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\n- paper is well structured, fairly easy to follow and the description of the method appears to be complete\n- 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.\n- 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.", "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.\n- 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].\n- 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.\n\n[1] https://arxiv.org/abs/2210.02019", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dILmsZ9DQU", "reviewer": "Reviewer_PJxf", "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.", "questions": "Refer to the questions above.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "- This paper is well-written. The structure of CSR is clearly explained, and the authors provide some technical details. \n\n- 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.\n\n- Including ablation studies and comparative evaluations for various expansion strategies (random, deterministic, self-adaptive) provides valuable insights into the effectiveness of different design choices.", "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.\n\n- 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.\n\n- Based on my understanding, $\\theta$ is trained on the data collected from the target domain. Therefore, how does CSR adapt to an unseen task?\n\n- 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?\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TeZnDRhJ8K", "reviewer": "Reviewer_jHns", "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:\n1. distribution shifts\n2. State/Action space expansion. \nThe 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.\n\nThus, the authors evaluate their baselines on different sets of environments such as the famous Atari and Procgen.", "questions": "* how does CSR compare with the other methods in terms of added parameters and training/inference time ?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n\n**Soundness of the method**. The method is overall sound.\n\n**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).\n\n**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.\n\n[1] Farebrother, J., Machado, M. C., & Bowling, M. Generalization and regularization in DQN. arXiv (2018).", "weaknesses": "Several points are unclear: \n* **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]). \n* **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.\n* **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)... \n\nA potentially striking point for the paper would be:\nCould 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 ?\nA discussion on this issue or its use in introduction/future work could demonstrate the potential benefit of causal WM in RL.\n\n------------\n* [1] Luo et al. \"INSIGHT: End-to-End Neuro-Symbolic Visual Reinforcement Learning with Language Explanations.\" ICML (2024).\n* [2] Delfosse et al. \"Interpretable concept bottlenecks to align reinforcement learning agents.\" Advances in Neural Information Processing Systems 37 (2024).\n* [3] Di Langosco et al. \"Goal misgeneralization in deep reinforcement learning.\" International Conference on Machine Learning. PMLR (2022).\n* [4] Delfosse et al. \"HackAtari: Atari Learning Environments for Robust and Continual Reinforcement Learning.\" arXiv (2024).\n* [5] Marton, et al. \"SYMPOL: Symbolic Tree-Based On-Policy Reinforcement Learning.\" arXiv (2024).\n\nI am willing to augment my scores if these are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G9Z681XXLV", "reviewer": "Reviewer_qQnp", "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. \n\nCSR 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. \n\nThree theorems are given which show that the world models and their sub-components are identifiable under certain assumptions. \n\nDetails are then given for how CSR operates:\n* How the world model is trained\n* How distributional shifts are detected\n* How the state-action space is modified when distributional shifts are insufficient to capture changes in observations\n\nLastly, experiments are done showing that, compared to various existing methods, CSR performs favorably following changes in the environment", "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? \n\n(276) - In the expression for $J_{\\rm rec}$, is there meant to be a log around the second probability term? \n\n(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? \n\n(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. \n\n(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? \n\n(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$.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, I liked this paper a lot, and thought it contained many good ideas. I was impressed by the experimental results. \n\nOriginality: \n- 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. \n- Moreover, the algorithm is non-trivial, and involves multiple steps. The authors have solved numerous problems in order to create a pipeline which works.  \n\nQuality: \n- The algorithm proposed by the paper, CSR, works well, and can outperform policy transfer with existing algorithms. \n- The authors perform numerous experiments across different domains to demonstrate the performance of CSR. \n- Solutions to various problems are not ad-hoc, and seem well-motivated. \n- Finally, the authors provide multiple theoretical results which justify their approach.  \n\nClarity: \n- The writing was in general good, and the paper had a nice overall flow.\n\nSignificance: \n- Transfer learning is an increasingly important topic in ML, and the authors do a good job of tackling the problem head-on.", "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. \n\n(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. \n\n(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? \n\n(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?\n\n(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. \n\n(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! \n\nA 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?  \n\nAdditionally, 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.\n\n---------------------------------------------------------------------\n\nAs 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)