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

Position within the literature. Yes, our current focus is on task adaptation in sequential settings. Nonetheless, CSR holds significant potential for applications in the continual learning (CL) domain. Actually, we have previously investigated several CRL benchmarks, such as Brax [5], CausalWorld [6], Continual World [7], and CORA [8]. However, these benchmarks generally assume a consistent environment space across tasks [9], which contrasts with the setting of SCR. Consequently, we have not conducted tests within the continual learning context. We plan to extend CSR to CRL scenarios and adapt these benchmarks for testing as part of our future work. Thank you very much for pointing this out!

[5] Freeman, C. Daniel, et al. "Brax--a differentiable physics engine for large scale rigid body simulation." arXiv preprint arXiv:2106.13281 (2021).

[6] Ahmed, Ossama, et al. "Causalworld: A robotic manipulation benchmark for causal structure and transfer learning." arXiv preprint arXiv:2010.04296 (2020).

[7] Wołczyk, Maciej, et al. "Continual world: A robotic benchmark for continual reinforcement learning." Advances in Neural Information Processing Systems 34 (2021): 28496-28510.

[8] Powers, Sam, et al. "Cora: Benchmarks, baselines, and metrics as a platform for continual reinforcement learning agents." Conference on Lifelong Learning Agents. PMLR, 2022.

[9] Khetarpal, Khimya, et al. "Towards continual reinforcement learning: A review and perspectives." Journal of Artificial Intelligence Research 75 (2022): 1401-1476.

Adapting CSR to domain-agnostic scenarios. Yes, CSR is currently designed for a domain-aware setting. However, it can easily handle domain-agnostic settings by incorporating domain shift detection techniques, as demonstrated in [10]. During the method design process, we also explored and experimented in this direction but ultimately chose not to include this component to avoid overcomplicating the CSR pipeline with excessive sub-components. A potential future research direction could be the seamless integration of distribution shift detection with domain shift detection. Thank you once again for your thoughtful and valuable insights!

[10] Wang, Zhenyi, et al. "Learning to learn and remember super long multi-domain task sequence." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

Expanding CSR beyond sequential settings. CSR can be naturally extended to settings where the agent leverages a replay buffer not only from the current task but also from previous tasks. However, as it represents an initial exploration of generalizable RL by learning causal world models, we currently focus on a simplified sequential setting where the agent concentrates on domain adaptation in a single target task. We look forward to further enhancing CSR in future work to unlock its potential in more complex scenarios like CRL.

(276). Yes, we have corrected this typo.

(276-283). $J$ are functions of $\psi$, $\beta$, $\alpha$, $\theta$, and $D$ because these variables are jointly optimized through $J$.

(293). Yes, we utilize $J$ to update $\theta$. To be specific, we empirically set the maximum training steps for distribution shift detection as: 1k in Simulation, 2k in CartPole, 100k in Atari, and 250k in CoinRun.

(303). Empirically, we observe that such incorrect detections do not occur in our experiments. This is because introducing a new state typically results in a substantial increase in prediction error (at least twice as much), even with minimal changes in dimensionality. Furthermore, we have also illustrated the evolution of $L_\text{pred}$ across tasks in Fig. 14, which further supports this observation.

We sincerely appreciate your recognition of our work and efforts in providing valuable feedback, which significantly help improve its quality. Following your suggestions, we have enriched the appendix and Algorithm 1 with sufficient implementation details to facilitate reproduction by interested researchers. Furthermore, we commit to release the code upon publication, accompanied by detailed instructions for installation. Please see below for our responses to your questions.

(315). We have corrected it as you suggested.

(319). The original use of "finetune" was inaccurate. What we intended to convey is: "Accordingly, we focus on learning the newly incorporated components with only a few samples, as the previous model has already captured the existing relationships between variables." We have revised the sentence to clarify this.

On the other hand, we believe that CSR can effectively address the issue of catastrophic forgetting by leveraging $\theta$. In other words, as long as the task-specific $\theta_{M_1}$ (e.g., $\theta_{M_1} = g = 9.8$ in Cartpole) has been learned and stored, even after encountering $M_2, \ldots, M_9$, where $\theta_{M_i}$ may take other values, policy adaptation can still be achieved by setting $\theta_M = 9.8$. This is because CSR updates only the components that have changed, without overwriting the original model.

