Metadata-Version: 2.4
Name: pymimir-rgnn
Version: 0.3.0b2
Summary: Relational Graph Neural Network (R-GNN) package for Mimir based on PyTorch.
Author-email: Simon Stahlberg <simon.stahlberg@gmail.com>
License-Expression: GPL-3.0-or-later
Project-URL: Homepage, https://github.com/simon-stahlberg/mimir-rgnn
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.6.0
Requires-Dist: pymimir==0.14.0b3
Provides-Extra: dev
Requires-Dist: pytest>=6.0; extra == "dev"
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"
Requires-Dist: mypy>=1.0; extra == "dev"
Dynamic: license-file

# Mimir-RGNN

[![PyPI Version](https://img.shields.io/pypi/v/pymimir-rgnn)](https://pypi.org/project/pymimir-rgnn/)
[![Python Versions](https://img.shields.io/pypi/pyversions/pymimir-rgnn)](https://pypi.org/project/pymimir-rgnn/)
[![License](https://img.shields.io/pypi/l/pymimir-rgnn)](https://github.com/simon-stahlberg/mimir-rgnn/blob/master/LICENSE)
[![Tests](https://github.com/simon-stahlberg/mimir-rgnn/actions/workflows/test.yml/badge.svg)](https://github.com/simon-stahlberg/mimir-rgnn/actions/workflows/test.yml)

**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:

```bash
pip install pymimir-rgnn
```

### Requirements

- Python 3.11+
- PyTorch 2.6.0+
- Pymimir 0.14.0b3

## Quick Start

```python
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.0b3 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, effect_relations, goal_condition)`; relation index
pairs refer to positions in the successor sequence.

### Decoder Classes

Inherit from `Decoder` base class to define custom output readout:

```python
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)

## Examples and Tutorials

For an comprehensive example, visit:

- [Example Project Repository](https://github.com/simon-stahlberg/relational-neural-network-python/)

## Contributing

We welcome contributions! Please see our [Contributing Guidelines](CONTRIBUTING.md) 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](LICENSE) file for details.

## Citation

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

```bibtex
@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](https://github.com/simon-stahlberg/mimir-rgnn/issues)
- 📧 **Contact**: simon.stahlberg@gmail.com
