SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Mean  $\pm$  std. error of ROC AUC (↑) and F1 (↑) of local dependency prediction
THAWING	CARWASH	KITCHEN
ROC AUC	F1	ROC AUC	F1	ROC AUC	F1
ELDEN	0.71 ± 0.01	0.57 ± 0.00	0.78 ± 0.02	0.66 ± 0.02	0.66 ± 0.01	0.25 ± 0.01
pCMI	0.55 ± 0.01	0.60 ± 0.00	0.73 ± 0.02	0.78 ± 0.01	0.60 ± 0.00	0.28 ± 0.00
Attn	0.65 ± 0.04	0.63 ± 0.01	0.66 ± 0.01	0.55 ± 0.03	0.51 ± 0.01	0.22 ± 0.02
Input Mask	0.50 ± 0.00	0.40 ± 0.00	0.50 ± 0.00	0.32 ± 0.01	0.50 ± 0.00	0.08 ± 0.00
Attn Mask	0.45 ± 0.03	0.47 ± 0.02	0.47 ± 0.07	0.43 ± 0.03	0.52 ± 0.01	0.13 ± 0.01
Implementation Details For discrete state or action spaces, the partial derivatives w.r.t.  $s_t / a_t$  are undefined. To address this issue, we use Mixup [34] to create synthetic inputs and labels by linearly combining pairs of inputs and labels, thus approximately changing the input space to be continuous. Compared to learning from discrete inputs only, dynamics models trained on such data generate partial derivatives that better reflect local dependencies, as shown in Sec 4.3.1. For the 2D MINECRAFT and KITCHEN environments where some local dependencies have complex preconditions and thus are hard to induce, we apply sample prioritization to dynamics learning, where the priority is measured as prediction error. In this way, the dynamics model gets aware of unknown interactions faster and guides the exploration more efficiently than not using prioritization. Further details are provided in the Appendix.
4.1 Evaluating the Detection of Local Dependencies
We compare the local dependencies extracted by ELDEN with the following baselines (see implementation details in Appendix Sec. C):
- pCMI (point-wise conditional mutual information) [27, 31]: defined as  $\log \frac{p(s_{t+1}^j | s_t, a_t)}{p(s_{t+1}^j | s_t \setminus s_t^i, a_t)}$ . It quantifies how likely it is that  $s_{t+1}^j$  depends on  $s_t^i$ .
- Attn (attention): Use the score between each entity pair computed by the attention modules inside the dynamics model to quantify local dependencies.
- Input Mask: we implement a learnable binary mask to the dynamics model that can zero out some inputs conditioned on  $(s_t, a_t)$ :  $f([s_t, a_t] \odot M(s_t, a_t))$ . During training, the mask is regularized to use as few inputs as possible with L1 regularization, leading to a quantification of minimal local dependencies.
- Attn Mask: we implement a learnable mask to the dynamics model similar to the one in Input Mask, but in this case, the mask is applied to the attention scores. The mask is regularized following the method by Weiss et al. [32].
We train the dynamics model of each method with three random seeds on pre-collected transition data and evaluate their performance by predicting the local causal graph  $\mathcal{G}_t$  for 50 unseen episodes based on the state-action pair  $(s_t, a_t)$ . We compare their predictions with the ground truth local dependencies extracted from the simulator. In the three environments, many potential local dependencies are inactive most of the time, and thus only a small portion ( $\leq 3\%$ ) of the ground truth labels indicate the existence of local dependencies for a given entity pair. To account for such imbalance, we use the area under the receiver operating characteristic curve (ROC-AUC) and the best achievable F-score (F1) as evaluation metrics.
The results of the evaluation on the detection of local dependencies are summarized in Table 1. ELDEN outperforms all baselines in terms of ROC-AUC consistently across all environments (Q1). For the F1 score, pCMI performs best in most environments (especially in the more complex CARWASH and KITCHEN), but ELDEN performs comparably or achieves the second-best F1 scores with much less computation: pCMI computation cost is  $N$  times higher than ELDEN, where  $N$  is the number of environment entities, and thus pCMI scales badly to environments with a large number of objects. Further evaluation details can be found in the Appendix.
Disagreement Dynamics Curiosity CAI Vanilla PPO ELDEN (ours)
(a) Thawing
Figure 3: Learning curve of ELDEN (ours) compared to baseline approaches. Each method uses three random seeds, and we show the mean  $\pm$  std dev of the number of stages completed toward task success. The stage count is normalized to [0, 1], where 1 corresponds to task completion. ELDEN learns successful policies in all four test environments, and is the only method that succeeds in the CarWash, 2D Minecraft, and Kitchen environments with complex chained dependencies.
(c) Craft
4.2 Evaluating Exploration in Sparse-Reward RL Tasks
The ultimate goal of our method is to improve exploration for RL in sparse-reward setups. In the second evaluation, we compare the performance of ELDEN against several state-of-the-art intrinsic-motivation exploration algorithms in reinforcement learning, including:
- Disagreement [21]: the intrinsic reward is computed based on the variance of the predictions from an ensemble of forward dynamics models.
- Dynamics Curiosity [5]: intrinsic reward is computed based on the prediction error of a trained forward dynamics model.
- CAI (Causal Influence Detection) [27]: an empowerment-based method, where the agent is given intrinsic reward for maximizing the number of state entities that depend on its action.
- Vanilla PPO [26]: baseline without intrinsic reward that serves as control signal.
While ELDEN can be used with any RL algorithm, in our experiments we use proximal policy optimization (PPO) [26], as well as with the baselines. To facilitate the introspection of the results, we define manually a set of semantic stages representing internal progress toward task completion. The stage definitions for each of the environments are described in detail in the Appendix. Notice that these stages are not used by the agents during training and do not provide any additional reward; they are only used to facilitate the analysis of the results.
Fig. 3 depicts the count of reached stages per episode during training for each task, normalized by the number of stages to complete the task. In the normalized stage count, a value of 1 corresponds to successfully completing the task and it is the only stage where the learning agents receive sparse task reward (not intrinsic). Fig. 3 indicates that ELDEN is able to learn successful policies in all four environments. Importantly, in CARWASH, 2D MINECRAFT and KITCHEN, ELDEN is the only method that successfully learns to complete the task, demonstrating the advantage of ELDEN over the baseline algorithms in tackling tasks with complex chained dependencies (Q2).
The normalized stage count of ELDEN in CARWASH, 2D MINECRAFT and KITCHEN does not converge to 1 (completing the entire task in all episodes) mainly due to two reasons: First, in both tasks, the locations of the objects are randomly initialized at the start of each episode. For some initialization (e.g. a target object is blocked by unmovable obstacles), the task is impossible to solve. Second, in both tasks, two out of the three ELDEN training procedures with different random seeds converge to succeeding most of the time, but the training process with one seed fails to find a good policy, dragging down the mean value of the normalized stage count. This large variance in success is a current limitation of ELDEN.
In the relatively simple THAWING environment, we found ELDEN does not provide a significant advantage over the other baseline methods. The Dynamics Curiosity baseline learns faster to achieve the task indicating a better sample efficiency. This was rather expected: as with any exploration heuristic, ELDEN is not universally better than previous intrinsic reward methods — instead, it is better suited for a specific type of environment, where there are many complex and chained object
Table 2: Ablation of ELDEN on local dependency prediction (mean ± std. error of ROC AUC and F1)
THAWING	CARWASH	KITCHEN
ROC AUC (↑)	F1 (↑)	ROC AUC (↑)	F1 (↑)	ROC AUC (↑)	F1 (↑)
no Mixup & no Reg	0.48 ± 0.01	0.42 ± 0.01	0.44 ± 0.00	0.27 ± 0.01	N/A	N/A
no Reg, i.e., λ = 0	0.57 ± 0.01	0.52 ± 0.01	0.54 ± 0.01	0.42 ± 0.01	0.64 ± 0.01	0.24 ± 0.01
λ = 10-1	0.68 ± 0.00	0.57 ± 0.00	0.73 ± 0.01	0.58 ± 0.02	0.55 ± 0.00	0.14 ± 0.00
λ = 10-2	0.71 ± 0.01	0.57 ± 0.00	0.76 ± 0.01	0.60 ± 0.00	0.60 ± 0.01	0.21 ± 0.01
λ = 10-3	0.64 ± 0.01	0.55 ± 0.01	0.78 ± 0.02	0.66 ± 0.02	0.65 ± 0.00	0.24 ± 0.01
λ = 10-4	0.65 ± 0.02	0.55 ± 0.01	0.75 ± 0.01	0.60 ± 0.01	0.66 ± 0.00	0.25 ± 0.01
λ = 10-5	0.63 ± 0.00	0.53 ± 0.01	0.72 ± 0.00	0.57 ± 0.00	0.65 ± 0.01	0.24 ± 0.01
dependencies, not the case for THAWING. We provide additional experimental evaluations of the failure cases of ELDEN in Appendix Sec E.
We ablate different components of ELDEN to examine their importance to the overall methods.
4.3.1 Ablations for Local Dependency Detection
In our ablation study on ELDEN for local dependency detection, we investigate the impact of each component with the following variations:
- No Mixup & No Reg: We disable the use of Mixup for discrete space prediction, and no partial derivative regularization is applied in this case.
- Different partial derivative regularization coefficients: we test with different  $\lambda$  values in  $\{0,10^{-1},10^{-2},10^{-3},10^{-4},10^{-5}\}$ .
As shown in Table. 2, in Thawing and CarWash environments, partial derivative regularization with appropriate coefficients significantly improves ELDEN's detection of local dependencies, compared to no regularization (i.e.,  $\lambda = 0$ ) or inappropriate  $\lambda$  values. Furthermore, in discrete-state environments, Mixup smooths the landscape of partial derivatives by providing synthesized continuous inputs as exemplified in Fig. 1(b) of Zhang et al. [34], thus facilitating local dependency prediction — even when compared to using Mixup without any regularization, not using Mixup leads to a noticeable degradation in the prediction performance.
4.3.2 Ablations for Task Learning
Next, we examine how different components and hyperparameters of ELDEN affect task learning:
Ablation of Local Dependency Metrics We compare the exploration performance when using different local dependency detection methods. Specifically, we compare with pCMI as it achieves the best local dependency detection in Sec. 4.1. We present the comparison results between ELDEN and pCMI in the Kitchen environment in Fig. 4(a) where both methods successfully learn to solve the task. However, it is important to notice that the computation cost of pCMI is  $N$  times more than that of ELDEN (where  $N$  is the number of environment entities), and thus may not scale to environments with a large number of entities.
Ablation of Dynamics Sample Prioritization We study the effectiveness of applying sample prioritization in dynamics model training. Specifically, we test ELDEN with and without prioritization in the Kitchen environment, and show the result in Fig. 4(b). We can see that ELDEN without prioritization fails to learn a useful policy. The reason is that some key entity interactions occur rather rarely before the agent masters them, e.g., frying meatball with butter. In such cases, the dynamics model needs to quickly learn that unknown dependencies appear so that it can bias the exploration toward reproducing such dependencies. Sample prioritization helps the dynamics model learn such
(a) Dependency detection
Figure 4: Ablation of ELDEN on task learning. Each curve uses three random seeds and shows the mean  $\pm$  std dev of the normalized stages. We found ELDEN to have moderate tolerance towards hyperparameters. We found sample prioritization in dynamics learning to be crucial to the performance of ELDEN.
(b) Sample prioritization
(c) Gradient thresholds
(d) Reward coefficients
infrequent dependencies quickly, making it critical in environments with novel and hard-to-induce local dependencies.
Ablation of Partial Derivative Threshold: The partial derivative threshold  $\epsilon$  determines the dependency predictions. A threshold that is too large / too small will make all dependency predictions negative / positive respectively, leading to deteriorated performance. In this section, we examine whether our method is sensitive to the choice of threshold in the CarWash environment, where the results are presented in Fig. 4(c). We observe that our method is relatively sensitive to the choice of threshold, and an inappropriate threshold could cause catastrophic failure. A potential next step for ELDEN is to automatically determine the partial derivative threshold.
Ablation of Intrinsic Reward Coefficient: The intrinsic reward coefficient controls the scale of the intrinsic reward relative to the task reward. We examine the effect of this coefficient by experimenting with different values in the CarWash environment, where the results are presented in Fig. 4(d). We find that our methods work well in a large range of the intrinsic reward coefficients (1 - 10), since the task only gives sparse rewards and the intrinsic rewards are the only learning signal most of the time. The only exceptions are (1) when the intrinsic reward coefficient is too large (e.g., 100), the intrinsic reward significantly surpasses the task reward, and (2) when the coefficient is too small (e.g., 0.1), the episode intrinsic reward is too small (e.g., 0.03) for PPO to learn any useful policy.
We introduce ELDEN, a method for improving exploration in sparse reward reinforcement learning tasks. ELDEN identifies local dependencies between environment entities and uses the uncertainty about such dependencies as an intrinsic reward to improve exploration. Experiments demonstrate that ELDEN uncovers local dependencies more accurately compared to related methods, and significantly outperforms previous exploration methods in tasks with complex chained dependencies.
However, ELDEN is not without limitations. First, ELDEN intentionally bias exploration towards "covering up the possible interactions between objects" rather than "becoming an expert at manipulating a particular object". While such an inductive bias works well in many practical domains, it may fail when facing tasks that require precise object interaction (e.g. rotating the meatball in the pot to a specific orientation). A future direction to alleviate this problem and expand the scope of solvable tasks is to combine ELDEN with dynamics curiosity and formulate a composite intrinsic reward. Second, as noted in the experiment section, the variance of ELDEN across different random seeds can be large, While the high variance is a general problem to Reinforcement Learning, finding ways to further stabilize ELDEN can be an important direction for future work.
Acknowledgements This work has taken place in the Robot Interactive Intelligence Lab (RobIn) and Learning Agents Research Group (LARG) at the Artificial Intelligence Laboratory, The University of Texas at Austin. LARG research is supported in part by the National Science Foundation (FAIN2019844, NRT-2125858), the Office of Naval Research (N00014-18-2243), Army Research Office (E2061621), Bosch, Lockheed Martin. Both LARG and RobIn are supported by Good Systems, a research grand challenge at the University of Texas at Austin. The views and conclusions contained in this document are those of the authors alone. Peter Stone serves as the Executive Director of Sony AI America and receives financial compensation for this work. The terms of this arrangement have been reviewed and approved by the University of Texas at Austin in accordance with its policy on objectivity in research. We thank Bo Liu and Caleb Chuck for their valuable feedback on the manuscript.
[1] Jacob Andreas, Dan Klein, and Sergey Levine. Modular multitask reinforcement learning with policy sketches. In International Conference on Machine Learning, pages 166-175. PMLR, 2017.
[2] OpenAI: Marcin Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research, 39(1):3-20, 2020.
[3] Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. Advances in neural information processing systems, 29, 2016.
[4] Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemysław Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.06680, 2019.
[5] Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A Efros. Large-scale study of curiosity-driven learning. arXiv preprint arXiv:1808.04355, 2018.
[6] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv:1810.12894, 2018.
[7] Maxime Chevalier-Boisvert, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Sahara, Thien Huu Nguyen, and Yoshua Bengio. Babyai: A platform to study the sample efficiency of grounded language learning. arXiv preprint arXiv:1810.08272, 2018.
[8] Jongwook Choi, Yijie Guo, Marcin Moczulski, Junhyuk Oh, Neal Wu, Mohammad Norouzi, and Honglak Lee. Contingency-aware exploration in reinforcement learning. arXiv preprint arXiv:1811.01483, 2018.
[9] Tuomas Haarnoja, Sehoon Ha, Aurick Zhou, Jie Tan, George Tucker, and Sergey Levine. Learning to walk via deep reinforcement learning. RSS, 2018.
[10] Jiaheng Hu, Peter Stone, and Roberto Martín-Martín. Causal policy gradient for whole-body mobile manipulation. arXiv preprint arXiv:2305.04866, 2023.
[11] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144, 2016.
[12] Emily Jin, Jiaheng Hu, Zhuoyi Huang, Ruohan Zhang, Jiajun Wu, Li Fei-Fei, and Roberto Martín-Martín. Mini-behavior: A procedurally generated benchmark for long-horizon decision-making in embodied ai. arXiv preprint 2310.01824, 2023.
[13] Alina Kloss, Stefan Schaal, and Jeannette Bohg. Combining learned and analytical models for predicting action effects. arXiv preprint arXiv:1710.04102, 11, 2017.
[14] Alexander S Klyubin, Daniel Polani, and Chrystopher L Nehaniv. All else being equal be empowered. In Advances in Artificial Life: 8th European Conference, ECAL 2005, Canterbury, UK, September 5-9, 2005. Proceedings 8, pages 744-753. Springer, 2005.
[33] Peter R Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas J Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, et al. Outracing champion gran turismo drivers with deep reinforcement learning. Nature, 602(7896): 223-228, 2022.
[34] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09412, 2017.
[35] Rui Zhao, Yang Gao, Pieter Abbeel, Volker Tresp, and Wei Xu. Mutual information state intrinsic control. arXiv preprint arXiv:2103.08107, 2021.
[36] Yuke Zhu, Josiah Wong, Ajay Mandlekar, and Roberto Martin-Martin. robosite: A modular simulation framework and benchmark for robot learning. arXiv preprint arXiv:2009.12293, 2020.
A ELDEN details
Assumptions We summarize our assumptions on the MDP as follows:
1. The state space can be factored as  $\mathcal{S} = \mathcal{S}^1 \times \dots \times \mathcal{S}^N$ .
2. The transition of each state factor is independent, i.e., the dynamics can be represented  $\mathcal{P}(s_{t + 1}|s_t,a_t) = \prod_{n = 1}^{N}p(s_{t + 1}^n |\mathrm{Pa}(\mathcal{S}_{t + 1}^n))$
3. There is no instantaneous dependency between state factors at the same time step  $t$ , i.e., no dependency such as  $s_t^i \rightarrow s_t^j$  for any  $i, j$ .
For assumption 1, factored state space is commonly employed in causality literature and applies to many simulated or robotics environments. In cases where low-level observations or partial observability are present, disentangled representation or causal representation methods can be utilized to learn a factored state space [17]. When a factored state space is available, assumptions 2 and 3 generally hold.
Network Architecture In Figure 5(a), the architecture of ELDEN for predicting each state factor  $s_{t+1}^j$  is illustrated. The process consists of the following steps:
1. Feature Extraction: For each input state factor  $s_t^i$ , ELDEN utilizes a separate multi-layer perception (MLP) to extract its corresponding feature  $g^i$ .
2. Entity Interaction: ELDEN employs a multi-head self-attention module to model entity interactions and generates a set of transformed features  $h^i$  that incorporate information from other state factors.
3. Prediction using Multi-Head Attention: With  $h^j$  as the query, ELDEN utilizes a multi-head attention module to compute the prediction  $\hat{p}(s_{t+1}^j | s_t, a_t)$  for each state factor. For continuous state factor,  $\hat{p}(s_{t+1}^j)$  is modeled as a normal distribution with the mean computed by the network and a fixed variance equal to 1. For discrete factor,  $\hat{p}(s_{t+1}^j)$  is a categorical distribution with network outputs as class probabilities.
Throughout the prediction process, there are a total of  $N$  such networks in ELDEN, with each network responsible for predicting a separate state factor  $s_{t + 1}^{j}$ .
The training loss for the dynamics model is:
$$
L = - \log \prod_ {j = 1} ^ {N} \hat {p} \left(s _ {t + 1} ^ {j} \mid s _ {t}, a _ {t}\right) + \lambda \sum_ {i, j} \left| \frac {\partial \hat {p} \left(s _ {t + 1} ^ {j}\right)}{\partial s _ {t} ^ {i}} \right|, \tag {3}
$$
where  $\lambda$  is the coefficient for partial derivative regularization.
In this section, we provide a detailed description of the environment, including its semantic stages representing internal progress toward task completion, state space, and action space. We also highlight that while each task consists of multiple semantic stages, agents do not have access to this information. The learning signal for agents is solely based on a sparse reward of 0 or 1, indicating whether the task has been completed or not. Additionally, in each environment, the poses of all environment entities are randomly initialized for each episode.
Meanwhile, as ELDEN focuses on exploring novel local dependencies between environment entities, in all environments, the action space consists of hard-coded skills to increase the probability of entity interactions and bypass navigation challenges under sparse rewards. Extending ELDEN to explore local dependency and learn such skills simultaneously would be an important direction for future work.
Thawing As shown in Fig. 6(a), the Thawing environment consists of a sink, a refrigerator, and a frozen fish. The task requires the agent to complete the following stages: (1) open the refrigerator, (2) take the frozen fish out of the refrigerator, and (3) put the fish into the sink to thaw it. The discrete state space consists of (i) the agent's position and direction, (ii) the positions of all environment entities,
(a) ELDEN (ours)
Figure 5: The dynamics model of each local dependency detection method. (a) The dynamics model of ELDEN for predicting  $s_{t+1}^j$ . Notice that each network predicts  $s_{t+1}^j$  only, and there are  $N$  such networks in total, each responsible for predicting one state factor in  $s_{t+1}$ . For visual simplicity, the “ $\times N$ ” symbol is only shown in (a). (b) pCMI computes  $p(s_{t+1}^j | s_t, a_t)$  and  $p(s_{t+1}^j | s_t \setminus s_t^i, a_t)$  by manually setting the binary mask  $M$  to different values, where  $\otimes$  represents element-wise multiplication. (c) For Input Mask,  $M$  is learned to condition on  $(s_t, a_t)$  and is regularized to use as few inputs as possible. (d) For Attn Mask,  $M$  also conditions on  $(s_t, a_t)$  but is applied to the attention score in the self-attention module.
(iii) the thawing status of the fish, and (iv) whether the refrigerator door is opened. The discrete action space consists of (i) moving to a specified environment entity, (ii) picking up / dropping down the fish, and (iii) opening / closing the refrigerator door.