Paper: Quantized Local Independence Discovery for Fine-Grained Causal Dynamics Learning in Reinforcement Learning (ICLR-2024)

Reviewer comment (Experiments):
The validation environment for the experiment is somewhat simple, and the assumption of full observables seems to have significant limitations.

Author response:
- Discovering fine-grained causal relationships (i.e., local causality) from the high-dimensional image observation is certainly an important future direction. However, it has not been explored even in a fully-observable setting. We would like to note that our work is the first step towards discovering and utilizing fine-grained causal relationships for enhancing model-based RL. We indeed agree with the reviewer that it is a promising future direction to further investigate this in more complex environments and in a partially-observable setting.

Reviewer comment (Experiments):
In Table 3, it seems that the larger K is, the better the performance is. What happens if K is increased to equal the number of samples?

Author response:
- In our experiments, the total number of samples we used is 15k and 20k for Chemical and Magnetic, respectively. We speculate that as $K \rightarrow \infty$, it would suffer from overfitting similar to NCD, a sample-specific inference method.

Reviewer comment (Theory):
Intuitively, an event should be a sequence of continuous states and actions. I would like to ask whether Eq. (7) is more inclined to aggregate neighboring states?

Author response:
- We speculate that Eq. 7 would be inclined to aggregate neighboring states because we use neural networks for an embedding function in practice.

Reviewer comment (Writing):
How can we understand the concept of 'event'? Is it akin to a specific state value or an additional variable, similar to the domain ID in domain generalization tasks? Could you provide some examples in the context of reinforcement learning scenarios?

Author response:
> **[Q1]** How can we understand the concept of 'event'? Is it akin to a specific state value or an additional variable, similar to the domain ID in domain generalization tasks? Could you provide some examples in the context of reinforcement learning scenarios?
> 
- Consider the mobile home robot. In general, it could navigate the rooms and interact with various objects. However, when the door is closed, it cannot interact with the objects in the other rooms; to do so, it would have to first opens the door. In the context of RL, the state variables (entities) consist of robot, door, and other objects. The specific event we are concerned with is the context of the door opened.

Reviewer comment (Theory):
Definition 1 may not fully describe an event-conditioned system. What if we consider a situation where changing the event alters the entire causal relationship? Definition 1 only discusses cases where the edges under certain conditions should be a subset of the original edge set.

Author response:
> **[Q2]** Definition 1 may not fully describe an event-conditioned system. What if we consider a situation where changing the event alters the entire causal relationship? Definition 1 only discusses cases where the edges under certain conditions should be a subset of the original edge set.
> 
- As the reviewer pointed out, the entire causal relationship could change in general (e.g., changes through time [1]). In our work, we consider the setting where the underlying causal relationships remain stationary. In this setting, LCGs are always subgraphs of the (global) causal graph because if a causal relationship (i.e., edge) is absent in the whole domain, then it is also absent in any arbitrary event.

Reviewer comment (Theory):
Can the score (equation 4) assist in identifying the causal graph? In Huang et al.'s paper, the score is a direct measure of conditional independence. Does equation 4 also measure conditional independence? In Brouillard et al.'s work, I found that the intervention data are required to support identifiability. Previous works have used score-based methods to discover causal graphs only under certain assumptions (e.g., linear and non-Gaussian Structural Causal Models). However, this paper lacks rigorous assumptions that substantiate identifiability.

Author response:
> **[Q3]** Can the score (equation 4) assist in identifying the causal graph? In Huang et al.'s paper, the score is a direct measure of conditional independence. Does equation 4 also measure conditional independence? In Brouillard et al.'s work, I found that the intervention data are required to support identifiability. Previous works have used score-based methods to discover causal graphs only under certain assumptions (e.g., linear and non-Gaussian Structural Causal Models). However, this paper lacks rigorous assumptions that substantiate identifiability.
> 
- Sorry for inconvenience. Our score function is the regularized maximum likelihood score, similar to [2]. In contrast to [2], which used interventional data to identify the graph up to $\mathcal{I}$-MEC, we use observational data to identify the graph up to Markov equivalence class (MEC). In our setting of factored MDP, each MEC contains a unique graph because temporal precedence determines the orientation of the edges. We also assume that our model using a neural network has sufficient capacity to model the ground truth density (Assumption 3 in Appendix B in the revised manuscript). Throughout the Appendix B in our revised manuscript, we have included rigorous assumptions and proofs that substantiate identifiability.

