Boosting Efficiency in Task-Agnostic Exploration through Causal Knowledge

Thanks for the review (part 1)

We greatly appreciate the reviewer's thorough and constructive comments. We attempt to address all the concerns in the following. **Weakness**$\quad$After carefully reading the mentioned works, it holds true that there exist several papers that are relevant to ours. We apologize for any oversight in our previous acknowledgment of related works and appreciate the opportunity to clarify this and will appropriately cite these papers and make comparisons. However, we also want to clarify that the exploration term used in [1,2] is actually the same as the ASR method, which has alreadby been taken as one of our baselines. In addition, the conditional independence test in [3] uses all of the coming data, while we selectively collect representative data points during exploration in an incremental way (equation 1). **Questions** **PC algorithm**$\quad$Thank you for point this out. We acknowledge that a more rigorous phrasing would be, "We extend PC to time series data with time-lagged causal relationships". We will promptly address this clarification in our latest version. **Score-based approach** $\quad$Thank you for your suggestion. We will incorporate the justification and discussions comparing our approach with score-based methods in the latest version. However, we cannot fully agree with the statement, "Score-based approach is more efficient compared to constraint-based methods for causal discovery." As mentioned in the Q3 section of the referenced work [3], under the MDP and Causal Faithfulness assumptions, "Based on the performance of the Score model summarized in Table 1, we also conclude that it (the score-based method) is not as good as our constraint-based method and has a large variance due to the unstable learning of the causal graph." This observation similarly applies to our work. **Depth in Fig 2 (b)**$\quad$"depth" here refers to the cardinality of the conditioning set required for conditional independence tests. **Difference between $w$ and $w^c$**$\quad$As stated in Section B of the supplementary materials, we " denote $w^c_{t}$ as the network parameters incorporating causal constraints with respect to $w_t$ that does not gather causal information". Sorry for any inconvenience caused in your reading and we will modify the corresponding expression in the main paper. **Active reward**$\quad$The method in [1] trains a strategy based on active reward to determine whether to annotate the incoming data for a supervised classifier. If the active reward value is positive, it indicates that the annotated data is helpful for the training process. Similarly, we use active rewards to encourage agents to explore toward data that is helpful for training. Therefore, active rewards are used to measure the training quality. To maintain consistency in the order of prediction loss, we need to ensure that the agent can receive prediction errors and active rewards at every time step $t$. Therefore, $t$ in equation 5 refers to the time step. However, the optimization of the world model $f$ doesn't occur at every time step $t$, so we need to separately extract a dynamics model $\phi_t$ for calculating the active reward. We discuss this in detail in Section C.1 of the supplementary materials. More information on the intuition of active reward can be found in [1]. [1] Meng Fang, Yuan Li, and Trevor Cohn. Learning how to active learn: A deep reinforcement learning approach. arXiv preprint arXiv:1708.02383, 2017.

Thanks for the review (part 2)

**Mujoco** Thank you for the thoughts. A detailed explanation of the performance gain in Figure 7 is outlined as follows: * Complexity of Causal Relationships in Mujoco: As shown in prior works [1-4], the dynamics in many real world problems is typically sparse. Therefore, it is necessary to clarify whether this is also the case in Mujoco. However, an extensive review of existing literature did not yield documentation or theoretical works describing the authentic causal relationships in Mujoco tasks. * Causal Structure Identified by PC Algorithm: In our work, the causal structure identified by the PC algorithm appears to be relatively not that dense, a trend supported by empirical results such as prediction error and performance gain. Unfortunately, due to the absence of benchmark works, we cannot guarantee the complete correctness of the identified causal graph. We appreciate your valuable feedback that points this out, and we will include a discussion on the correctness of identified causal relationships in Mujoco tasks in our latest version, as demonstrated in simulated data (Figure 5) and the traffic light control task (Figure 6). * Two-Stage Approach and Model Learning Impact: Our method consists of two stages: model learning and policy learning. Policy learning is based on the model learned in the first stage. Therefore, the accuracy of model learning significantly influences policy learning. In the first stage, we use a limited amount of data for model learning, which might be sufficient for causal models but inadequate for non-causal models, as reflected in Figure 7(a). This explains the results in Figure 7(b). **The value of $\beta$**$\quad$We set $\beta$ as $3$ according to the nonlinear results in Section C.1 of the supplementary materials. **Minor comments and typos** Thank you very much for pointing out these typos. We will carefully review and address them in the latest version.

Thanks for the review

