Quantized Local Independence Discovery for Fine-Grained Causal Dynamics Learning in Reinforcement Learning

We thank the reviewer for the positive assessment and insightful remarks. We respond to each of your comments below:

[W1] The validation environment for the experiment is somewhat simple, and the assumption of full observables seems to have significant limitations.

• 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.

[Q1] 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?

• 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.

[Q2] 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?

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

We again appreciate the reviewer for the valuable comments to improve our manuscript. We tried our best to address the reviewer’s concerns and questions, and we incorporated our responses into our revised manuscript. If our responses are insufficient or if there is a misunderstanding of the reviewer’s questions, we are happy to engage in further discussions.

We sincerely appreciate for reviewing our paper and providing us with constructive feedback. We respond to each of your comments below:

it seems that for Magnetic, there is a uniform causal graph which does not change locally. [..] what are the ground truth decomposition and the corresponding LCGs in both Chemical and Magnetic environment?

• 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.

There is no quantitative result that supporting that identifiability is or is not achieved in the experiment. [..] 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.

• 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).

Minor presentation issues about the notations. […] In Proposition 3. why $\mathbb{E}\left[\lvert\mathcal{G}_z\rvert\right] \leq \mathbb{E}\left[\lvert \hat{\mathcal{G}}_z\rvert \right]$? Can you provide detailed derivation?

• 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.

This method can only handle the case when different parts in the decomposition have different LCG. [..] 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'\mid s, a)$. […] is the proposed method flexible enough to generalize to those case?

• 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.

For ID setting, the proposed method is not the best, is there any explanation for this?

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

How to decide the number of codebook vector K?

• 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).

We again deeply thank the reviewer for the constructive comments. We tried our best to address the reviewer’s concerns and questions, and we incorporated our responses into our revised manuscript. If our responses are insufficient or if there is a misunderstanding of the reviewer’s questions, we are happy to engage in further discussions.

We sincerely appreciate the reviewer’s time and efforts to provide constructive feedback to improve our paper. We respond to each of your comments below.

[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.

[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.

[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.

[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}_{`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.

Dear reviewer eMfR,

As the discussion period is ending soon, we are wondering if our responses and revision have addressed your concerns. We find the constructive comments from the reviewer extremely helpful in improving our initial manuscript, and we would be happy to engage in further discussions if you have any additional questions or suggestions. If our responses have addressed your concerns, we would highly appreciate it if you could re-evaluate our work based on our responses and revised manuscript.

Best regards,

Authors of the submission 4522.

We thank the reviewer for the positive assessment and insightful remarks. We respond to each of your comments below:

[W1] Citation of VQ-VAE

• 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.

[Q1] Why an event-specific graph has to be constrained to be a subgraph of the underlying G?

• 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.

[Q2] Why vector quantization can correctly identify the event?

• 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.

[Q3] Does the identifiability in Theorem 1 rely on the specific choice of K?

• 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.

We again appreciate the reviewer for the valuable comments to improve our manuscript. We tried our best to address the reviewer’s concerns and questions, and we incorporated our responses into our revised manuscript. If our responses are insufficient or if there is a misunderstanding of the reviewer’s questions, we are happy to engage in further discussions.

Thank you for your response. After reading all the review comments I tend to keep my rating unchanged.

We sincerely appreciate the reviewer’s time and efforts to provide constructive feedback to improve our paper. We respond to each of your comments below.

Questions for the proposed method

[Q1-2] Backpropagation in Eq. 5-7. [..] Exactly what is being updated by $\mathcal{L}\_{`pred`}$ and $\mathcal{L}_{`quant`}$? What are exactly computed for the differentiation?

• Gradient of $\mathcal{L}_{`pred`}$ (eq. 6) does update the latent embedding $h$, and consequently, affects the codebook $C$.
  • First, recall that $A\sim g_{`dec`}(e)$, backpropagation from $A$ in $\mathcal{L}\_{`pred`}$ updates the quantization decoder $g_{`dec`}$through $e$.
  • Second, we copy gradients from $e$ (= input of $g_{`dec`}$) to $h$ (= output of $g_{`enc`}$) during the backward path in eq. 5, following a popular trick used in VQ-VAE [1]. By doing so, $\mathcal{L}\_{`pred`}$ updates the quantization encoder $g_{`enc`}$ and $h$.
  • Finally, this also affects the learning of the codebook $C$ since $C$ is updated with $\mathcal{L}_{`quant`}$ (eq. 7) which involves in $h$.
• To sum up, $\mathcal{L}\_{`pred`}$ updates the encoder $g_{`enc`}(s,a)$, decoder $g_{`dec`}(e)$, and the dynamics model $\hat{p}$. $\mathcal{L}\_{`quant`}$ updates $g_{`enc`}$ and the codebook $C$.
• Since the learning of $C$ and $h$ are also based on $\mathcal{L}_{`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.

Questions for the experiment

How is the data sampled?

• We use a random policy for the initial data collection.

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.

Questions for the theory

[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.

Dear reviewer RHPG,

As the discussion period is ending soon, we are wondering if our responses and revision have addressed your concerns. We find the constructive comments from the reviewer extremely helpful in improving our initial manuscript, and we would be happy to engage in further discussions if you have any additional questions or suggestions. If our responses have addressed your concerns, we would highly appreciate it if you could re-evaluate our work based on our responses and revised manuscript.

Best regards,

Authors of the submission 4522.

Dear reviewers,

Thank you again for your detailed feedback! As the discussion period will close in a day, we are wondering if you have any other questions. We are happy to engage in further discussions and address any remaining questions or concerns.

Authors of the submission 4522.