Reviewer comment (Theory):
In equation 5, both h and e_j are learned. How can we ensure a nontrivial result? Specifically, if both h and e_j are learned as the static value 0, how can we avoid trivial outcomes?

Author response:
> **[Q4]** In equation 5, both h and e_j are learned. How can we ensure a nontrivial result? Specifically, if both h and e_j are learned as the static value 0, how can we avoid trivial outcomes?
> 
- In theory, $\mathcal{L}_{\texttt{pred}}$ in eq. 6, i.e., regularized log likelihood loss, also updates $h$, and thus affects the learning of latent codes $e_j$. Thus this prevents the degenerate solution, e.g., constant static $h$, because such trivial solution that learns a single causal graph over the entire domain has higher regularization loss.
- In practice, the training of vector quantization still often becomes unstable and yields a degenerate solution even though it achieves higher loss: a phenomenon called *codebook collapsing* in VQ-VAE literature. A popular technique to prevent such collapsing and stabilize the training is to use exponential moving average (EMA) to update the codebook, which is also adopted in our implementation. We discuss in Appendix C.3.3 recent relevant techniques and tricks to further stabilize the training.

Reviewer comment (Theory):
Although there are some positive results in Table 3 shows that considering the subgraph can help to enhance the causal discovery score. I’m still curious about why breaking the full graph into subgraphs by removing unnecessary edges can help enhance the robustness. If the edges are unnecessary, the causal function on the full graph will not consider the causal effect of it.

Author response:
> **[Q5]** Although there are some positive results in Table 3 shows that considering the subgraph can help to enhance the causal discovery score. I’m still curious about why breaking the full graph into subgraphs by removing unnecessary edges can help enhance the robustness. If the edges are unnecessary, the causal function on the full graph will not consider the causal effect of it.
> 
- Consider the above example of mobile home robot. If the door is closed, the objects on the other side of the door become irrelevant, thus the robot should keep function well even if there are unseen objects on the other side of the door. In other words, these edges are *locally unnecessary* only under the specific events. They are relevant in general, thus (learned) causal function using the full graph would be affected by these edges. Consequently, this harms their prediction accuracy and RL performance as shown in Table 1-3.

Reviewer comment (Novelty):
The quantization technique in eq 5 is not novel. VQVAE and related papers should be cited there.

Author response:
- Indeed, we adapt the quantization technique of VQ-VAE, where we cited in eq. 7, and provide discussions regarding the recent techniques related to VQ-VAE in Appendix C.3.3. Following the reviewer’s suggestion, we have also added citation in eq. 5 in the revised manuscript. As a side note, our utilization of vector quantization is quite different from most of the works on learning discrete latent codebook with vector quantization whose main focus is the reconstruction of the observation.

Reviewer comment (Theory):
Why an event-specific graph has to be constrained to be a subgraph of the underlying G?

Author response:
- In certain events, causal relationships may no longer exist. LCGs are always subgraphs of the (global) causal graph because if a causal relationship (i.e., edge) is absent in the whole domain, then it is also absent in any arbitrary event.

Reviewer comment (Theory):
Why vector quantization can correctly identify the event?

Author response:
- We use vector quantization to partition the state-action space. By imposing the dynamics model to use sparse subgraphs for each corresponding partition, vector quantization tends to cluster samples that causal relationships are explained by the same LCG.

Reviewer comment (Theory):
Does the identifiability in Theorem 1 rely on the specific choice of K?

Author response:
- Thm. 1 holds for any choice of $K\geq 2$. If $K=1$, it is indeed impossible to identify particular event $\mathcal{E}\subsetneq \mathcal{X}$ which induces fine-grained causal relationships.