We greatly appreciate your insightful comments and time devoted to our work. We attempt to address all the concerns as follows. **Contribution**$\quad$We believe that our key contributions are summarized at the end of the introduction section, which are: * In order to enhance the sample efficiency and reliability of model training with causal knowledge, we introduce a novel concept: causal exploration, and focus particularly on the domain of task-agnostic reinforcement learning. * To efficiently learn and use structural constraints during world model learning, we design a novel weight-sharing-decomposition schema that can avoid additional computational burden. * We demonstrate the effectiveness of causal exploration across a range of demanding reinforcement learning environments. Theoretically, we show that, given strong convexity and smoothness assumptions, our approach attains a superior convergence rate compared to non-causal methods. Empirically, experimental results on synthetic data demonstrate the ability to learn accurate world models with exceptional data utilization, surpassing existing exploration methods in complex scenarios. Notably, our approach also produces outstanding performance in challenging tasks such as traffic light control and MuJoCo, highlighting its practicality in real-world applications. Sorry for any inconvenience caused in your reading and we will make a better organization in the revised version. **Causal identifiability**$\quad$We sincerely appreciate your thoughtful comments. Consider that we focus on a Markov decision process (MDP) where the underlying states $\boldsymbol{s}_t$ are observed, the causal assumptions for our method are the Markov condition and faithfulness assumption. Besides, the identifiability theory of MDP has been established in prior works [1-3]. [2] also provides theoretical guarantees for the correctness of the causal graph we identified. We acknowledge the need for providing a theoretical analysis for this in the main paper and we will add this in the revised manuscript. [1] 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.} [2] Ding et al., “Generalizing goal-conditioned reinforcement learning with variational causal reasoning”, NeurIPS 2022. [3] Wang et al., “Causal dynamics learning for task-independent state abstraction”, ICML 2022.

Thanks for the review

We appreciate your thoughtful comments and time devoted to reviewing this paper, and hope the following response properly addresses your concerns. **Equation 5**$\quad$The method in [1] trains a strategy based on active reward to determine whether to annotate the incoming data for a supervised classifier. If the active reward value is positive, it indicates that the annotated data is helpful for the training process. Similarly, we use active rewards to encourage agents to explore toward data that is helpful for training. Therefore, active rewards are used to measure the training quality. To maintain consistency in the order of prediction loss, we need to ensure that the agent can receive prediction errors and active rewards at every time step $t$. That is, $t$ in equation 5 refers to the time step. [1] Meng Fang, Yuan Li, and Trevor Cohn. Learning how to active learn: A deep reinforcement learning approach. arXiv preprint arXiv:1708.02383, 2017. **Identifiability**$\quad$The identification of causal relationships is completed by the PC algorithm. In addition, the input of the world model only contains information about the parent nodes, so the shared schema does not affect causal identification. **Enhancing efficiency**$\quad$In our paper, "boosting efficiency" entails that agents can achieve comparable or superior performance with fewer sample size. As has been studied in [2,3,4], causal knowledge improves sample efficiency during training but calls for factorization. The shared schema reduces the computation burden caused by causal factorization. Besides, the scale of both the shard and decomposed layers can be regulated so there is no introduction of an additional layer even compared with non-causal methods. What's more, introducing active reward does indeed lead to an increase in computational complexity. But such a slight increase in computational complexity can be ignored compared to the improvement in model performance (Figure 9 in the supplementary materials). Furthermore, the additional computational costs can be reduced by controlling the scale of model $\phi$ and the test set $\mathcal {D}_h $. [2] Maximilian Seitzer, et al., Causal influence detection for improving efficiency in reinforcement learning. Advances in Neural Information Processing Systems, 34: 22905–22918, 2021. [3] Wang et al., “Causal dynamics learning for task-independent state abstraction”, ICML 2022. [4] Ding et al., “Generalizing goal-conditioned reinforcement learning with variational causal reasoning”, NeurIPS 2022. **Convergence analysis**$\quad$As is depicted in section B of the supplementary materials, "Denote $w_t^c$ as the network parameters taking causal constraints with respect to $w_t$ that dose not gather causal information. Theorem 3.1 shows that causal exploration gets a prediction error bound $\delta_t$ times lower at time $t$, where $\delta_t$ is a density measurement of causal matrix $D$." We believe this is consistent with what you said about “whether the convergence rate is better than what was previously required”. **Data collection**$\quad$Actually, we learn the parameters based on the collected data $B_t$ during exploration instead of $\mathcal{D}_h$ which is used for the calculation of active reward (line 9 and 12 in Algorithm 1).

Thanks for the review (part 1)

