Technical Summary: Learning to Search and Searching to Learn for Generalization in Planning

Problem Statement

Combinatorial generalization remains a central challenge in Deep Reinforcement Learning (DRL), where policies or value functions must generalize to out-of-distribution scenarios defined by a common problem structure. Classical planning offers an ideal testbed for this problem because the underlying relational structure (predicates and action schemas) is fixed and explicit, while instances vary in object counts, initial states, and goals.

However, standard DRL approaches struggle in sparse-reward planning domains. Real-time search (iterative state-to-state transitions), typical in RL algorithms like LRTA* or RTDP, is ineffective for finding plans in these settings. Conversely, classical planners use best-first search (e.g., A*, WA*) but rely on hand-crafted heuristics. While recent work has attempted to learn general policies or heuristics, existing methods often rely on real-time search, supervised learning from expert demonstrations, or hindsight relabeling, failing to establish a self-improving loop where search and learning mutually reinforce each other to achieve zero-shot generalization across varying problem sizes.

Methodology: Generalized Search for Planning (GSP)

The authors propose GSP, a self-improving framework that combines best-first search with Q-learning to learn general action-value functions ($Q_\theta$).

Core Loop

1. Search-Guided Learning: Instead of real-time search, GSP uses Weighted A (WA)** to solve training instances. The search is guided by the current learned $Q_\theta(s, a)$ values.
2. Data Generation: As WA* expands nodes, it generates a replay buffer of transitions. Crucially, the search provides supervision signals that standard RL lacks:
   • Goal Paths: When a goal is reached, the algorithm backtracks to assign actual return-to-go values to the pairs on the solution path.
   • Dead-Ends: If a state is a dead-end (no applicable actions), the parent pair is assigned a fixed penalty return.
   • Lower Bounds: The search provides a lower bound on the optimal return for any state-action pair encountered on a solution path.
3. Q-Learning Update: The $Q_\theta$ network is updated using these search-derived targets. For goal-path pairs, the target is $\max(R, \hat{y})$, where $R$ is the search-derived lower bound and $\hat{y}$ is the standard bootstrap target. This prevents the value estimates from dropping below what the search has already proven achievable.
4. Instance Selection: To optimize learning efficiency, the framework maintains three pools of instances (unsolved, solved, and satisficed) and samples from them with exponentially increasing weights, prioritizing instances that provide the most informative updates (e.g., those with suboptimal solutions).

Representation: Relational GNNs

To ensure generalization across different numbers of objects and goals, $Q_\theta$ is represented by a Relational Graph Neural Network (R-GNN).

• Input: A state $s$ and goal $g$ are encoded as a set of ground atoms. Applicable actions are explicitly represented by introducing auxiliary "action objects" and atoms, allowing the network to evaluate specific action choices within the context of the state.
• Architecture: The network uses message passing over the relational structure. Predicate-specific MLPs generate messages based on argument roles, which are aggregated via permutation-invariant functions.
• Output: A shared readout MLP combines the embedding of the specific action object with a pooled summary of the state and goal to produce the Q-value. This architecture ensures the model learns a common scoring principle for grounded actions rather than separate predictors for each schema.

Key Contributions

1. Self-Improving Search-Learning Cycle: The paper introduces a framework where best-first search (WA*) generates high-quality training data (including lower bounds and dead-end signals) to update a learned heuristic, which in turn guides more efficient future searches.
2. Zero-Shot Generalization: The learned Q-functions function as general policies capable of solving instances with significantly more objects and different goals than seen during training, without requiring search at test time in many cases.
3. Unified Framework: The approach applies uniformly to diverse domains, from the 2023 International Planning Competition (IPC) benchmarks to combinatorial puzzles (Sokoban, The Witness) and the PushWorld domain, using a single architecture and hyperparameter set.
4. Superiority over Real-Time Search: The work demonstrates that replacing real-time search with best-first search in the learning loop is critical for success in sparse-reward, high-branching-factor planning domains.

Experimental Results

The authors evaluated GSP on three categories of benchmarks:

• 2023 IPC Learning Track:

  • Greedy Execution ($GSP_\pi$): Achieved 100% coverage on Blocksworld, Miconic, and Spanner, outperforming baselines like Lifted HER, LAMA, and AlphaZero-style policies.
  • Search-Guided ($GSP_{WA^*}$): Showed mixed results; while it improved coverage on puzzles like Sokoban (14% $\to$ 32%), it sometimes reduced coverage on domains with massive branching factors (e.g., Blocksworld) where the learned heuristic, while good for local ranking, lacked global calibration.
  • Scaling: Heuristics trained on Blocksworld instances with $<30$ blocks successfully solved instances with 488 blocks without search.

• Combinatorial Puzzles (Sokoban, The Witness, 24-Puzzle):

  • On Sokoban, $GSP_{WA^*}$ solved 100% of instances with an average of 207 expansions, significantly outperforming WA* with standard heuristics (3298 expansions) and DeepCubeA (1050 expansions).
  • On The Witness, the method solved all instances, demonstrating the value of training on dead-end states.
  • On 24-Puzzle, the method failed to find solutions during training, highlighting that the approach requires the ability to find at least one solution to bootstrap learning, which is difficult in extremely hard fixed-size puzzles without initial informed heuristics.

• PushWorld:

  • Trained on Level 0, the model generalized to Level 1. $GSP_{WA^*}$ solved 48/63 Level 1 instances with shorter plans and ~30$\times$ fewer expansions than the classical planner LAMA.
  • As a greedy policy, it significantly outperformed model-free RL baselines (DQN, PPO).

Significance and Claims

The paper claims that best-first search is a superior exploration mechanism for generalized reinforcement learning when the transition model is known, particularly in sparse-reward domains where real-time search fails to discover informative trajectories.

The primary significance lies in the demonstration of a self-reinforcing cycle: the search provides the training targets (lower bounds, dead-ends) that allow the heuristic to improve, and the improved heuristic makes subsequent searches more efficient. This allows the system to learn general policies that can solve problems of arbitrary size and complexity without search, or highly efficient heuristics that guide search to optimal solutions.

The authors modestly note limitations: the learned Q-function often acts as a strong local action-ranking mechanism but may lack global calibration across the entire search frontier, which can lead to suboptimal performance in WA* on domains with extremely large branching factors. They suggest that future work could focus on improving value generalization across states and combining best-first search with complementary strengths of real-time search and hindsight experience relabeling.