Learning to Search and Searching to Learn for Generalization in Planning

Combinatorial generalization remains a central challenge in Deep Reinforcement Learning (DRL). Classical planning provides a simple yet challenging setting to study this problem through explicit relational descriptions, without requiring learning from perception. In sparse-reward domains, standard RL exploration via real-time search is ineffective, and learning-based planning methods often rely on expert demonstrations, hindsight relabeling, or random walks from the goal state. In contrast, planners rely on best-first search methods such as $\mathrm{A}^\star$ to solve problems from scratch. We propose a self-improving $\mathrm{WA}^\star$ learning framework in combination with a value heuristic represented by a Relational Graph Neural Network: the heuristic guides search, and the resulting search data updates the heuristic via $Q$-learning. This loop yields heuristics that can function as general policies and solve new instances even without search, where DRL otherwise fails, as we show on puzzles such as Sokoban, PushWorld, The Witness, and the 2023 International Planning Competition benchmarks. Notably, we demonstrate strong zero-shot generalization: For example, heuristics trained on Blocksworld instances with fewer than $30$ blocks successfully solve instances with $488$ blocks without search.

Many AI systems can learn to solve small examples of a task, but fail when the same task becomes larger or structurally different. This is a major challenge in planning problems, where an AI system must find a sequence of actions to reach a goal, such as solving puzzles, moving objects, or arranging blocks. Our work studies how to make learning systems generalize better across such problems. Instead of relying on trial-and-error exploration one step at a time, we let the system use a planning-style search to find solutions. The results of this search are then used to train a neural network to better estimate which actions are promising. As the network improves, it guides future searches more effectively, creating a self-improving loop between searching and learning. Because the model represents problems in terms of objects and their relations, the same learned strategy can be applied to larger problem instances with many more objects. In experiments on planning benchmarks and puzzle domains, this approach solves many unseen and larger problems, sometimes even without search at test time. This suggests that combining structured search with learning can make AI systems more reliable on tasks that require systematic generalization.