Thanks a lot for your thoughtful comments and suggestions. Please see below for our responses to your specific comments. **Lack of problem-setting section**$\quad$Sorry for any inconvenience caused in your reading and we will add a Preliminary section before section 2 following your invaluable suggestions. **Causal assumptions and identification theory**$\quad$We sincerely appreciate your thoughtful comments. Consider that we focus on a Markov decision process (MDP) where the underlying states $\boldsymbol{s}_t$ are observed, the causal assumptions for our method are the Markov condition and faithfulness assumption. Besides, the identifiability theory of MDP has been established in prior works [7]. [8] provides theoretical guarantees for the correctness of the causal graph we identified. We acknowledge the need for providing a theoretical analysis for this in the main paper and we will add this in the revised manuscript. Moreover, we would like to clarify that the PC algorithm is also a common method for causal discovery in SCM models through the application of conditional independence tests, as demonstrated in [8,9]. Furthermore, [8] has proven that the use of independence tests for causal discovery tends to outperform and exhibit greater stability compared to other causal techniques like score-based methods for SCM in our settings. In addition, the ASR method precisely employs an end-to-end approach for embedding causal structure and world model learning. We have included it as a baseline method for experimental comparison (Figure 5). [7] 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. [8] Ding et al., “Generalizing goal-conditioned reinforcement learning with variational causal reasoning”, NeurIPS 2022. [9] Wang et al., “Causal dynamics learning for task-independent state abstraction”, ICML 2022. **Disconnection from SOTA** (1)We appreciate the reviewer's attention to this specific point. We would like to clarify that while the exploration policy in our work shares a conceptual similarity with the notion of experiment design in active causal discovery, the practical implications and challenges are different in the context of reinforcement learning. In traditional active causal discovery, experiment design often involves interventions to perturb the system and observe the causal effects. However, in reinforcement learning, interventions are typically challenging due to the sequential, interactive nature of decision-making. Agents in reinforcement learning environments operate in an ongoing interaction with the environment, making controlled interventions more complex compared to static observational settings. Our exploration policy, while conceptually related to experiment design, is designed to efficiently collect data for model learning rather than explicit interventions. This adaptation is necessary to accommodate the unique challenges posed by sequential decision-making in reinforcement learning settings. We will highlight this distinction more explicitly in the revised manuscript. (2)We greatly appreciate your provided references. However, after carefully reviewing them, we did not find a similar weight-sharing schema as employed in our work. In [1], the causal discovery is formulated as a continuous optimization problem, and the authors use the Lagrangian method for solving. This approach differs from ours, as we employ deep neural network structures. Additionally, [1] focuses solely on linear Structural Equation Models (SEMs), while our deep network has the capacity to approximate a wide range of linear and nonlinear functions. In [2], the authors merge the decomposed networks into a unified one by learning a training embedding $\boldsymbol{u}_i$ to serve as index information for each dimension, i.e., $f_i(\cdot)=f(\boldsymbol{u}_i,\cdot)$. However, it's crucial to note that $\boldsymbol{u}_i$ is present in the form of an input, whereas our structure involves passing inputs through a shared layer before transitioning to respective decomposed networks. Moreover, $\boldsymbol{u}_i$ is a d-dimensional vector in [2], whereas our shared network's layer count is dynamically adjustable. (3)Similar to the answer for causal assumptions, [8] has provided both theoretical and experimental evidence supporting the superiority and greater stability of using independence tests for causal discovery compared to other causal techniques, such as score-based methods for Structural Causal Models (SCM) in our settings. Furthermore, the ASR method precisely employs an end-to-end approach for embedding causal structure and world model learning. We have included it as a baseline method for experimental comparison, as shown in Figure 5.

Thanks for the review (part 2)

**Experiment evaluation methods**$\quad$We appreciate your valuable insights. But we would like to clarify that the prediction error serves as a suitable metric for assessing model training efficacy as is highlighted in [10] and has been used in several works like [8,9]. Its application to exploration data precisely reflects that our model achieves shorter training times and utilizes fewer samples. Furthermore, all baseline methods, with the exception of CID, employ intrinsic rewards for training. Points of distinction include: firstly, our intrinsic reward formulation encompasses active reward, which cannot be considered a specialized design for using prediction loss. Secondly, we provide additional causal structural information to the model. Moreover, CID employs a strategy of random sampling for model training, whereas utilizing intrinsic reward forms a more appropriate exploration approach. Certainly, we also utilize the discrepancy between the learned model and the ground truth model as an evaluation criterion. In contrast to a direct model-to-model comparison, we choose to compare the learned causal graph with the ground truth one (Figure 5). Additionally, the downstream performance for traffic signal control and the MuJoCo task, as depicted in Table 1 and Figure 7(b), adheres to the concept of offline equivalence in real-world scenarios as you mentioned. [10] Ke, Nan Rosemary et al. "Systematic evaluation of causal discovery in visual model-based reinforcement learning." arXiv preprint arXiv:2107.00848 (2021).