Reviewer comment (Theory):
I am not sure about the soundness of the proposed method. The paper claims that they are jointly optimizing both the event decomposition/partitioning scheme and the causal models, but I doubt if the learning of the decomposition scheme is indeed effectively signaled by the impact of the decomposition on the quality of the resulted causal models. Moreover, I doubt if the learning objective designed for optimizing decomposition (Eq.7) may lead to degenerate results. See my question 1 and 2 below for the detailed concerns. Since joint optimization is a main selling point of this paper (without it the theory part becomes largely irrelevant to the experiment part, for example), this is a serious issue that must be clarified.

Author response:
> **[Q1-2]** Backpropagation in Eq. 5-7. [..] Exactly what is being updated by $\mathcal{L}\_{\texttt{pred}}$ and $\mathcal{L}_{\texttt{quant}}$? What are exactly computed for the differentiation?
> 
- Gradient of $\mathcal{L}_{\texttt{pred}}$ (eq. 6) **does** update the latent embedding $h$, and consequently, affects the codebook $C$.
    - First, recall that $A\sim g_{\texttt{dec}}(e)$, backpropagation from $A$ in $\mathcal{L}\_{\texttt{pred}}$ updates the quantization decoder $g_{\texttt{dec}}$through $e$.
    - Second, we copy gradients from $e$ (= input of $g_{\texttt{dec}}$) to $h$ (= output of $g_{\texttt{enc}}$) during the backward path in eq. 5, following a popular trick used in VQ-VAE [1]. By doing so, $\mathcal{L}\_{\texttt{pred}}$ updates the quantization encoder $g_{\texttt{enc}}$ and $h$.
    - Finally, this also affects the learning of the codebook $C$ since $C$ is updated with $\mathcal{L}_{\texttt{quant}}$ (eq. 7) which involves in $h$.
- To sum up, $\mathcal{L}\_{\texttt{pred}}$ updates the encoder $g_{\texttt{enc}}(s,a)$, decoder $g_{\texttt{dec}}(e)$, and the dynamics model $\hat{p}$. $\mathcal{L}\_{\texttt{quant}}$ updates $g_{\texttt{enc}}$ and the codebook $C$.
- Since the learning of $C$ and $h$ are also based on $\mathcal{L}_{\texttt{pred}}$, this prevents yielding a degenerate solution such as constant vector $h$.
- We have included our responses and further details on our method and its implementations in our revised manuscript (Appendix C.2.1). We promise to publicly open-source the experimentation code after the paper being published.

### Response to “regarding the rationale of the proposed method (Q 1-2)”

- Following the reviewer’s suggestion, we will include background on VQ-VAE and training details in the Preliminary section of the final version.

> The above reflection suggests that $\mathcal{L}_{\texttt{pred}}$ does not have well-defined gradient with respect to $h$ at all.
> 

$$
e = e_z, \quad \text{where} \quad  z = \text{argmin}\_{j \in \left[K\right]} \\|h - e_j\\|_2.  \quad \text{(Eq.5)}
$$

- While there is no real gradient defined for $h=g_{\texttt{enc}}(s, a)$ in Eq. 5, we approximate the gradient using the trick of VQ-VAE, which copies the gradient $\nabla_e \mathcal{L}\_{\texttt{pred}}$ to update the encoder $g_{\texttt{enc}}$, similar to straight-through estimator.
- The rationale behind this trick is that the gradients contain useful information for how the encoder $g_{\texttt{enc}}$ has to change its output $h=g_{\texttt{enc}}(s, a)$ to lower the prediction loss $\mathcal{L}_{\texttt{pred}}$.
    - This gradient $\nabla_e \mathcal{L}_{\texttt{pred}}$ will update the encoder to change its output $h$, which could alter the quantization (i.e., assignment of the cluster) in the next forward pass.
    - A larger prediction loss (which implies that this sample $(s, a)$ is assigned to the wrong cluster) induces a bigger change on $h$ and consequently it would be more likely to cause a re-assignment of the cluster.

