# PowerZooJax: A JAX-based Power System Benchmark for Reinforcement Learning

arXiv: https://arxiv.org/abs/2609.36052 · HTML: https://mlsanity.com/paper/2609.36052/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL for energy, networks and operations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 19 of its 40 references that are in the map.
- First posted on arXiv 28 Sep 2026.

## Placed in

- Deep reinforcement learning methods / [RL for energy, networks and operations](https://mlsanity.com/deeprl/rl-for-energy-networks-and-operations/)

## Builds on (most cited first, 19 of 19)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Is Independent Learning All You Need in the StarCraft Multi-Agent Challenge?](https://mlsanity.com/paper/2011.09533/) (2020)
- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [pandapower - an Open Source Python Tool for Convenient Modeling, Analysis and Optimization of Electric Power Systems](https://mlsanity.com/paper/1709.06743/) (2017)
- [JaxMARL: Multi-Agent RL Environments and Algorithms in JAX](https://mlsanity.com/paper/2311.10090/) (2023)
- [Learning to run a Power Network Challenge: a Retrospective Analysis](https://mlsanity.com/paper/2103.03104/) (2021)
- [Multi-Agent Reinforcement Learning for Active Voltage Control on Power Distribution Networks](https://mlsanity.com/paper/2110.14300/) (2021)
- [Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning](https://mlsanity.com/paper/2303.17503/) (2023)
- [Reinforcement Learning for Selective Key Applications in Power Systems: Recent Advances and Future Challenges](https://mlsanity.com/paper/2102.01168/) (2021)
- [PowerGym: A Reinforcement Learning Environment for Volt-Var Control in Power Distribution Systems](https://mlsanity.com/paper/2109.03970/) (2021)
- [GridLearn: Multiagent Reinforcement Learning for Grid-Aware Building Energy Management](https://mlsanity.com/paper/2110.06396/) (2021)
- [pymgrid: An Open-Source Python Microgrid Simulator for Applied Artificial Intelligence Research](https://mlsanity.com/paper/2011.08004/) (2020)
- [RL2Grid: Benchmarking Reinforcement Learning in Power Grid Operations](https://mlsanity.com/paper/2503.23101/) (2025)
- [Small vulnerable sets determine large network cascades in power grids](https://mlsanity.com/paper/1804.06432/) (2018)
- [Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX](https://mlsanity.com/paper/2510.01764/) (2025)
- [Andes_gym: A Versatile Environment for Deep Reinforcement Learning in Power Systems](https://mlsanity.com/paper/2203.01292/) (2022)
- [CommonPower: A Framework for Safe Data-Driven Smart Grid Control](https://mlsanity.com/paper/2406.03231/) (2024)

## Abstract (arXiv)

Power system operation is a safety-critical sequential decision-making problem, making it a natural testbed for reinforcement learning (RL). However, existing RL environments for power systems are often narrow in scope and computationally limited by CPU-based simulation workflows, making large-scale evaluation difficult. We introduce PowerZooJax, a JAX-based benchmark suite for RL in power system operation. It provides five constrained Markov decision process tasks spanning generation, transmission, distribution, distributed energy resources, and data center microgrid. By rewriting power flow, economic dispatch, market clearing, and device dynamics as JAX computation graphs, PowerZooJax keeps the entire training and evaluation loop on the GPU. Experiments show substantial speedups over CPU-based simulations and demonstrate standardized evaluation of policy returns, safety violations, and out-of-distribution stress conditions. Our open-source benchmark is available at: https://github.com/powerzoojax/PowerZooJax.