SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Before comparing different methods based on task performance, we specifically test the quality of initiation sets learned in MINIGRID-FOURROOMS [Chevalier-Boisvert et al., 2018] and the first screen of MONTEZUMA'S REVENGE [Bellemare et al., 2013]. In both domains, observations are  $84 \times 84$  images and the action-space is discrete. For this experiment, we design start states  $S_0$  in each domain—in MINIGRID-FOURROOMS,  $S_0$  is the set of all states (since this is a small tabular domain) and in MONTEZUMA'S REVENGE,  $S_0$  is a series of 100 states scattered across the first screen (more details in the appendix). For the purpose of evaluating our initiation set learning algorithms, at every episode, we sweep through states  $s \in S_0$  and reset the simulator to  $s$ . Option termination conditions are also hand-defined: in FourRooms, we create options that target the center of each room; in Montezuma's Revenge, we define five options: those that navigate the player to each of the two doors, one that navigates the player to the bottom-right and one to the bottom-left of the first screen and finally an option that attempts to get the key.
Initiation Set Accuracy. At each state  $s \in S_0$ , we record the initiation decision made by the learning algorithm as  $\hat{\mathcal{I}}_o(s; \theta) \in \{0, 1\}$ . We then execute the option policy  $\pi_o$  from that state and record whether or not the agent reached the option's subgoal as  $Y_s \in \{0, 1\}$ . Accuracy at state  $s$ , for option  $o$  is then given by  $\mathbb{1}(\hat{\mathcal{I}}_o(s; \theta) = Y_s)$ . This process is repeated several times for all options  $o \in \mathcal{O}$  and start states  $s \in S_0$ . To be faithful to how options are learned in the online RL setting, we only use the trajectory obtained by running the option for updating the initiation classifier and the option policy if the learned classifier returned true at the start state, i.e.,  $\hat{\mathcal{I}}_o(s; \theta) = 1$  (in which case  $\pi_o$  is updated using Rainbow [Hessel et al., 2018]; pseudocode in Appendix C.1).
Initiation Set Size. Although accuracy is a natural evaluation metric, it does not fully capture the quality of the learned initiation set. For example, if the predicted initiation function returns false
Figure 2: Measuring the quality of initiation sets in MiniGrid-FourRooms (left) and the first screen of Montezuma's Revenge (right). Solid lines denote mean accuracy and initiation set size, shaded regions denote standard error. All curves are averaged over all state-option pairs and 5 random seeds.
everywhere (i.e, if  $\mathcal{I}_o(s;\theta) = 0,\forall s\in S_0$ ), then the option policy would not get new data to improve. If the policy never improves, its true initiation set could also collapse; while such an initiation learner would register high accuracy, the option would not be useful. As a result, we additionally measure the normalized size of the "true" initiation set  $|Y_{s}|$ : this is the fraction of start states  $S_0$  from which Monte Carlo rollouts of the policy succeeds. A well-learned initiation set is not only accurate, it is also as large as possible, reflecting a large region of option policy competence.
Competing methods. We compare the accuracy and size of the initiation sets learned by:
- Baseline binary. Binary classifier used to learn the initiation set. This is what is used in essentially all prior work and is the only baseline method in this experiment.
- IVF. Threshold the Initiation Value Function as discussed in Section 3.1.
- Weighted. This is a weighted binary classifier discussed in Section 3.2.
Both the IVF and Weighted approaches use the competence progress exploration bonus described in Section 3.3, its impact is ablated in Appendix E.
Discussion of results. Figure 2 shows that our proposed methods significantly outperform the baseline binary classification approach, which learns classifiers that are less accurate and smaller than those learned using other techniques. Furthermore, the accuracy of the baseline binary classifier falls over time; this is because even though the policy improves (and hence the size of the true initiation set increases), the learned initiation set remains small and unable to adapt to an improving policy. By modifying the binary classifier using the weighting scheme described in Section 3.2, we are able to learn more accurate initiation sets that reflect highly competent option policies. Interestingly, the optimistic version of the IVF is necessary to learn good initiation sets in Montezuma's Revenge, but it underperforms compared to the plain IVF in FourRooms; we hypothesize that this is because FourRooms is a small and simple enough domain that an exploration bonus is not strictly necessary.
4.2 Robot Manipulation: Identifying Promising Grasps
In the previous section, we showed that our methods can lead to higher quality initiation sets; now we evaluate whether that can in turn improve option learning as a whole. We adopt robot manipulation as a testbed because it is a natural problem setting for evaluating initiation set learning algorithms—for example, consider a robot manipulating a hammer: not only should it choose a feasible, stable grasp on the hammer, but it should also choose one that would allow it to subsequently drive a nail. This problem, referred to as task-specific grasping in robotics, and has been studied extensively for a fixed policy [Kokic et al., 2020, Fang et al., 2020, Zhao et al., 2021, Wen et al., 2022, Schiavi et al., 2022]; here we use RL to deduce which grasps are most likely to afford success while simultaneously learning the skill policies. Furthermore, dense reward functions are often necessary to train RL policies in
This is similar to the considerations of precision/recall in supervised learning.
Figure 3: Task success rate for manipulation domains aggregated over 4 random seeds.
robot manipulation domains [Gu et al., 2017]; this section highlights that initiation probabilities can be learned even when the reward function used to train  $\pi_{o}$  is not sparse like the initiation cumulant.
We use three constrained manipulation tasks in ROBOSUITe [Zhu et al., 2020]: opening a door, flipping a lever, and manipulating a sliding mechanism. Each task is modeled as a single option and the policy and initiation set are learned simultaneously. Each environment has 250 possible start states  $S_0$ . At the start of each episode, a configuration is sampled from the initiation set learner:
$$
s _ {0} \sim \frac {\mathcal {I} _ {o} (s)}{\sum_ {s ^ {\prime} \in \mathcal {S} _ {0}} \mathcal {I} _ {o} (s ^ {\prime})}, \forall s \in \mathcal {S} _ {0};
$$
the manipulator is then reset to  $s_0$ . The continuous observation space is 52-dimensional, the continuous action space is 13-dimensional; we use TD3 for policy learning [Fujimoto et al., 2018].
Our algorithms are compared against two baselines: Random, which selects an arm configuration at random at the start of each episode<sup>5</sup>, and Binary which treats initiation learning as vanilla binary classification. Our algorithms, IVF and Weighted, employ the optimistic bias discussed in Section 3.3 by adding a count-based bonus to the predicted initiation probability.
Task success rates are shown in Figure 3. The DOOR task is the simplest: most of the grasp candidates afford task success. As a result, the Random and Binary baselines are competitive with our algorithms. The LEVER and SLIDE tasks are more challenging as the number of promising grasps is significantly lower. In these domains, only the Weighted variant is able to consistently recover a policy with success rate above  $50\%$ .
Figure 4 shows that our methods can identify good grasp poses without task-specific engineering or human demonstrations. This is impressive given that human data [Mandikal and Grauman, 2021, 2022] or heuristics like image segmentation [Kokic et al., 2017, Rosen et al., 2022] are typically required to learn these affordances efficiently, even in the case of a fixed policy. Additional visualizations and ablations can be found in Appendix D and E.
4.3 Improving Option Discovery
Finally, to evaluate whether better option learning can improve option discovery, we integrate our approach into an existing state-of-the-art algorithm: deep skill chaining (DSC) [Konidaris and Barto, 2009, Bagaria and Konidaris, 2020]. DSC learns a collection of options so that the subgoal region  $\beta_{o_i}$  of an option  $o_i$  is the initiation region  $\mathcal{I}_{o_j}$  of another option  $o_j$ ; by learning options that funnel into each other's initiation sets, DSC learns how to sequentially compose options to reliably achieve a goal. It does so by first learning an option that reaches the task goal, then another option that targets the first option's initiation set and so on, until the start-state is inside the initiation set of some option; for more details, please refer to Appendix C.3 and the original papers [Konidaris and Barto, 2009, Bagaria and Konidaris, 2020]. We chose to integrate our techniques with DSC because of its focus on learning initiation sets jointly with options in an online and incremental reinforcement learning setting and because the quality of initiation sets is crucial to its performance and stability [Bagaria et al., 2021a].
5This baseline is equivalent to an initiation set that is uniformly true everywhere.
Figure 4: Examples of promising grasp poses (top row) in DOOR (left), LEVER (middle) and SLIDE (right): these are high probability samples from our initiation function; contrasted with bad grasp poses (bottom row), which are low probability samples and are ruled out by the initiation function.
Figure 5: (left) Comparing the performance of baseline DSC [Bagaria et al., 2021a] to versions of DSC with our initiation learning techniques. Solid lines denote average success rate over 5 random seeds, shaded regions denote standard error. (right) A visualization of initiation sets learned in ANT MEDIUM-MAZE, where the ant robot has to navigate from the bottom-right to the top-left. Each color denotes the initiation set of a different option; although the plot only shows the location of the ant in the maze, the initiation set is learned using the full 30-dimensional state.
Experimental setup. We compare baseline DSC to versions that modify the way in which option initiation sets are learned. We use the ANT MEDIUM MAZE environment where the agent gets a sparse terminating reward of 0 for reaching the goal and  $-1$  every other step; each episode lasts a maximum of 1000 steps [Fu et al., 2020, Todorov et al., 2012]. This is a challenging problem that cannot be solved with flat RL algorithms in a reasonable amount of time. Baseline DSC was able to solve this problem by learning initiation sets defined over a subset of the state  $(x,y)$  location of the ant) [Bagaria et al., 2021a], but here, we do not assume access to such privileged information—all methods learn initiation sets using a dense neural network that maps the full 30-dimensional continuous state to an initiation probability.
We used the author implementation [Bagaria et al., 2021a] where option policies are parameterized using goal-conditioned value functions [Schaul et al., 2015]. When using the IVF directly as the initiation set, we mimic the structure of main agent and use a goal-conditioned value function to represent the IVF; the final initiation set is defined as
$$
\mathcal {I} _ {o} = \left\{s: \max  _ {g \in \beta_ {o}} \mathcal {V} _ {\theta} (s, g) > T, \forall s \in \mathcal {S} \right\}.
$$
More details about the task and implementation details can be found in Appendix A and C.3.
Figure 5 (left) shows that baseline DSC is unable to solve Medium Ant-Maze; but by changing the way options learn their initiation sets, we are able to solve the problem quickly and reliably. The weighted classifier slightly underperforms the pure IVF approach in terms of mean return, but it has lower variance across runs. Figure 5 (right) visualizes the initiation sets learned by the weighted classification approach; the figure makes it clear that different learned options specialize in different parts of the maze.
Learning initiation sets is a critical component of skill discovery, but treating it as binary classification misses key characteristics of the problem. Specifically, it does not address the non-stationarity that results from a continually changing policy, the pessimistic bias of learning initiation sets online, or the temporal structure in the problem. To address these challenges, we proposed the Initiation Value Function (IVF), a general value function tailored specifically to learning initiation sets. We used the IVF directly via thresholding and also via a weighted binary classifier which adapts to changing option policies. Through our experiments, we showed that our proposals lead to higher quality initiation sets, can lead to faster learning of a single option and boost the performance of an existing skill discovery algorithm.
A limitation of the IVF cumulant proposed in Section 3.1 is that only applies for goal-reaching options, which although is quite general, it is not universal. For example, if the option's task is to maximize the velocity of a robot and there is no specific target velocity, then we could not write down a  $0/1$  cumulant that faithfully describes that subtask. Designing cumulants that result in initiation probabilities for general option subtasks is an important avenue for future work.
We thank Mark Rowland and Rafael Rodriguez-Sanchez for useful discussions. We also thank our anonymous reviewers for making thoughtful suggestions that improved the quality of our paper. This research was supported in part by NSF 1955361, NSF CAREER 1844960 to Konidaris and ONR N00014-21-1-2200. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. The content is solely the responsibility of the authors and does not necessarily represent the official views of NSF or the ONR.
Barrett Ames and George Konidaris. Bounded-error LQR-Trees. In 2019 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 144-150. IEEE, 2019.
Barrett Ames, Jeremy Morgan, and George Konidaris. Ikflow: Generating diverse inverse kinematics solutions. IEEE Robotics and Automation Letters, 7(3):7177-7184, 2022.
Bowen Wen, Wenzhao Lian, Kostas Bekris, and Stefan Schaal. Catgrasp: Learning category-level task-relevant grasping in clutter from simulation. In 2022 International Conference on Robotics and Automation (ICRA), pages 6401-6408. IEEE, 2022.
Zheng Wen, Doina Precup, Morteza Ibrahimi, Andre Barreto, Benjamin Van Roy, and Satinder Singh. On efficiency in hierarchical reinforcement learning. In Advances in Neural Information Processing Systems, volume 33, pages 6708-6718, 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/4a5cfa9281924139db466a8a19291aff-Paper.pdf.
Adam White. *Developing a predictive approach to knowledge*. PhD thesis, University of Alberta, 2015.
Martha White. Unifying task specification in reinforcement learning. In International Conference on Machine Learning, pages 3742-3750. PMLR, 2017.
Danfei Xu, Ajay Mandlekar, Roberto Martin-Martin, Yuke Zhu, Silvio Savarese, and Li Fei-Fei. Deep affordance foresight: Planning through what can be done in the future. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 6206-6213. IEEE, 2021.
Jialiang Zhao, Daniel Troniak, and Oliver Kroemer. Towards robotic assembly by predicting robust, precise and task-oriented grasps. In Conference on Robot Learning, pages 1184-1194. PMLR, 2021.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martin-Martin, Abhishek Joshi, Soroush Nasiriany, and Yifeng Zhu. Robosuite: A modular simulation framework and benchmark for robot learning. arXiv preprint arXiv:2009.12293, 2020.
Figure 6: Domains used for our experiments in Section 4. Enumerating left to right from the top-left: first screen of Montezuma's Revenge, MiniGrid FourRooms, Ant Medium Maze, Robosuite Door, Robosuite Lever, and Robosuite Slide.
A Task Descriptions
Four Rooms. This task, which is part of the MINIGRID suite [Chevalier-Boisvert et al., 2018], is an adaptation of the four rooms problem presented in the paper introducing the Options framework [Sutton et al., 1999]. Observations are  $84 \times 84$  images of the fully-observable underlying state-space, which contains approximately  $19 \times 19$  grid locations and 4 possible orientations of the player. We defined 4 options whose termination conditions (subgoals) were to navigate to the center of the 4 rooms.  $S_0$  was the set of all empty grid locations.
Montezuma's Revenge. As is standard in ALE [Bellemare et al., 2013], observations are  $84 \times 84$  images, action space is a set of 18 discrete actions [Machado et al., 2018]. We defined start states scattered across the first room where the player was on the ground (not jumping or falling), was not too close to the skull and did not already have the key in its possession. We also defined 5 options whose termination conditions are reaching the left door, right door, bottom-right of the first screen, bottom-left of the first screen and getting the key.
Robosuite manipulation tasks. Three constrained manipulation tasks were used to study the task-oriented grasping performance of our initiation set learning algorithms: opening a door, flipping a lever, and sliding a knob. The door task was originally implemented in ROBOSUITE [Zhu et al., 2020]; the others were implemented and described in the work of Rosen et al. [2022]. All three are 1-DoF articulated objects which require making and sustaining contact to manipulate. The 52-dimensional observation space consists of the robot's proprioceptive state (joint position, joint velocity, end-effector position, gripper state, tactile data) as well as the object state (object pose, joint position, handle position). The action space employed is operational space control with variable impedance [Martín-Martín et al., 2019]: the agent controls the 6-DoF change in position and orientation of the end-effector, the 6-DoF change in stiffness, and 1-DoF gripper state. Episodes have a maximum length of 250 steps. In each task,  $S_0$  was a set of arm configurations establishing contact with the object; see Section C.2.
Ant Medium Maze. The goal location is small region around (20, 20). A state is considered to satisfy a goal if the two have a euclidean distance of 0.5 units or less  $R(s, g) = ||s - g||_2 < 0.5$ . The agent is evaluated by rolling out the learned policy once every 10 episodes; during evaluation,
Figure 7: Illustration of the pessimistic bias in a tabular MDP. Action  $o_{t}$  has non-stationary effects, just like an option in HRL. As  $o_{t}$  executed more, the probability of landing in  $s_{3}$  increases, otherwise it leaves the agent in  $s_{1}$ .
the agent starts from a small region around  $(0,0)$ , during training it starts at a location randomly sampled from the open locations in the maze. The task reward function is  $-1$  for non-goal transitions and a terminating reward of 0 for reaching the goal. Episodes last a maximum of 1000 steps. The training and evaluation protocol is identical to Bagaria et al. [2021a], except for the fact that we learn initiation sets over the full state.
B Pessimistic Bias: Connections to Non-Hierarchical MDPs
MDPs with actions that have non-stationary effects mirror hierarchical MDPs in which the option policies are being learned online. Figure 7 shows such an MDP—there are three states:  $s_1, s_2, s_3$ , two actions  $a$  and  $o_t$ , and discount factor  $\gamma < 1$ . Action  $a$  takes the agent one state to the right, state  $s_3$  yields a terminating reward of 1, resetting the state to  $s_1$ . Action  $o_t$  has a non-stationary effect: it either leaves the agent at state  $s_1$ , or takes the agent to  $s_3$ . The probability with which  $o_t$  takes the agent to  $s_3$  evolves over time—the more it is executed, the more consistently it takes the agent to  $s_3$ , similar to an option in HRL whose policy improves with execution.
A Q-learning agent quickly learns  $Q(s_{1},a) = \gamma$  and  $Q(s_{1},o_{t}) < \gamma$  leading to greedy policy  $\pi (s_1) = a$ . But if the agent executes  $o_t$  enough times, it will get higher return  $Q(s_{1},o_{t}) = 1 > Q(s_{1},a) = \gamma$  leading to the greedy policy  $\pi (s_1) = o_t$ . Vanilla Q-learning alone would not discover the higher value solution because it does not account for non-stationary effect of  $o_t$ ; some form of exploration is needed for that.
C Details about Learning Algorithm
C.1 Accuracy Experiment
Algorithm 1 is the pseudocode used for the experiments described in Section 4.1. Every episode, every option is executed from every start state in  $S_0$ . The result of that execution is recorded as ground-truth  $Y_{s,o}(t)$  and stored to later compute the size of the true initiation set  $|Y| = \frac{\sum_{s,o,t} Y_{s,o}(t)}{|Y_{s,o}(t)|}$ . If the start was predicted to be inside the initiation set by the learning algorithm, then the trajectory generated by rolling out the option policy is used to update the policy and the initiation learner (e.g., IVF, classifier). We report the agreement between the predicted initiations and the ground-truth as an accuracy measurement for that state-option pair. The fraction of start states in  $S_0$  that lead to success is reported as a measurement of the size of the true initiation set.
C.2 Robot Manipulation
The task-specific grasping problem is typically phrased as identifying grasp poses  $g \in \mathbb{S}\mathbb{E}(3)$  that afford task success. In practice, the difficulty of this problem is compounded by the fact that, for redundant manipulators, each grasp pose  $g$  yields an infinite number of corresponding arm
Inputs: Option termination conditions  $\beta_{o},\forall o\in \mathcal{O}$  , start states  $S_0$  , number of episodes n Episodes.
Outputs: Accuracy table  $A$  and Initiation Size table  $S$ ; both map state-option pairs to a list of booleans.
Initialize goal-conditioned policy  $\pi_{\theta}:\mathcal{S}\times \mathcal{G}\to a\in \mathcal{A}$
Initialize Initiation Value Function (IVF)  $\mathcal{V}_{\phi}:\mathcal{S}\times \mathcal{G}\to \mathbb{R}$
Initialize binary classifier for each option  $\mathcal{I}_o(s;\psi)\to \{0,1\}$
Initialize replay buffers for Rainbow  $B_{R}$  and IVF  $B_{I}$ .
Initialize buffers to store positive and negative examples for each option's initiation classifier.
Initialize tables  $A$  and  $S$  as mapping each state-option pair to an empty list.
Algorithm 1 Accuracy/Size Experiment Procedure
for episode  $\in$  range(n Episodes) do
for start state  $s_0\in S_0$  do   
for option  $o\in \mathcal{O}$  do Reset simulator to  $s_0$  Record option o's initiation decision  $X = \mathcal{I}_o(s_0;\psi)$  Rollout option policy  $\pi_o(s_0,g\sim \beta_o)$  to get trajectory  $\tau$  and next state  $s^{\prime}$  Record whether the option policy reached the goal  $Y = s^{\prime}\in \beta_{o}$  Record accuracy  $A[s_0][o]$  . append(1(X=Y)). Update ground-truth size table  $S[s_0][o]$  . append(Y). if predicted initiation  $X = 1$  then Add trajectory  $\tau$  to policy's replay  $B_{R}$  Relabel trajectory  $\tau$  with initiation cumulant  $c_{0}:S\to \{0,1\}$  Add relabeled trajectory to IVF's replay  $B_{I}$  Add trajectory  $\tau$  to o's positive/negative example buffer. end if end for   
end for   
end for   
Sample minibatch and update  $\pi_o$  using Rainbow. Sample minibatch and update  $\mathcal{V}_{\phi}$  using TD(0).   
for option  $o\in \mathcal{O}$  do Compute weights  $w(s)$  for all training examples using Equation 1. Update o's initiation classifier by minimizing weighted cross-entropy loss.   
end for   
end for