Reviewer comment (Reproducibility):
I have concern on the reproducibility of the experiment part. No pseudo-code is provided for the proposed method, and the method description in Section 3.3, 4.1.2, and C.2, missed some important pieces. For example, how is the data sampled (what's the exploration policy being used)? Exactly what is being updated by L_pred and L_quant (the gradient is respective to what)? What are exactly computed for the differentiation?

Author response:
> **[Q1-2]** Backpropagation in Eq. 5-7. [..] Exactly what is being updated by $\mathcal{L}\_{\texttt{pred}}$ and $\mathcal{L}_{\texttt{quant}}$? What are exactly computed for the differentiation?
> 
- Gradient of $\mathcal{L}_{\texttt{pred}}$ (eq. 6) **does** update the latent embedding $h$, and consequently, affects the codebook $C$.
    - First, recall that $A\sim g_{\texttt{dec}}(e)$, backpropagation from $A$ in $\mathcal{L}\_{\texttt{pred}}$ updates the quantization decoder $g_{\texttt{dec}}$through $e$.
    - Second, we copy gradients from $e$ (= input of $g_{\texttt{dec}}$) to $h$ (= output of $g_{\texttt{enc}}$) during the backward path in eq. 5, following a popular trick used in VQ-VAE [1]. By doing so, $\mathcal{L}\_{\texttt{pred}}$ updates the quantization encoder $g_{\texttt{enc}}$ and $h$.
    - Finally, this also affects the learning of the codebook $C$ since $C$ is updated with $\mathcal{L}_{\texttt{quant}}$ (eq. 7) which involves in $h$.
- To sum up, $\mathcal{L}\_{\texttt{pred}}$ updates the encoder $g_{\texttt{enc}}(s,a)$, decoder $g_{\texttt{dec}}(e)$, and the dynamics model $\hat{p}$. $\mathcal{L}\_{\texttt{quant}}$ updates $g_{\texttt{enc}}$ and the codebook $C$.
- Since the learning of $C$ and $h$ are also based on $\mathcal{L}_{\texttt{pred}}$, this prevents yielding a degenerate solution such as constant vector $h$.
- We have included our responses and further details on our method and its implementations in our revised manuscript (Appendix C.2.1). We promise to publicly open-source the experimentation code after the paper being published.

> How is the data sampled?
> 
- We use a random policy for the initial data collection.

Reviewer comment (Reproducibility):
The experimentation code is not given, and the description of the two MDP environments are not detailed enough for third-parties to reproduce the experiment. For example, how exactly do the nodes in the Chemical environment affect each other? How exactly is the noise injected at testing time? In the robot-arm environment, how exactly is the "unseen locations" of the box determined at testing time? In Section C.1.3, it is said that at training time the positions of the box are "randomly sampled within the range of the table", does this mean every location on the table has a positive probability density to be selected at training time (If so, the location of the box at testing time is not really out-of-distribution, even though it may be unseen)?

Author response:
> Detailed description of two environments.
> 
- In Chemical, the underlying color-change mechanism of each node given its parents follows a predefined conditional probability table implemented with a randomly initialized neural network. To create OOD states, we use a noise sampled from $\mathcal{N}(0, \sigma^2)$, similar to [2]. In Chemical, the noise is multiplied to the one-hot encoding representing color during the test.
- In Magnetic, the box position is sampled from $\mathcal{N}(0, \sigma^2)$ during the test. We note that the box can be located outside of the table, which never happens during training. We use $\sigma=100$ for both environments.

Reviewer comment (Theory):
I have concerns on the theory part too. On one hand, the main theorem seems to only apply to a rather special case (where the causal model can only have two irreducible phases). On the other hand, the lemmas (Prop. 1-3) used for proving the main theorems seem to have flawed statements and proofs. Some important assumptions are not explicitly stated. See my question 3~8 for the details in this regard.

Author response:
> **[Q3]** Requirement on the value of $\lambda$.
> 
- Prop. 2 holds for small enough $\lambda$, i.e., $0<\forall\lambda\leq \eta(\\{\mathcal{E}\_z\\})$, where we have included detailed derivations in the proof of Prop. 2 in our revised manuscript. We also note that $\eta(\\{\mathcal{E}\_z\\})>0$ is a value specific to a decomposition. Thus, for Prop. 3, we introduce Assumption 5 (on page 18) which states that there exist a lambda that works for all decompositions, i.e., $\inf_{\\{\mathcal{E}\_z\\}\in \mathcal{T}} \eta(\\{\mathcal{E}\_z\\}) > 0$ where $\mathcal{T}$ is a set of all decompositions. We also discuss the violation of this assumption: in that case, the arguments hold for decompositions $\mathcal{T}\_\lambda=\\{\\{\mathcal{E}\_z\\}\mid \eta(\\{\mathcal{E}\_z\\})\geq \lambda \\}$, where $\mathcal{T}\_\lambda \rightarrow \mathcal{T}$ as $\lambda \rightarrow 0$ (see page 18 in our revised manuscript).

> **[Q4]** The last equality in the proof of Prop. 3.
> 
- We have also included more detailed derivations for the proof of Prop. 3 in our revised manuscript. In short, for the arbitrary decomposition $\\{\mathcal{E}\_z\\}$ and its corresponding *true* LCGs $\\{\mathcal{G}\_z\\}$, $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K)=\mathbb{E}\_{p(s, a, s')} \log p(s'\mid s, a) - \lambda \sum_z p(\mathcal{E}_z) \cdot \lvert \mathcal{G}_z \rvert$, and thus the first term is canceled out in the last equality in the proof of Prop. 3.

> **[Q5]** Factored MDP setting.
> 
- The probability factorization $p(s'\mid s, a) = \prod_j p(s'_j\mid s, a)$ implies that each $S'_j$ conditioned on $S, A$ are independent to each other. This is the setting of factored MDP where our paper and this line of related works commonly considers. Our propositions and theorem are developed under this setting since we assume a bipartite causal graph. As the reviewer pointed out, a causal graph is not necessarily bipartite without this assumption, and we would no longer have unique identifiability anymore. In our experiments, Magnetic satisfies the factored MDP assumption but not Chemical, which contains instantaneous effects.

> **[Q6]** Please explicitly and formally state all the external propositions
> 
- Sorry for the inconvenience. We have made the external proposition explicit and improved the clarity by adding detailed derivations throughout Appendix B in our revised manuscript.

> **[Q7]** how is $Pa(j; \mathcal{E})$ defined? Does the global causal model entail the local models through the induced local independence relationships? In that case a proper *definition* of local causal model would be entirely conditioned on the E-faithfulness assumption, am I right?
> 
- We have added formal definition of $Pa(j; \mathcal{E})$ in Def. 5 on page 15 in our revised manuscript. In words, $Pa(j; \mathcal{E})$ is a minimal subset of $Pa(j)$ in which the local independence on $\mathcal{E}$ holds. And Yes, the global causal model entails the local models through the induced local independence relationships. To clarify, LCG always exists because $Pa(j; \mathcal{E})$ always exists regardless of the $\mathcal{E}$-faithfulness assumption. $\mathcal{E}$-faithfulness implies that $Pa(j; \mathcal{E})$ is unique, and thus implies the uniqueness of LCG $\mathcal{G}_\mathcal{E}$.

> **[Q8]** Assumptions.
> 
- Sorry for the inconvenience. Following the reviewer’s suggestion, we have made our assumptions explicit, and they are now given in the statement of the involved propositions. Due to the space constraints, precise statement of the assumptions are currently presented before the proof, in the Appendix B in our manuscript.

Reviewer comment (Novelty):
The paper may need to better contrast to prior work, especially to (Hwang et al., 2023). It seems (Hwang et al., 2023) is not limited to sample-specific decomposition, but also discusses event-level or event-set-level decomposition. Proposition 1 in the paper under review is essentially a rephrasing of Proposition 4 in (Hwang et al., 2023), for example.

Author response:
> The paper may need to better contrast to prior work [..] It seems (Hwang et al., 2023) is not limited to sample-specific decomposition, but also discusses event-level or event-set-level decomposition.
> 
- To clarify, Hwang et al., (2023) indeed discussed event-level decomposition, which our work draws inspiration from. However, their proposed method does not explicitly discovers such decomposition, but only infers LCG for each sample. In contrast, our method explicitly discovers $\\{\mathcal{G}_z, \mathcal{E}_z\\}$, decomposition and corresponding LCGs. We have included related discussions in Appendix C.3.3 in our revised manuscript to better contrast to prior work.

### Response to ''regarding contrast to (Hwang et al., 2023) (Q 9)''

- To clarify, such decomposition is indeed first studied in the prior work (Hwang et al., 2023), but their method (i) lacks theoretical guarantees to discover true LCGs, and (ii) does not **explicitly** identify events.
    - The method introduced in (Hwang et al., 2023) does not *explicitly* identify events since it only learns the function that maps each sample $(s, a)$ to LCG (i.e., sample-specific inference). Here, the events are *implicitly* identified: the samples having the same inferred LCG belong to the same event. Due to its limitation, important attributes of the decomposition such as the number of discovered events (i.e., size of the decomposition) are unclear, and it lacks theoretical guarantee.
    - In contrast, our method learns the function that maps each code $e$ to LCG (i.e., event-specific inference) and events are *explicitly* identified by the latent clustering. Thus, the identified events are more interpretable and our method is also theoretically-grounded.

Reviewer comment (Experiments):
In theory there are ground truth decomposition which the LCG are the same within partition and different otherwise, however, it seems that for Magnetic, there is a uniform causal graph which does not change locally.

Author response:
- In Magnetic, the context $D=$ {at least one of the objects is non-magnetic, i.e., colored black} induces the fine-grained causal relationships, and its corresponding LCG is shown in Fig. 7(c). The (global) CG is shown in Fig. 7(a) in our revised manuscript.  
- In Chemical, the context $\mathcal{D}=$ {color of the root node is red or blue} induces the fine-grained causal relationships and its corresponding LCG is shown in Fig. 4(c). The (global) CG is the lower triangular matrix.

Reviewer comment (Evaluation):
Identifiability results are not reflected in the method. There is no quantitative result that supporting that identifiability is or is not achieved in the experiment.

Author response:
- The evaluation of local causal discovery is provided in Appendix C.3.2 and Fig. 16. We use the structural hamming distance (SHD), a metric used to quantify the dissimilarity between two graphs (i.e., the distance between the learned LCG and true LCG). We also provide an analysis of whether our method successfully identifies meaningful events through the codebook histograms and learned LCGs (e.g., Fig. 9).

Reviewer comment (Theory):
The assumption on the decomposition is not mild. In real world data, the transition function is usually highly complex. This method can only handle the case when different parts in the decomposition have different LCG. However, in real case, all parts may share the same LCG due to the highly complex transition function. Even though the conditional independence relations are the same (same LCG), the transition function itself may be used to further partition the sample space into further small parts. That is saying same LCG but with different \(p(s'\vert s,a)\).

Author response:
- This is a great question! Recall eq. 3 that $p(s'_j\mid s, a)=p(s'_j\mid Pa(j; \mathcal{E}_z), z)$ for $(s, a)\in \mathcal{E}_z$, the dynamics prediction model takes not only $Pa(j; \mathcal{E}_z)$, but also $z$ as an input. This is because the transition function could be indeed different among parts with the same LCG, as the reviewer pointed out. Here, $z$ guides the network to learn (possibly) different transition functions even if the LCG is the same. Recall that each latent code $e_z\in C=\\{e_z\\}^K\_{z=1}$ denotes the partition, our model takes a one-hot encoding of size $K$ according to the latent code as the additional input to deal with such cases.

Reviewer comment (Presentation):
Minor presentation issues about the notations.

Author response:
- Sorry for the inconvenience. We have uploaded the revised manuscript to improve the clarity. We have included a precise statement of the score function and derivation in Eqs. 8-12 in Appendix B in our revised manuscript. Specifically,  
    - Event $\mathcal{E}$ is a subset of the joint space action space, i.e., $\mathcal{E}\subseteq \mathcal{X}= \mathcal{S}\times \mathcal{A}$.  
    - Score function is $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K):=\text{sup}\mathbb{E}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}) - \lambda \vert \mathcal{G}_z\vert\right]:=\text{sup}\_\phi \mathbb{E}\_{p(s, a, s')}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}, \phi) - \lambda \vert \mathcal{G}_z\vert\right]$. (Eq. 8-12)  
- We have also included more detailed proof of Prop. 3 with derivations in our revised manuscript.  
    - In short, for an arbitrary decomposition $\\{\mathcal{E}_z\\}$ and its corresponding true LCGs $\\{\mathcal{G}_z\\}$, $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K)=\mathbb{E}\_{p(s, a, s')}  \log p(s'\mid s, a) - \lambda \sum_z p(\mathcal{E}_z) \cdot \lvert \mathcal{G}_z \rvert$, and thus the first term is canceled out in the last equality in the proof of Prop. 3.

Reviewer comment (Writing):
In Sec 3.1, what is the definition of event \(\mathcal{E}\)?

Author response:
- Sorry for the inconvenience. We have uploaded the revised manuscript to improve the clarity. We have included a precise statement of the score function and derivation in Eqs. 8-12 in Appendix B in our revised manuscript. Specifically,  
    - Event $\mathcal{E}$ is a subset of the joint space action space, i.e., $\mathcal{E}\subseteq \mathcal{X}= \mathcal{S}\times \mathcal{A}$.  
    - Score function is $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K):=\text{sup}\mathbb{E}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}) - \lambda \vert \mathcal{G}_z\vert\right]:=\text{sup}\_\phi \mathbb{E}\_{p(s, a, s')}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}, \phi) - \lambda \vert \mathcal{G}_z\vert\right]$. (Eq. 8-12)  
- We have also included more detailed proof of Prop. 3 with derivations in our revised manuscript.  
    - In short, for an arbitrary decomposition $\\{\mathcal{E}_z\\}$ and its corresponding true LCGs $\\{\mathcal{G}_z\\}$, $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K)=\mathbb{E}\_{p(s, a, s')}  \log p(s'\mid s, a) - \lambda \sum_z p(\mathcal{E}_z) \cdot \lvert \mathcal{G}_z \rvert$, and thus the first term is canceled out in the last equality in the proof of Prop. 3.

Reviewer comment (Writing):
In terms of the score function \(\mathcal{S}(\{\mathcal{G}_z,\mathcal{E}_z\}_1^K):=\text{sup }\mathbb{E}[\text{log}\hat{p}(s' \vert s,a;\mathcal{G}_z) - \lambda \vert \mathcal{G}_z\vert]\), which variables are for sup and which variables are for \(\mathbb{E}\)?

Author response:
- Sorry for the inconvenience. We have uploaded the revised manuscript to improve the clarity. We have included a precise statement of the score function and derivation in Eqs. 8-12 in Appendix B in our revised manuscript. Specifically,  
    - Event $\mathcal{E}$ is a subset of the joint space action space, i.e., $\mathcal{E}\subseteq \mathcal{X}= \mathcal{S}\times \mathcal{A}$.  
    - Score function is $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K):=\text{sup}\mathbb{E}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}) - \lambda \vert \mathcal{G}_z\vert\right]:=\text{sup}\_\phi \mathbb{E}\_{p(s, a, s')}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}, \phi) - \lambda \vert \mathcal{G}_z\vert\right]$. (Eq. 8-12)  
- We have also included more detailed proof of Prop. 3 with derivations in our revised manuscript.  
    - In short, for an arbitrary decomposition $\\{\mathcal{E}_z\\}$ and its corresponding true LCGs $\\{\mathcal{G}_z\\}$, $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K)=\mathbb{E}\_{p(s, a, s')}  \log p(s'\mid s, a) - \lambda \sum_z p(\mathcal{E}_z) \cdot \lvert \mathcal{G}_z \rvert$, and thus the first term is canceled out in the last equality in the proof of Prop. 3.

Reviewer comment (Theory):
In Proposition 3, why \(\mathbb{E}[\vert\mathcal{G}_z\vert] \le \mathbb{E}[\vert\mathcal{\hat{G}}_z\vert]\)?

Author response:
- Sorry for the inconvenience. We have uploaded the revised manuscript to improve the clarity. We have included a precise statement of the score function and derivation in Eqs. 8-12 in Appendix B in our revised manuscript. Specifically,  
    - Event $\mathcal{E}$ is a subset of the joint space action space, i.e., $\mathcal{E}\subseteq \mathcal{X}= \mathcal{S}\times \mathcal{A}$.  
    - Score function is $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K):=\text{sup}\mathbb{E}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}) - \lambda \vert \mathcal{G}_z\vert\right]:=\text{sup}\_\phi \mathbb{E}\_{p(s, a, s')}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}, \phi) - \lambda \vert \mathcal{G}_z\vert\right]$. (Eq. 8-12)  
- We have also included more detailed proof of Prop. 3 with derivations in our revised manuscript.  
    - In short, for an arbitrary decomposition $\\{\mathcal{E}_z\\}$ and its corresponding true LCGs $\\{\mathcal{G}_z\\}$, $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K)=\mathbb{E}\_{p(s, a, s')}  \log p(s'\mid s, a) - \lambda \sum_z p(\mathcal{E}_z) \cdot \lvert \mathcal{G}_z \rvert$, and thus the first term is canceled out in the last equality in the proof of Prop. 3.

Reviewer comment (Theory):
Specifically, in appendix B.2 proof of Proposition 3 why the last equality holds? Can you provide detailed derivation involving the definition of score function \(\mathcal{S}(\cdot)\)?

Author response:
- Sorry for the inconvenience. We have uploaded the revised manuscript to improve the clarity. We have included a precise statement of the score function and derivation in Eqs. 8-12 in Appendix B in our revised manuscript. Specifically,  
    - Event $\mathcal{E}$ is a subset of the joint space action space, i.e., $\mathcal{E}\subseteq \mathcal{X}= \mathcal{S}\times \mathcal{A}$.  
    - Score function is $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K):=\text{sup}\mathbb{E}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}) - \lambda \vert \mathcal{G}_z\vert\right]:=\text{sup}\_\phi \mathbb{E}\_{p(s, a, s')}\left[\text{log}\hat{p}(s' \vert s,a;\\{\mathcal{G}_z, \mathcal{E}_z \\}, \phi) - \lambda \vert \mathcal{G}_z\vert\right]$. (Eq. 8-12)  
