Subgoal Search For Complex Reasoning Tasks

1 Introduction . Reasoning is often regarded as a defining property of advanced intelligence [ 40 , 18 ] . When confronted with a complicated task , humans ’ thinking process often moves from one idea to a related idea , and the progress is made through milestones , or subgoals , rather than through atomic actions that are necessary to transition between subgoals [ 15 ] . During this process , thinking about one subgoal can lead to a possibly diverse set of subsequent subgoals that are conceptually reachable and make a promising step towards the problem ’ s solution . This intuitive introspection is backed by neuroscience evidence [ 20 ] , and in this work , we present an algorithm that mimics this process . Our approach couples a deep learning generative subgoal modeling with classical search algorithms to allow for successful planning with subgoals . We showcase the efficiency of our method on the following complex reasoning tasks : two popular puzzle games Sokoban and the Rubik ’ s Cube , and an inequality theorem proving benchmark INT [ 55 ] , achieving the state-of-the-art results in INT and competitive results for the remaining two . The deep learning revolution has brought spectacular advancements in pattern recognition techniques and models . Given the hard nature of reasoning problems , these are natural candidates to provide ∗equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . search heuristics [ 4 ] . Indeed , such a blend can produce impressive results [ 44 , 45 , 37 , 1 ] . These approaches seek solutions using elementary actions . Others , e.g . [ 30 , 34 , 24 ] , utilize variational subgoals generators to deal with long-horizon visual tasks . We show that these ideas can be pushed further to provide algorithms capable of dealing with combinatorial complexity . We present Subgoal Search ( kSubS ) method and give its practical implementations : MCTS-kSubS and BF-kSubS . kSubS consists of the following four components : planner , subgoal generator , a low-level policy , and a value function . The planner is used to search over the graph induced by the subgoal generator and is guided by the value function . The role of the low-level policy is to prune the search tree as well as to transition between subgoals . In this paper , we assume that the generator predicts subgoals that are k step ahead ( towards the solution ) from the current state , and to emphasize this we henceforth add k to the method ’ s abbreviation . MCTS-kSubS and BF-kSubS differ in the choice of the search engine : the former uses Monte-Carlo Tree Search ( MCTS ) , while the latter is backed by Best-First Search ( BestFS ) . We provide two sets of implementations for the generator , the low-level policy , and the value functions . The first one uses transformer architecture [ 49 ] for each component , while the second utilizes a convolutional network for the generator and the value function , and the classical breadth-first search for the low-level policy . This lets us showcase the versatility and effectiveness of the approach . The subgoal generator lies at the very heart of Subgoal Search , being an implementation of reasoning with high-level ideas . To be useful in a broad spectrum of contexts , the generator should be implemented as a learnable ( generative ) model . As a result , it is expected to be imperfect and ( sometimes ) generate incorrect predictions , which may turn the search procedure invalid . Can we thus make planning with learned subgoals work ? In this paper , we answer this question affirmatively : we show that the autoregressive framework of transformer-based neural network architecture [ 49 ] leads to superior results in challenging domains . We train the transformer with the objective to predict the k-th step ahead . The main advantages of this subgoal objective are simplicity and empirical efficiency . We used expert data to generate labels for supervised training . When offline datasets are available , which is the case for the environments considered in this paper2 , such an approach allows for stable and efficient optimization with highquality gradients . Consequently , this method is often taken when dealing with complex domains ( see e.g . [ 42 , 52 ] ) or when only an offline expert is available3 . Furthermore , we found evidence of out-of-distribution generalization . Finally , we formulate the following hypothesis aiming to shed some light on why kSubS is successful : we speculate that subgoal generation may alleviate errors in the value function estimation . Planning methods based on learning , including kSubS , typically use imperfect value function-based information to guide the search . While traditional low-level search methods are susceptible to local noise , subgoal generation allows for evaluations of the value functions at temporally distant subgoals , which improves the signal-to-noise ratio and allows a “ leap over ” the noise . To sum up , our contributions are : 1 . We propose Subgoal Search method with two implementations : MCTS-kSubS , BF-kSubS . We demonstrate that our approach requires a relatively little search or , equivalently , is able to handle bigger problems . We also observe evidence of out-of-distribution generalization . 2 . We provide evidence that a transformer-based autoregressive model learned with a simple supervised objective to predict states k-th step ahead is an effective tool to generate valid and diverse subgoals . 3 . We show in our experiments that using subgoal planning help to might mitigate the negative influence of value function errors on planning . We provide the code of our method and experiment settings at https : //github.com/ subgoal-search/subgoal-search , and a dedicated website https : //sites.google.com/ view/subgoal-search . 2The dataset for INT or Sokoban can be easily generated or are publicly available . For the Rubik ’ s Cube , we use random data or simple heuristic ( random data are often sufficient for robotic tasks and navigation . ) 3For example , the INT engine can easily generate multiple proves of random statements , but can not prove a given theorem . 2 Related work . In classical AI , reasoning is often achieved by search ( [ 40 ] ) . Search rarely can be exhaustive , and a large body of algorithms and heuristics has been developed over the years , [ 40 , Section 3.5 ] . It is hypothesized that progress can be achieved by combining search with learning [ 4 ] . Among notable successful examples of this approach are Alpha Zero [ 43 ] , or solving Rubik ’ s cube using a learned heuristic function to guide the A∗ algorithm ( see [ 1 ] ) . An eminent early example of using goal-directed reasoning is the PARADISE algorithm ( [ 53 ] ) . In deep learning literature , [ 25 ] was perhaps the first work implementing subgoal planning . This was followed by a large body of work on planning with subgoals in the latent space for visual tasks ( [ 24 , 31 , 34 , 30 , 21 , 9 , 35 ] ) or landmark-based navigation methods ( [ 41 , 27 , 14 , 46 , 56 , 23 ] ) . The tasks considered in the aforementioned studies are often quite forgiving when it comes to small errors in the subgoal generation . This can be contrasted with complex reasoning domains , in which even a small variation of a state may drastically change its meaning or render it invalid . Thus , neural networks may struggle to generate semantically valid states ( [ 4 , Section 6.1 ] ) . Assuming that this problem was solved , a generated subgoal still remains to be assessed . The exact evaluation may , in general , require exhaustive search or access to an oracle ( in which case the original problem is essentially solved ) . Consequently , it is unlikely that a simple planner ( e.g. , one unrolling independent sequences of subgoals [ 9 ] ) will either produce an informative outcome , could be easily improved using only local changes via gradient descent [ 31 ] , or cross-entropy method ( CEM ) [ 30 , 34 , 35 ] . Existing approaches , which are based on more powerful subgoal search methods , have their limitations , on the other hand . [ 13 ] is perhaps the closest to our method and uses MCTS to search the subgoal-induced graph . However , it uses a predefined ( not learned ) predicate function as a subgoal generator , limiting applicability to the problems with available high-quality heuristics . Learned subgoals are used in [ 32 ] , a method of hierarchical planning . That said , the subgoal space needs to be relatively small for this method to work ( or crafted manually to reduce cardinality ) . To the best of our knowledge , our algorithm is the first domain agnostic hierarchical planner for combinatorially complex domains . More generally , concepts related to goals and subgoals percolated to reinforcement learning early on , leading , among others , to prominent ideas like hindsight [ 22 ] , hierarchical learning [ 48 , 7 ] or the Horde architecture [ 47 ] . Recently , with the advent of deep reinforcement learning , these ideas have been resurfacing and scaled up to deliver their initial promises . For example , [ 51 ] implements ideas of [ 7 ] and a very successful hindsight technique [ 2 ] is already considered to be a core RL method . Further , a recent paper [ 36 ] utilizes a maximum entropy objective to select achievable goals in hindsight training and [ 50 ] uses meta-learned subgoals to discover options in multi-task RL settings . Apart from the aforementioned hierarchical planning , the idea to use neural networks for subgoal generation was used to improve model-free reinforcement learning agents [ 11 , 6 ] and imitation learning algorithms [ 33 ] . 3 Method . Our method , Subgoal Search ( kSubS ) , is designed for problems , which can be formulated as a search over a graph with a known transition model . Formally , let G = ( S , E ) be a directed graph and S̃ ⊂ S be the set of success states . We assume that , during the solving process , the algorithm can , for a given node g , determine the edges starting at g and check if g ∈ S̃ . Subgoal Search consists of four components : planner , subgoal generator , low-level policy , and a value function . The planner , coupled with a value function , is used to search over the graph induced by the subgoal generator . Namely , for each selected subgoal , the generator allows for sampling the candidates for the next subgoals . Only these reachable by the low-level policy are used . The procedure continues until the solution is found or the computational budget is exhausted . Our method searches over a graph G̃ = ( S , Es ) , with the edges Es given by the subgoal generator . Provided a reasonable generator , paths to S̃ are shorter in G̃ than in G and thus easier to find . In this paper we provide BestFS- and MCTS- backed implementations kSubS . Algorithm 1 presents BF-kSubS ; see Algorithm 4 in Appendix A.1 for MCTS-kSubS . For INT and Rubik ’ s Cube , we use transformer models ( see Appendix B.1 ) in all components other than the planner . For Sokoban , we use convolutional networks ( see Appendix B.2 ) . While transformers could also be used in Sokoban , we show that a simplified setup already achieves strong results . This showcases that Subgoal Search is general enough to work with different design choices . In Section 4.2 we describe datasets used train these neural models . Algorithm 1 Best-First Subgoal Search ( BF-kSubS ) Require : C1 max number of nodes V value function network SOLVED predicate of solution function SOLVE ( s0 ) T.PUSH ( ( V ( s0 ) , s0 ) ) . T is priority queue paths [ s0 ] ← [ ] . paths is dict of lists seen.ADD ( s0 ) . seen is set while 0 < LEN ( T ) and LEN ( seen ) < C1 do _ , s← T.EXTRACT_MAX ( ) subgoals← SUB_GENERATE ( s ) . see Algorithm 3 for s′ in subgoals do if s′ in seen then continue seen.ADD ( s′ ) actions← GET_PATH ( s , s′ ) . see Alg 2 or Alg 9 if actions.EMPTY ( ) then continue T.PUSH ( ( V ( s′ ) , s′ ) ) paths [ s′ ] ← paths [ s ] + actions if SOLVED ( s′ ) then return paths [ s′ ] return False Algorithm 2 Low-level conditional policy Require : C2 steps limit π low-level conditional policy network M model of the environment function GET_PATH ( s0 , subgoal ) step← 0 s← s0 action_path← [ ] while step < C2 do action← π.PREDICT ( s , subgoal ) action_path.APPEND ( action ) s←M.NEXT_STATE ( s , action ) if s = subgoal then return action_path step← step + 1 return [ ] Subgoal generator Formally , it is a mapping ρ : S → P ( S ) , where P ( S ) is a family of probability distributions over the environment ’ s state space S. More precisely , let us define a trajectory as a sequence of state-action pairs ( s0 , a0 ) , ( s1 , a1 ) , . . . , ( sn , an ) , with ( si , ai ) ∈ S ×A , whereA stands for the action space and n is the trajectory ’ s length . We will say that the generator predicts k-step ahead subgoals , if at any state s ` it aims to predict smin ( ` +k , n ) . We show , perhaps surprisingly , that this simple objective is an efficient way to improve planning , even for small values of k , i.e . k ∈ { 2 , 3 , 4 , 5 } . Operationally , the subgoal generator takes as input an element of the s ∈ S and returns subgoals , a set of new candidate states expected to be closer to the solution and is implemented by Algorithm 3 . It works as follows : first , we generate C3 subgoal candidates with SUB_NET_GENERATE . Then , we prune this set to obtain a total probability greater than C4 . For INT and Rubik ’ s Cube , we represent states as sequences modeled with a transformer . Following the practice routinely used in language modeling , [ 49 ] , SUB_NET_GENERATE employs beam search to generate a set of high likelihood outputs.4 With the same goal in mind , for Sokoban , we use another method described Appendix C.1 . SUB_NET_GENERATE uses the subgoal generator network , which is trained in a supervised way on the expert data , with training examples being : s ` ( input ) and smin ( ` +k , n ) ( label ) , see Appendix D for details . Low-level conditional policy Formally , it is a mapping π : S × S → A∗ . It is used to generate a sequence of actions on how to reach a subgoal starting from a given initial state . Operationally , it may return an empty sequence if the subgoal can not be reached within C2 steps , see Algorithm 2 . This is used as a pruning mechanism for the subgoals set in Algorithm 1 . 4In language modeling , typically , only one beam search output is used . In our case , however , we utilize all of them , which turns out to be a diverse set of subgoals . In INT and Rubik ’ s Cube , we use Algorithm 2 , which utilizes low-level policy network π . Similarly to the subgoal generator , it is trained using expert data in a supervised fashion , i.e . for a pair ( s ` , smin ( ` +i , n ) ) , with i ≤ k , its objective is to predict a ` . Algorithm 3 Subgoal generator Require : C3 number of subgoals C4 target probability ρ subgoal generator network function SUB_GENERATE ( s ) subgoals← ∅ states , probs← SUB_NET_GENERATE ( ρ , s ; C3 ) . ( states , probs ) is sorted wrt probs total_p← 0 for state , p ∈ ( states , probs ) do if total_p > C4 then break subgoals.ADD ( state ) total_p← total_p + p return subgoals When the branching factor is small , the low-level policy can be realized by a simple breadth-first search mechanism , see Algorithm 9 in Appendix , which we illustrate on Sokoban . Value function Formally , it is a mapping V : S → R , that assigns to each state a value related to its distance to the solution , and it is used to guide the search ( see Algorithm 1 and Algorithm 4 ) . For its training , we use expert data . For each state s ` the training target is negated distance to the goal : ` − n , where n denotes the end step of a trajectory that s ` belongs to . Planner This is the engine that we use to search the subgoal-induced graph . In this paper , we use BestFS ( Algorithm 1 ) and MCTS ( Algorithm 4 in Appendix A.1 ) . The former is a classic planning method , which maintains a priority queue of states waiting to be explored , and greedily ( with respect to their value ) selects elements from it ( see , e.g. , [ 40 ] ) . MCTS is a search method that iteratively and explicitly builds a search tree , using ( and updating ) the collected node statistics ( see , e.g. , [ 5 ] ) . In this paper , we use an AlphaZero-like [ 42 ] algorithm for single-player games . We note that the subgoal generator can be combined with virtually any search algorithm and can benefit from an additional structure . For example , for domains providing a factored state representation , the width-based methods [ 26 , 12 ] would likely be stronger search mechanisms .

