Learning to Search and Searching to Learn for Generalization in Planning

TL;DR: Take the classic AlphaZero-style "search generates training data, learning improves search" loop, but swap Monte Carlo tree search for weighted A* — the workhorse of classical planning — and represent the Q-function with a relational GNN over PDDL predicates instead of a convnet over pixels. The result (GSP, "Generalized Search for Planning") learns heuristics that transfer across problem sizes, not just states: trained on Blocksworld with ≤29 blocks, the greedy policy solves all 90 IPC test instances including one with 488 blocks and a 1,786-step plan, with no search at test time. On Sokoban it solves all 1000 standard test puzzles with ~16× fewer node expansions than the best prior search-and-learn systems. The most interesting finding is a subtle one: the learned Q-function is a much better local action ranker than a globally calibrated value function, which is why the searchless policy sometimes beats the search that trained it.

The problem: three kinds of generalization, and why RL fails at the third

Classical planning gives you a clean laboratory for studying combinatorial generalization. A domain (say, Blocksworld) fixes a relational vocabulary — predicates like \mathit{on}(x,y), \mathit{clear}(x) — and a handful of action schemas like \mathit{stack}(x,y). An instance fills in the objects, the initial state, and the goal. Because the vocabulary is shared across all instances, "generalize to a 488-block problem after training on 29-block problems" is a well-posed question, not a hope. The paper distinguishes three levels: generalization across states (what any RL value function must do), across states and goals (goal-conditioned RL), and across states, goals, and problem size — the target here.

The catch is that planning problems are deterministic goal-reaching MDPs with reward −1 per step and nothing else: the sparsest possible reward. Standard deep RL explores via real-time search — the agent occupies one state, takes an action, occupies the next. In a Sokoban instance with a 10,000-step optimal plan, random real-time exploration will never see the goal, so there's no learning signal. The field has worked around this with expert demonstrations, hindsight relabeling (HER), or random walks backward from the goal — but hindsight relabeling assumes goals decompose into relabelable subgoals, which fails on many puzzles, and backward walks need an invertible representation.

