SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiayuan Mao $^{1}$  Tomás Lozano-Pérez $^{1}$  Joshua B. Tenenbaum $^{1,2,3}$  Leslie Pack Kaelbling $^{1}$
<sup>1</sup> MIT Computer Science & Artificial Intelligence Laboratory
$^{2}$  MIT Department of Brain and Cognitive Sciences
3 Center for Brains, Minds and Machines
Goal-conditioned policies are generally understood to be "feed-forward" circuits, in the form of neural networks that map from the current state and the goal specification to the next action to take. However, under what circumstances such a policy can be learned and how efficient the policy will be are not well understood. In this paper, we present a circuit complexity analysis for relational neural networks (such as graph neural networks and transformers) representing policies for planning problems, by drawing connections with serialized goal regression search (S-GRS). We show that there are three general classes of planning problems, in terms of the growth of circuit width and depth as a function of the number of objects and planning horizon, providing constructive proofs. We also illustrate the utility of this analysis for designing neural networks for policy learning.
Goal-conditioned policies are generally understood to be "feed-forward" circuits, in the form of neural networks such as multi-layer perceptrons (MLPs) or transformers [Vaswani et al., 2017]. They take a representation of the current world state and the goal as input, and generate an action as output. Previous work has proposed methods for learning such policies for particular problems via imitation or reinforcement learning [Wang et al., 2018, Dong et al., 2019, Li et al., 2020]; recently, others have tried probing current large-language models to understand the extent to which they already embody policies for a wide variety of problems [Liang et al., 2022, Carta et al., 2023].
However, a major theoretical challenge remains. In general, we understand that planning problems are PSPACE-hard with respect to the size of the state space [Bylander, 1994], but there seem to exist efficient (possibly suboptimal) policies for many problems such as block stacking [Dong et al., 2019] that generalize to arbitrarily sized problem instances. In this paper, we seek to understand and clarify the circuit complexity of goal-conditioned policies for classes of "classical" discrete planning problems: under what circumstances can a polynomial-sized policy circuit be constructed, and what is its size? Specifically, we highlight a relationship between a problem hardness measure (regression width, related to the notion of "width" in the forward planning literature [Chen and Giménez, 2007]) and circuit complexity. We concretely provide upper bounds on the policy circuit complexity as a function of the problem's regression width, using constructive proofs that yield algorithms for directly compiling a planning problem description into a goal-conditioned policy. We further show that such policies can be learned with conventional policy gradient methods from environment interactions only, with the theoretical results predicting the necessary size of the network.
There are several useful implications of these results. First, by analyzing concrete planning domains, we show that for many domains, there do exist simple policies that will generalize to problem instances with an arbitrary number of objects. Second, our theory predicts the circuit complexity of policies for these problems. Finally, our analysis suggests insights into why certain planning
Project page: https://concepts-ai.com/p/goal-regression-width/
Figure 1: (a) Illustration of the Blocks World domain that we will be using as the main example. (b) The action schema definition in Blocks World. clear(x) means there is no object on x. (c) A backward search tree for solving the goal clear(B). (d) A serialized goal regression search tree for the same goal. (e) The form of a goal-conditioned policy for this problem.
problems, such as Sokoban and general task and motion planning (TAMP) problems, are hard, and likely cannot be solved in general by fixed-sized MLP or transformer-based policies [Culberson, 1997, Vega-Brown and Roy, 2020]. In the rest of the paper, Section 2 provides problem definitions, Sections 3 and 4 provides complexity definitions, theoretical results, and the policy compilation algorithm; Section 5 discusses the practical implications of these results.
2 Preliminaries: Planning Domain and Problem
Throughout this paper, we focus on analyzing the search complexity and policy circuit complexity of classical planning problems. Importantly, these planning problems have an object-centric representation and sparsity in the transition models: the state of the world is represented as a set of entities and their properties and relationships, while each action only changes a sparse set of properties and relations of a few objects. These features will contribute to search efficiency.
Formally, we consider the problem of planning and policy learning for a space  $S$  of world states. A planning problem is a tuple  $\langle S, s_0, \mathcal{G}, \mathcal{A}, \mathcal{T} \rangle$ , where  $s_0 \in S$  is the initial state,  $\mathcal{G} \subseteq S$  is a goal specification,  $\mathcal{A}$  is a set of actions that the agent can execute, and  $\mathcal{T}$  is a (possibly partial) environmental transition model  $\mathcal{T}: S \times \mathcal{A} \to S$ . The task of planning is to output a sequence of actions  $\bar{a}$  in which the terminal state  $s_T$  induced by applying  $a_i$  sequentially following  $\mathcal{T}$  satisfies  $s_T \in \mathcal{G}$ . The task of policy learning is to find a function  $\pi(s, \mathcal{G})$  that maps from the state and a goal specification to the next action so that applying  $\pi$  recurrently on  $s_0$  eventually yields a state  $s_T \in \mathcal{G}$ .
We use atomic STRIPS [Fikes and Nilsson, 1971, Lifschitz, 1987] to formalize the object-centric factorization of our planning problems. Specifically, as illustrated in Fig. 1a, the environmental state can be represented as a tuple  $(\mathcal{U}_s,\mathcal{P}_s)$ , where  $\mathcal{U}_s$  is the set of objects in state  $s$ , denoted by arbitrary unique names (e.g., A, B). The second component,  $\mathcal{P}_s$ , is a set of atoms: handfree(), on-table(A), on(C, B), etc. Each atom contains a predicate name (e.g., on) and a list of arguments (e.g., C, B). All atoms in  $\mathcal{P}_s$  are true at state  $s$ , while any atoms not in  $\mathcal{P}_s$  are false. Since we do not consider object creation and deletion, we simply use  $\mathcal{P}_0$  to denote the set of all possible atoms for the object universe  $\mathcal{U}_{s_0}^*$ .
Illustrated in Fig. 1b, a (possibly partial) transition model  $\mathcal{T}$  is specified in terms of a set of object-parameterized action schemas  $\langle \text{name}, \text{args}, \text{precond}, \text{effect} \rangle$ , where name is the name of the action, args is a list of symbols, precond and effect are descriptions of the action's effects. In the basic STRIPS formulation, both precond and effect are sets of atoms with variables in args, and effect is further decomposed into  $e f_{+}$  and  $e f_{-}$ , denoting "add" effects and "delete" effects. An action schema (e.g., unstack) can be grounded into a concrete action  $a$  (e.g.,  $unstack(D, C)$ ) by binding all variables in args.
*For simplicity, throughout the paper, we will be focusing on Boolean state variables.
Algorithm 1 Plain backward search.   
function  $\mathrm{BWD}(s_0,\mathcal{A},\mathrm{goal\_set})$  if goal_set  $\subseteq s_0$  then return empty_list() possible_t  $=$  empty_set() for  $a\in \{a\in \mathcal{A}\mid$  goal_set  $\cap$  eff_  $(a) =$  0 and goal_set  $\cap$  eff  $^+\neq \emptyset$  do new_goals  $=$  goal_set  $\cup$  pre(a)  $\backslash$  eff_  $(a)$  if new_goals  $\in$  goal stack then continue  $\pi = \mathrm{BWD}(s_0,\mathrm{new\_goals})$  if  $\pi \neq \bot$  then possible_t.add(  $\pi +\{a\}$  ) return shortest path in possible_t
Algorithm 2 Serialized goal regression search.  
function S-GRS(s0, R, g, cons)  
if  $g \in s_0$  then return (s0)  
possible_t = empty_set()  
for  $r \in \{r \in \mathcal{R} \mid \text{cons} \cap \text{eff}_{-}(\text{action}(r)) = \emptyset$  and goal(r) = g do  
p1, p2, ..., pk = subgoals(r)  
if ∃pi s.t. pi ∈ goal stack then continue  
for i in 1, 2, ..., k do  
new_c = cons ∪ {p1, ..., pi-1}  
 $\pi_i = \text{s - GRS}(s_{i-1}, p_i, \text{new}_c)$   
if  $\pi_i == \bot$  then break  
si = T(si-1,  $\pi_i[-1]$ )  
if  $\pi_k \neq \bot$  then  
possible_t.add( $\pi_1 + \ldots + \pi_k + \{a\}$ )  
return shortest path in possible_t
The formal semantics of STRIPS actions is:  $\forall s, \forall a. pre(a) \subseteq s \Rightarrow (\mathcal{T}(s, a) = s \cup eff_{+}(a) \setminus eff_{-}(a))$ . That is, for any state  $s$  and any action  $a$ , if the preconditions of  $a$  are satisfied in  $s$ , the state resulting from applying  $a$  to  $s$  will be  $s \cup eff_{+}(a) \setminus eff_{-}(a)$ . Note that  $\mathcal{T}$  may not be defined for all  $(s, a)$  pairs. Furthermore, we will only consider cases where the goal specification is a single atom, although more complicated conjunctive goals or goals that involve existential quantifiers can be easily supported by introducing an additional "goal-achieve" action that lists all atoms in the original goal as its precondition. From now on, we will use  $g$  to denote the single goal atom of the problem.
3 Goal Regression Search and Width
A simple and effective way to solve STRIPS problems is backward search†. Illustrated in Algorithm 1‡, we start from the goal  $\{g\}$ , and search for the last action  $a$  that can be applied to achieve the goal (i.e.,  $g \in eff_{+}(a)$ ). Then, we add all the preconditions of action  $a$  to our search target and recurse. To avoid infinite loops, we also keep track of a "goal stack" variable composed of all goal_set's that have appeared in the depth-first search stack. The run-time of this algorithm has critical dependencies on (1) the number of steps in the resulting plan and (2) the number of atoms in the intermediate goal sets in search. In general, it has a worst-case time complexity exponential in the number of atoms.
3.1 Serialized Goal Regression Search
One possible way to speed up backward search is to serialize the subgoals in the goal set: that is, to commit to achieving them in a particular order. In the following sections, we study a variant of the plain backward search algorithm, namely serialized goal regression search (S-GRS, Algorithm 2), which uses a sequential decomposition of the preconditions: given the goal atom  $g$ , it searches for the last action  $a$  and also an order in which the preconditions of  $a$  should be accomplished. Of course, this method cannot change the worst-case complexity of the problem (actually, it may slow down the algorithm in cases where we need to search through many orders of the preconditions), but there are useful subclasses of planning problems for which it can achieve a substantial speed-up.
A serialized goal regression rule in  $\mathcal{R}$ , informally, takes the following form:  $g \not\sim p_1, p_2, \dots, p_k \parallel a$ . It reads: in order to achieve  $g$ , which is a single atom, we need to achieve atoms  $p_1, \dots, p_k$  sequentially, and then execute action  $a$ . To formally define the notion of achieving atoms  $p_1, \dots, p_k$  "sequentially," we must include constraints in goal regression rules. Formally, we consider the application of a rule under constraint  $c$  to have the following definition,  $g^c \not\sim p_1^c, p_2^{c \cup \{p_1\}}, \dots, p_k^{c \cup \{p_1, \dots, p_{k-1}\}} \parallel a$  which reads: in order to achieve  $g$  without deleting atoms in set  $c$ , we can first achieve  $p_1$  while maintaining  $c$ , then achieve  $p_2$  while maintaining  $c$  and  $p_1, \dots$ , until we have achieved  $p_k$ . Finally, we perform  $a$  to achieve  $g$ . For example, in Blocks World, we have the rule: clear(B) $^0$ $\not\sim$  on(A,B), clear(A) ${on(A,B)}$ , handsfree()  ${on(A,B),clear(A)}$ $\parallel$  unstack(A,B). Note that many rules might be infeasible for a given state, because it is simply not possible to achieve  $p_i$  while
$\dagger$  Fundamentally, backward search has the same worst-case time complexity as forward search. We choose to analyze backward search because goal regression rules are helpful in determining fine-grained search complexity.
$^{\ddagger}$ We show the recursive version for clarity. This algorithm can also be implemented as a breadth-first search.
maintaining  $c$  and  $\{p_1, \dots, p_{i-1}\}$ . In this case, achieving handsfree() before clear(A) is infeasible because handsfree() may be falsified while achieving clear(A), if we need to move other blocks that are currently on A. Similarly, on(A, B) must be achieved before clear(A) because moving A onto B will break clear(A) (because we need to pick up A).
The set of all possible goal regression rules can be instantiated based on all ground actions  $\mathcal{A}$  in a planning problem. Specifically, for a given atom  $g$  (e.g., holding(A)) and a set of constraints  $c$ , and for each action  $a \in \mathcal{A}$  (e.g., pick-table(A)), if  $g \in eff_{+}(a)$  while  $c \cap eff_{-}(a) = \emptyset$ , then for any permutation of  $pre(a)$ , there will be a goal regression rule:  $g^{c} \Lleftarrow p_{1}, \dots, p_{k} \| a$ , where  $p_{1}, \dots, p_{k}$  is a permutation of  $pre(a)$ . We call this regression rule set  $\mathcal{R}_0$ .
Given a set of regression rules (e.g.,  $\mathcal{R}_0$ ), we can apply the serialized goal-regression search (S-GRS) algorithm, shown in Algorithm 2. S-GRS returns a shortest path from state to achieve goal while maintaining cons. Unfortunately, this algorithm is not optimal nor even complete (e.g., a sequence of preconditions may not be "serializable," which we will define formally later), in the general case. To make it complete, it is necessary to backtrack through different action sequences to achieve each subgoal, which increases time complexity. We include a discussion in Appendix A.1.
3.2 Serialization of Goal Regression Rules
Although Algorithm 2 is not complete in general, it provides insights about goal regression. In the following, we will introduce two properties of planning problems such that S-GRS becomes optimal, complete, and efficient. For brevity, we define OptSearch(state, goal, cons) as the set of optimal trajectories that achieve goal, from state, while maintaining cons. Here, goal can be a conjunction.
Definition 3.1 (Optimal serializability). A goal regression rule  $g^{cons} \Leftrightarrow p_1, \dots, p_k \parallel a$  is optimally serializable w.r.t. a state  $s$  if and only if, for all steps  $i$ , if  $\forall \pi_{<i} \in OptSearch(s, p_{<i} \wedge \dots \wedge p_{i-1}, cons)$  and  $\forall \pi_i \in OptSearch(\mathcal{T}(s, \pi_1), p_i, cons \cup \{p_1, \dots, p_{i-1}\})$  then  $\pi_{<i} \oplus \pi_i \in OptSearch(s, p_1 \wedge \dots \wedge p_i, cons)$ . Furthermore,  $\forall \pi \in OptSearch(s, p_1 \wedge \dots \wedge p_k, cons)$ ,  $\pi \oplus a \in OptSearch(s, g, cons)^{\S}$ .
Intuitively, a rule is optimally serializable if any optimal plan for the length  $i - 1$  prefix of its preconditions can be extended into an optimal plan for achieving the length  $i$  prefix. For example, in Blocks World, the rule holding(A)  $\Leftarrow$  on(A,B), clear(A), handsfree() || unstack(A,B) is optimally serializable when on(A,B) is true for  $s$ . We define the set of optimally serializable rules for a state  $s$  as  $OSR(s)$ , and the set of single-literal goals that can be solved with  $OSR(\cdot)$  from  $s_0$  as  $OSG(s_0)^{\sharp}$ .
Theorem 3.1. For any goal  $g \in OSG(s_0)$ , S-GRS is optimal and complete. See proof in Appendix A.2.
3.3 Width of Search Problems
Chen and Giménez [2007] introduced the notion of the width of a planning problem and showed that the forward-search complexity is exponential in the width of a problem. We extend this notion to regression-based searches and introduce a generalized version of regression rules in which not all previously achieved preconditions must be explicitly maintained.
Definition 3.2 (Generalized regression rules). A generalized, optimally-serializable regression rule is  $g^{c} \Leftrightarrow p_{1}^{c \cup c_{1}}, p_{2}^{c \cup c_{2}}, \dots, p_{k}^{c \cup c_{k}} \| a$ , where  $c_{i} \subseteq \{p_{1}, p_{2}, \dots, p_{i - 1}\}$ . Its width is  $|c| + \max \{|c_{1}|, \dots, |c_{k}|\}$ .
For example, for any initial state  $s_0$  that satisfies on(A,B), the following generalized rule is optimally serializable: clear(B) $\varnothing$ $\rightsquigarrow$  on(A,B) $\varnothing$ , clear(A) $\varnothing$ , handsfree()  ${}^{[clear(A)]}$ $\parallel$  unstack(A,B). When we plan for the second precondition clear(A), we can ignore the first condition on(A,B), because the optimal plan for clear(A) will not change on(A,B). Using such generalized rules to replace the default rule set  $\mathcal{R}_0$  improves search efficiency by reducing the number of possible subgoals. To define existence conditions for highly efficient search algorithms and policies, we need a stronger notion:
Definition 3.3 (Strong optimally-serializable (SOS) width of regression rules). A generalized regression rule has strong optimally-serializable width  $k$  w.r.t. state  $s$  if (a) it is optimally serializable, (b) its width is  $k$ , and, (c)  $\forall c \text{OptSearch}(s, p_i \cup c_i \cup c, \emptyset) \subseteq \text{OptSearch}(s, \{p_1, \dots, p_i\} \cup c, \emptyset)$ .
In the Blocks World domain, given a particular state  $s$ , if  $on(A,B) \in s$ , the generalized goal regression rule  $clear(B) \Leftarrow on(A,B)$ ,  $clear(A)$ ,  $handsfree()^{[clear(A)]}$ $\parallel$  unstack(A,B) is strong optimally-serializable, because the optimal trajectory to achieve  $clear(A)$  will not move A.
$^{\S}a\oplus b$  denotes concatenation of sequences.  $\pi_{< i}$  denotes the subsequence  $\{a_0,\dots ,a_{i - 1}\}$
It is possible to make a non-optimally-serializable rule into an optimally serializable one by introducing "super predicates" as the conjunction of existing predicates, at the cost of increasing the state and action space sizes. We include discussions in Appendix A.3.
Definition 3.4 (SOS width of problems). A planning problem  $P = \langle S, s_0, g, \mathcal{A}, \mathcal{T} \rangle$  has strong optimally-serializable width  $k$  if there exists a set of strong optimally-serializable width  $k$  rules  $\mathcal{R}$  w.r.t.  $s_0$ , such that  $S$ -GRS can solve  $P$  using only rules in  $\mathcal{R}$ .
Theorem 3.2. If a problem  $P$  is of SOS width  $k$ , it can be solved in time  $O(N^{k + 1})$  with s-GRS, where  $N$  is the number of atoms. Here we have omitted polynomials related to enumerating all possible actions and their serializations (which are polynomial w.r.t. the number of objects in  $\mathcal{U}$ ).
Proof idea. We start by assuming knowing  $\mathcal{R}$ : the proof can be simply done by analyzing function calls: there are at most  $O(N^{k + 1})$  possible argument combinations of (goal, cons). Next, since the enumeration of all possible width-  $k$  rules can be done in polynomial time. Therefore, the algorithm runs in polynomial time and it does not need to know  $\mathcal{R}$  a priori. Note that the SOS condition cannot be removed because even if regression rules have a small width, there will be a possibly exponential branching factor caused by "free variables" in regression rules that do not appear in the goal atom. See the full proof in Appendix A.4. We also describe the connections with classical planning width and other related concepts, as well as how to generalize to  $\forall$ -quantified preconditions in Appendix A.5.
Importantly, planning problems such as Blocks World and Logistics have constant SOS widths, independent of the number of objects. Therefore, there exist polynomial-time algorithms for finding their solutions. Our analysis is inspired by and closely related to classical (forward) planning width Chen and Giménez [2007] and Lipovetzky and Geffner [2012]. Indeed, we have:
Theorem 3.3. Any planning problem that has SOS-width  $k$  has a forward width of at most  $k + 1$  and, hence, can be solved by the algorithm  $IW(k + 1)$ . See the full proof in Appendix A.6.
Remark. Unfortunately, SOS width and forward width are not exactly equivalent to each other. There exist problems whose forward width is smaller than  $k + 1$ , where  $k$  is the SOS width. Appendix A.6 presents a concrete example. However, there are two advantages of our SOS width analysis over the forward width analysis. In particular, proofs for (forward) widths of planning problems were mostly done by analyzing solution structures. By contrast, our constructions form a new perspective: a problem has a large width if the number of constraints to track during a goal regression search is large. This view is helpful because our analyses can now focus on concrete regression rules for each individual operator in the domain — how its preconditions can be serialized. Second and more importantly, our "backward" view allows us to reason about circuit complexities of policies.
4 Policy Realization
We have studied sequential backward search algorithms for planning; now, we address the questions of how best to "compile" or "learn" them as feed-forward circuits representing goal-conditioned policies, and how the size of these circuits depends on properties of the problem. Here, we will focus on understanding the complexity of a certain problem class; that is, a set of problems in the same domain (e.g., Blocks World) with a similar goal predicate (e.g., having one particular block clear). In particular, we will make use of relational neural networks (RelNN, such as graph neural networks [Morris et al., 2019] and Transformers [Vaswani et al., 2017]). They accept inputs of arbitrary size, which is critical for building solutions to problems with arbitrary-sized universes. Another critical fact is that these networks have "parameter-tying" in the sense that there is a constant-sized set of parameters that is re-used according to a given pattern to realize a network.
In particular, as illustrated in Fig. 2, we will show how to take the set of operator descriptions for a planning domain and construct a relational neural network that represents a goal-conditioned policy: it takes a state and a goal (coded as a special predicate, e.g., on-goal) and outputs a ground action. Here, the state is represented as a graph (objects are nodes, and relations are edges; possibly there will be hyperedges for high-arity relations). For example, the input contains nullary features (a single vector for the entire environment), unary features (a vector for each object), binary features (a vector for each pair of objects), etc. The goal predicate can be represented similarly. These inputs are usually represented as tensors. The complexity of this neural network will depend on the SOS width of problems in the problem class of interest.
We will develop two compilation strategies. First, in Section 4.2, we directly compile the BWD and the S-GRS algorithms into RelNN circuits. For problems with constant SOS width that is independent of the size of the universe, there will exist finite-breadth circuits (but the depth may depend on the size of the universe). In Section 4.3, we discuss a new property of planning problems—the existence of regression rule selector circuits—which results in policy circuits that are smaller than those generated
by previous strategies, and they are potentially of finite depth. The construction of the regression rule selector is not automatic and will generally require human priors or machine learning.
4.1 Relational Neural Network Background