Mimir-RGNN

Mimir-RGNN is a Python library that implements Relational Graph Neural Networks (R-GNN) for AI planning applications. Built on PyTorch and Mimir, it provides a powerful and flexible interface for learning on structured relational data, particularly PDDL planning domains.

Key Features

• 🧠 Relational Graph Neural Networks: R-GNN implementation for structured reasoning
• 📋 PDDL Integration: Seamless integration with PDDL planning domains and problems via Mimir
• ⚡ PyTorch Backend: Built on PyTorch for GPU acceleration
• 🔧 Flexible Configuration: Declarative configuration system for input/output specifications
• 🎯 Planning-Focused: Designed specifically for AI planning and reinforcement learning applications
• 📊 Multiple Aggregation Functions: Support for various message aggregation strategies
• 🏗️ Typed API: Clean and type-safe interface

Installation

Install Mimir-RGNN from PyPI:

pip install pymimir-rgnn

Requirements

• Python 3.11+
• PyTorch 2.6.0+
• Pymimir 0.14.0b6

Quick Start

import pymimir as mm
import pymimir_rgnn as rgnn

# Load a PDDL domain
domain = mm.Domain.from_file('path/to/domain.pddl')

# Configure the R-GNN hyperparameters
hparam_config = rgnn.HyperparameterConfig(
    domain=domain,
    embedding_size=64,
    num_layers=30,
)

# Define input and output specifications using encoder/decoder classes
input_spec = (rgnn.StateEncoder(), rgnn.GroundActionsEncoder(), rgnn.GoalEncoder())
output_spec = [('q_values', rgnn.ActionScalarDecoder(hparam_config))]

# Configure the R-GNN modules (aggregation, message, and update functions)
module_config = rgnn.ModuleConfig(
    aggregation_function=rgnn.MeanAggregation(),
    message_function=rgnn.PredicateMLPMessages(hparam_config, input_spec),
    update_function=rgnn.MLPUpdates(hparam_config)
)

# Create and initialize the model
model = rgnn.RelationalGraphNeuralNetwork(hparam_config, module_config, input_spec, output_spec)

# Use the model for inference
# problem = mm.Problem.from_file(domain, 'path/to/problem.pddl')
# state = problem.initial_state
# actions = state.applicable_actions()
# goal = problem.goal
#
# inputs = [(state, actions, goal)]  # Input tuple matching input_spec order
# outputs = model(inputs)
# q_values = outputs.readout('q_values')

Pymimir 0.14.0b6 provides the native pymimir.learning extraction used by all built-in RGNN encoders. One native encoding context owns the complete input batch: RGNN begins and ends each instance around the existing instance-major, polymorphic encoder loop, then materializes one packed int32 relation buffer and the node metadata once. PyTorch creates one CPU tensor from that buffer; each native relation is a view, and accelerator encoding transfers the packed tensor only once. State, goal, ground-action, transition-effect, virtual-node, and expressive encoders delegate their traversal and relation assembly directly to Mimir. Custom encoders remain freely mixable in specification order and allocate IDs through that same context. The public EncodedTensors interface and model checkpoint layout is unchanged for domains without nullary predicates.

Nullary predicates are represented as unary relations over every canonical problem object, including domain constants. Goal and expressive encoders use the same lifting. A nullary transition effect is instead represented by the transition node alone, as P(transition). Consequently, relation arities—and therefore checkpoint structure—change for domains that contain nullary predicates.

Numeric PDDL remains intentionally unsupported.

API Overview

Core Components

HyperparameterConfig

Configuration class for R-GNN model hyperparameters:

• Domain: The PDDL domain for the planning problem
• Model Parameters: Embedding size, number of layers
• Training Settings: Normalization, global readout options

ModuleConfig

Configuration class for R-GNN neural network modules:

• Aggregation Function: How messages are aggregated (mean, sum, max, etc.)
• Message Function: How messages are computed between related nodes
• Update Function: How node embeddings are updated with aggregated messages

Encoder/Decoder Classes

Extensible class-based system for defining inputs and outputs:

• Input Specification: Tuple of encoder instances (StateEncoder, GoalEncoder, etc.)
• Output Specification: List of named decoder instances with custom readout logic

RelationalGraphNeuralNetwork

The main R-GNN model class that:

• Takes hyperparameter config, module config, input specification, and output specification
• Processes relational graph structures from PDDL problems
• Supports extensible encoder/decoder system for custom input/output handling
• Handles batched inference efficiently

Encoder Classes

Inherit from Encoder base class to define custom input processing:

• StateEncoder: Current state of the planning problem
• GoalEncoder: Goal specification
• GroundActionsEncoder: Available ground actions
• TransitionEffectsEncoder: Ordered successor states and their realized transition effects

StateEncoder, GoalEncoder, and TransitionEffectsEncoder expose get_relation_descriptors(domain) for semantic relation discovery. Each EncoderRelation identifies its source symbol and EncoderRelationKind, so callers do not need to infer meaning from relation names or list positions. TransitionEffectsEncoder consumes (ordered_successor_states, actions, effect_relations, goal_condition), where actions[i] produced successors[i]; relation index pairs refer to positions in the successor sequence. Each transition node also gets the unary relation action_name_<schema><suffix> for the action that produced it.

Decoder Classes

Inherit from Decoder base class to define custom output readout:

input_spec = (StateEncoder(), GroundActionsEncoder(), GoalEncoder())
output_spec = [
    ('actor', ActionScalarDecoder(hparam_config)),
    ('critic', ObjectsScalarDecoder(hparam_config)), 
    ('embeddings', ActionEmbeddingDecoder())
]

Aggregation Functions

Available in the ModuleConfig:

• MeanAggregation(): Mean aggregation
• SumAggregation(): Sum aggregation
• HardMaximumAggregation(): Hard maximum
• SmoothMaximumAggregation(): Smooth maximum (LogSumExp)

Replayed Layers

On CUDA, a forward pass without gradients of a network with at least four layers records its first message-passing layer as a CUDA graph and replays it for every layer, which is several times faster for deep networks. The outputs are those of running every layer. This applies when the message, aggregation and update functions all set supports_replay = True, no hooks are registered, and messages and updates are not quantized; otherwise every layer runs as usual.

PredicateMLPMessages, PredicateLinearMessages, all aggregation functions, MLPUpdates and LinearUpdates support replay. A custom function may set supports_replay = True if its forward runs the same operations in every layer, independent of the layer, never synchronizes with the host (no .item(), no .cpu(), no indexing with Python lists, no shapes that depend on tensor values), and draws no random numbers. A recording that synchronizes fails with an error; one that depends on the layer would silently repeat the first layer's behavior.

Examples and Tutorials

For an comprehensive example, visit:

• Example Project Repository

Contributing

We welcome contributions! Please see our Contributing Guidelines for details on:

• Development setup
• Coding standards
• Testing requirements
• Pull request process

License

This project is licensed under the GNU General Public License v3.0 or later. See the LICENSE file for details.

Citation

If you use Mimir-RGNN in your research, please cite:

@inproceedings{stahlberg-bonet-geffner-icaps2022,
  author       = {Simon St{\aa}hlberg and Blai Bonet and Hector Geffner},
  title        = {Learning General Optimal Policies with Graph Neural Networks: Expressive Power, Transparency, and Limits},
  booktitle    = {Proceedings of the Thirty-Second International Conference on Automated Planning and Scheduling, {ICAPS} 2022, Singapore (virtual), June 13-24, 2022},
  pages        = {629--637},
  year         = {2022}
}

Support

• 🐛 Bug Reports: GitHub Issues
• 📧 Contact: simon.stahlberg@gmail.com