SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Akhil Bagaria*
Providence, RI, USA.
akhil_bagaria@brown.edu
Ben Abbatematteo*
Brown University,
abba@brown.edu
Omer Gottesman†
Amazon,
New York, NY, USA.
omergott@gmail.com
Matt Corsaro
matthew_corsaro@brown.edu
Sreehari Rammohan
sreehari_rammohan@brown.edu
George Konidaris
gdk@cs.brown.edu
An agent learning an option in hierarchical reinforcement learning must solve three problems: identify the option's subgoal (termination condition), learn a policy, and learn where that policy will succeed (initiation set). The termination condition is typically identified first, but the option policy and initiation set must be learned simultaneously, which is challenging because the initiation set depends on the option policy, which changes as the agent learns. Consequently, data obtained from option execution becomes invalid over time, leading to an inaccurate initiation set that subsequently harms downstream task performance. We highlight three issues—data non-stationarity, temporal credit assignment, and pessimism—specific to learning initiation sets, and propose to address them using tools from off-policy value estimation and classification. We show that our method learns higher-quality initiation sets faster than existing methods (in MINIGRID and MONTEZUMA'S REVENGE), can automatically discover promising grasps for robot manipulation (in ROBOSUITE), and improves the performance of a state-of-the-art option discovery method in a challenging maze navigation task in MuJoCo.
Temporal abstraction, which is crucial for scaling RL to long-horizon problems [Konidaris, 2019, Sutton et al., 2022], is elegantly captured using the options framework [Sutton et al., 1999]. Unlike primitive actions that only last for a single timestep, options execute a policy until some termination condition is met. General purpose agents must possess many such options to solve a wide variety of tasks [White, 2017], but increasing the size of the action set raises both the complexity of learning and the branching factor of search [Littman et al., 1995]. This necessitates a mechanism for locally pruning irrelevant options: most options do not apply in a given state and the agent can lower its effective branching factor by ignoring them [Precup et al., 1998, Wen et al., 2020]. Fortunately, the options framework includes the concept of initiation sets which captures the idea of local applicability: an option can only be executed from states inside the initiation set.
It is natural to view an initiation set as the set of states from which option execution is likely to succeed; in that case, learning it is straightforward when the option's policy is fixed: run the
$^{\dagger}$ Work done while at Brown University.
option from different states, record the outcome, and learn a classifier that predicts whether option execution will succeed at a given state [Konidaris and Barto, 2009, Kaelbling and Lozano-Perez, 2017]. However, this is complicated by non-stationarity during learning: as the option policy changes, previously recorded outcome labels become invalid. How should the initiation classifier change in response? Another challenge is that initiation set learning suffers from a pessimistic bias: the option policy mostly improves for states inside the initiation set because option execution only occurs from such states. Therefore, if a state is excluded from the initiation set early during learning, it will likely remain excluded even as the policy improves over time. This pessimistic bias causes initiation sets to mostly shrink and rarely expand [Bagaria et al., 2021a], even in the face of policy improvement; this is potentially catastrophic for learning since the agent prematurely gives up on options in favor of primitive actions [Harb et al., 2018].
To address the issue of non-stationarity, we introduce the Initiation Value Function (IVF), a general value function (GVF) [Sutton et al., 2011, White, 2015] that predicts the probability that option execution will succeed from a given state. Since the IVF is purely predictive, it can be learned using tools from off-policy policy evaluation (OPE) [Voloshin et al., 2021] and does not interfere with the option's main task of maximizing its internal reward function [Sutton et al., 2023]. Unlike the classification approach described above, the IVF adapts to a changing policy: as the policy improves over time, so does the IVF's estimates of its probability of success. We show how the IVF can be used as a direct estimate of the option's initiation set or as input to a weighted classifier that accounts for changes in the option policy.
To address the pessimistic bias of learning initiation sets, we expand the criterion for including a state in the option's initiation set: in addition to states from which option execution is likely to succeed, we also include states for which the option policy is most likely to improve. By identifying and adding these states to the initiation set, we mitigate the pessimistic bias and prevent initiation sets from collapsing as the option is learned.
For evaluation, we first measure how accurately we can deduce an option's initiation set in MINIGRID-FOURROOMS and the first screen of MONTEZUMA'S REVENGE. Then, we demonstrate that our proposed methods can effectively identify promising grasps in challenging robot manipulation problems in ROBOSUITE. Finally, by integrating our method for learning initiation sets into an existing option discovery algorithm, we solve a maze navigation problem in MuJoCo that baseline agents are unable to.
As is standard in HRL [Barto and Mahadevan, 2003], we consider problems that can be modeled as Semi Markov Decision Processes (SMDPs)  $\mathcal{M} = (\mathcal{S},\mathcal{A},\mathcal{R},\mathcal{T},\gamma)$ , where  $\mathcal{S}$  is the state space,  $\mathcal{A}$  is the action space (which contains options and primitive actions),  $\mathcal{R}$  is the reward function,  $\mathcal{T}$  is the transition function and  $\gamma$  is the discount factor. An option  $o\in \mathcal{O}\subseteq \mathcal{A}$  models temporally extended behaviors;  $o = (\mathcal{I}_o,\pi_o,\beta_o)$ , where  $\mathcal{I}_o$ , the initiation set of the option, describes the states from which the option can be executed,  $\beta_{o}$  is the termination or subgoal region and  $\pi_o$  is the policy [Sutton et al., 1999]. When using RL to learn the option policy, an internal option subgoal reward function  $\mathcal{R}_o$  and timescale  $\gamma_{o}$  completes the subtask description of the option and is used to train  $\pi_o$  [Precup, 2000, Barto and Mahadevan, 2003, White, 2017].
An important problem in HRL is that of option discovery, where algorithms focus on identifying option termination conditions (subgoals) and usually assume that all options apply everywhere [Dayan and Hinton, 1993, Bacon et al., 2017, Eysenbach et al., 2019, Machado et al., 2017]; as previously discussed, this is not a scalable strategy for developing general-purpose agents. We hypothesize that well-learned initiation sets will improve the quality of options, regardless of how they are discovered, and will eventually lead to more performant hierarchical RL agents.
General Value Functions. Value-based RL typically predicts and maximizes a single scalar reward, but a value function can be generalized to predict (and sometimes control [Jaderberg et al., 2017]) the discounted sum of any real-valued "cumulant" that can be computed from the agent's observations [Sutton et al., 2011, White, 2015]. Such Generalized Value Functions (GVFs) efficiently solve the temporal prediction tasks and can represent rich knowledge about the world [Schaul and Ring, 2013].
Initiation Sets in Robotics and Planning. Initiation sets are crucial in classical planning [Fikes et al., 1972] as well as task and motion planning (where they are referred to as preconditions) [Garrett et al., 2021], but they are usually designed by a domain expert [Kaelbling and Lozano-Pérez, 2017]. In robotics (affordances), they are usually programmed to avoid collisions [Xu et al., 2021] or to correspond to relevant object attributes [Şahin et al., 2007, Huang et al., 2023]. In control theory (regions of attraction), they denote states where policies are provably stable [Tedrake, 2009, Ames and Konidaris, 2019]. While our proposed methods do not come with finite-time convergence guarantees, they are arguably more general and scalable because they do not assume pre-programmed policies or additional structure such as objects and robot dynamics.
Techniques for learning initiation sets fall into 3 categories: classification, value-based and end-to-end.
Classification approach. When initiation sets are learned, it is most often framed as binary classification [Konidaris and Barto, 2009, Bagaria and Konidaris, 2020, Khetarpal et al., 2020, Bagaria et al., 2021b]. This approach is sound when the option policy is fixed [Konidaris et al., 2018] or when the affordances correspond to primitive actions [Khetarpal et al., 2020]. But, this approach does not scale to the continual learning setting [Ring, 1995, Mcgovern, 2002] where new, temporally extended options keep getting discovered and their policies and initiation sets must be learned simultaneously.
Value function approach. SayCan grounds large language models using value functions of skills trained using offline RL [Brohan et al., 2023]. Since the policies are pre-trained, they do not confront the challenges of jointly learning policies and their initiation sets discussed in this paper. Nica et al. [2022] learn affordances via Monte Carlo estimation (which is equivalent to binary classification under a restrictive  $0/1$  reward function), but do not account for the pessimistic bias or the temporal structure in the initiation set learning problem. Furthermore, both SayCan and Nica et al. [2022] design the option reward function so that the resulting value function can be interpreted as an affordance/initialization probability; this is a strong restriction because RL is famously sensitive to reward design [Randlov and Alström, 1998]. Our method allows the designer to pick the reward function appropriate for their task and the initiation function is learned using the appropriate cumulant using policy evaluation. Relay Networks [Kumar et al., 2018] learn initiation regions by thresholding the critic, but have difficulty picking a threshold because it is not possible to interpret arbitrary value functions as initiation probabilities.
End-to-end approaches. IoC [Khetarpal and Precup, 2019] learns initiation sets in the Option-Critic (OC) framework; GrASP [Veeriah et al., 2022] learns affordances that are useful for Monte Carlo Tree Search (MCTS). These are promising gradient-based approaches to learning affordances, but are specific to OC and MCTS respectively. Furthermore, they both learn affordances that maximize task reward, we instead focus on learning options that each represent their own subtask [White, 2017] and could later be composed to solve many downstream tasks [Barreto et al., 2019].
3 Effectively Learning Initiation Sets
The problem of learning initiation sets can be naturally framed as training a classifier where states on trajectories that achieve the subgoal are labeled as positive examples, and states on trajectories that fail to achieve the subgoal are labeled as negative examples. A probabilistic classifier trained on this data will predict the probability that an option will succeed from a given state, which is exactly the desired semantics of an initiation set. However, when initiation sets and policies have to be learned together, such a classifier is no longer a good estimator, for three reasons:
Data non-stationarity. The first reason is data non-stationarity, which refers to the phenomenon that previously collected training examples become invalid during the course of learning, but continue to impact classifier training. Consider a trajectory  $\tau_0$  obtained by rolling out an option policy  $\pi_o^{t_0}$  at time  $t_0$ . If the policy is in the early stages of training,  $\pi_o^{t_0}$  will likely fail to reach the subgoal region  $\beta_o$  and  $\tau_0$  will be labeled as a negative example  $(Y(s) = 0, \forall s \in \tau_0)$ . At some future time  $t_1$ , the policy might improve, but the states along  $\tau_0$  will continue to have a negative label, which means that they will continue to be outside the initiation set. As a result, we will fail to capture the growing competence/reachability of  $\pi_o^{t > t_0}$ . A naive solution to this problem would be to discard old training
Figure 1: Learning an initiation set for an option targeting subgoal  $\beta_{o}$  (star). (left) Temporal structure in initiation set learning: even though  $\tau_{1}$  (black) fails to reach the goal, states along it should get credit because of the portion shared with the successful trajectory  $\tau_{2}$  (red); this does not happen using a simple classification approach. (middle) the initiation value function (IVF)  $\mathcal{V}_{o}$  (Sec 3.1) is shown in green: darker the shade, higher the value; classification examples should be re-weighted to account for a changing policy. (right) Pessimistic bias causes initiation sets to shrink (blue  $\rightarrow$  red) over time.
examples, but not only would that be sample-inefficient, it is also unclear at what point a training example is "old enough" to be discarded.
Pessimistic bias. The second reason, pessimism, refers to the fact that once states are deemed to be outside the initiation set, they are unlikely to switch labels because the option has low probability of being executed from them, and so policy improvement from these states is unlikely (Figure 1 (right)).
Temporal Structure. The third reason is that classification (or equivalently, Monte Carlo learning [Sutton and Barto, 2018]) does not exploit the temporal structure in the initiation set learning problem. Figure 1 (left) shows two trajectories sampled from the same option's policy:  $\tau_{1}$  reaches the goal and  $\tau_{2}$  does not. Any form of Monte Carlo estimation using this data would assign an initiation probability of 0 to most states in trajectory  $\tau_{1}$ . However, this does not assign enough credit to most of the good decisions made along  $\tau_{1}$ —the only way to achieve that would be via temporal bootstrapping. This form of generalization is not possible when using classification.
To address non-stationarity and to exploit temporal structure, we formulate a new general value function, the Initiation Value Function (IVF; Section 3.1), learned using temporal difference (TD) methods [Sutton, 1988]. To address pessimism, we augment the initiation set with states from which the option policy is most likely to improve (Section 3.3).
3.1 Initiation Value Function (IVF)
An option  $o$ 's initiation set  $\mathcal{I}_o$  is the set of states from which the option policy  $\pi_o$  can succeed with high probability:  $\mathcal{I}_o = \{s : \mathbb{P}(s' \in \beta_o | S = s, A = o) > T\}$  ( $T$  is a predefined threshold). Treating the success condition of the option as a cumulant  $c_o : s \to \{0, 1\}$  and setting the timescale  $\gamma_{c_o} := 1$ , the corresponding GVF
$$
\mathcal {V} ^ {\pi_ {o}} (s _ {t}) = \mathbb {E} _ {\pi_ {o}} \Big [ \sum_ {t = 0} ^ {t = H _ {o}} c _ {o} (s _ {t + 1}) | s _ {t + 1} \sim \mathcal {T} (s _ {t}, \pi_ {o} (s _ {t})) \Big ] = \mathbb {P} (c _ {o} = 1 | S = s _ {t}, O = o)
$$
represents the initiation probability at state  $s_t$  ( $H_{o}$  is the option horizon). Note that  $\mathcal{V}^{\pi_o}$  is different from  $V_{o}$ , the value function used by the option policy for control; this is because  $\pi_{o}$  maximizes an arbitrary reward function  $\mathcal{R}_{o}$  and so the resulting value function  $V_{o}$ , which approximates the value of the optimal option policy  $V^{\pi_{o}}^{*}$ , cannot be interpreted as an initiation probability.
Using the IVF as an initiation set. To directly use the IVF as the initiation set, we need to pick a threshold  $T$  above which a state is deemed to be in the option's initiation set. Previous work [Kumar et al., 2018] reused the option value function  $V_{o}$  as the IVF, but had to develop heuristic schemes to pick a threshold. But since  $\mathcal{V}_{\phi}^{\pi_o}$  outputs an interpretable number between 0 and 1 (we use a sigmoid layer at the end of  $\phi$  to enforce this), it is easy to threshold, regardless of  $\mathcal{R}_o$ .
Learning the IVF. Computing the probability of success of a given policy  $\pi_o^t$  is equivalent to the problem of off-policy policy evaluation. Our approach to policy evaluation must be sample-efficient if it is not, the option policy will change before we are able to sufficiently evaluate it [Sutton et al.,
2007]. Furthermore, the initiation cumulant  $c_{o}$  is a sparse binary function, which makes learning with the Monte Carlo estimator high variance and sample inefficient. This once again underscores the importance of exploiting temporal structure of the problem using TD-learning (we use TD(0)) to estimate  $\mathcal{V}_{\phi}^{\pi_o}$ , which unlike a classifier, is able to propagate value from partial trajectories (as illustrated in Figure 1 (left)). In this view, using a classifier is equivalent to using a Monte Carlo estimate of the IVF: high-variance, sample inefficient, and unable to bootstrap.
3.2 Combining Classification and the Initiation Value Function (IVF)
In Section 3.1 we connected the initiation set to a GVF, which allows us to use all the tools of value-based RL [Sutton and Barto, 2018] to learn the initiation set. However, the classification approach mentioned earlier has some attractive qualities: specifically, classification is easier than regression [Bishop and Nasrabadi, 2006] and the supervised learning community has developed powerful tools for learning classifiers using neural networks and cross-entropy loss [Goodfellow et al., 2016]. To get the best of both worlds, we propose an additional method that combines value estimation and classification.
Recall that binary classification can be described as the process of finding parameters  $\theta$  that minimize the cross-entropy loss  $\mathcal{L}(\theta)$  in the training data  $\mathcal{D} = \{(s_i,Y_i)\}_{i = 0}^{|\mathcal{D}|}$  where  $s_i$  are the inputs (states in our case) and  $Y_{i}$  are the classification labels (whether option execution was successful from  $s$ ). To deal with non-stationarity, we instead use the weighted binary cross-entropy loss  $\mathcal{L}_w(\theta)$  which minimizes loss over a dataset  $\mathcal{D} = \{s_i,Y_i,w_t(s_i)\}$  where  $w_{t}(s_{i})$ , a state-dependent weight, represents the desired contribution of the training example to the classification loss [Natarajan et al., 2013].
How should we set the weights  $w_{t}(s)$  of a training example  $(s_i, Y_i)$  in a way that reflects the evolving competence of the option policy? Since the IVF evaluates the success probability of the current option policy  $\pi_{o}^{t}$ , we can heuristically use it as a reflection of our certainty in the training example  $(s_i, Y_i)$ —positively (negatively) labeled states contribute more to  $\mathcal{L}_w$  when the IVF prediction is high (low). Specifically, the weighting function is defined via the following simple equation:
$$
w _ {t} (s) = Y _ {s} \mathcal {V} (s) + \left(1 - Y _ {s}\right) \left(1 - \mathcal {V} (s)\right). \tag {1}
$$
As the IVF estimates that the option policy is becoming more competent at state  $s$ , the classifier's uncertainty about a negative example at  $s$  will increase, thereby causing that negative example to contribute less to the classifier's decision boundary. Similarly, if the policy degrades at  $s$ , then the contribution of a positive label at  $s$  to  $\mathcal{L}_w$  will go down. By weighing training examples in this way, a binary classifier is able to adapt to data non-stationarity that results from a changing option policy.
At the end of every option execution, we therefore recompute weights using Equation 1 and re-train option initiation classifiers using a weighted cross-entropy loss  $\mathcal{L}_w$ .
Policy Improvement Prior. We can use the fact that policies in RL usually improve over time to guide classifier training. To do this, we set the weight of positive examples  $w_{t}(s)$  as 1 throughout training. The underlying idea is that a state which leads to success once is likely to do so again with an even better policy. The weights of negative examples are allowed to vary with the IVF as in Eq 1.
3.3 Overcoming Pessimistic Bias
Unsuccessful option trajectories cause the initiation set to shrink. When a state is outside the initiation set, the option can no longer be executed from there. So even if the option could succeed from that state in the future, it remains outside the initiation set and the option ends up with a smaller region of competence than it needs to. This issue, which we call the pessimistic bias of learning initiation sets, can prevent the learning of useful options and lowers the effectiveness of hierarchical RL agents<sup>3</sup>.
To mitigate this pessimistic bias, we expand the initiation set to also include states from which policy improvement is most likely:
$$
\mathcal {I} _ {o} = \{s: \mathcal {V} _ {o} (s) + \mathcal {B} (s) > T, s \in \mathcal {S} \}.
$$
Effectively identifying such states in high-dimensional spaces is the topic of bonus-based exploration in deep RL [Taiga et al., 2019]. We propose two simple approaches:
This issue is analogous to non-stationary dynamics in flat MDPs; see Section B in the Appendix.
1. Competence progress attempts to capture regions where a policy is either improving or regressing [Şimşek and Barto, 2006, Stout and Barto, 2010, Baranes and Oudeyer, 2013]. This can be computed as changes in the IVF over time:  $\mathcal{B}_1(s) = \left|\mathcal{V}_o^t (s) - \mathcal{V}_o^{t - K}(s)\right|$ , where  $\mathcal{V}_o^t$  is the current IVF and  $\mathcal{V}_o^{t - K}$  is the IVF estimate  $K$  timesteps ago (obtained using the target network).
2. Count-based bonus approach keeps track of the number of times  $N(s,o)$  option  $o$  has been executed from state  $s$ . This is then converted into an uncertainty measure:  $\mathcal{B}_2(s) = c / \sqrt{N(s,o)}$ , where  $c$  is a scalar hyperparameter [Strehl and Littman, 2008].
We use count-based bonuses in problems where tabular counts are readily available; otherwise, we use competence progress.
We aim to evaluate whether our methods can improve hierarchical RL. First, we evaluate whether they result in better, more efficiently learned initiation sets. Second, we test if better initiation set learning improves option learning as a whole. Finally, we incorporate our changes into a state-of-the-art skill discovery method and check if resulting agent is able to outperform the baseline in a sparse-reward continuous control problem.
Implementation Details. Option policies are learned using Rainbow [Hessel et al., 2018] when the action-space is discrete and TD3 [Fujimoto et al., 2018] when it is continuous. Following Bagaria et al. [2021a], all options share the same UVFA [Schaul et al., 2015] but condition it using their own subgoals. The IVF is learned using Fitted Q-Evaluation [Le et al., 2019], prioritized experience replay [Schaul et al., 2016] and target networks [Mnih et al., 2015]. The IVF Q-function and initiation classifier are parameterized using neural networks that have the same architecture as the Rainbow/TD3. Each option has a "gestation period" of 5 [Konidaris and Barto, 2009], which means that before the option sees 5 successful trajectories, its initiation set is optimistically initialized to be true everywhere. Since the training data for classification can be severely imbalanced, we use the standard technique of upweighting the minority class [Japkowicz and Stephen, 2002].
4.1 Measuring the Quality of Initiation Sets