SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sotetsu Koyamada $^{1,2*}$  Shinri Okano $^{1}$  Soichiro Nishimori $^{3}$
Yu Murata<sup>1</sup> Keigo Habara<sup>1</sup> Haruka Kita<sup>1</sup> Shin Ishii<sup>1,2</sup>
$^{1}$ Kyoto University  $^{2}$  ATR  $^{3}$ The University of Tokyo
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 https://github.com/sotetsuk/pgx.
Shogi
Figure 1: Example games included in Pgx.
Developing algorithms for solving challenging games is a standard artificial intelligence (AI) research benchmark. Especially building AI, which can defeat skilled professional players in complex games like chess, shogi, and Go, has been a crucial milestone. Though reinforcement learning (RL) algorithms, which combine deep learning and tree search, are successful in obtaining such high-level strategies [1, 2, 3], complex games like chess and Go are still in the interest of AI research for developing RL and search algorithms in discrete state environments.
On the other hand, studying algorithms for solving large state-space games such as Go requires a huge sample size. MuZero [4], employing a learned model, has been successful in the domain of
\*Correspondence to: Sotetsu Koyamada <koyamada@atr.jp>
game AI and can significantly reduce the number of interactions with the real environment. However, this does not render research on approaches without a learned model like AlphaZero unnecessary. Notably, it is mentioned that AlphaZero's learning is faster than MuZero's in the case of chess [5].
Thus, in RL research, a fast simulator of the environment, which achieves high throughput is often required. Simulators that possess practical speed and performance are often implemented in  $\mathrm{C + + }$ . In the machine learning community, however, Python serves as a lingua franca. Therefore, libraries like OpenSpiel [6] wrap the core  $\mathrm{C + + }$  implementation and provide a Python API. However, this approach presents several challenges. Efficient parallelization of the environment is crucial for generating a large number of samples, but efficient parallel environments utilizing  $\mathrm{C + + }$  threading, such as EnvPool [7], may not always be accessible from Python. In fact, to our knowledge, there are no libraries publicly available in Python that allow the use of efficient parallel environments for important game AI benchmarks like chess and Go from the identical Python API. Also, the RL algorithm often runs on accelerators (such as GPUs and TPUs), whereas simulation runs on CPUs, which makes additional data transfer costs between CPUs and accelerators.
In continuous state space environments, Brax [8] and Isaac Gym [9] demonstrate that environments that work on accelerators can dramatically improve the simulation throughput and RL training speed, resolving the data transfer and parallelization problems. Brax, written in JAX [10], a Python library, provides hardware-accelerated environments that leverage JAX's auto-vectorization, parallelization over accelerators, and Just-In-Time (JIT) compilation optimized for individual hardware platforms.
In this study, we offer the hardware-accelerated environments of complex, discrete state domains like chess, shogi and Go. Specifically, we introduce Pgx, a suite of efficient game simulators developed in JAX. Thanks to JAX's auto-vectorization and parallelization across multiple accelerators, Pgx can achieve high throughput on GPUs/TPUs. As of writing this paper, to our best knowledge, there is no other comprehensive game environment library written in JAX. Highlighted features of Pgx include:
- Fast simulation: Pgx provides high-performance simulators written in JAX that run fast on GPU/TPUs, similar to Brax. We demonstrated that Pgx is 10-100x faster than existing Python libraries such as PettingZoo [11] and OpenSpiel [6] on a DGX-A100 workstation (see Fig. 3).
- Diverse set of games: Pgx offers over 20 games, ranging from perfect information games like chess to imperfect information games like bridge (see Table 1). Pgx also offers miniature versions of game environments (e.g., miniature chess) to facilitate research cycles.
- **Baseline models:** Evaluating agents in multi-agent games is relative, requiring baseline opponents for evaluation. Since it is not always easy to have appropriate baselines available, Pgx provides its own baseline models. In Sec. 5, we demonstrate the availability of them with AlphaZero training.
Pgx is open-sourced and freely available at https://github.com/sotetsuk/pgx².
Games in AI research. An early study that combined neural networks (NNs) with RL to build world-class agents in a complex board game was TD-Gammon [12]. After the breakthrough of deep learning [13], RL agents combined with NNs performed well in the video game domain [14] and large state fully-observable board games, including chess, shogi, and Go [1, 2, 3]. RL agents with NNs also performed well in large-scale, partially observable games like mahjong [15]. However, these RL agents in complex board games require a huge number of self-play samples.
Games as RL environment. Game AI studies often have to pay high engineering costs, and there are a variety of libraries behind the democratization of game AI research. Arcade learning environment (ALE) made using Atari 2600 games as RL environments possible [16]. Several RL environment libraries provide classic board game suits [6, 11, 17]. Pgx aims to implement (classic) board game environments with high throughput utilizing GPU/TPU acceleration.
Hardware-accelerated RL environments. While hardware acceleration is a more specific approach compared to methods that run on CPUs, such as EnvPool [7], it has a major advantage of its
See App. A for the license.
import jax   
import pgx   
env  $=$  pgx.make("go_19x19")   
init_fn  $=$  jax jit(jax.vmap(env.init))   
step_fn  $=$  jax jit(jax.vmap(env_STEP))   
batch_size  $= 1024$    
# Pseudo-random number generator keys determine the first players   
rng_keys  $=$  jax.random.split(jax.random.PRNGKey(9999)，batch_size)   
state  $=$  init_fn(rng_keys） # Vectorized initial states   
while not(stateterminated|state.truncated).all(): action  $=$  model(state.current_player，stateobservation，state.lega_l_action_mask) state  $=$  step_fn(state，action） # state.rewards with shape (1024,2)
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/.
ability to leverage accelerators for parallel execution, enabling high-speed simulations. Also, NN training is often performed on GPU/TPU accelerators, and there is an advantage that there is no data transfer cost between CPU and GPU/TPU accelerators. There is a wide range of environments available through various open-source software. In particular, JAX-based environments have gained popularity due to their high scalability over accelerators. These include:
- gymnax [18], which re-implements popular RL environments, including classic control (e.g., Pendulum and mountain car), bsuite [19] and a selection of environments from MinAtar [20].
- Brax [8], which provides a variety of continuous control tasks such as Ant and Humanoid.
- Jumanji [21], which offers RL environments for combinatorial optimization problems in routing (e.g., Travelling Salesman Problem; TSP), packing (e.g., bin-packing), and logic (e.g., 2048).
Pgx complements these environments by offering a (classic) board game suite for (multi-agent) RL research. Other hardware-accelerated environments include Isaac Gym [9] for continuous control, CuLE [22] as a GPU-based Atari emulator, and WarpDrive [23] for multi-agent RL research.
Algorithms and architectures that can leverage Pgx. The Anakin architecture [24] is an RL architecture that enables efficient utilization of accelerators and fast learning under the constraint that both the algorithm and environment are written as pure JAX functions. The architecture is capable of scaling up to (potentially) thousands of TPU cores with a simple configuration change. Since all Pgx environments are implemented using pure JAX functions, the Anakin architecture is applicable to any Pgx environment. Gumbel AlphaZero [5] improves the performance of AlphaZero when the number of simulations is small by employing the Gumbel-Top-k trick for search. They provide JAX-based Gumbel AlphaZero implementation[25], which allows batch planning on accelerators. In Sec. 5, we use this implementation to show the example of Pgx usage in AlphaZero training.
<sup>3</sup>https://github.com/google-deepmind/mctx
Table 1: Available games in Pgx (as of v1.4.0).
Env Name	# Players	Obs. shape	# Actions	Tag	Ref.
2048	1	4 × 4 × 31	4	perfect info. (w/ chance)	[27, 28]
Animal shogi	2	4 × 3 × 194	132	perfect info.	[29]
Backgammon	2	34	156	perfect info. (w/ chance)	[12, 27]
Bridge bidding	4	480	38	imperfect info.	[30, 31]
Chess	2	8 × 8 × 119	4672	perfect info.	[3]
Connect Four	2	6 × 7 × 2	7	perfect info.	[32]
Gardner chess	2	5 × 5 × 115	1225	perfect info.	[33]
Go 9x9	2	9 × 9 × 17	82	perfect info.	[1, 2]
Go 19x19	2	19 × 19 × 17	362	perfect info.	[1, 2]
Hex	2	11 × 11 × 4	122	perfect info.	[34, 35]
Kuhn poker	2	7	4	imperfect info.	[36, 6]
Leduc hold'em	2	34	3	imperfect info.	[37, 6]
MinAtar Asterix	1	10 × 10 × 4	5	Atari-like	[20, 38]
MinAtar Breakout	1	10 × 10 × 4	3	Atari-like	[20, 38]
MinAtar Freeway	1	10 × 10 × 7	3	Atari-like	[20, 38]
MinAtar Seaquest	1	10 × 10 × 10	6	Atari-like	[20, 38]
MinAtar Space Invaders	1	10 × 10 × 6	4	Atari-like	[20, 38]
Othello	2	8 × 8 × 2	65	perfect info.	[35]
Shogi	2	9 × 9 × 119	2187	perfect info.	[3]
Sparrow mahjong	3	11 × 15	11	imperfect info.
Tic-tac-toe	2	3 × 3 × 2	9	perfect info.	[39]
3 Pgx overview
In this section, we will provide an overview of the basic usage of Pgx, along with the fundamental principles behind the design of the Pgx API. Additionally, we will explain the overview of the game environments currently offered by Pgx as of this publication.
3.1 Pgx API design
Pgx does take inspiration from existing APIs, but it provides its own custom API for game environments. Specifically, two existing Python RL libraries highly inspired Pgx API:
- Brax [8], a physics engine written in JAX, providing continuous RL tasks, and
- PettingZoo [11], a multi-agent RL environment library available through Gym-like API [26].
Fig. 2 describes an example usage of Pgx API. The main difference from the Brax API comes from that the environments targeted by Pgx are multi-agent environments. Therefore, in Pgx environments, the next agent to act is specified as the current player, as in the PettingZoo API. On the other hand, the significant difference from the PettingZoo API stems from the fact that Pgx is a vectorization-oriented library. Therefore, Pgx does not use an agent iterator like PettingZoo and does not allow the number of agents to change. Concrete code examples comparing the Pgx API with the Brax and PettingZoo APIs are provided in App. B.
While this study does not focus on the design of the API, the Pgx API is sufficiently generic. At present, all game environments implemented in Pgx can be converted to the PettingZoo API and called from the PettingZoo API through the Pgx API (see App. B). This fact demonstrates the practical generality of the Pgx API. However, there is a limitation of the Pgx API in that it cannot handle environments where the number of agents dynamically changes. This limitation arises from the fact that the Pgx API is specialized for efficient vectorized simulation.
3.2 Available games in Pgx
The Pgx framework offers a diverse range of games, as summarized in Table 1. While Pgx primarily emphasizes multi-agent board games, it also includes some single-agent environments and Atari-
Figure 3: Simulation throughputs. Policies are random without learning processes. Error bars are not visible in this scale.
like environments to assist comprehensive RL research. Thus, as we will describe below, games implemented in Pgx span various categories, including two-player perfect information games, games with chance events, imperfect information games, and Atari-like games.
Two-player perfect information games Pgx provides two-player perfect information games that range from simple games like Tic-tac-toe and Connect Four to complex strategic games like chess, shogi, and Go 19x19. While these traditional board games offer rich gameplay and strategic depth, they can be computationally demanding for many RL researchers. To address this, Pgx also includes smaller versions of shogi and chess: Animal shogi and Gardner chess. Although they have smaller board sizes compared to their original counterparts, these games are not mere toy environments. They retain enough complexity to provide engaging gameplay experiences for humans. Animal shogi, in particular, was specifically designed for children, while Gardner chess has a notable history of active play in Italy [33]. Additionally, Pgx offers other medium-sized two-player perfect information games such as Go 9x9, Hex and Othello.
Games with (stochastic) chance events Pgx supports perfect information games with chance events, including backgammon and 2048, which are popular benchmarks for RL algorithms in stochastic state transitions [27, 28]. These games introduce elements of randomness and uncertainty, adding a layer of complexity and decision-making under uncertainty to the gameplay.
Imperfect information games In the realm of imperfect information games, Pgx provides several environments. These include Kuhn poker, Leduc hold'em, Sparrow mahjong (a miniature version of mahjong), and bridge bidding. These games involve hidden information, requiring agents to reason and strategize based on imperfect knowledge of the current game state.
Atari-like games While the primary focus of Pgx is on board games, to make Pgx comprehensive and versatile, Pgx also implements all five environments from the MinAtar game suite: Asterix, Breakout, Freeway, Seaquest, and Space Invaders [20]. These Atari-like environments offer a more visually oriented and dynamic gameplay experience compared to traditional board games. Researchers often utilize MinAtar to conduct comprehensive ablation studies on RL methods in environments with visual inputs [38, 40]. Of these games, Freeway and Seaquest are highlighted as significant benchmarks for assessing the exploration capabilities of RL algorithms [38]. However, as of this writing, gymnax has not incorporated Seaquest into its suite.
For detailed descriptions of each environment, please refer to App. C.
4 Performance benchmarking: simulation throughput
Pgx excels in efficient and scalable simulation on accelerators thanks to JAX's auto-vectorization, parallelization over accelerators, and JIT-compiler. In this section, we validate it through experiments on an NVIDIA DGX-A100 workstation.
Figure 4: Simulation throughputs of all Pgx environments. Error bars are not visible in this scale.
4.1 Comparison to existing Python libraries
Experiment setup. We compare the pure simulation throughput of Pgx with existing popular RL libraries available in Python: PettingZoo [11] and OpenSpiel [6]. For the evaluation, we specifically selected Tic-tac-toe, Connect Four, chess, and Go, as these games are included in all three libraries. To our knowledge, neither PettingZoo nor OpenSpiel provides official parallelized environments from the Python API. Therefore, we prepared two implementations for each library:
1. For-loop (DummyVecEnv): sequentially executes and does not parallelize actually.
2. Subprocess (SubprocVecEnv): parallelize using the multiprocessing module in Python.
We modified and used SubprocVecEnv provided by Tianshou [41] for all competitor libraries. We performed all experiments on an NVIDIA DGX-A100 workstation with 256 cores; Pgx simulations used a single A100 GPU or eight A100 GPUs. We used random policies to evaluate the pure performance of simulators without agent learning. In all implementations, the environment automatically resets to the initial state upon reaching termination. The Pgx version used here was v0.8.0.
Results. Fig. 3 shows the results. We found that Pgx achieves at least  $10\mathrm{x}$  faster throughputs than other existing Python libraries when the number of vectorized environments (i.e., batch size) is large enough (e.g., 1024) on a single A100 GPU. Furthermore, when utilizing eight A100 GPUs, Pgx achieves throughputs of approximately  $100\mathrm{x}$  higher. This trend was identical from the simplest environment, Tic-tac-toe, to complex environments such as  $19\mathrm{x}19$  Go.
4.2 Throughputs of other Pgx environments
The throughput of each environment is influenced by several factors, including the complexity and nature of the game, as well as the quality of its implementation. For instance, OpenSpiel demonstrates throughput of the same order for chess and Go 19x19. In contrast, PettingZoo's chess implementation exhibits a throughput approximately 10x slower than its Go 19x19 counterpart. This suggests that there might be room for optimization in PettingZoo's chess implementation regarding execution speed. To ensure that all Pgx environments achieve reasonable throughput and scalability like the four environments shown in Fig. 3, we measured the sample throughputs of all other Pgx environments on a DGX-A100 workstation, following the same approach as in the previous section. The number of vectorized environments was 1024 for a single A100 GPU and 8192 for eight A100 GPUs. The results are shown in Fig. 4. From these results, we can see that even in the slowest environment, Pgx achieves a throughput of approximately  $10^{5}$  samples/second with a single A100 GPU. This demonstrates the efficiency of Pgx, considering that achieving such throughput with other Python libraries, as shown in the previous section, is challenging even in simpler environments like Tic-tac-toe. Furthermore, we observe a significant improvement in throughput when using eight A100 GPUs compared to a single A100 GPU in all environments. The throughput increased by an average of 7.4x across all
environments (at least 6.6x in the MinAtar Breakout environment). This highlights the excellent parallelization performance of Pgx across multiple accelerators.
5 AlphaZero on Pgx environments
In this section, we showcase the effectiveness of RL training on accelerators using the Pgx environments, with a specific focus on two-player perfect information games such as Go. For demonstrations of RL training in other game types, please refer to App. D. We begin by providing a brief overview of the Gumbel AlphaZero. Next, we describe the Pgx environments where we apply the Gumbel AlphaZero. Subsequently, we explain the experimental setup and discuss the selection of baseline models, which serve as anchor opponents for evaluation purposes. Finally, we present the results.
AlphaZero [3] and Gumbel AlphaZero [5]. AlphaZero is an RL algorithm that leverages a combination of NNs and Monte Carlo Tree Search (MCTS). It has achieved state-of-the-art performance in chess, shogi, and Go, employing a unified approach. Through self-play, AlphaZero integrates MCTS with the current parameters of NNs, continually updating them through training on the samples generated during self-play. Gumbel AlphaZero is an adaptation of AlphaZero that removes several heuristics present in AlphaZero, enabling it to function even with a reduced number of simulations. In particular, Gumbel AlphaZero addressed the issue in the original AlphaZero where utilizing Dirichlet noise at the root node during tree search did not guarantee policy improvement. To overcome this limitation, Gumbel AlphaZero introduced an enhancement by utilizing the Gumbel-Top-k trick to perform sampling actions without replacement. They also proposed a MuZero version of the algorithm but we focus on the evaluation of AlphaZero in this paper. They released the Mctx library, which includes a JAX implementation of Gumbel AlphaZero. We utilized this library in our experiments. From now on, unless otherwise specified, when referring to "AlphaZero," it refers to Gumbel AlphaZero, not the original AlphaZero.
**Environments.** While Pgx enables fast simulations on accelerators, large-scale environments such as chess, shogi, and 19x19 Go pose significant challenges for researchers in terms of computational resources for learning. To address this, Pgx provides several small-scale environments, including miniature versions of shogi (Animal shogi) and chess (Gardner chess), as well as Hex (11x11) and Othello (8x8). In addition to these environments, we included the 9x9 Go environment to create a set of five environments for training AlphaZero. Here, we describe these environments briefly:
- Animal shogi is a 4x3 miniature version of shogi designed originally for children. Like shogi, players can reuse captured pieces. The small board size allows researchers to conduct research in an environment where planning ability is important with minimal computational resources.
- Gardner chess is a 5x5 variant of minichess that uses the leftmost five columns of the standard chessboard. It has a history of active play by human players in Italy [33].
- Go 9x9 maintains the essential aspects of full-sized Go while being the smallest playable board size in the game. The advantage of 9x9 Go is that several full-sized Go AI models are also capable of playing on the 9x9 board, allowing us to use 9x9 Go as a reliable benchmark (e.g., [5]).
- Hex is a game played on an 11x11 board where two players take turns placing stones, and the player who forms a connected path from one side of the board to the other with their stones wins. Its rules are simple, making it relatively easy to interpret for researchers.
- Othello, also known as Reversi, is played on an 8x8 board. Players take turns placing stones and flip the opponent's discs that are sandwiched. The game concludes when neither player can make a valid move, and the player with the most discs on the board wins.
For more detailed information about each environment, please refer to App. C.
Training setup. We trained the models using the same network architecture and hyperparameters across all five environments. The network architecture is 6 ResNet blocks with policy head and value head, following the structure outlined in the original AlphaZero study [3] basically but with a smaller model size. During the self-play, 32 simulations were performed at each position for policy improvement. In each iteration, we generated data for 256 steps with a self-play batch size of 1024 (i.e., the number of vectorized environments). We then divided this data into mini-batches of