Summary: The paper considers the problem of finding a sequence of actions that, starting at an initial state, achieve a goal. The search graph is known in this problem: transition model is known, states are fully observable and actions are deterministic. That means any search algorithm could solve the problem for graphs like Breadth First Search, Best First Search or A*.  The paper proposes to abstract the transition function by mapping a state into a set of steps $k$ steps ahead. When $k=1$, the abstraction corresponds to a subset of the states that can be achieved by applying one action.  The use of lookahead well-established idea in the search community, in different forms, but it seems under-explored for data-driven methods for search. So, the line of work is original.  I celebrate that the paper focuses on deterministic observable known models because this setting allows understanding better the relationship between "search" and "learning". That relationship is a the core of the tension between data-driven and model/reasoning-driven methods.  The research question about the value of predicting $k$ steps ahead is interesting. On the one hand, data-driven methods are combined with search because they perform better than following predictions greedily. AlphaZero is "exhibit A". On the other hand, bounded-length predictions are more likely to remain within the practical scope of data-driven methods.  I do have a serious problem with the use of the word "**analog**" across the paper. I think we cannot write papers about "analog of human intuition" because "analog" implies similarity or comparison. It's a very strong word. I could accept "inspired by an aspect of human intuition", but I find it unacceptable because it sounds like we know what human intuition is. I might be wrong. If that's the case, please justify that using the relevant literature. Otherwise, please just state how the manuscript would be revised so this expressions have a lower tone.  The findings are very interesting. Fig 2 is the most insightful.  However, I think the manuscript have some confusions that make the contributions less clear.  Only three domains are considered. Two use the same setting with a shared family of hyperparameters. A third one uses a different setting.