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: This work proposes a method for search, considering the transition model is given. The proposed framework works as follows: given a state, a subgoal predictor is used to predict subgoals k-steps ahead. A policy network is used to get the actions to reach those subgoals and a value function is used to judge the value or benefit of being in a state. All of these are trained on an expert set of demos. Empirical results show that this algorithm performs better than BestFS or MTCS.