🤖 AI

Differentiable Learning of Lifted Action Schemas for Classical Planning

This paper introduces a novel neural network architecture that robustly learns lifted action schemas and infers unobserved action arguments from fully observed state traces, serving as a differentiable component for neuro-symbolic planning models.

Original authors: Jonas Reiter, Jakob Elias Gebler, Hector Geffner

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jonas Reiter, Jakob Elias Gebler, Hector Geffner

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to play a game like Sokoban (pushing boxes) or Tower of Hanoi (moving disks), but you have a very specific constraint: you can see the board before and after every move, and you know what move was made (e.g., "move the block"), but you cannot see the robot's internal instructions on which specific blocks it chose to move.

This is the puzzle the paper "Differentiable Learning of Lifted Action Schemas for Classical Planning" tries to solve. The authors, Jonas Reiter, Jakob Elias Gebler, and Hector Geffner, have built a new type of AI called DIAS (Differentiable Induction of Action Schemas) to figure out the hidden rules of the game just by watching the board change.

Here is a simple breakdown of how they did it, using everyday analogies.

1. The Problem: The "Black Box" Chef

Imagine a master chef cooking a complex meal. You can see the ingredients on the counter before they start (State A) and the finished dish on the plate afterward (State B). You also know the chef used the "Chop" action.

However, you don't know which specific carrot the chef chopped. Did they chop the big one? The small one? The one that was already peeled?

Old methods usually required you to tell the AI exactly which carrot was chopped.
DIAS has to figure out: "Based on how the pile of carrots changed, the chef must have chopped the big one."

The goal is to learn the general rule (the "Schema") that applies to any carrot, not just the specific one in the video. This is like learning the recipe "Chop any vegetable" rather than just "Chop this specific carrot."

2. The Solution: The "Detective" Neural Network

The authors created a neural network that acts like a detective. It works in two main stages:

Stage 1: The "Who Did It?" Detective (Selection)

The AI looks at the "before" and "after" pictures of the game board. It uses a Graph Neural Network (GNN)—think of this as a super-smart eye that sees the relationships between objects (e.g., "Block A is on top of Block B").

The AI creates a "fingerprint" (embedding) for every object on the board.
It then asks: "Which of these fingerprints matches the 'slot' for the action 'Move'?"
It uses a mathematical trick called Sinkhorn (imagine a very efficient way of sorting cards into piles) to assign the right objects to the right roles in the action. It's like the AI saying, "I'm 90% sure the robot moved the red block, not the blue one."

Stage 2: The "What Happened?" Detective (Effects)

Once the AI guesses which objects were involved, it tries to learn the rules of the game:

Preconditions: What had to be true before the move? (e.g., "The block must be clear on top").
Effects: What changed because of the move? (e.g., "The block is now on the table").

The AI writes these rules down in a symbolic format (like a computer code called PDDL). It then simulates the move using these guessed rules to see if it predicts the "after" picture correctly. If the simulation matches the real "after" picture, the AI gets a "good job" signal. If not, it tweaks its guesses and tries again.

3. The "Magic" Ingredient: Differentiable Learning

