Learning Generalizable Environment Models via Discovering Superposed Causal Relationships

Thank you for your evaluation of our work and for your valuable suggestions. We have incorporated your feedback into the revised paper, which is available via the link in the General Response.

W1: The writing is not strong at a point that prevents getting easily the information out of the paper.

A: We appreciate your insightful comments and thank you for highlighting these important issues. In response, we have carefully addressed each point and made revisions in line with your suggestions. Specifically, we have refined several expressions throughout the manuscript to improve clarity and precision. Additionally, we have reorganized Section 3 (Preliminaries) and Section 4.1 (Problem Description) to enhance the logical flow and readability. Furthermore, we have included two new diagrams: a Data Generation Process diagram (Figure 2) and a Network Architecture diagram (Figure 3), to provide clearer visual representations of the process and model structure.

W2: Some words are not clearly defined.

A: Thank you for your comment. We apologize for any lack of clarity and carefully revise our expression . In line 182, the term "decomposition" refers to the fact that the C environments share the same underlying state-action space structure, but each environment may exhibit different causal relationships. We have revised the manuscript to provide a clearer definition of "decomposition" to ensure it is better understood in this context.

W3: The formalization is somewhat unclear at places.

A: We have revised the formalization to improve its clarity and rigor. Specifically, we 1) standardized the presentation to ensure a more formal structure[L 143-150, 162-176, 207-215], and 2) added several diagrams(figure 2,3) to aid in the understanding of the formalization.

W4: The key contributions are not fully clear

A: Thank you for sharing your valuable perspectives. This is the first study to incorporate superposed causal relations in decision-making, where multiple causal relationships are considered for a single state-action pair. Noting that while the significance of this topic has been recognized in the field of causality (e.g., Varambally et al., 2024), studies in reinforcement learning (RL) remain limited. Some existing works in RL, particularly those on local causal graphs (LCGs), are related but differ fundamentally from our approach. For instance, a local causal model aims to learn a mask function $M: S \times A \to \{L_i\}$, where $L_i \subset S \times A$, that maps each state-action pair $(s, a)$ to the adjacency matrix of $\mathcal{G}_L$ (Pitis et al., 2022).This assumes a unique causal graph for each state-action pair, whereas our approach considers multiple causal relationships simultaneously. We have also expanded the discussion of related work and highlighted how our method contrasts with existing approaches, as detailed in [L103-107]. Notably, our method is distinguished by its simplicity and efficiency.

Q1: the specific architecture of the different NN/transformers components.

A: We are grateful for your constructive feedback. Please refer to Figure 3 for the updated diagram. The causal graph section illustrates how trajectory data is utilized to predict the causal graph, while the dynamics model section highlights the components of the network responsible for generating predictions for the next time step based on the causal graph and the data at the current time step. We hope it will contribute to a more thorough and nuanced understanding of our method.

Thank you for your positive evaluation of our work and for your valuable suggestions. We have incorporated your feedback into the revised paper, which is available via the link in the General Response.

W1: Novelty of approach and discussion of relevant prior work

A: We would like to emphasize that, while the significance of this topic has been recognized in the field of causality (e.g., Varambally et al., 2024), studies in reinforcement learning (RL) are relatively limited, with only a few works (specifically addressing the LCG setting). Notably, one of the baselines in our study is based on the LCG, as detailed in FCDL. We have also supplemented a discussion on the related work [L102-107]

Local causal models are indeed related, but they differ from superposed causal relations. A local causal model aims to learn a mask function $M: S \times A \to \{L_i\}$, where $L_i \subset S \times A$, that maps each state-action pair $(s, a)$ to the adjacency matrix of $\mathcal{G}_L$ (Pitis et al., 2022). This approach assumes that the causal graph is unique for the current state-action pair. However, in practice, states may be only partially observed, which may lead to multiple causal relationships for a single state-action pair. As a result, we aim to infer the causal graph by leveraging causal transition information from historical trajectories.

w2，Q3: More precise formal presentation of the framework

A: We appreciate the valuable feedback. We have supplemented the data generating process (figure 2) and the network diagram(figure 3) .

An expression for the masks and how the masks fit into this product are shown in network diagram.

Regarding the point "whether the environment is described by one mask or several ." we have revised the manuscript by presenting the modeling of the single mask environment as a separate subsection within the preliminary section, and in our method section, we now only introduce the dataset collected from multiple environments thus including several masks. We believe this revision makes the explanation clearer.

In response to the query on ”the relation between masks and trajectories” as well as “is \\mathcal{G}\ a random variable”, we have added a diagram depicting the data generation and inference process (see figure2). The masks are trajectory-independent, and we use the trajectory to infer the mask.

Learning a mixture of causal mechanisms refers to the process of learning from a dataset that contains multiple causal relationships[L167-169].

Q2: Misleading claims and requiring more related work

A: Thank you for your insightful comments. We have revised the misleading claim and included the relevant related work as suggested.

Regarding LCG, we have added further details in [L103-107].

For the Bayesian approach, additional information has been included in [L113-114 ].

Regarding block MDPs, we have clarified that a substantial number of the related works we discuss are implemented under the factored MDP framework, as outlined in [L097-102].

Q4: control for the length of trajectory used for structure inference/identification

A: Thank you for your valuable input. This is an interesting question. We will address it in future discussions.

Thank you for your positive evaluation of our work and for your valuable suggestions. We have incorporated your feedback into the revised paper, which is available via the link in the General Response.

W1.1 : distinguishing between causal and predictive aspects of the model.

A: We appreciate the insightful feedback. In response, we have enhanced the network architecture diagram to clearly distinguish the causal and predictive components of the model. Please refer to Figure 3 for the updated diagram. Specifically, we have marked the sections corresponding to causal graph prediction and dynamics model prediction to provide a clearer visual distinction. The causal graph section illustrates how trajectory data is utilized to predict the causal graph, while the dynamics model section highlights the components of the network responsible for generating predictions for the next time step based on the causal graph and the data at the current time step. We aim for this differentiation to contribute to a more thorough and nuanced understanding of our method.

W1.2: Further clarification and rigorous testing are needed to explain how the method identifies and validates causal relationships.

A: Thank you for your valuable feedback. We have included a detailed explanation of the Structural Hamming Distance (SHD), which is used to evaluate the performance of our causal discovery methods. The SHD is a widely recognized metric in the causal inference community, quantifying the dissimilarity between the estimated causal graph and the true underlying causal structure.

W2 : repeat the experiments

A: We appreciate the valuable feedback. The model was tested across multiple seeds, but trained using only a single seed. Due to time constraints, we will improve this approach in future work.

W3 : clearly summarize the contributions of the proposed method compared to existing approaches：

A: We have included additional related work [L103-107][L113-114].

Noting that while the significance of this topic has been recognized in the field of causality (e.g., Varambally et al., 2024), studies in reinforcement learning (RL) are relatively scarce, with only a few notable papers. In the RL domain, although there are some existing multi-graph works under the setting of local causal graph, our work focus on a different situation which is not investigate. A local causal model aims to learn a mask function $M: S \times A \to \{L_i\}$, where $L_i \subset S \times A$, that maps each state-action pair $(s, a)$ to the adjacency matrix of \mathcal ${G}_L$ (Pitis et al., 2022). This approach assumes that the causal graph is unique for the current state-action pair. However, in practice, states may be only partially observed, which may lead to multiple causal relationships for a single state-action pair that previous method can’t handle. We aim to infer the causal graph in this situation by leveraging causal transition information from historical trajectories.