- We have also included more detailed proof of Prop. 3 with derivations in our revised manuscript.  
    - In short, for an arbitrary decomposition $\\{\mathcal{E}_z\\}$ and its corresponding true LCGs $\\{\mathcal{G}_z\\}$, $\mathcal{S}(\\{\mathcal{G}_z,\mathcal{E}_z\\}\_{z=1}^K)=\mathbb{E}\_{p(s, a, s')}  \log p(s'\mid s, a) - \lambda \sum_z p(\mathcal{E}_z) \cdot \lvert \mathcal{G}_z \rvert$, and thus the first term is canceled out in the last equality in the proof of Prop. 3.

Reviewer comment (Evaluation):
How to measure if the identifiability is achieved during learning process? Some measure of the distance between the learned LCG and the true LCG may be used here.

Author response:
- The evaluation of local causal discovery is provided in Appendix C.3.2 and Fig. 16. We use the structural hamming distance (SHD), a metric used to quantify the dissimilarity between two graphs (i.e., the distance between the learned LCG and true LCG). We also provide an analysis of whether our method successfully identifies meaningful events through the codebook histograms and learned LCGs (e.g., Fig. 9).

Reviewer comment (Evaluation):
For ID setting, the proposed method is not the best, is there any explanation for this?

Author response:
- During training (i.e., ID setting), all methods achieve similar near-optimal performances as shown in the learning curve (Fig. 18 on page 27).

Reviewer comment (Experiments):
How to decide the number of codebook vector K?

Author response:
- In theory, the identifiability result holds for any choice of $K\geq 2$. In practice, we find that any choice of $K$ works reasonably well, except for $K=2$ where the training was often unstable (see relevant discussions in Appendix C.3.3).