Zenith Grid Rerouter ✨: Autonomous Self-Healing Power Grid

Zenith Grid Rerouter is an imitation and reinforcement learning platform for distribution system service restoration. This repository includes:

• A custom Gymnasium-compatible environment for power grid restoration
• PPO training and benchmarking scripts
• A Streamlit command-center UI (Project Zenith) for real-time visualization
• Optimization-backed simulation with Pyomo + HiGHS

Table of Contents

• Overview
• Key Features
• Project Structure
• Requirements
• Quick Start (Windows)
• Install From Source
• Run the Streamlit App
• Train a PPO Model
• Benchmark the Trained Model
• Verify Environment Setup
• Environment IDs and Registration Notes
• How the System Works
• Troubleshooting
• Documentation
• Roadmap Ideas
• Citation
• License

Overview

Distribution networks face outages caused by faults, weather events, and cascading disturbances. Zenith Grid Rerouter learns restoration policies that can quickly decide tie-line switching actions to recover service.

Project Zenith is the operational dashboard layer on top of this system:

• Displays a digital twin of the IEEE 33-bus feeder
• Simulates disturbances and autonomous recovery actions
• Visualizes voltage and load restoration telemetry
• Streams SCADA-style restoration logs

Key Features

• Gymnasium-based restoration environment in gym_power_res/envs
• Reinforcement learning workflow using Stable-Baselines3 PPO
• Optimization-informed environment dynamics via Pyomo
• HiGHS solver integration for fast MILP/LP backends
• Streamlit app with interactive controls and live Plotly visuals

Project Structure

.
|-- app.py                               # Streamlit command center UI
|-- train_main.py                        # PPO training entrypoint
|-- benchmark_zenith.py                  # Benchmark/evaluation script
|-- verify_grid.py                       # Environment verification script
|-- Zenith_33node_v1.zip                 # Pretrained PPO model artifact
|-- setup.py                             # Package setup
|-- gym_power_res/
|   |-- __init__.py                      # Environment registration
|   |-- envs/
|   |   |-- GYM_ENV_restoration_distribution.py
|   |   |-- DS_pyomo.py
|   |   `-- data_test_case.py
|   `-- sorces/
|-- project_dis_restoration/             # Research/training scripts and assets
|-- docs/                                # Sphinx docs build and source

Requirements

• Python 3.10+ (tested in Python 3.12 virtual environment)
• pip
• Internet access for first-time dependency installation

Core Python dependencies used by the app and scripts:

• streamlit
• plotly
• gymnasium
• stable-baselines3
• pyomo
• highspy
• torch
• numpy

Quick Start (Windows)

From the project root:

python -m venv .venv
.\.venv\Scripts\Activate.ps1
python -m pip install --upgrade pip
python -m pip install -e . (if this command throws error, dont worry about it and go to the next command.)
python -m pip install streamlit plotly gymnasium stable-baselines3 pyomo highspy torch numpy

Run the UI:

python -m streamlit run app.py

Install From Source

Editable install registers local package imports while you develop:

python -m pip install -e .

If you need a clean reinstall:

python -m pip install --force-reinstall -e .

Run the Streamlit App

python -m streamlit run app.py

What you get:

• A digital twin network plot for the 33-bus system
• Voltage monitor and grid stability gauges
• Disturbance injection controls
• Autonomous healing simulation loop
• SCADA-style event log

The app expects the pretrained model file Zenith_33node_v1.zip in the repository root.

Train a PPO Model

python train_main.py

Default behavior in train_main.py:

• Registers environment ID power-res-v1
• Trains PPO with MlpPolicy
• Saves model as Zenith_33node_v1.zip

Benchmark the Trained Model

python benchmark_zenith.py

The benchmark script:

• Registers environment ID zenith-v1
• Loads Zenith_33node_v1
• Runs multiple episodes
• Reports success rate and average restoration performance

Verify Environment Setup

Use this for a quick sanity check:

python verify_grid.py

Expected checks include:

• Environment registration success
• Observation space availability
• One successful step transition

Environment IDs and Registration Notes

Multiple scripts register different IDs for the same environment class. That is valid, but you should keep ID usage consistent per workflow.

Current IDs in this repository include:

• RestorationDisEnv-v1 (package-level registration)
• power-res-v1 (training script)
• zenith-v1 (benchmark script)
• zenith-cmd-v1 (Streamlit app)

If you see an error like "Environment ... does not exist", ensure:

1. gym_power_res can be imported in your active environment
2. The script performing gym.make(...) also registers that ID first (or imports a module that does)
3. You are running with the same Python environment where dependencies were installed

How the System Works

1. Fault injection: The environment introduces random line disturbances.
2. Observation: The agent receives a fixed-dimensional grid-state vector.
3. Action: PPO selects tie-line switching actions.
4. Physics/optimization update: Pyomo + HiGHS computes resulting feasible grid state transitions.
5. Reward and telemetry: Load restoration and system stability drive rewards and displayed metrics.

Troubleshooting

1) Environment not found

Symptom:

• gymnasium.error.NameNotFound for zenith-cmd-v1 or similar

Fix:

python -m pip install -e .

Then ensure the script registers the same ID before gym.make(...).

2) Pyomo solver not available (NotFound)

Symptom:

• Solver ... is not available (NotFound)

Fix:

python -m pip install highspy pyomo

3) Stable-Baselines3 import or model loading errors

Fix:

python -m pip install stable-baselines3 torch

Also verify Zenith_33node_v1.zip exists in the repository root.

4) Streamlit duplicate Plotly element ID

Symptom:

• streamlit.errors.StreamlitDuplicateElementId

Fix:

• Pass unique key values to each st.plotly_chart call.

5) Streamlit warning about use_container_width

Symptom:

• Deprecation warning recommending width='stretch'

Status:

• Non-blocking warning. App still runs.

Documentation

Sphinx docs are included in docs/source.

To build documentation locally (from docs directory):

pip install -r requirements-rtd.txt
make html

On Windows, if make is unavailable, use:

.\make.bat html

Roadmap Ideas

• Consolidate environment IDs to a single naming convention
• Add unit tests for environment registration and reset/step contracts
• Add CI workflow for linting, tests, and docs build
• Export benchmark metrics to structured JSON/CSV for experiment tracking
• Add dockerized runtime for one-command startup

Citation

If you use Zenith Grid Rerouter in research, cite:

Y. Zhang, F. Qiu, T. Hong, Z. Wang, and F. Li, Hybrid imitation learning for real-time service restoration in resilient distribution systems, IEEE Transactions on Industrial Informatics, 2021.

@article{Zhang2021_ZenithGridRerouter,
  author = {Zhang, Yichen and Qiu, Feng and Hong, Tianqi and Wang, Zhaoyu and Li, Fangxing Fran},
  journal = {IEEE Trans. Ind. Informat.},
  pages = {1--11},
  note = {early access},
  title = {Hybrid imitation learning for real-time service restoration in resilient distribution systems},
  year = {2021}
}

License

Released under the modified BSD license. See LICENSE.