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

2026-09-28

AI Summary

This study addresses the limitations of existing power system reinforcement learning (RL) environments, namely their task specificity and reliance on CPU-based simulations that hinder large-scale evaluation. To overcome these bottlenecks, this work proposes a JAX-based RL benchmark suite for power systems. Methodologically, core modules such as power flow calculation and economic dispatch are reformulated as JAX computational graphs, achieving, for the first time, end-to-end GPU-parallel acceleration across five tasks spanning generation to microgrids. Furthermore, constrained Markov decision processes are introduced to model safety constraints. By breaking through traditional computational barriers, the proposed framework significantly accelerates simulation efficiency and supports standardized, large-scale evaluation of policy returns, safety compliance, and out-of-distribution stress scenarios. Ultimately, this work provides an efficient, open-source benchmark to advance RL research in power systems.

Abstract

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.

Problem

• Reinforcement Learning
• Power System
• Benchmark
• Sequential Decision-Making
• GPU Acceleration

Innovation

• JAX
• Reinforcement Learning
• Power System Benchmark
• GPU Acceleration
• Constrained Markov Decision Process