SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kalle Kujanpää<sup>1,3*</sup>; Joni Pajarinen<sup>2,3</sup>, Alexander Ilin<sup>1,3,4</sup>
$^{1}$ Department of Computer Science, Aalto University
$^{2}$ Department of Electrical Engineering and Automation, Aalto University
$^{3}$ Finnish Center for Artificial Intelligence FCAI
4System 2 AI
{kalle.kujanpaa,joni.pajarinen,alexander.ilin}@aalto.fi
Solving complex planning problems has been a long-standing challenge in computer science. Learning-based subgoal search methods have shown promise in tackling these problems, but they often suffer from a lack of completeness guarantees, meaning that they may fail to find a solution even if one exists. In this paper, we propose an efficient approach to augment a subgoal search method to achieve completeness in discrete action spaces. Specifically, we augment the high-level search with low-level actions to execute a multi-level (hybrid) search, which we call complete subgoal search. This solution achieves the best of both worlds: the practical efficiency of high-level search and the completeness of low-level search. We apply the proposed search method to a recently proposed subgoal search algorithm and evaluate the algorithm trained on offline data on complex planning problems. We demonstrate that our complete subgoal search not only guarantees completeness but can even improve performance in terms of search expansions for instances that the high-level could solve without low-level augmentations. Our approach makes it possible to apply subgoal-level planning for systems where completeness is a critical requirement.
Combining planning with deep learning has led to significant advances in many fields, such as automated theorem proving [33], classical board games [39], puzzles [1], Atari games [38], video compression [24], robotics [5], and autonomous driving [43]. However, deep learning-based methods often plan in terms of low-level actions, such as individual moves in chess or commands in Atari games, whereas humans tend to plan by decomposing problems into smaller subproblems [11]. This observation has inspired a lot of recent research into methods that solve complex tasks by performing hierarchical planning and search. In the continuous setting, hierarchical planning has been applied to, for instance, robotic manipulation and navigation [31, 25, 12, 6, 15]. Discrete hierarchical planning methods, when trained entirely on offline data, can solve complex problems with high combinatorial complexity. These methods are efficient at long-term planning thanks to the hierarchy reducing the effective planning horizon [4, 18, 45]. Furthermore, they reduce the impact of noise on planning [4] and show promising generalization ability to out-of-distribution tasks [45]. Training RL agents on diverse multi-task offline data has been shown to scale and generalize broadly to new tasks, and being compatible with training without any environment interaction can be seen as an additional strength of these methods [20, 5, 42].
(a) low-level search
(b) high-level search
Figure 1: (a): Low-level search systematically visits all states reachable from the root node and, therefore, is guaranteed to find a path to the terminal state (shown in red). (b): High-level search can find a solution with much fewer steps due to its ability to operate with high-level actions that span several time steps. (c): High-level search may fail to find a solution due to an imperfect subgoal generation model. (d): Complete subgoal search can find a solution with few steps thanks to high-level actions, and it has completeness guarantees due to also considering low-level actions.
(c) high-level search
(d) complete search
Despite their excellent performance, the discrete subgoal search methods share a weakness: neither AdaSubS [45], kSubS [4], nor HIPS [18] are guaranteed to find a solution even if it exists. All these methods rely on a learned subgoal generator that proposes new subgoals to be used with some classical search algorithms for planning. If the generator fails to perform adequately, the result can be a failure to discover solutions to solvable problems. We call this property the lack of completeness. The guarantee of the algorithm finding a solution is critical for both theoretical science and practical algorithms. In the context of subgoal search, completeness guarantees discovering the solution, even if the generative model is imperfect. Completeness allows extensions and new incremental algorithms that require building on top of exact solutions. One example of that could be curriculum learning. Furthermore, we show that completeness significantly improves the out-of-distribution generalization of subgoal search. Achieving completeness also makes applying high-level search as an alternative to low-level search possible in safety-critical real-world systems. Zawalski et al. [45] mentioned that AdaSubS can be made complete by adding an exhaustive one-step subgoal generator that is only utilized when the search would otherwise fail, whereas Kujanpää et al. [18] proposed combining high- and low-level actions to attain completeness of HIPS. However, to the best of our knowledge, this idea has not been formalized, analyzed, or evaluated in prior work.
We present a multi-level (hybrid) search method, complete subgoal search, that combines hierarchical planning and classical exhaustive low-level search. The idea is illustrated in Figure 1. We apply our complete search approach to HIPS [18] which 1) has been shown to outperform several other subgoal search algorithms in challenging discrete problem domains, 2) does not require training multiple subgoal generators in parallel to perform adaptive-length planning, and 3) is particularly suited for being combined with low-level actions due to using Policy-guided Heuristic Search (PHS) as the main search algorithm [18, 27]. We use the name HIPS- $\varepsilon$  for the proposed enhancement of HIPS with complete subgoal search. We argue that the idea suggested in [45] can be seen as a specific case of our search method that works well in some environments but can be improved upon in others. We evaluate HIPS- $\varepsilon$  on challenging, long-term discrete reasoning problems and show that it outperforms HIPS, other subgoal search methods, and strong offline reinforcement learning (RL) baselines. Not only does HIPS- $\varepsilon$  attain completeness, but our complete subgoal search also demonstrates improved performance on problem instances that subgoal search alone was sufficient for solving.
We propose to augment the high-level search used in learning-based hierarchical planning with low-level actions to achieve completeness while even improving the search performance in terms of node expansions. First, we discuss our work in the general context of hierarchical planning and then focus on the closest methods in discrete subgoal search, to which we compare the proposed approach.
Continuous Hierarchical Planning Hierarchical planning has been widely applied in the continuous setting. The Cross-Entropy Method (CEM) [34] is often used as the planning algorithm for different hierarchical planning methods [22, 25, 12, 32]. However, the cross-entropy method is a numerical optimization method, unsuitable for addressing discrete planning problems. There is also prior work on hierarchical planning for continuous and visual tasks, including control and navigation,
where CEM is not used as the optimization algorithm [21, 26, 6, 31, 36, 46]. Kim et al. [14] plan how to select landmarks and use them to train a high-level policy for generating subgoals. Planning with search is a common approach in classical AI research [35]. Combining search with hierarchical planning has been proposed [23, 8], but these approaches are generally inapplicable to complex, discrete reasoning domains, where precise subgoal generation and planning are necessary [45].
Recursive Subgoal-Based Planning In this work, we focus on search algorithms where subgoals are generated sequentially from the previous subgoal. Recursive subgoal generation is an orthogonal alternative to our approach. In certain domains, given a start and a goal state, it is possible to hierarchically partition the tasks into simpler ones, for instance, by generating subgoals that are approximately halfway between the start and the end state. This idea can be applied repeatedly for an efficient planning algorithm [30, 31, 13]. However, the recursive approach can be difficult to combine with search and limits the class of problems that can be solved.
Discrete Subgoal Search Discrete subgoal search methods can solve complex reasoning problems efficiently with limited search node expansions. S-MCTS [7] is a method for MCTS-based subgoal search with predefined subgoal generators and heuristics, which significantly limits its usability to novel problems. Allen et al. [2] improve the efficiency of search through temporal abstraction by learning macro-actions, which are sequences of low-level actions. Methods that learn subgoal generators, on the other hand, generally suffer from a lack of completeness guarantees. kSubS [4] learns a subgoal generator and performs planning in the subgoal space to solve demanding reasoning tasks. AdaSubS [45] builds on top of kSubS and proposes learning multiple subgoal generators for different distances and performing adaptive planning. HIPS [18] learns to segment trajectories using RL and trains one generator that can propose multiple-length subgoals for adaptive planning.
We present HIPS- $\varepsilon$ , an extension of Hierarchical Imitation Planning with Search (HIPS), a search-based hierarchical imitation learning algorithm for solving difficult goal-conditioned reasoning problems [18]. We use the Markov decision process (MDP) formalism, and the agent's objective is to enter a terminal state in the MDP. However, we deviate from the reward maximization objective. Instead, the objective is to minimize the total search loss, that is, the number of search nodes expanded before a solution is found. A similar objective has been adopted in prior work on subgoal search and, in practice, the solutions found are efficient [4, 45].
We assume the MDPs to be deterministic and fully observable with a discrete state space  $S$  and action space  $\mathcal{A}$ . We also assume the standard offline setting: the agent cannot interact with the environment during training, and it learns to solve tasks only from  $D$ , an existing dataset of expert demonstrations. These trajectories may be heavily suboptimal, but all lead to a goal state. Therefore, the agent should be capable of performing stitching to discover the solutions efficiently [40].
HIPS learns to segment trajectories into subgoals using RL and a subgoal-conditioned low-level policy  $\pi (a|s,s_g)$ . The segmented trajectories are used to train a generative model  $p(s_{g}|s)$  over subgoals  $s_g$  that is implemented as a VQVAE with discrete latent codes [41], a strategy inspired by Ozair et al. [29]. The VQVAE learns a state-conditioned prior distribution  $p(e|s)$  over latent codes  $e$  from a codebook  $\mathbb{E}$  and a decoder  $g(s_{g}|e,s)$  acting as a generator that outputs a subgoal  $s_g$  given the latent code  $e$  and state  $s$  deterministically as  $s_g = g(e,s)$ . Each code  $e_k\in \mathbb{E}$  can be considered a state-dependent high-level action that induces a sequence of low-level actions  $A_{k}(s) = (a_{1},\dots ,a_{n_{k}})$ ,  $a_i\in \mathcal{A}$ , that takes the environment from state  $s$  to  $s_g$ . The low-level action sequence is generated deterministically by the subgoal-conditioned policy  $\pi (a|s,s_g)$ , when the most likely action is used at each step. The prior  $p(e|s)$  can be interpreted as a high-level policy  $\pi_{\mathrm{SG}}(e|s)$  that assigns a probability to each latent code. HIPS also learns a heuristic that predicts the number of low-level actions  $a\in \mathcal{A}$  needed to reach a terminal state from  $s$  and a single-step dynamics model  $f_{\mathrm{dyn}}(s_{i + 1}|a_i,s_i)$ . The heuristic is denoted by  $V(s)$  in [18], but we denote it by  $h$  in this paper.
HIPS performs high-level planning in the subgoal space with PHS [27], although other search methods can be used as well. PHS is a variant of best-first search in which the algorithm maintains
a priority queue and always expands the node  $n$  with the lowest evaluation function value. There is a non-negative loss function  $L$ , and when the search expands node  $n$ , it incurs a loss  $L(n)$ . The objective of PHS is to minimize the total search loss. The evaluation function of PHS is
$$
\varphi (n) = \eta (n) \frac {g (n)}{\pi (n)}, \tag {1}
$$
where  $\eta(n) \geq 1$  is a heuristic factor that depends on the heuristic function  $h(n)$ .  $g(n)$  is the path loss, the sum of the losses from the root  $n_0$  to  $n$ .  $\pi(n) \leq 1$  is the probability of node  $n$  that is defined recursively:  $\pi(n') = \pi(n'|n)\pi(n)$ ,  $\pi(n_0) = 1$ ,  $n$  is the parent of  $n'$  and  $n \in \mathrm{desc}_*(n_0)$ , that is,  $n$  is a descendant of the root node  $n_0$ . When HIPS plans, each search node  $n$  corresponds to a state  $s$ ,  $g(n)$  to the number of low-level steps needed to reach  $n$  from the root  $n_0$ , and  $h(n)$  to the output of the value function  $V(s)$ . The children  $\mathcal{C}(n)$  of node  $n$  are the valid subgoals generated by the VQVAE decoder  $g(e_k, s)$  for all  $e_k \in \mathbb{E}$ . The search policy  $\pi(n)$  is induced by the conditional policy  $\pi(n'|n)$ , which is represented by the state-conditioned prior  $p(e|s)$ .
3.2 Complete Subgoal Search
HIPS performs search solely in the space of subgoals proposed by the trained generator network, which may cause failures to discover solutions to solvable problems due to possible imperfections of the subgoal generation model (see Fig. 1c). This problem can be tackled in discrete-action MDPs by augmenting the search with low-level actions (see Fig. 1d). In this work, we apply this idea to HIPS to guarantee solution discovery when the solution exists while still utilizing the subgoal generator for temporally abstracted efficient search.
Formally, we propose modifying the search procedure of HIPS such that in addition to the subgoals  $\{s_H = g(e,s),\forall e\in \mathbb{E}\}$  produced by the HIPS generator, the search also considers states  $\{s_L = T(s,a),\forall a\in \mathcal{A}\}$  that result from the agent taking every low-level action  $a\in \mathcal{A}$  in state  $s$ , where  $T(s,a)$  is the environment transition function. We assume that  $T$  is known but we also evaluate the method with a learned  $T$ . The augmented action space is then  $\mathcal{A}^{+} = \mathcal{A}\cup \mathbb{E}$ . We use PHS as the search algorithm and compute the node probabilities used for the evaluation function (1) as
$$
\pi (n) = \prod_ {j = 1} ^ {d} \pi (s _ {j} (n) | s _ {j - 1} (n)),
$$
where  $s_d(n)$  is the state that corresponds to node  $n$  at depth  $d$ ,  $s_{d - 1}(n), s_{d - 2}(n), \ldots, s_0(n)$  are the states of the ancestor nodes of  $n$  and  $\pi (s_j(n)|s_{j - 1}(n))$  is the search policy. We propose to compute the probabilities  $\pi (s_j|s_{j - 1})$  in the following way:
$$
\pi \left(s _ {j} \mid s _ {j - 1}\right) = \left\{ \begin{array}{l l} (1 - \varepsilon) \pi_ {\mathrm {S G}} \left(e _ {k} \mid s _ {j - 1}\right) & \text {i f} s _ {j} = g \left(e _ {k}, s _ {j - 1}\right) \text {i s p r o p o s e d b y H I P S} \\ \varepsilon \pi_ {\mathrm {B C}} \left(a \mid s _ {j - 1}\right) & \text {i f} s _ {j} = T \left(s _ {j - 1}, a\right) \text {i s p r o p o s e d b y l o w - l e v e l s e a r c h} \end{array} \right. \tag {2}
$$
where  $\pi_{\mathrm{SG}}(e_k|s_{j-1}) = p(e_k|s_{j-1})$  is the prior over the high-level actions learned by HIPS and  $\pi_{\mathrm{BC}}(a|s_{j-1})$  is a low-level policy that we train from available expert demonstrations with BC. We use hyperparameter  $\varepsilon$  to balance the probabilities computed with the high and low-level policies: higher values of  $\varepsilon$  prioritize more low-level exploration. We call this complete search HIPS- $\varepsilon$ .
The proposed complete search approach can be combined with any policy-guided search algorithm, but we focus on PHS due to its state-of-the-art performance. Complete search can also be used with policyless search algorithms, but in that case,  $\varepsilon$  cannot directly control the frequency of low and high-level actions.
Zawalski et al. [45] suggested using low-level actions only when the high-level search would otherwise fail. In the context of PHS, this can be interpreted as having an infinitesimal  $\varepsilon$  such that every high-level action has a lower evaluation function value (1) and, therefore, a higher priority than every low-level action. With some abuse of notation, we refer to this approach as  $\varepsilon \rightarrow 0$ .
3.3 Complete Subgoal Search Heuristic
The objective of PHS is to minimize the search loss, and Orseau and Lelis [27] propose using a heuristic to estimate the  $g$ -cost of the least costly solution node  $n^*$  that is a descendant of the current
search node  $n$ . They propose to approximate the probability  $\pi(n^{*})$  for node  $n^{*}$  that contains a found solution as
$$
\pi (n ^ {*}) = \left[ \pi (n) ^ {1 / g (n)} \right] ^ {g (n) + h (n)} \tag {3}
$$
and use this approximation to arrive at a heuristic factor  $\hat{\eta}_h(n) = \frac{1 + h(n) / g(n)}{\pi(n)^{h(n) / g(n)}}$  to be used in (1). PHS* is the variant of PHS that uses this heuristic factor. The derivation of the heuristic factor  $\hat{\eta}_h(n)$  assumes that  $g(n)$ , the path loss of node  $n$ , and  $h(n)$ , the heuristic function approximating the distance to the least-cost descendant goal node  $n^*$  are expressed in the same scale. This does not hold for the proposed search algorithm. The search objective of HIPS- $\varepsilon$  is to minimize the number of node expansions, which implies that  $g(n)$  is equal to the depth of node  $n$ , whereas the heuristic  $h(n)$  is trained to predict the number of low-level actions to the goal. One solution is to let  $g(n)$  be equal to the number of low-level steps to node  $n$ . However, this would distort the search objective. The cost of expanding a node corresponding to a subgoal would equal the length of the corresponding low-level action sequence, whereas the loss of low-level actions would be one. However, the cost should be the same for every node, as the objective is to minimize the number of node expansions.
Instead, we propose re-scaling the heuristic factor  $h$  to be equal to the estimated number of search nodes on the path from  $n$  to the goal node  $n^*$  by dividing the expected number of low-level actions from node  $n$  to the terminal node by the average low-level distance between the nodes on the path from  $n_0$  to node  $n$ . Let  $g(n)$  be equal to  $d(n)$ , the depth of the node, and let  $l(n)$  be the number of low-level actions from the root node. Then, we define a scaled heuristic  $\hat{h}(n) = \frac{h(n)g(n)}{l(n)}$  and by using it instead of  $h(n)$  in (3), we get a new approximation  $\pi(n^*) = \pi^{(1 + h(n) / l(n))}$ , which yields the following heuristic factor and evaluation function:
$$
\hat {\eta} _ {\hat {h}} = \frac {1 + h (n) / l (n)}{\pi (n) ^ {h (n) / l (n)}}, \quad \hat {\varphi} _ {\hat {h}} (n) = \frac {g (n) \cdot (1 + h (n) / l (n))}{\pi (n) ^ {1 + h (n) / l (n)}}. \tag {4}
$$
For the full derivation, please see Appendix C. Note that  $\hat{\eta}_{\tilde{h}}$  may not be PHS-admissible, even if  $h(n)$  were admissible. In HIPS, there are no guarantees of the learned heuristic  $h(n)$  being admissible.
3.4 Analysis of HIPS-ε
Orseau and Lelis [27] showed that the search loss of PHS has an upper bound:
$$
L \left(\mathrm {P H S}, n ^ {*}\right) \leq L _ {\mathrm {U}} \left(\mathrm {P H S}, n ^ {*}\right) = \frac {g \left(n ^ {*}\right)}{\pi \left(n ^ {*}\right)} \eta^ {+} \left(n ^ {*}\right) \sum_ {n \in \mathcal {L} _ {\varphi} \left(n ^ {*}\right)} \frac {\pi (n)}{\eta^ {+} (n)}, \tag {5}
$$
where  $n^*$  is the solution node returned by PHS,  $g(n)$  is the path loss from the root to  $n$ ,  $\pi$  the search policy, and  $\eta^+$  the modified heuristic factor corresponding to the monotone non-decreasing evaluation function  $\varphi^{+}(n)$ , and  $\mathcal{L}_{\varphi}(n^{*})$  the set of nodes that have been expanded before  $n^*$ , but the children of which have not been expanded. Now, given the recursively defined policy  $\pi (n)$  that depends on  $\pi_{\mathrm{BC}}$  and  $\pi_{\mathrm{SG}}$  as outlined in Subsection 3.2, the upper bound derived in (5) is unchanged.
However, the upper bound of (5) can be seen as uninformative in the context of complete search, as it depends on the high-level policy  $\pi_{\mathrm{SG}}$ . In general, we cannot make any assumptions about the behavior of the high-level policy as it depends on the subgoals proposed by the generator network. However, we can derive the following upper bound, where the probability of the terminal node  $n^*$  is independent of the behavior of the generator, and consequently, the high-level policy  $\pi_{\mathrm{SG}}$ :
Corollary 1. (Subgoal-PHS upper bound without  $\pi_{SG}$  in the denominator). For any non-negative loss function  $L$ , non-empty set of solution nodes  $\mathcal{N}_{\mathcal{G}} \subseteq \text{desc}_*(n_0)$ , low-level policy  $\pi_{BC}$ , high-level policy  $\pi_{SG}$ ,  $\varepsilon \in (0,1]$ , policy  $\pi(n'|n)$  as defined in (2), and heuristic factor  $\eta(\cdot) \geq 1$ , Subgoal-PHS returns a solution node  $n^* \in \arg \min_{n^* \in \mathcal{N}_{\mathcal{G}}} \varphi^+(n^*)$ . Then, there exists a terminal node  $\hat{n} \in \mathcal{N}_{\mathcal{G}}$  such that state( $\hat{n}$ ) = state( $n^*$ ), where  $\hat{n}$  corresponds to low-level trajectory ( $s_0, a_0, \ldots, a_{N-1}, \text{state}(n^*)$ ) and the search loss is bounded by
$$
L \left(P H S - h, n ^ {*}\right) \leq L \left(P H S - h, \hat {n}\right) \leq \frac {g (\hat {n})}{\varepsilon^ {N} \prod_ {i = 0} ^ {N - 1} \pi_ {B C} \left(a _ {i} \mid s _ {i}\right)} \eta^ {+} (\hat {n}) \sum_ {n \in \mathcal {L} _ {\varphi} (\hat {n})} \frac {\pi (n)}{\eta^ {+} (n)} \tag {6}
$$
Proof. Subgoal-PHS does not change the underlying search algorithm, so the assumption of Subgoal-PHS returning the minimum score solution node follows from Theorem 1 in [27]. The first inequality
Sliding Tile Puzzle
Box-World
Figure 2: We evaluate our complete search approach, HIPS- $\varepsilon$ , in Sokoban, Sliding Tile Puzzle (STP), Box-World (BW), and Travelling Salesman Problem (TSP). In Sokoban, the agent must push the yellow boxes onto the red target locations. In the Sliding Tile Puzzle, the agent must slide the tiles to sort them from 1 to 24. In Box-World, the task is to collect the gem (marked with $) by opening locks with keys of the corresponding color. In Travelling Salesman, the agent (marked with a circle) must visit all unvisited cities (red squares) before returning to the start (black square).
follows from the fact that the loss function  $L$  is always non-negative and the node  $n^*$  will be expanded before any other node  $n$  with  $\varphi^{+}(n) > \varphi^{+}(n^{*})$ . In particular, note that it is possible that  $n^* = \hat{n}$ . The second inequality follows from the Theorem 1 in [27] and  $\pi(\hat{n}) = \varepsilon^{N} \prod_{i=0}^{N-1} \pi_{\mathrm{BC}}(a_{i}|s_{i})$ .