Usually, figuring out "which object was moved" is a binary choice (it's either the red block or the blue block). This is hard for AI to learn because you can't easily "nudge" the answer halfway.

The paper's breakthrough is making this process differentiable.

Analogy: Imagine you are trying to tune a radio to a clear station. Instead of jumping from station 1 to station 2, you can slowly slide the dial.
DIAS doesn't just guess "Red Block." It guesses "80% Red Block, 20% Blue Block." This allows the AI to use gradient descent (a standard machine learning technique) to slowly slide its guesses toward the perfect answer, rather than getting stuck in a loop of wrong guesses.

4. What They Found (The Results)

The team tested DIAS on 13 different classic planning domains (like Blocksworld, Logistics, and Hanoi).

Perfect Scores: When they gave the AI the full list of arguments (telling it exactly which blocks moved), it learned the rules perfectly every single time.
The Hard Mode: When they hid the arguments (only giving the action name, like "move"), it still learned the rules perfectly in 8 out of 13 domains. In the others, it was very close.
Noise Resistance: They tested it with "noisy" data (where some facts about the board were flipped randomly, like saying a block is clear when it's actually covered). DIAS handled this surprisingly well, though too much noise eventually confused it.
Comparison: They compared their method to a traditional symbolic method (L1). DIAS was much better at finding the correct rules, especially in complex domains where the old method failed to find all the necessary conditions.

5. Why This Matters (Without Overpromising)

The paper claims this is a simplification of a much harder problem: learning planning rules directly from images (like watching a video of a robot arm moving blocks).

The Current Achievement: They solved the "middle step" perfectly. They proved that if you can see the symbolic state (the list of facts) but not the arguments, you can still learn the rules using a neural network.
The Future Goal: The authors hope to eventually plug this "DIAS" module into a larger system that looks at raw images and learns the rules directly, without needing a human to translate the image into a list of facts first.

In short: The paper presents a new AI detective that can watch a game, guess which pieces moved, and deduce the universal rules of the game, all by using a smooth, math-based learning process that avoids getting stuck in dead ends.

Technical Summary: Differentiable Learning of Lifted Action Schemas for Classical Planning

Problem Formulation

The paper addresses the challenge of learning lifted action schemas for classical planning domains from observational traces. Specifically, the authors tackle a setting where states are fully observed (represented as grounded STRIPS atoms) but action arguments are unobserved. The input consists of triplets $\langle s, a, s' \rangle$ , where $s$ and $s'$ are grounded states and $a$ is either a full action name, a STRIPS+ action (with some arguments observed), or an action name only.

The core difficulty lies in the isomorphism-hard problem of simultaneously inferring:

Action Arguments: Identifying which objects in the state bind to the parameters of an action.
Lifted Schemas: Learning the preconditions and effects (add/delete lists) for each action schema in a lifted, symbolic form.

This problem is framed as a simplification of the more realistic task of learning from image sequences and action labels, aiming to solve the symbolic abstraction step with near-perfect accuracy to enable integration into larger neuro-symbolic systems.

Methodology: DIAS Architecture

The authors propose DIAS (Differentiable Induction of Action Schemas), a neural architecture that learns action schemas via gradient descent. The approach operates in two main stages, utilizing soft, differentiable operations to allow end-to-end training.

1. Object Selection via Graph Neural Networks (GNN)

Graph Construction: The state $s$ and the transition to $s'$ are encoded as a directed graph $G$ . Nodes represent objects, and edges represent predicates (atoms) true in the state or added/deleted during the transition.
Embedding: A Relational Graph Convolutional Network (R-GCN) processes the graph to generate $d$ -dimensional embedding vectors (keys $K$ ) for each object.
Slot Assignment: For each action name $a$ , the model maintains learnable query vectors $Q_a$ corresponding to a fixed number of "slots" (an upper bound on action arity).
Soft Matching: A correspondence matrix is computed via scaled dot-product attention between object keys and action queries. A rectangular Sinkhorn algorithm is then applied to produce a soft assignment matrix $\tilde{S}_a$ , probabilistically mapping objects to action argument slots. Learnable slot weights $w_a$ determine the activation probability of each slot, allowing the model to infer the action's arity.

2. Learning Preconditions and Effects

Lifted Parameters: The model maintains learnable probability tensors for every lifted atom $p(x_i, x_j)$ to determine if it is a positive/negative precondition, an add effect, or a delete effect.
Grounding: The soft assignment $\tilde{S}_a$ projects these lifted probabilities into the grounded adjacency matrix space, yielding probabilities for specific atoms being preconditions or effects in the current state.
State Prediction:
- Precondition Check: The model computes the probability $p_{pre}$ that all preconditions are met by evaluating the conjunction of individual atom probabilities. To avoid vanishing gradients from random initialization, a geometric mean scheduling strategy is used, gradually transitioning to the true product.
- Effect Application: The predicted successor state $\hat{A}'$ is computed by applying the grounded add/delete probabilities to the current state adjacency matrix $A$ , scaled by $p_{pre}$ .

3. Loss Function and Optimization

The training objective minimizes a composite loss:

Main Loss ( $L_{adj}$ ): Binary cross-entropy between the predicted successor state $\hat{A}'$ and the true successor state $A'$ . This drives the model to correctly select arguments and predict state changes.
Auxiliary Loss ( $L_{aux}$ ): Encourages minimal effects and maximal preconditions. This prevents the model from learning trivial solutions (e.g., assuming all atoms are effects or no preconditions exist) by penalizing inactive effects and inactive preconditions.
Gradient Projection: To ensure the primary goal of state prediction is not compromised by the auxiliary objectives, the authors employ PCGrad (Gradient Surgery). The auxiliary gradient is projected onto the hyperplane orthogonal to the main gradient if they conflict, ensuring $L_{main}$ optimization takes precedence.

Key Results

The authors evaluated DIAS on 13 standard IPC planning domains (e.g., Blocks, Logistics, Hanoi, Sokoban) using three observation variants: full STRIPS actions, STRIPS+ actions, and action names only.

Perfect Recovery: In 8 out of 13 domains, DIAS learned the ground-truth lifted action schemas perfectly in all 10 training runs across all three observation variants.
Robustness:
- Noise: The model demonstrated robustness to observation noise (randomly flipped atoms) up to 2 flips per transition for most domains. Performance degraded significantly at 8 flips.
- Comparison to Symbolic Methods: Compared against the L1 algorithm (a SAT-based approach), DIAS achieved superior or comparable performance, particularly in domains where L1 failed to recover preconditions involving variables not present in effects.
- Failure Analysis: Failures were primarily attributed to:
  - Dead ends: Domains like Sokoban where random walks fail to explore critical state regions.
  - State-space structure: The Hanoi domain's Sierpiński triangle structure limits random walk coverage.
  - Redundant actions: The Satellite domain, where actions like take_image often result in no state change, confounding the learning of effects.
MLP Ablation: An ablation replacing STRIPS effects with a Multi-Layer Perceptron (MLP) showed that while the architecture can learn state-dependent dynamics, it fails to produce symbolic schemas suitable for classical planning, highlighting the necessity of the lifted STRIPS formulation for generalization.

Significance and Claims

The paper claims that DIAS represents a significant step toward learning planning domains from raw observations (such as images) by solving the intermediate symbolic abstraction problem with high fidelity.

Differentiability: The primary contribution is a robust differentiable component that can learn lifted action schemas without requiring explicit action arguments, bridging the gap between neural observation processing and symbolic planning.
Generalization: Unlike slot-based dynamic models (e.g., RIMs) which struggle with varying numbers of objects, DIAS learns lifted representations that generalize to arbitrary numbers of objects, a property inherent to STRIPS but difficult to learn from data.
Modest Scope: The authors explicitly frame this work as a simplification of the broader challenge of learning from images. They do not claim to solve the image-to-planning problem directly but argue that achieving "near-perfect" performance on the symbolic abstraction step is a necessary and challenging prerequisite for such systems.
Limitations: The approach assumes pure STRIPS domains (no conditional effects or numerical features) and relies on the assumption that action arguments are injectively bound to objects. The method also struggles with domains where random walk sampling provides insufficient coverage of the state space.