SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zizhao Wang*
University of Texas at Austin  
zizhao.wang@utexas.edu
Peter Stone†
University of Texas at Austin, Sony Aipstone@cs.utexas.edu
Jiaheng Hu*
University of Texas at Austin jhu@cs.utexas.edu
Roberto Martín-Martín†
University of Texas at Austin robertomm@cs.utexas.edu
Tasks with large state space and sparse rewards present a longstanding challenge to reinforcement learning. In these tasks, an agent needs to explore the state space efficiently until it finds a reward. To deal with this problem, the community has proposed to augment the reward function with intrinsic reward, a bonus signal that encourages the agent to visit interesting states. In this work, we propose a new way of defining interesting states for environments with factored state spaces and complex chained dependencies, where an agent's actions may change the value of one entity that, in order, may affect the value of another entity. Our insight is that, in these environments, interesting states for exploration are states where the agent is uncertain whether (as opposed to how) entities such as the agent or objects have some influence on each other. We present ELDEN, Exploration via Local DepENDencies, a novel intrinsic reward that encourages the discovery of new interactions between entities. ELDEN utilizes a novel scheme — the partial derivative of the learned dynamics to model the local dependencies between entities accurately and computationally efficiently. The uncertainty of the predicted dependencies is then used as an intrinsic reward to encourage exploration toward new interactions. We evaluate the performance of ELDEN on four different domains with complex dependencies, ranging from 2D grid worlds to 3D robotic tasks. In all domains, ELDEN correctly identifies local dependencies and learns successful policies, significantly outperforming previous state-of-the-art exploration methods.
Reinforcement learning (RL) has achieved remarkable success in recent years in tasks where a well-shaped dense reward function is easy to define, such as playing video games [33, 18, 4] and controlling robots [9, 2, 15, 16]. However, for many real-world tasks, defining a dense reward function is non-trivial, yet a sparse reward function based on success or failure is directly available. For such reward functions, learning good policies is often challenging, as it requires efficient exploration of the state space.
To address this challenge, RL researchers proposed the use of an intrinsic reward, an additional task-agnostic signal given to the agent for visiting interesting states. Intrinsic reward methods can be roughly classified into two main paradigms: curiosity [20, 25, 6] and empowerment [29, 27, 14], where the agent is rewarded either for visiting novel states or for obtaining maximal control over the environment, respectively.
Equal supervision
Figure 1: (Left) In a kitchen task with multiple potential agent-object and object-object interactions, (Middle) for a curiosity-based agent interested in hard-to-predict entity motion, it will initially focus on exploring arm movement, then on pot and meatball manipulation, and finally keep rolling the meatball whose outcomes are challenging to predict. On the other hand, for an empowerment-based agent interested in maximizing the action's influence, it begins with controlling the arm and then learning to move the pot and meatball simultaneously, but it ignores the potential interaction between the stove and the meatball. (Right) ELDEN avoids those issues by identifying whether dependencies between entities happen and focusing the exploration on novel ones. After the agent learns that it can control the pot and meatball, it will move on to explore other potential interactions, e.g., whether the stove can influence the meatball. Hence it has a larger opportunity to learn this task, compared with a curiosity or empowerment-based agent.
While these methods significantly improve exploration in some domains, there are cases where the aforementioned methods fail. Consider, for example, a kitchen environment with several objects where there are multiple potential agent-object and object-object interactions, and an agent is tasked with putting a meatball in a pot and cooking it on the stove (Fig. 1). On the one hand, curiosity-driven methods will encourage the agent to explore the environment by visiting states where the exact outcome of an action is uncertain. Consequently, for each interactable object, the agent will exhaust any possible interaction until it can accurately predict every change in the object's state. As a result, such an exploration strategy can be inefficient, especially for environments with many objects. In the kitchen example, it is hard to predict how the meatball rolls in the pot, and thus the curiosity-driven agent would keep rolling it. On the other hand, for empowerment methods, the agent is encouraged to remain in states where it can influence as many states (objects) simultaneously as possible (e.g. holding the pot with the meatball inside). By doing so, however, it ignores object-object interactions that the action cannot directly control but indirectly induce, which can be the key to task completion. In the kitchen case, an empowerment-driven agent will therefore not be interested in placing the pot and the meatball on the stove, as it forfeits control of them by doing so, even though this action enables the stove to heat the meatball. Our main insight is that, in this type of environment, an intelligently exploring agent should be able to learn that it can use the pot to move the meatball after a few trials. Then, instead of spending time learning the complex meatball movement or different styles to manipulate the pot, it would move on to explore other modes of interacting with other objects, e.g., putting the pot on the stove and switching on the stove.
Following this motivation, we propose a new definition of interesting states — focusing on whether the environment entities (consisting of the agents and objects) can interact, rather than how exactly they interact. We present ELDEN, Exploration via Local DepENDencies, a novel intrinsic reward mechanism that models the local dependencies between entities in the scene (agent-object, object-object) and uses the uncertainty about the dependencies to guide exploration. By relaxing the curiosity signal from dynamics prediction to dependencies prediction, ELDEN implicitly biases the exploration toward states where novel interaction modes happen rather than states where the state value is novel but dependencies remain the same. Specifically, ELDEN trains an ensemble of dynamics models. In each model, the local dependencies between objects are modeled by the partial derivatives of state predictions w.r.t. the current state and action. Then, the local dependency uncertainty is measured as the variance across all dynamic models.
We evaluate ELDEN on discrete and continuous domains with multiple objects leading to many interaction modes and tasks with chained dependencies. Our results show that using a partial
derivative-based extractor on dynamic models training allows us to accurately identify local connectivities between environment entities. Furthermore, the intrinsic reward derived from the identified local connectivities allows ELDEN to outperform state-of-the-art exploration methods (curiosity and empowerment-driven) on the tested domains.
The idea of using intrinsic reward to facilitate exploration in reinforcement learning is a long-studied topic that can be dated back to Schmidhuber [25]. In this section, we first discuss two main classes of intrinsic reward related to ELDEN. Since ELDEN requires reasoning about the local dependencies between environment entities, we also discuss works that involve utilizing dependencies/causality in reinforcement learning.
2.1 Curiosity-Driven Exploration
Curiosity-driven exploration rewards an agent for visiting "novel" states, where different methods define the "novelness" of a state in different ways. For methods that utilize visit count to define state novelty, Bellemare et al. [3] utilized a density model to estimate pseudo-count for each state, which is then used to derive intrinsic reward for the agent; Tang et al. [30] uses locality-sensitive hashing code to convert high-dimensional states to hash codes such that the visit count for each hash code can be explicitly kept track of. For methods that utilize predictiveness to define state novelty, Stadie et al. [29] learns a forward dynamics model that operates on a learned latent space of the observation, and uses the prediction error as the intrinsic reward for exploration; Burda et al. [6] uses randomly initialized networks to extract state features, where the agent is encouraged to visit states where predictions about the state features are inaccurate; Pathak et al. [21] utilizes disagreements within an ensemble of dynamic models as a signal for intrinsic reward, and directly backpropagates from the reward signal to the policy parameters to improve exploration. Pathak et al. [20] incorporates empowerment into the curiosity-based method by learning an inverse dynamics model that maps from state to action. The inverse dynamics model defines a feature space for the states, where the prediction error of a forward dynamic model is measured in this feature space.
However, learning accurate dynamics can be difficult and require a significant coverage of a (possibly large) state-action space; e.g., when a robot manipulates a block, it would have to experience multiple action-reaction pairs to be able to make accurate predictions [13]. Furthermore, prior curiosity-driven exploration methods can be derailed by the stochasticity in the dynamics, i.e., when the outcome of an action has large entropy (e.g., tossing a coin), so the agent would keep repeating it. ELDEN can be considered a type of curiosity-driven exploration method. However, unlike previous works which only consider improving knowledge about the dynamics, ELDEN explicitly considers knowledge about local dependencies between environment entities, and utilizes it to encourage exploration. Thus, it avoids the need to learn an accurate dynamic model and is less sensitive to environmental noise and stochasticity.
2.2 Empowerment-Driven Exploration
Empowerment-based exploration methods are based on a different understanding of what states should be encouraged for task-agnostic exploration [35, 27, 8, 19]. Their main idea is that the most interesting states to explore for any task are states where the agent has the most controllable diversity about what the next state will be, i.e., states where there are multiple possible next states that can be chosen by the agent. From those states, it is easier to fulfill any downstream task that requires purposefully changing the state. Empowerment-based exploration methods reason about the controllable elements in the environment (states that can be influenced by agent's actions), and encourage the agent to find states where this controllability is large, typically through some form of mutual information maximization between the agent's actions and the next states. In particular, Zhao et al. [35] uses domain knowledge to divide up the state space into internal state and external state, and maximize the mutual information between them. Seitzer et al. [27] measures the local dependencies between action and environment entities by estimating their conditional mutual information (CMI) and using it as intrinsic reward signal to encourage the agent to maximize the influence of action over environment entities.
However, due to the difficulty in measuring the mutual information across a multi-step trajectory, existing empowerment-based methods only measure 1-step empowerment, e.g., how much the agent directly influences variables. Thus they cannot detect dependencies between objects (e.g., indirect tool use such as using the stove to cook meals) and the environment's influence on the agent (e.g., the elevator can take the agent to the desired floor). Furthermore, since the objective of empowerment-based methods is to maximize controllability, they tend to only control the easiest-to-manipulate object to maximize its empowerment when there are multiple controllable variables.
ELDEN is closely related to Seitzer et al. [27], but differs from it in three main aspects: first, unlike Seitzer et al. [27], which only considers the interaction between action and environment entities, ELDEN also considers the interaction between environment entities that are not locally dependant on the action, which allows curiosity about indirect interaction between the agent and the environment entities to propagate through time during the RL training. Second, unlike Seitzer et al. [27], ELDEN tries to visit states with novel interactions, instead of maximizing controllability, thus avoiding the tendency to only interact with easy-to-control objects. Lastly, ELDEN estimates local dependencies through reasoning about the partial derivatives of learned dynamic models, which we empirically show to be more accurate and computationally efficient compared to the CMI-based estimation in Seitzer et al. [27].
2.3 Causality in Reinforcement Learning
The concept of incorporating causality in the training of a reinforcement learning agent has been utilized in many different forms. Wang et al. [31] demonstrates that incorporating causal dependencies between environment entities can greatly improve generalization to out-of-distribution states. Hu et al. [10] exploited causal dependencies between action dimensions and reward terms to reduce variance in the gradients and facilitate policy learning of mobile manipulation tasks. Pitis et al. [24] shows that knowing the local causal dependencies between objects can facilitate policies to generalize to unseen states and actions. Pitis et al. [23] uses local dependencies to generate counterfactual samples in order to facilitate sample efficiency. Sontakke et al. [28] discovers causal factors in the dynamics of a given environment through a causal curiosity reward term.
Like in Pitis et al. [23], ELDEN learns to predict the local connectivities between environment entities depending on the state and the action values. However, we do not generate counterfactuals but create an intrinsic reward based on the local dependency that facilitates exploration with RL in sparse-reward setups.
3 ELDEN: Exploration via Local Dependencies
In the section, we introduce ELDEN, which infers the local dependencies between environment entities and uses the uncertainty of dependencies as an intrinsic reward for tackling hard-exploration problems. In Sec. 3.1, we formally define the problem setup of ELDEN. In Sec. 3.2, we discuss how ELDEN uncovers local dependencies. In Sec. 3.3, we describe how ELDEN improves exploration with the intrinsic reward.
3.1 Problem Statement
We consider decision-making as a discrete-time Markov Decision Process  $(S, \mathcal{A}, \mathcal{P}, \mathcal{R})$ , where  $S$  is a state space which we assume can factored as  $S = S^1 \times \dots \times S^N$ ,  $\mathcal{A}$  is an action space,  $\mathcal{P}$  is a Markovian transition model, and  $\mathcal{R}$  is a reward function. The goal of the RL agent is to optimize the parameters  $\theta$  of a policy  $\pi_{\theta}$  such that the total expected return under  $\pi_{\theta}$  is maximized. Specifically, we focus on cases where  $\mathcal{R}$  is sparse, and therefore intelligent exploration is crucial to the discovery of optimal policies.
Local Causal Graph Model We can model the transition at time step  $t$  as a Causal Graphical Model (CGM) [22] consisting of (1) nodes  $(\mathcal{S}_t, \mathcal{A}_t, \mathcal{S}_{t+1})$ , (2) a directed graph  $\mathcal{G}$  describing global dependencies between nodes, and (3) a conditional distribution  $p$  for each state variable at the next time step,  $S_{t+1}^n$ . We assume the transition can be factorized as  $\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))$ , where the  $\mathrm{Pa}(v)$  are parents of a node  $v$  in the causal graph  $\mathcal{G}$ . For many environments,  $\mathcal{G}$  can be dense or even fully connected, because whenever it is possible for  $S^j$  to depend on  $S^i$ , no matter how
Algorithm 1 Training of ELDEN (on-policy)
1: Initialize the dynamics ensemble  $\{f\}^{M}$ , policy  $\pi_{\theta}$ , and replay buffer  $D$ .  
2: for number of training iterations do  
3: Collect K environment transitions  $\{s_t, a_t, r_{t,\text{task}}, s_{t+1}\}^{K}$  with current policy  $\pi_{\theta}$   
4: for  $k = 1\dots K$  do  
5:  $\mathcal{G}^{k,m} = f^{m}$ .compute  $\mathcal{G}(s_t^k, a_t^k)$  (Sec. 3.2)  $\triangleright$  Compute the local dependency graph  
6:  $r_{t,\text{intrinsic}}^k = \text{variance}(\{\mathcal{G}^{k,m}\}_{m=1}^M)$   
7:  $r_t^k = r_{t,\text{task}}^k + \beta \cdot r_{t,\text{intrinsic}}^k$   
8: end for  
9: Update the policy  $\pi_{\theta}$  with  $\{s_t, a_t, r_t, s_{t+1}\}^{K}$  (Sec. 3.3)  
10: Add  $\{s_t, a_t, s_{t+1}\}^{K}$  into the replay buffer  $D$   
11: Sample a mini-batch  $B = \{s_t, a_t, s_{t+1}\}_{\text{buff}}^{K}$  from  $D$   
12: Train the dynamics ensemble  $\{f\}^{M}$  with  $B$  (Sec. 3.2)  
13: end for
unlikely, it is necessary to include the edge  $S^i \to S^j$ . However, in the real world, even if possible to interact, most entities are independent of each other most of the time. Following this observation, we are interested in inferring local dependencies that are specific to  $(s_t, a_t)$ , represented by a local causal graph  $\mathcal{G}_t$  that is minimal by removing inactive edges in  $\mathcal{G}$ .
3.2 Identifying Local Dependencies with Dynamics Partial Derivatives
Based on the definition in Sec. 3.1, the key component of ELDEN is to accurately evaluate which potential dependencies between environment entities are locally active, i.e., identify the local causal graph  $\mathcal{G}_t$  given  $(s_t, a_t)$ . This identification requires answering a challenging question — whether entity  $i$ 's value,  $S_t^i = s_t^i$  is the cause of entity  $j$  to have value  $S_{t+1}^j = s_{t+1}^j$ . ELDEN approaches it with the inspiration from the but-for test: the local dependency exists if  $S_{t+1}^j = s_{t+1}^j$  would not happen but for  $S_t^i = s_t^i$ . In other words, we assess whether  $S_{t+1}^j$  would change if  $S_t^i$  has a different value. Notice that, since we focus on local dependencies, we only want to vary  $S_t^i$  near its actual value  $s_t^i$  rather than trying all its possible values.
To this end, ELDEN utilizes partial derivatives to identify local dependencies, as they naturally capture the extent of change in  $S_{t + 1}^j$  with respect to  $S_t^i$ . Specifically, assuming the access of ground truth transition probability  $p$  (which we will relax later), ELDEN considers  $S_{t + 1}^{j} = s_{t + 1}^{j}$  locally depends on  $S_{t}^{i} = s_{t}^{i}$  if
$$
\left| \frac {\partial p \left(s _ {t + 1} ^ {j} \mid s _ {t} , a _ {t}\right)}{\partial s _ {t} ^ {i}} \right| \geq \epsilon , \tag {1}
$$
where  $\epsilon$  is a predefined threshold. A large partial derivative indicates that a slight change in  $S_{t}^{i}$  will lead to a substantial change in  $S_{t + 1}^{j}$ , thus satisfying the but-for test.
To evaluate partial derivatives without the ground truth transition probability, ELDEN approximates  $p$  with a dynamics model  $f$  parameterized by a neural network  $\hat{p}(s_{t+1}^j) = f(s_t, a_t)$  and trains  $f$  by maximizing the log-likelihood of  $\hat{p}(s_{t+1}^j)$  (for notational simplicity, we omit the conditionals  $s_t, a_t$  in  $p$  in this section). Due to limited data and training errors, even when two environment entities are not locally dependent, there occasionally exists a large partial derivative between them. To reduce such false positives, ELDEN further applies regularization to suppress partial derivatives w.r.t. inputs that are not necessary for predicting  $s_{t+1}^j$ . The overall loss of dynamics training is
$$
L _ {f} = - \log \hat {p} \left(s _ {t + 1} ^ {j}\right) + \lambda \sum_ {i, j} \left| \frac {\partial \hat {p} \left(s _ {t + 1} ^ {j}\right)}{\partial s _ {t} ^ {i}} \right|, \tag {2}
$$
where  $\lambda$  is the regularization coefficient.
(a) Thawing
(b) CarWash
(c) Minecraft 2D
Figure 2: We test ELDEN on three domains and four environments. (a) (b) Mini-behavior and (c) Minecraft 2D with discrete state spaces, where the agent has to achieve a series of temporally extended tasks with complex object interactions. (d) Robosuite, a robot table-top manipulation simulation environment with continuous state spaces, where the robot needs to perform multiple interdependent subtasks to finish the cooking task.
(d) Kitchen
3.3 ELDEN Policy Learning
ELDEN utilizes the local dependency identification described in Sec. 3.2 to improve exploration for model-free RL. The key idea behind ELDEN is to encourage an agent to visit states where new local dependencies are likely to emerge. When the ground truth local dependencies are available, the novelty of local dependencies between state variables can be measured by the magnitude of error of the dependencies identified by our method. Unfortunately, in many cases, it is hard to manually specify the ground truth local dependencies. Instead, ELDEN trains an ensemble of dynamics models, and measures dependency novelty by the variance of the local dependency graphs extracted independently from each of the dynamics models in the ensemble. Specifically, ELDEN first computes the variance of each edge in the graph and then uses the mean of the edge variance as the graph variance. Finally, the calculated variance is used as the intrinsic reward and is scaled with a coefficient  $\beta$  that controls the magnitude of the exploration bonus, and added to the task reward  $r_{\mathrm{task}}$ .
ELDEN is applicable to both on-policy and off-policy settings. We show the pseudo-code for on-policy ELDEN in Algorithm 1. The off-policy version of ELDEN can be easily derived by updating the policy with transitions sampled from the replay buffer in line 9 of Algorithm 1. Importantly, any model-free RL algorithm can be used for the policy update step.
In our experiments, we aim to answer two main questions:  $Q1$ : Is ELDEN able to accurately detect local dependencies between environment entities in factored state spaces (Sec. 4.1)?  $Q2$ : Does ELDEN improve the performance of RL algorithms in sparse-reward environments with chained dependencies (Sec. 4.2)?
**Environments** As shown in Fig. 2, we evaluate ELDEN in four simulated environments with different objects that have complex and chained dependencies: (1) CARWASH, (2) THAWING, (4) 2D MINECRAFT and (3) KITCHEN. Both CARWASH and THAWING are long-horizon household tasks in discrete gridworld from the Mini-BEHAVIOR Benchmark [12]. MINECRAFT 2D is an environment modified from the one used by Andreas et al. [1], where the agent needs to master a complex technology tree to finish the task. KITCHEN is a continuous robot table-top manipulation domain implemented in RoboSuite [36]. To complete tasks in these environments and receive the sparse reward, the agent has to conduct a series of actions that change not only the state of the interacted entities but also induce further interaction between interacted entities and others (e.g., interacting with the stove switch that enact interaction between the stove and the cooking the meatball). The agents in all environments can select between a set of action primitives, a set of discrete actions that can be applied on each object, e.g., goTo(obj) or pick(obj). Notice that even with the action primitives, these domains are still very hard to solve due to the presence of many interaction modes and the difficulty in finding the correct (potentially long) sequence of interactions among many options that will lead to task success. We provide further descriptions of each environment in Appendix Sec B.