(361). In causal representation learning, $D$ represents binary masks that indicate the structural relationships among variables, and thus cannot take negative values. Upon implementation, we adopt the Gumbel-Softmax [1] and Sigmoid methods to approximate the binary masks and optimize them with $J$, which is a commonly used approach. Additionally, we do not threshold $D$, and the results demonstrate that the obtained values are effectively near 0 or 1.

[1] Ng, Ignavier, et al. "Masked gradient-based causal structure learning." Proceedings of the 2022 SIAM International Conference on Data Mining (SDM). Society for Industrial and Applied Mathematics, 2022.

(371), (377) and (276-283). As mentioned in line 271, CSR is built upon Dreamer, which iteratively performs exploration, model estimation, and policy learning (see Algorithm 1 of [2]). To maintain conciseness, some training details were not repeated in the main text, but we have now included them in Appendix D and the CSR algorithm. Below are our responses to your concerns regarding implementation details:

• (371) We do not adopt Thompson sampling. Instead, for model estimation, latent state variables are inferred from pixel observations using an RSSM [3]. To balance the exploration-exploitation trade-off, we apply the $\epsilon$-greedy strategy. For policy learning, the trajectories are purely obtained through imagination that starts at the true model states $s_t$ of observation sequences and predicts both the actions and state values using the learned world model.

• (377) The policy is learned using an actor-critic approach instead of DDPG. We use reinforce gradients to optimize the policy.

• (276-283) The world model is trained by randomly sampling from the memory replay buffer that contains sequences of observation $o_t$, actions $a_t$ and rewards $r_t$. We use the representation model $q_\alpha$ to encode categorical latent state $s_t$ from the sampled sequences, and take straight-through gradients through the sampling step (see DreamerV2 [4]).

[2] Hafner, Danijar, et al. "Dream to control: Learning behaviors by latent imagination." arxiv preprint arxiv:1912.01603 (2019).

[3] Hafner, Danijar, et al. "Learning latent dynamics for planning from pixels." International conference on machine learning. PMLR, 2019.

[4] Hafner, Danijar, et al. "Mastering Atari with discrete world models." arXiv preprint arXiv:2010.02193, 2020.

Dear Reviewer qQnp,

Thank you for your thoughtful feedback and for highlighting the importance of detailed implementation in our paper, which has greatly helped us improve our work. Following your suggestions, we have included additional implementation details in Appendix D and also in Algorithm 1 to facilitate reproducibility and to inspire further advancements in this area.

As the discussion phase is coming to an end, we kindly ask whether our response has resolved your main concerns. If there are any remaining issues, please let us know, and we are willing to address them.

If your concerns have been addressed, we would deeply appreciate it if you could consider updating your recommendation accordingly. We value your time and insights and look forward to your response.

Thank you again for your attention to our work.

Sincerely,

Submission620 Authors

We appreciate your thoughtful comments and time devoted to reviewing this paper. We have added the citation as suggested and hope the following response properly addresses your concerns.

Implicit world models. Thank you for raising this point. But we would like to clarify that CSR indeed encodes a world model that is designed to be interpretable. By discovering and utilizing causal graphs and their corresponding causal matrices $D$, CSR can effectively identify the underlying causal relationships in each task, thereby assisting in interpretable decision-making, as in [1, 2, 5]. Besides, in contrast to neuro-symbolic approaches, one of CSR’s key advantages is that it does not require predefined symbols, enabling it to be applied more generally across various domains.

For example, consider Fig. 5(a), where the causal graph represents a simple MDP scenario. In Cartpole, $s_1$ could represent the cart position, and $\theta^o$ denotes the background color. Based on the causal graph, we observe that $D^{s_1 \to r} = 1$ and $D^{\theta^o \to r} = 0$, revealing that "policy learning needs to consider the cart position as it influences the reward, but does not need to account for the background color, as it is an irrelevant variable." The agent uses this rule to guide its decisions, which is precisely the approach employed in Sections 3.3 and 3.4.

