Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning

TL;DR

Pgx presents a JAX-based, hardware-accelerated suite of multi-agent board-game environments designed for high-throughput RL experiments. By leveraging auto-vectorization and accelerator parallelism, Pgx achieves 10-100x faster simulation than Python libraries and scales across multiple GPUs, enabling rapid AlphaZero-style training on Go, chess, shogi, and related games. The paper demonstrates Pgx's effectiveness with Gumbel AlphaZero across five compact environments and shows the practical benefits of multi-accelerator scaling for faster learning. It also provides a diverse game catalog (including MinAtar-like Atari variants) and baseline models to accelerate research cycles, while acknowledging current limitations and outlining paths for future extension.

Abstract

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.

Figures (10)

• Figure 1: Example games included in Pgx.
• Figure 2: Basic usage of Pgx. The init function generates the initial state object. The state object has an attribute current player that indicates the agent which acts next. In this case, since we are using a batch size of 1024 for a 2-player game, current player is a binary vector whose size is 1024. Note that current player is independent of the colors (i.e., first player or second player). Here, current player is determined randomly using a pseudo-random number generator. The step function takes the previous state and an action vector, whose size is 1024, as input and returns the next state. The observation of the current player can be accessed through state. In the case of Go, for example, each observation has a shape of $1024 \times 19 \times 19 \times 17$. The available actions at the current state can be obtained through the boolean vector legal action mask. Here, it has a shape of $1024 \times 362$. For more detailed API description and usage, refer to the Pgx documentation at https://sotetsuk.github.io/pgx/.
• Figure 3: Simulation throughputs. Policies are random without learning processes. Error bars are not visible in this scale.
• Figure 4: Simulation throughputs of all Pgx environments. Error bars are not visible in this scale.
• Figure 5: AlphaZero training results. Black line represents the Elo rating of baseline models provided by Pgx (1000 Elo). The shaded area represents the standard errors of two runs.