Code for the paper: > Learning Topology Actions for Power Grid Control: A Graph-Based Soft-Label Imitation Learning Approach Mohamed Hassouna, Clara Holzhüter, Malte Lehna, Matthijs de Jong, Jan Viebahn, Bernhard Sick, Christoph Scholz At ECML PKDD 2025. Lecture Notes in Computer Science(), vol 16022. Springer, Cham. https://doi.org/10.1007/978-3-032-06129-4_8 [Pre-print Paper PDF]
This repository contains the official implementation of the SoftGNN agent, a novel Graph Neural Network (GNN)-based imitation learning framework for power grid topology control. The approach improves over traditional hard-label imitation learning by learning from soft labels that capture multiple viable actions for grid congestion mitigation. The agent operates in the Grid2Op L2RPN WCCI 2022 environment, outperforming both the expert and state-of-the-art RL agents.
- Soft-Label Generation: Learn from a distribution over viable actions rather than a single expert action.
- GNN-Based Architecture: Use Graph Attention Networks (GAT) to encode power grid topology.
- Action Feasibility Enhancements: Improved substation reconfiguration support and line-disconnection handling.
- N-1 Security Evaluation: Post-hoc contingency-aware action selection for increased robustness.
- Benchmarking: Evaluation against greedy expert, and SOTA DRL agents in the L2RPN WCCI 2022 environment.
To run the experiments and use the agents in this repository, follow the steps below to set up the environment and dependencies.
- Python >= 3.8
- Grid2Op
- LightSim2Grid (Highly recommended for performance)
- PyTorch & PyTorch Geometric
- Create a Virtual Environment (Recommended)
conda create -n softgnn python=3.9 conda activate softgnn
- Install Dependencies Install the core dependencies required for Grid2Op and GNN training:
# Install Grid2Op and LightSim2Grid pip install grid2op lightsim2grid # Install PyTorch (adjust for your CUDA version) pip install torch torchvision torchaudio # Install PyTorch Geometric pip install torch_geometric # Install other utilities pip install optuna pandas scikit-learn h5py
The repository is organized as follows:
soft_label_gnn/
├── data/ # Stores generated soft-label datasets and scaler files
│ ├── actions.npy # Action space definitions
│ └── scaler_all.pkl # Data scaler for normalization
├── evaluation/ # Scripts for evaluating agent performance
│ ├── general_tutor.py # Logic for the expert tutor (Taken and extended from curriculumagent)
│ ├── score_agent.py # Scoring script for Grid2Op agents (Taken from curriculumagent)
│ └── utilities.py # Helper functions for simulations and metrics (Taken and extended from curriculumagent)
├── gnn/ # GNN model definitions and training logic
│ ├── gnn_models.py # Implementation of GAT architecture
│ ├── gnn_prediction.py # Training loop and prediction logic
│ ├── obs_converter.py # Converts Grid2Op observations to PyG graphs
│ └── torch_geometric_datasets.py # Custom PyG dataset loader
├── saved_models/ # Directory for saving trained model checkpoints
├── data_generation.py # Script to generate soft-label experience data
├── Data_Processing.ipynb # Code to process the collected data into graphs ready to use for GNN training
├── soft_target_optuna_distributed.py # Distributed hyperparameter optimization (Training Script)
├── GNNAgent.py # Main agent classes (GNNAgent & GNNAgentN1)
├── get_seed_gnn_array.py # Evaluation on 20 seeds
└── README.md # Project documentation
The Soft-Label dataset is generated by an expert tutor (Greedy) that explores viable topology actions for specific grid states.
To generate new experience data:
python data_generation.py
-
Input: Uses data/actions.npy and the l2rpn_wcci_2022 environment.
-
Output: Saves experience (states, action IDs, and rho values) to the data/ directory in .h5 or .npy format.
-
Run the Notebook Data_Processing.ipynb to process the data into graphs and save them ready for training.
We use Optuna to optimize the GNN architecture (e.g., number of GAT layers, heads, hidden dimensions) and training hyperparameters.
To start the distributed training/optimization process:
python soft_target_optuna_distributed.py- Configuration: Defines the search space for learning rate, weight decay, GAT layers, and dropout.
- Storage: Results are stored in a local SQLite database (study.db) in saved_models/.
The repository provides two agent variants in GNNAgent.py:
GNNAgent: The standard GNN-based soft-label agent.
GNNAgentN1: An advanced agent that performs "N-1" security checks (post-hoc contingency analysis) on the top candidate actions.
from GNNAgent import GNNAgent
from pathlib import Path
import grid2op
# Initialize environment
env = grid2op.make("l2rpn_wcci_2022")
# Load the trained agent
agent = GNNAgent(
action_space=env.action_space,
model_path=Path("./saved_models/best_model_path"),
action_space_path=Path("./data/actions.npy"),
run_with_tf=False, # Set False for PyTorch models
topo=True # Enable topological actions
)
# Run interaction
obs = env.reset()
done = False
while not done:
action = agent.act(obs, reward=0, done=False)
obs, reward, done, info = env.step(action)To robustly evaluate your agent, make sure to run your agent on 20 seeds or more using:
python get_seed_gnn_array.pyPlease make sure to create a suitable validation environment to evaluate on (see. https://grid2op.readthedocs.io/en/latest/user/environment.html#splitting-into-raining-validation-test-scenarios)
The core model is a Graph Attention Network (GAT) defined in gnn/gnn_models.py. It processes the power grid as a graph where:
- Nodes: Represent loads, generators, line ends
- Edges: Represent electrical connectivity.
The model outputs a probability distribution over the discrete action space (defined in actions.npy), which acts as the soft label for the agent to imitate. The architecture of our best model can be found in data/best_model/paper_train_config.pkl
If you use this code or dataset in your research, please cite our paper:
@InProceedings{hassouna25_softgnn,
author="Hassouna, Mohamed
and Holzh{\"u}ter, Clara
and Lehna, Malte
and de Jong, Matthijs
and Viebahn, Jan
and Sick, Bernhard
and Scholz, Christoph",
editor="Dutra, In{\^e}s
and Pechenizkiy, Mykola
and Cortez, Paulo
and Pashami, Sepideh
and Pasquali, Arian
and Moniz, Nuno
and Jorge, Al{\'i}pio M.
and Soares, Carlos
and Abreu, Pedro H.
and Gama, Jo{\~a}o",
title="Learning Topology Actions for Power Grid Control: A Graph-Based Soft-Label Imitation Learning Approach",
booktitle="Machine Learning and Knowledge Discovery in Databases. Applied Data Science Track and Demo Track",
year="2026",
publisher="Springer Nature Switzerland",
address="Cham",
pages="129--146",
isbn="978-3-032-06129-4"
}