Moreover, CSR can also explain its adaptation mechanism (i.e., why a distribution shift detection or space expansion is chosen). Referring to Fig. 5(b), after re-estimating the values of $\theta^{r'} = \{g\}$ in Task 2, CSR clarifies that the reason for not selecting space expansion is: "Only the gravity ($g$) have changed in the environment, with no new variables introduced, thus no expansion is necessary."

We further illustrate in Fig. 10 the estimated causal matrices from our experiments, which reflect which variables are significant and explain why they are chosen for policy learning. It is important to note that unlike the MDP example above, all of our experiments involve POMDP scenarios, where latent state variables are learned from pixel images. Thus, directly determining the representation of each variable is not possible. However, a more intuitive approach is to observe the activation or deactivation of causal variables and analyze the corresponding changes in the reconstructed image, specifically noting which parts of the image appear or disappear. Fig. 13 provides an example from the Crazy Climber game, showing the meaning of expanded variables and explaining why the expansion is necessary. A detailed analysis can be found in Appendix D.3.4.

The difference between the two types of change is not clear. Thank you for your feedback. We have added a simple MDP example with a graphical illustration in Appendix C to clarify the difference between the two types of change.

What are the answered scientific questions? Thank you for the valuable suggestion. We have reorganized the experimental evaluations around the following scientific questions:

• Q1: Can CSR effectively detect and adapt to 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?

Could a causal world model help mitigate misgeneralisation/misalignment problems? Yes! CSR can indeed help mitigate misgeneralization problems. The misalignment you mention can be traced back to the distinction between correlation and causality. For instance, in the LazyEnemy variation of Pong, the causal graph is: Y (enemy position) $\leftarrow$ X (ball position) $\rightarrow$ Z (agent position). It is clear that there is no causal edge between Y and Z, as the agent position is not directly caused by the enemy’s position, but by the ball position, which aligns with the goal. In this structure, X is a common cause of both Y and Z, creating a correlation that leads to the goal misalignment problem [6, 7]. CSR can address this by learning a causal world model that identifies causal edges rather than correlations. We have discussed this further in Appendix E. Thank you very much for this insightful point!

[6] Pearl, J. , Mackenzie, D. . (2018). The book of why : the new science of cause and effect. Science, 361(6405), 855.2-855.

[7] Spirtes, P.; Glymour, C.; and Scheines, R. 2001. Causation, Prediction, and Search. Cambridge, MA: MIT Press, 2nd edition.

Thank you for your thorough review of our rebuttal and for your appreciation of our efforts! Below are the responses to the remaining concerns you have raised.

The captions of the figures and table can be improved to again improve clarity. Thank you for your valuable suggestions. We have revised the captions as you recommended and highlighted the changes in orange for easier identification.

An implicit world model. Yes, your understanding is correct. The CSR agent cannot directly explain the reason for each action choice because it is not feasible to determine precisely what each latent state, inferred from pixel images, represents. In light of your comment, moving forward, we will explore how to enable the model to autonomously identify the meanings of each state, thereby revealing the reasoning behind each decision. Another potential benefit of this approach is that it allows us to determine the minimal sufficient state representation for each task, rather than empirically choosing a suitable value for the state dimension and then pruning them as we currently do, which could help reduce the network size of the model. This is indeed an invaluable point; we thank you again for highlighting these considerations.

The learned causal WM on Pong. Yes. We have illustrated the state space in Pong and the evolution of the graphical structure $D_k^{z \to h}$ across tasks in Fig. 10(c), where $z$ represents the latent state variable and $h$ is the deterministic state variable in our implementation based on a RSSM [8] world model.

[8] Hafner, Danijar, et al. "Learning latent dynamics for planning from pixels." International conference on machine learning. PMLR, 2019.

Test CSR in the LazyEnemy version. It may not be as straightforward as it seems. After reviewing the HackAtari code, we identified a number of differences between HackAtari and our implementation in handling the Gym Atari environment, particularly in the step(), reset(), and certain gym.Wrapper classes. These differences make it challenging to directly test CSR in the LazyEnemy version. We are actively exploring solutions and will do our best to complete the testing within the discussion period, though we cannot fully guarantee this timeline. Updates will be shared as soon as the results are available. Thank you for your understanding and for this valuable suggestion.

Please let us know if you have any further concerns. Thank you once again for your time and consideration!

We greatly appreciate your insightful comments and time devoted to our work. We attempt to address all the concerns as follows.

Visualization of correlation between a task and the adaptation mechanism. Thank you for your valuable suggestion. We have visualized the evolution of $L_\text{pred}$ and the ratio between the two modes of CSR across all four experimental scenarios in Fig. 14. A detailed analysis of these correlations is provided in Appendix D.4.

Empirical results on Atari. The selection of Atari games in our experiments is guided by a key principle: the tasks must exhibit varying modes and difficulty levels. However, not all games meet this criterion. For example, among the Atari-10 games [1] you referenced, only "Name the Game" and "Bowling" satisfy this requirement. To address this, we carefully selected five games, each with at least two distinct modes and two difficulty levels. Additionally, since our focus is on sequential tasks, we aim to limit training time per task to control costs. This is why we chose games from the Atari 100K set (requiring 1 GPU day per task) instead of Atari 200M (requiring 16 days).

Regarding the comparison with Dreamer, the superior adaptation ability of CSR is particularly evident in its more efficient sample usage and faster convergence to higher scores. As shown in Fig. 16, CSR significantly outperforms other baselines at 100K steps. While this performance gap narrows with more training steps, this trend is entirely expected: with enough training time, any world model (even one trained from scratch) can achieve good scores. This observation actually highlights CSR's higher transfer efficiency, which is one of our key contributions. Furthermore, we have also conducted experiments on Atari-10 games, including "Amidar" (1 mode, 2 difficulties), "Battle Zone" (3 modes, 1 difficulty), "Frostbite" (2 modes, 1 difficulty), and "Qbert" (1 mode, 2 difficulties). The results are presented below.

Distinction between the two scenarios. We have included more detailed examples in Appendix C, along with graphical illustrations, to better clarify the distinction between the two scenarios. Thank you once again for your thoughtful comments!

Necessity of the expansion mechanism. In our setting, simply learning $D$ and $\theta$ is not always sufficient to ensure the identification of underlying changes. This is because tasks arrive sequentially, and we cannot predict in advance how many variables each task will involve. In other words, it is not feasible to predefine the number of variables, as this approach often turns out to be inadequate.

To illustrate this more clearly, consider an extreme example in the CartPole task. Suppose we set the number of state variables to 4, where these 4 learned variables correspond exactly to the real environment variables: $s_1 = x$ (cart position), $s_2 = \dot{x}$ (cart velocity), $s_3 = \psi$ (pole angle), and $s_4 = \dot{\psi}$ (pole angle velocity). Additionally, the parameters $\theta_1 = m_c$ (cart mass) and $\theta_2 = g$ (gravity) are learned. In this case, for Tasks 1 and 2, the model can easily capture the changes by learning $D$ and $\theta$. However, when a new variable—such as friction $\mu_c$—appears in Task 3, we must expand the state space from $\mathcal{S}_1 \times \mathcal{S}_2 \times \mathcal{S}_3 \times \mathcal{S}_4$ to $\mathcal{S}_1 \times \mathcal{S}_2 \times \mathcal{S}_3 \times \mathcal{S}_4 \times \mathcal{S}_5$, where $s_5 = \mu_c$, in order to adapt to this change. This highlights the necessity of the expansion mechanism in our method.

Parameter consistency. Yes, we keep the parameters consistent across all baselines for a fair comparison.

Thanks a lot for your thoughtful comments. Please see below for our responses to your concerns.

CSR's potential in noisy and complex environments. Causal representation learning methods, such as CSR, are capable of performing well under complex and noisy conditions for the following reasons:

• Causal methods, when appropriately designed, hold the potential to outperform non-causal approaches in managing noisy environments [1].
  For instance, CIRL [2] employs a causal intervention module, causal factorization module, and adversarial mask module to effectively transform noisy and dependent representations into clean and independent ones, resulting in superior generalization ability. Moreover, learning a causal model uncovers the underlying mechanisms of the data generation process, potentially addressing challenges that non-causal methods may struggle with. For instance, in the case of misalignment [3]—a scenario where an agent’s performance appears aligned with the target goal but is actually influenced by a side goal due to a hidden common cause—causal modeling offers a way to mitigate this issue by identifying causal relationships rather than correlations.

• Theoretical and empirical evidence supports CSR's robustness.
  Theorems 1-3 in our paper provide identifiability guarantees for causal world models under conditions where latent variables or processes have i.i.d. noise terms. Additionally, our experimental results on POMDP tasks, including Atari and CoinRun games, highlight CSR's ability to handle complex environments more effectively than non-causal approaches.

• Advances in identifiability extend CSR's applicability to challenging scenarios.
  Recent developments in causal identifiability have tackled issues such as nonstationary processes [4], time-delayed dynamics [5], heteroscedastic noises [6], and even arbitrary noise distributions [7]. These advancements reinforce the potential of CSR to adapt to diverse and complex environments.

We hope the explanations provided clarify how CSR can leverage causal structures to address such challenges effectively. Thank you again for highlighting this important aspect.

[1] Sanyal et al., 2024. Accuracy on the wrong line: On the pitfalls of noisy data for out-of-distribution generalisation. arXiv preprint arXiv:2406.19049.

[2] Lv et al., 2022. Causality inspired representation learning for domain generalization. CVPR, 8046-8056.

[3] Di Langosco et al., 2022. Goal misgeneralization in deep reinforcement learning. ICML.

[4] Yao et al., 2021. Learning temporally causal latent processes from general temporal data. arXiv preprint arXiv:2110.05428.

[5] Yao et al., 2022. Temporally disentangled representation learning. NeurIPS, 35, 26492-26503.

[6] Lin et al., 2024. A skewness-based criterion for addressing heteroscedastic noise in causal discovery. arXiv preprint arXiv:2410.06407.

[7] Montagna et al., 2023. Causal discovery with score matching on additive models with arbitrary noise. CLeaR, 726-751.

Novelty differs from AdaRL. CSR differs significantly from AdaRL in the following ways:

• Much More General Application Scenarios:
  Unlike AdaRL, which focuses solely on simple distribution shifts for policy adaptation (as stated in line 418), CSR extends to broader scenarios where the spaces may also change across tasks. Details about the distinctions between the two types of scenarios are further given in the Introduction and Appendix C.

• A Novel Three-step Self-adaptation Strategy:
  CSR introduces a three-step strategy specifically designed to address the dynamic nature of the environments it is applied to: (1) Distribution Shifts Detection and Characterization, (2) State/Action Space Expansions, and (3) Causal Graph Pruning. This strategy enables CSR to effectively capture and adapt to environmental changes, a capability not addressed by AdaRL.

• Theoretical Guarantees:
  Prior works, such as AdaRL, lack theoretical support for the identifiability of world models in Eq. (2), whereas CSR not only addresses this gap but also establishes the identifiability of $\theta$ in linear cases, supported by empirical results in nonlinear scenarios. Furthermore, CSR provides rigorous proofs for the identifiability of newly introduced variables in space expansion settings.

How does CSR adapt to an unseen task? CSR can be viewed as a method for few-shot policy adaptation. As mentioned in lines 140–141, when a new target task is introduced, CSR collects data from the target domain, specifically $\{\langle o_t, a_t, r_t \rangle \}$, and performs adaptation accordingly. Algorithm 1 in lines 1098-1133 presents the pseudocode for CSR.

Dear Reviewer PJxf,

Thank you once again for your time and thoughtful review of our paper. We have addressed your comments, providing clarification on the potential of causal models to handle noisy and complex environments, as well as on the novelty of CSR compared to AdaRL.

Could you please check whether they properly addressed your concern? We would greatly appreciate your feedback. Should there be any further issues, we would be happy to address them. Thank you very much!

Sincerely,

Submission620 Authors