# Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning

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

## Key facts

- Cited by 51 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 50 references that are in the map.
- First posted on arXiv 29 Mar 2023.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Identity Mappings in Deep Residual Networks](https://mlsanity.com/paper/1603.05027/) (2016)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Thinking Fast and Slow with Deep Learning and Tree Search](https://mlsanity.com/paper/1705.08439/) (2017)
- [PettingZoo: Gym for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2009.14471/) (2020)
- [OpenSpiel: A Framework for Reinforcement Learning in Games](https://mlsanity.com/paper/1908.09453/) (2019)
- [Tianshou: a Highly Modularized Deep Reinforcement Learning Library](https://mlsanity.com/paper/2107.14171/) (2021)
- [Behaviour Suite for Reinforcement Learning](https://mlsanity.com/paper/1908.03568/) (2019)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Revisiting Rainbow: Promoting more Insightful and Inclusive Deep Reinforcement Learning Research](https://mlsanity.com/paper/2011.14826/) (2020)
- [Real World Games Look Like Spinning Tops](https://mlsanity.com/paper/2004.09468/) (2020)
- [Spectral Normalisation for Deep Reinforcement Learning: an Optimisation Perspective](https://mlsanity.com/paper/2105.05246/) (2021)
- [You Can't Count on Luck: Why Decision Transformers and RvS Fail in Stochastic Environments](https://mlsanity.com/paper/2205.15967/) (2022)
- [EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine](https://mlsanity.com/paper/2206.10558/) (2022)
- [Suphx: Mastering Mahjong with Deep Reinforcement Learning](https://mlsanity.com/paper/2003.13590/) (2020)
- [Podracer architectures for scalable Reinforcement Learning](https://mlsanity.com/paper/2104.06272/) (2021)
- [Joint Policy Search for Multi-agent Collaboration with Imperfect Information](https://mlsanity.com/paper/2008.06495/) (2020)
- [Human-Agent Cooperation in Bridge Bidding](https://mlsanity.com/paper/2011.14124/) (2020)
- [Competitive Bridge Bidding with Deep Neural Networks](https://mlsanity.com/paper/1903.00900/) (2019)
- [Gardner's Minichess Variant is solved](https://mlsanity.com/paper/1307.7118/) (2013)

## Built on by (most cited first, 30 of 51)

- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [JaxMARL: Multi-Agent RL Environments and Algorithms in JAX](https://mlsanity.com/paper/2311.10090/) (2023)
- [Gymnasium: A Standard Interface for Reinforcement Learning Environments](https://mlsanity.com/paper/2407.17032/) (2024)
- [Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks](https://mlsanity.com/paper/2410.23208/) (2024)
- [NAVIX: Scaling MiniGrid Environments with JAX](https://mlsanity.com/paper/2407.19396/) (2024)
- [Games for Artificial Intelligence Research: A Review and Perspectives](https://mlsanity.com/paper/2304.13269/) (2023)
- [The Impact of On-Policy Parallelized Data Collection on Deep Reinforcement Learning Networks](https://mlsanity.com/paper/2506.03404/) (2025)
- [JaxUED: A simple and useable UED library in Jax](https://mlsanity.com/paper/2403.13091/) (2024)
- [Benchmarking Massively Parallelized Multi-Task Reinforcement Learning for Robotics Tasks](https://mlsanity.com/paper/2507.23172/) (2025)
- [PC-Gym: Benchmark Environments For Process Control Problems](https://mlsanity.com/paper/2410.22093/) (2024)
- [SocialJax: An Evaluation Suite for Multi-agent Reinforcement Learning in Sequential Social Dilemmas](https://mlsanity.com/paper/2503.14576/) (2025)
- [BuilderBench: The Building Blocks of Intelligent Agents](https://mlsanity.com/paper/2510.06288/) (2025)
- [Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX](https://mlsanity.com/paper/2510.01764/) (2025)
- [Emergence of Exploration in Policy Gradient Reinforcement Learning via Retrying](https://mlsanity.com/paper/2606.00151/) (2026)
- [PuzzleJAX: A Benchmark for Reasoning and Learning](https://mlsanity.com/paper/2508.16821/) (2025)
- [Discovering Minimal Reinforcement Learning Environments](https://mlsanity.com/paper/2406.12589/) (2024)
- [Investigating Memory in Model-Free RL with POPGym Arcade](https://mlsanity.com/paper/2503.01450/) (2025)
- [gfnx: Fast and Scalable Library for Generative Flow Networks in JAX](https://mlsanity.com/paper/2511.16592/) (2025)
- [Spyx: A Library for Just-In-Time Compiled Optimization of Spiking Neural Networks](https://mlsanity.com/paper/2402.18994/) (2024)
- [NashPG: A Policy Gradient Method with Iteratively Refined Regularization for Finding Nash Equilibria](https://mlsanity.com/paper/2510.18183/) (2025)
- [Multi-Agent Craftax: Benchmarking Open-Ended Multi-Agent Reinforcement Learning at the Hyperscale](https://mlsanity.com/paper/2511.04904/) (2025)
- [OrderGrad: Optimizing Beyond the Mean with Order-Statistic Policy Gradient Estimation](https://mlsanity.com/paper/2606.06096/) (2026)
- [Finding the Time to Think: Learning Planning Budgets in Real-Time RL](https://mlsanity.com/paper/2606.26463/) (2026)
- [A Batch Sequential Halving Algorithm without Performance Degradation](https://mlsanity.com/paper/2406.00424/) (2024)
- [EconoJax: A Fast & Scalable Economic Simulation in Jax](https://mlsanity.com/paper/2410.22165/) (2024)
- [EvoRL: A GPU-accelerated Framework for Evolutionary Reinforcement Learning](https://mlsanity.com/paper/2501.15129/) (2025)
- [NiceWebRL: a Python library for human subject experiments with reinforcement learning environments](https://mlsanity.com/paper/2508.15693/) (2025)
- [The Emergence of Complex Behavior in Large-Scale Ecological Environments](https://mlsanity.com/paper/2510.18221/) (2025)
- [Revisiting Regularized Policy Optimization for Stable and Efficient Reinforcement Learning in Two-Player Games](https://mlsanity.com/paper/2602.10894/) (2026)
- [AlphaExploitem: Going Beyond the Nash Equilibrium in Poker by Learning to Exploit Suboptimal Play](https://mlsanity.com/paper/2605.09150/) (2026)

## Abstract (arXiv)

We propose Pgx, a suite of board game reinforcement learning (RL) environments written in JAX and optimized for GPU/TPU accelerators. By leveraging JAX's auto-vectorization and parallelization over accelerators, Pgx can efficiently scale to thousands of simultaneous simulations over accelerators. In our experiments on a DGX-A100 workstation, we discovered that Pgx can simulate RL environments 10-100x faster than existing implementations available in Python. Pgx includes RL environments commonly used as benchmarks in RL research, such as backgammon, chess, shogi, and Go. Additionally, Pgx offers miniature game sets and baseline models to facilitate rapid research cycles. We demonstrate the efficient training of the Gumbel AlphaZero algorithm with Pgx environments. Overall, Pgx provides high-performance environment simulators for researchers to accelerate their RL experiments. Pgx is available at http://github.com/sotetsuk/pgx.