Meanwhile, classical planners never search this way. They use best-first search — A*, weighted A*, greedy best-first — which maintains a global frontier of candidate nodes and always expands the most promising one anywhere in the tree. Best-first search is systematic, complete, and immune to dead ends. When you have the transition model (and in planning you always do), there is simply no reason to shackle exploration to a single physical agent walking one step at a time. Prior work exploited this for single environments (DeepCubeA for Rubik's cube, policy-guided tree search for Sokoban), but those systems learn a value function for one fixed state space. The other camp — learned general heuristics for planning — trains GNNs supervised on optimal costs precomputed by planners, which caps them at what the teacher planner can solve and forfeits any self-improvement loop.

GSP's proposition: use best-first search as the exploration mechanism inside the RL loop, over an entire family of instances of varying size, and let a size-agnostic relational network close the loop.

The mechanism

Search episode. On a sampled training instance, GSP runs weighted A* over state–action pairs, scoring each frontier pair by f(s,a) = g(s) + w\,Q_\theta(s,a), where g(s) is the (negative) depth of s in the search tree, Q_\theta is the learned action-value, and w=2 throughout. Rewards are −1 per step, undiscounted, so higher Q means fewer steps to goal.

Three flavors of training target. This is where search pays off over plain DQN. When the search reaches a goal, it backtracks along the solution path and assigns each pair on the path its actual return-to-go \underline{R}. Since a shorter plan might exist, this is a lower bound on the optimal return, so the target becomes y = \max\{\underline{R},\ \hat{y}\} where \hat{y} is the usual one-step Bellman bootstrap. The max prevents bootstrapped targets from regressing below what search has already demonstrated — the ablations show removing this hurts most of any component. Dead-end successors get a fixed penalty target R_\bot. Everything else gets standard Q-learning. The whole thing is a clean marriage of certificate-based supervision (search proves a path exists) with bootstrapping (which propagates value beyond found paths).

Curriculum by bookkeeping. Instances are sorted into three pools — unsolved, satisficed (a plan was found), and solved (the search expanded exactly as many nodes as the plan is long, i.e., the heuristic guided it perfectly) — and sampled with exponentially increasing weight toward satisficed instances. The intuition: perfectly-solved instances teach nothing new, unsolved ones give no signal, and sloppily-solved ones are exactly where the heuristic is wrong in a fixable way.

The size-agnostic network. Q_\theta is a relational GNN operating directly on the set of true ground atoms of the state plus the goal atoms. The clever architectural bit: every applicable ground action a = A(\bar{o}) is injected as an extra action object o_a with an atom A(o_a, \bar{o}) linking it to its arguments. Message passing then runs over objects and action-objects jointly for L layers (predicate-specific message MLPs, permutation-invariant smoothmax aggregation, shared residual updates), and a single shared readout MLP scores each action from its embedding concatenated with a pooled state summary. Nothing in this construction references the number of objects, so the same 32-dimensional-embedding network evaluates a 3-block state or a 488-block state. One set of hyperparameters across every experiment in the paper.

The evidence

IPC 2023 learning track (10 domains, 90 test instances each, deliberately scaled to break planners — one Childsnack instance has 46.7 million applicable actions in the initial state). The greedy searchless policy GSP_\pi hits 100% coverage on Blocksworld, Miconic, and Spanner, beating every baseline including the LAMA planner:

The AlphaZero baseline — same relational architecture, MCTS instead of WA* as the learning driver — collapses almost everywhere, mostly 0–31% coverage. That's the paper's central comparative claim vindicated: for single-goal, sparse-reward pathfinding with a known model, MCTS is the wrong search engine, and swapping it for best-first search is what makes the loop turn over.

Sokoban and The Witness (fixed-size puzzles, same benchmark as Orseau & Lelis 2021). GSP with WA* solves all 1000 Sokoban test instances with dramatically fewer expansions than anything prior:

The greedy policy alone solves 681/1000 Sokoban instances (vs. 309 for Lifted HER's policy) and 667/1000 Witness instances — a searchless policy beating the searching GBFS baseline (290) on The Witness, which the authors attribute to explicit dead-end supervision, something goal-relabeling methods can't easily provide.

PushWorld shows the same pattern against deep RL: GSP_\pi solves 93/200 Level-0 test instances vs. 20 for DQN and 11 for PPO; with WA*, 200/200 with plans shorter than LAMA's (18 vs. 24 steps), and Level-1 transfer with ~30× fewer expansions than LAMA on jointly-solved instances.

The interesting wrinkle: ranking beats calibration

Look back at the IPC table and notice something odd: on Blocksworld the searchless policy gets 100%, but using the same Q-function as a WA* heuristic drops coverage to 79%. Search makes things worse. The authors' diagnosis (Section 6.1, worth reading carefully) is the most transferable insight in the paper.

Because all applicable actions are embedded into one joint relational graph, the GNN scores actions in context with each other — it learns to answer "which of these moves is best here?" extremely reliably, and that relative ranking survives extrapolation to instances 16× larger than anything seen in training. But WA* needs more: it must compare f-values across the frontier — a node at depth 40 in one subtree against a node at depth 3 in another — which requires the Q-values to be globally calibrated as actual costs-to-go. That calibration degrades out of distribution much faster than local ranking does. So on huge instances the greedy policy glides along on relative judgments while WA* thrashes among incomparably-mis-scaled frontier scores (amplified by the benchmark's enormous branching factors).

The flip side holds too: on intractable domains where no compact perfect policy exists — Sokoban, Floortile — the policy is weak (14%, 20%) and search rescues it (32%, 28%). Policy and search are complementary readouts of one learned object, and which one to trust depends on whether the domain admits a compact general policy at all.

What to be skeptical about

The headline "30 blocks → 488 blocks" is real but chosen kindly: Blocksworld is exactly the sort of domain where a simple relational policy generalizes indefinitely once found, and Lifted HER already gets 98% there. The genuinely hard evidence is the Sokoban efficiency numbers and the AlphaZero contrast.

Honest failures, which the paper reports plainly: on the 24-puzzle, GSP never finds a single goal during training — the fixed 60-second search budget with an uninformed initial heuristic can't crack any training instance, so the loop never starts, while baselines that adaptively grow budgets (and use cheap fixed-size MLPs rather than GNNs) bootstrap successfully. This is a real cold-start fragility: self-improvement requires at least one solvable rung on the ladder. On Rovers and Satellite, GSP loses badly to LAMA for a known theoretical reason — these domains require features beyond 1-WL expressivity, a ceiling shared by all GNN-based methods. Childsnack's 46M-action branching factors defeat everyone except a symmetry-pruning method. And the whole framework presupposes clean PDDL-style relational state descriptions and a known transition model — no perception, no stochasticity — so this is a statement about the structure of the learning loop, not a deployable robot planner.

Also note the model selection: best checkpoint across five seeds by validation coverage, and the ablation table shows non-trivial seed variance (e.g., Blocksworld GSP_\pi at 83%±24% across seeds vs. the reported 100% best run). The headline numbers are best-of-five.

Why it matters

If the result holds up, the recipe is broadly reusable: when the model is known, exploration should be best-first, not real-time, and search should feed the learner not just trajectories but certificates — proven lower bounds and identified dead ends — that anchor bootstrapped values. That reframes the AlphaZero template for the large class of single-agent, sparse-reward, known-model problems where MCTS underperforms. The ranking-vs-calibration diagnosis also suggests a concrete research direction the authors flag: hybrid readouts, or training objectives that enforce cross-state calibration rather than only within-state ordering (and, from the appendix, multi-queue search combining Q_\theta with classical h_{\mathrm{ff}}, whose failures are complementary).

If you read one part of the paper, make it the results discussion at the end of Section 6.1 (the local-ranking vs. global-calibration analysis) plus Table 3 — that's where both the capability jump and the conceptual takeaway live. Section 4 is a quick read if you want the algorithm precisely; the appendix's Spanner discussion (a greedy suboptimal strategy that generalizes better than the optimal one, destabilizing training) is a nice bonus puzzle about what "the right" general policy even is.