SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
size 4096 for gradient estimation and parameter updates. We performed training for 400 iterations ( $\approx$  105M frames) in each environment. The choice of accelerator varied across the environments, but for example, in the case of  $9 \times 9$  Go, we trained the model using a single A100 GPU, and the training process took approximately 8 hours. For more detailed information about the network architecture, hyperparameters, and accelerator specifications used in the training process, refer to App. E.
Evaluation and baseline model selection. We trained agents using the AlphaZero algorithm on the five environments described above and performed evaluations. However, in multi-agent games, the performance of the trained agents is relative, so we need a reference agent for comparison. However, finding a suitable baseline model for any environment is difficult, or even if it exists, it may not be computationally efficient. Therefore, for researchers and practitioners using Pgx, we created our own baseline models. It is important to note that our baseline models are not designed to be state-of-the-art or oracle models, but rather serve the purpose of examining the learning process within the Pgx environment. In the given learning setup, we selected the 200-iteration ( $\approx$  52M frames) model for 9x9 Go and the 100-iteration ( $\approx$  26M frames) model for other environments, considering their lower complexity compared to 9x9 Go. To evaluate the agents, we estimated the Elo rating through their pairwise matches. We used the BayesElo program to calculate the Elo rating [42]. We adjusted the Elo rating to ensure that our baseline models had 1000 Elo. During the evaluation matches, the agents conducted 32 simulations for each move like during the training.
Results. Fig. 5 presents the learning results of AlphaZero in the five environments. We can observe that the agents successfully learn in all five environments starting from a random policy with the same network architecture and hyperparameters. Furthermore, for the baseline model in  $9 \times 9$  Go, we evaluated its performance by playing against Pachi [43] with 10K simulations per move, which was used as a baseline opponent in prior study [5]. The baseline model conducted 800 simulations for each move. Our baseline model outperformed Pachi with a record of 62 wins and 38 losses out of 100 matches, confirming its reasonable strength as a baseline model. Although no comparisons were made with other AIs in environments other than  $9 \times 9$  Go, we trained them using the same network architecture and hyperparameters as in  $9 \times 9$  Go. Given that the baseline model obtained in  $9 \times 9$  Go exhibited reasonable strength throughout the learning process, we suppose that the baseline models in other environments, which were trained with exactly the same settings, have also learned reasonably. Therefore, we believe that researchers can accelerate their research cycles using the five environments and baseline models presented here, instead of relying on full-scale chess, shogi, and  $19 \times 19$  Go, while exploring RL algorithms such as AlphaZero.
4https://www.remi-coulom.fr/Bayesian-Elo
6 Training scalability to multiple accelerators
In Sec. 4, we demonstrated a significant improvement in the pure throughput of Pgx when increasing the number of accelerators using a random agent. Here, we will show that increasing the number of cores also improves the learning speed in the AlphaZero training on Pgx environments.
Experiment setup. To demonstrate the improvement in learning speed by increasing the number of cores in AlphaZero training with Pgx, we conducted experiments in the 9x9 Go environment. In the experiments of Sec. 5, we performed training using a single A100 GPU in the 9x9 Go environment. Here, we conducted the exact same number of training frames but with an increased number of GPUs and batch size during self-play. In the experiment using a single A100 GPU, the batch size during self-play was set to 1024. However, in training with eight A100 GPUs, we increased the self-play batch size to 8192, which is eight times larger. It is important to note that the training batch size, learning rate, and the other hyperparameters were kept the same as in the single GPU case, ensuring that these hyperparameters did not affect the learning speed. Similar to Sec. 5, we used the baseline model as an anchor and adjusted Elo ratings so that the rating of the baseline model is 1000. Furthermore, in this setup, we want to mention that the time spent on self-play was dominant (more than  $90\%$ ) compared to the time spent on training (gradient calculation and parameter updates).
Results. Fig. 6 shows the learning curves for the 9x9 Go environment using one A100 GPU and eight A100 GPUs. The shaded regions represent the standard error of runs with two different seeds. Based on the figure, it is evident that when both models are trained with the same number of training frames, the model trained with eight A100 GPUs achieves the same level of performance approximately four times faster than the model trained with a single GPU. This experimental outcome highlights
Figure 6: Multi-GPU training.
the fact that Pgx not only enhances throughputs in random play but also accelerates the learning process when training RL algorithms such as AlphaZero. These results underscore the practical utility of Pgx in the field of RL, providing researchers and practitioners with a powerful tool for efficient experimentation utilizing multiple accelerators.
There are several notable limitations users should take care of regarding Pgx, including:
- Lack of support for Atari: One of the limitations of Pgx is that it does not support Atari, which is an important benchmark in RL research. This limitation arises from the difficulty of implementing the Atari emulator and re-implementing the dynamics of each game in JAX.
- Pgx API limitation: While the games currently implemented in Pgx can be exported to the PettingZoo API, which is regarded as a general API for multi-agent games, Pgx API is not well-suited for handling certain types of games. These game types include those with a varying number of agents and those that involve chance players (nature players) such as poker.
- JAX lock-in: Although Pgx provides a convenient way to implement fast algorithms in Python without directly working with  $\mathrm{C}++$ , it has a reliance on JAX, which may require users to be familiar with JAX. This can make it less straightforward to utilize other frameworks like PyTorch [44].
Our future work for Pgx includes the following:
- Expansion of baseline algorithms and models: Currently, we are unable to provide learning examples or models for large-scale games like chess, shogi, and Go  $19 \times 19$ , making it an important area for future work. We plan to expand the availability of strong models through proprietary training and connect with other strong AI systems to enhance the baselines.
- Diversification of game types: The current game collection in Pgx is biased towards (two-player) perfect information games. We plan to implement games with imperfect information, such as Texas hold'em and mahjong, to broaden the range of supported game types.
- Verification on TPUs: While we validated Pgx performance on an NVIDIA DGX-A100 workstation, it is important to conduct verification using Google TPUs as well. This will provide valuable insights into the performance and scalability of Pgx on different hardware architectures.
- Human-vs-agent UI: Developing a user interface that enables human-versus-agent gameplay is important future work. This will allow researchers with domain knowledge to conduct high-quality evaluations and experiments, fostering improved research and assessment.
By addressing these areas in our future work, we aim to enhance the capabilities and applicability of Pgx in the field of RL research and game AI for both researchers and practitioners.
We proposed Pgx, a library of hardware-accelerated game simulators that operate efficiently on accelerators, implemented in JAX. Pgx achieves 10-100x higher throughput compared to other libraries available in Python and demonstrates its ability to scale and train using multiple accelerators in the context of AlphaZero training. By providing smaller game environments, along with their baselines, Pgx facilitates the development and research of RL algorithms and planning algorithms that can operate at faster speeds. We anticipate that Pgx will contribute to advancing the field in terms of developing efficient RL algorithms and planning algorithms in these accelerated environments.
This paper is based on results obtained from a project, JPNP20006, subsidized by the New Energy and Industrial Technology Development Organization (NEDO).
[1] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484-489, 2016.
[2] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap, Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of Go without human knowledge. Nature, 550(7676):354-359, October 2017.
[3] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Simonyan, and Demis Hassabis. A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play. Science, 362(6419):1140-1144, December 2018.
[4] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering Atari, Go, chess and shogi by planning with a learned model. Nature, 588(7839):604-609, 2020.
[5] Ivo Danihelka, Arthur Guez, Julian Schrittwieser, and David Silver. Policy improvement by planning with Gumbel. In *ICLR*, 2022.
[6] Marc Lanctot, Edward Lockhart, Jean-Baptiste Lespiau, Vinicius Zambaldi, Satyaki Upadhyay, Julien Pérolat, Sriram Srinivasan, Finbarr Timbers, Karl Tuyls, Shayegan Omidshafiei, Daniel Hennes, Dustin Morrill, Paul Muller, Timo Ewalds, Ryan Faulkner, János Kramár, Bart De Vylder, Brennan Saeta, James Bradbury, David Ding, Sebastian Borgeaud, Matthew Lai, Julian Schrittwieser, Thomas Anthony, Edward Hughes, Ivo Danihelka, and Jonah Ryan-Davis. OpenSpiel: A Framework for Reinforcement Learning in Games. arXiv:1908.09453, 2019.
[41] Jiayi Weng, Huayu Chen, Dong Yan, Kaichao You, Alexis Duburcq, Minghao Zhang, Yi Su, Hang Su, and Jun Zhu. Tianshou: A Highly Modularized Deep Reinforcement Learning Library. JMLR, 23(267):1-6, 2022.
[42] Rémi Coulom. Whole-History Rating: A Bayesian Rating System for Players of Time-Varying Strength. In CG, 2008.
[43] Petr Baudiš and Jean-loup Gailly. PACHI: State of the Art Open Source Go Program. In ACG, 2012.
[44] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In NeurIPS, 2019.
[45] Gabriele Cirulli. 2048. https://gabrielecirulli.github.io/2048/, 2014.
[46] Jiang Rong, Tao Qin, and Bo An. Competitive Bridge Bidding with Deep Neural Networks. In AAMAS, 2019.
[47] John Tromp. The game of Go. https://tromp.github.io/go.html, 1995.
[48] Tadao Yamaoka. dlshogi. https://github.com/TadaoYamaoka/DeepLearningShogi, 2017.
[49] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Policy Optimization Algorithms. arXiv:1707.06347, 2017.
[50] Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered Policy Optimisation. In NeurIPS, 2022.
[51] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity Mappings in Deep Residual Networks. In ECCV, 2016.
A License
The source code of Pgx is available at https://github.com/sotetsuk/pgx under the Apache-2.0 license. However, since the original source code of the MinAtar game suite is released under the GPL 3.0 license, a separate extension for Pgx called pgx-minatar is provided at https://github.com/sotetsuk/pgx-minatar, which adheres to the GPL 3.0 license.
B Comparison to Brax and PettingZoo APIs
The Pgx API draws inspiration from the Brax and PettingZoo APIs, and we illustrate this with actual code examples. Fig. 7 and Fig. 8 present the code examples for the Pgx and Brax APIs, respectively. While both APIs share similarities, the Pgx API handles multi-agent environments and the current player vector specifies which agents are to act. Similar to Pgx, the PettingZoo API (Fig. 9) designates the next agent to act using the agent iterator concept. However, since Pgx is a library centered on environment vectorization, we prefer a vectorized current player to a flexible but dynamic iterator.
1 import jax  
2 import pgx  
3  
4 env = pgx.make("go_19x19")  
5 init_fn = jax jit(jax.vmap(env.init))  
6 step_fn = jax jit(jax.vmap(env_STEP))  
7  
8 batch_size = 1024  
9rng_keys = jax.random.split(jax.random.PRNGKey(9999), batch_size)  
10  
11 state = init_fn(rng_keys)  
12 while not(stateterminated | state.truncated).all():  
13 action = model(state.current_player, stateobservation, state/legal_action_mask)  
14 state = step_fn(state, action)
Figure 7: Example usage of Pgx API.
1 import jax  
2 from brax import envs  
3  
4 env = envs.get_environment("ant")  
5 reset_fn = jax jot(jax.vmap(env.reset))  
6 step_fn = jax jot(jax.vmap(env.step))  
7  
8 batch_size = 1024  
9rng_keys = jax.random.split(jax.random.PRNGKey(9999), batch_size)  
10  
11 state = reset_fn(rng_keys)  
12 while not state_done.all():  
13 action = model(state)  
14 state = step_fn(state, action)
Figure 8: Example usage of Brax API.
1 from pettingzoo.classic import go_v5   
2   
3 env  $=$  go_v5/env()   
4 env.reset()   
5 for agent in env(agent_iter(): observation, reward, terminated, truncated, info  $=$  env.last() action  $=$  model(agent, observation["observation"], observation["action_mask"]); env.step(action)   
9 env.close()
Figure 9: Example usage of PettingZoo API.
To prove the practical generality of the Pgx API, we provide a demonstration illustrating the conversion from Pgx API to PettingZoo API, available at https://github.com/sotetsuk/pgx/blob/main/colab/pgx2pettingzoo.ipynb.
C Game explanations
This section describes the short summary of each game implemented in Pgx (as of v1.4.0), except the MinAtar suite. See [20] for the description of the MinAtar suite. For the full description of each game, please refer to the Pgx documentation from https://github.com/sotetsuk/pxg.
2048
- Animal shogi
- Backgammon
- Bridge bidding
- Chess
Connect Four
Gardner chess
Go
Hex
Kuhn poker
Leduc hold'em
Othello
Shogi
- Sparrow mahjong
Tic-tac-toe
4	4	2	2
8	4	16
2
4
Game	2048
Number of players	1
Observation shape	4 × 4 × 31
Number of actions	4
Tag	perfect information with chance
The game of 2048 [45] is a single-player perfect information game with chance events on a  $4 \times 4$  board. The rules are simple, but to play well, agents need planning ability in stochastic dynamics.
Rules. The objective is to create larger-numbered tiles by merging tiles. The player can take four actions: left, up, right, or down. All tiles move in the chosen direction, and when tiles with the same number collide, they merge, forming a new tile with twice the value (e.g.,  $4 + 4 = 8$ ). When a new tile is created, the new tile number is rewarded. After the tiles have moved, a new tile whose value is either two or four is randomly placed in empty positions. The probability of a two appearing is 0.9, and the probability of a four appearing is 0.1. The game ends when no legal move exists.
Observation. The observation design follows [27]. Each plane represents a  $4 \times 4$  board, and each tile number is encoded in a 31-bit binary representation.
Action. Left (0), up (1), right (2), or down (3).
Reward. The sum of merged tiles.
Animal shogi
Game	Animal shogi
Number of players	2
Observation shape	4 × 3 × 194
Number of actions	132
Tag	perfect information
Animal shogi is a miniature version of the traditional Japanese board game, shogi, played on a small 4x3 board. It is a two-player perfect information game. Originally developed for children, it possesses sufficient complexity for human play, making it more than just a toy environment.
Rules. Two players take turns moving their pieces, aiming to achieve checkmate – a situation where the king is attacked and cannot make a legal move to escape the threat. The pieces used in Animal shogi are Lion (King), Giraffe (Rook), Elephant (Bishop), and Chick (Pawn). The available directions for movement are indicated by circular dots for each piece. Note that each piece can move only one square at a time, even if the piece is a Giraffe (Rook). The Chick (Pawn) can be promoted to a Hen (Gold) by entering the opponent's territory (the farthest rank). Similar to ordinary shogi, captured opponent pieces can be reused by dropping them on the board. A player can win by checkmating the opponent's Lion (King) or by having their Lion (King) enter the opponent's territory (Try rule). If the same position occurs three times, it results in a repetition draw.
Observation. The observation design follows AlphaZero [3]. Each plane represents a  $4 \times 3$  board and encodes the position-dependent features in Table 2. P1 represents the current player, and P2 represents the opponent player. As in AlphaZero, the last 8-step history is encoded ( $24 \times 8 = 192$
Table 2: Animal shogi position-dependent features.
Feature	# Planes
P1 piece	5
P2 piece	5
P1 prisoner piece count	6
P2 prisoner piece count	6
Repetitions	2
Total	24
planes in total). Also, Table 3 shows the position-independent features. Each plane has the same values for all positions.
Table 3: Animal shogi position-independent features.
Feature	# Planes
Color	1
Elapsed timesteps (normalized to [0, 1])	1
Action. Action consists of (1) the source position of the piece to move and the direction to move  $(12 \times 8 = 96)$ , and (2) the position of the piece to drop and the type of piece to drop  $(12 \times 3 = 36)$ , for a total of 132 discrete actions.
Reward. Each player gets  $+1$  (win),  $-1$  (lose), or 0 (draw).
Game	Backgammon
Number of players	2
Observation shape	34
Number of actions	156
Tag	perfect information with chance
Backgammon is a two-player game with perfect information that also incorporates chance events. It serves as an important benchmark for RL in stochastic environments [12, 27]. To excel, agents require a high planning capability within stochastic environments.
Rules. Players, represented by white and black colors, aim to move their set of 15 checkers across a board consisting of 24 points. The objective is to be the first to bear off all their checkers, moving in the opposite direction. Each turn involves rolling two dice, determining the number of points a player can move their checkers. If both dice show the same number, the player can make four moves of that number. The game has several constraints on legal actions:
- Checkers cannot bear off until all of them have reached the home board (an area one-quarter of the distance from the goal).
- A point with two or more opponent's checkers stacked on it is blocked, and the player cannot move their checkers onto that point.
- By moving a checker to a point where the opponent has only one checker, the player can hit the opponent's checker and send it to the central bar. A player with checkers on the bar must first move those checkers before moving any other checkers.
The game ends when one of the players has borne off all their checkers. Victory rewards differ:
- A gammon win (2 points) is when the opponent has not borne off any checkers.
- A backgammon win (3 points) is a gammon win where the opponent has checkers left on the bar or within the winner's home board.
- All other victories are termed single wins (1 point).
Observation. Table 4 shows the backgammon features. The first 28 features are the same as in [27]. The last 6 features encode the number of playable die number.
Table 4: Backgammon features.
Feature	Size
Checkers on points	24
Checkers on bar	2
Checkers borne off	2
Number of available moves for each die number	6
Total	34
Action. Action in Pgx follows [27]. There are  $26 \times 6 = 156$  discrete actions. Each action consists of 26 source positions and die number. The first source position is a no-op when there is no movable checker, the second source is the bar, and the remaining 24 sources represent each point on board.
Reward. Each player gets the game payoff as a reward.
Bridge bidding
Game	Bridge bidding
Number of players	4
Observation shape	480
Number of actions	38
Tag	imperfect information
Contract bridge is an imperfect information game played by four people in teams of two. Cooperation within the team is required to win.
Game flow. In the contract bridge, after each player is dealt 13 cards, there are two phases: (1) bidding and (2) playing.
- Bidding: Each player bids in an auction format to determine the contract, the target number of tricks their team will try to achieve in the next playing phase.
- Playing: Each player plays a card, and the player who plays the "strongest" card wins the trick. This process is repeated 13 times, aiming to maximize the number of tricks won by the team and achieve the target number of tricks for their team or prevent the opponent team from achieving their target number of tricks.
Bidding is considered more challenging than playing and is believed to have a significant impact on the outcome of the game. Therefore, previous studies [46, 30, 31] have often focused only on the bidding part by replacing the playing part results using a double dummy solver<sup>5</sup>. Pgx also follows this setting.
Rules of bidding. The four players take turns to act. The available actions for each player are as follows: (1) Bid "higher" than the previous bid, (2) Pass, (3) Double the opponent's last bid, or (4) Redouble in response to the opponent team's double. There are 35 possible bid combinations, including the target trick number (1-7) and the suit (club, diamond, heart, spade, no trump). If three consecutive passes occur after someone's last bid, the bidding phase ends, and the bidding team, target number of tricks, and trump suit are determined (if no one bids and there are four consecutive passes, the game ends in a draw). The rewards are based on the achievement of the target, the magnitude of the target trick number, and whether a double or redouble is present.
Observation. Table 5 shows the bridge bidding features. Observation design follows [31, 6].
Table 5: Bridge bidding features.
Feature	Size
Vulnerability	4
Pass before the opening bid	4
Bidding history (35 × 4 × 3)	420
Current player's hand	52
Total	480
Action. There are 38 discrete actions: pass, double, redouble, and 35 bids.
<sup>5</sup>https://github.com/dds-bridge/dds
Game	Chess
Number of players	2
Observation shape	8 × 8 × 119
Number of actions	4672
Tag	perfect information
Chess is a two-player perfect information game on an 8x8 board.