SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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. Each piece has its own specific movement rules. While capturing the opponent's pieces is allowed, unlike in shogi, the captured pieces cannot be reused. In addition to checkmate, there are other terminal conditions for a draw, such as a threefold repetition of the same position or a stalemate (when the king is not in check but has no legal moves). There are also special moves called pawn promotion, en passant, and castling.
Observation. The observation design follows AlphaZero [3]. Each plane represents an  $8 \times 8$  board and encodes the position-dependent features (Table 6). As in AlphaZero, the last 8-step history is
Table 6: Chess position-dependent features.
Feature	# Planes
P1 piece	6
P2 piece	6
Repetitions	2
Total	14
encoded  $(14 \times 8 = 112$  planes in total). Also, Table 7 shows the 7 position-independent feature planes.
Table 7: Chess position-independent features.
Feature	# Planes
Color	1
Total move count	1
P1 castling	2
P2 castling	2
No progress count	1
Total	7
Action. Action design also follows AlphaZero [3]. There are  $64 \times 73 = 4672$  discrete actions. Each action consists of 64 source positions and 73 moves, where 56 moves are queen moves, 8 moves are knight moves, and 9 moves are underpromotions.
Game	Connect four
Number of players	2
Observation shape	6 × 7 × 2
Number of actions	7
Tag	perfect information
Connect Four is a two-player perfect information game played on a 7x6 board.
Rules. Players take turns dropping discs into any of the seven columns. The objective is to create a line of four of their own discs either vertically, horizontally, or diagonally. The player who achieves this first is declared the winner. If all the spaces on the board are filled, and neither player has managed to create a line of four discs, the game ends in a draw.
Observation. Observation consists of two  $6 \times 7$  feature planes (Table 8).
Table 8: Connect Four features.
Feature	# Planes
P1 discs	1
P2 discs	1
Action. Each action represents the column index into which the player drops the token.
Gardner chess
Game	Gardner chess
Number of players	2
Observation shape	5 × 5 × 115
Number of actions	1225
Tag	perfect information
Gardner chess is a two-player perfect information game played on a 5x5 board. It uses the pieces corresponding to the leftmost 5 columns of standard chess.
Rules. The rules are the same as in regular chess, except that double pawn moves, en passant, and castling are not allowed.
Observation. The observation design follows AlphaZero [3]. Each plane represents a  $5 \times 5$  board and encodes position-dependent features (Table 9). As in AlphaZero, the last 8-step history is
Table 9: Gardner chess position-dependent features.
Feature	# Planes
P1 piece	6
P2 piece	6
Repetitions	2
Total	14
encoded  $(14 \times 8 = 112$  planes in total). Also, Table 10 shows position-independent feature planes.
Table 10: Gardner chess position-independent features.
Feature	# Planes
Color	1
Total move count	1
No progress count	1
Total	3
Action. Action design also follows AlphaZero [3]. There are  $25 \times 49 = 1225$  discrete actions. Each action consists of 25 source positions and 49 moves, where 32 moves are queen moves, 8 moves are knight moves, and 9 moves are underpromotions.
Game	Go 9x9	Go 19x19
Number of players	2
Observation shape	9 × 9 × 17	19 × 19 × 17
Number of actions	82	362
Tag	perfect information
Go is a two-player perfect information game played on a 19x19 board. The essential strategic aspects of Go can be preserved even on smaller boards like 9x9. There are variations of the Go rules, such as Chinese rules and Japanese rules. In computer Go, the Tromp-Taylor rules [47] are commonly used, and Pgx also follows them. To address the inherent advantage of the first player (black), it is common to add a scoring adjustment called komi (e.g., 6.5) to the final score of the second player (white). This improves fairness and helps to avoid a draw. Pgx uses a komi of 6.5 by default.
Observation. The observation design follows the AlphaGo Zero observation design [2]. Each plane represents a  $9 \times 9$  (or  $19 \times 19$ ) board and encodes the position-dependent features (Table 11). As
Table 11: Go features.
# Planes	Description
P1 stones	1
P2 stones	1
is the case with AlphaGo Zero, the last 8-step history is encoded ( $2 \times 8 = 16$  planes in total). An additional plane encodes the color of the current player. This is necessary for the agent to know the komi information.
Action. Each action represents the position on the board to place a stone. The last action represents pass.
Game	Hex
Number of players	2
Observation shape	11 × 11 × 4
Number of actions	122
Tag	perfect information
Hex is a two-player perfect information game played on an  $11 \times 11$  board.
Rules. Players take turns placing stones on a board, aiming to connect one side of the board to the opposite side with their stones. A draw does not occur because both players cannot simultaneously connect their sides of the board. To balance the first-player advantage, the swap rule is implemented. This rule allows the second player, instead of placing a stone, to replace the color of the first player's stone with their own color at the mirrored position.
Observation. Each plane represents an  $11 \times 11$  board and encodes position-dependent features (Table 12). The last two planes encode the color of the current player and whether the swap is a legal
Table 12: Hex position-dependent features.
Feature	# Planes
P1 stones	1
P2 stones	1
action (Table 13). Color information is necessary for the agent to know the side to connect.
Table 13: Hex position-independent features.
Feature	# Planes
Color	1
Swap	1
Action. The first 121 actions represent placing a stone on each cell of the board. The final action 121 is the swap action available only at the second turn.
Reward. Each player gets  $+1$  (win) or  $-1$  (lose).
Kuhn poker
Game	Kuhn poker
Number of players	2
Observation shape	7
Number of actions	4
Tag	imperfect information
Kuhn poker is a two-player imperfect information game designed for research purposes [36].
Rules. The deck has three cards: Jack, Queen, and King. Each player is dealt one card, and the remaining card is unused. Players have four actions available: check (pass without placing a bet when no one else has bet), call, bet, and fold. The following scenarios can occur, with player A being the first to play and player B being the second:
- bet (A) - call (B): Showdown, and the winner takes  $+2$ .
- bet (A) - fold (B): A takes  $+1$ .
- check (A) - check (B): Showdown, and the winner takes  $+1$ .
- check (A) - bet (B) - call (A): Showdown, and the winner takes  $+2$ .
- check (A) - bet (B) - fold (A): B takes  $+1$ .
As Kuhn poker is a zero-sum game, the loser of the game receives the negative of the winner's payoff.
Observation. The observation consists of a binary vector of size 7 (Table 14).
Table 14: Kuhn poker features.
Feature	Size
P1 hand	3
P1 chip	2
P2 chip	2
Total	7
Action. Call (0), Bet (1), Fold (2) or Check (3).
Reward. Each player gets  $+2, +1, -1$ , or  $-2$ , depending on the game payoff.
Leduc hold'em
Game	Leduc hold'em
Number of players	2
Observation shape	34
Number of actions	3
Tag	imperfect information
Leduc hold'em is a two-player imperfect information game designed for research purposes [37].
Rules. The deck consists of six cards: two Jacks, two Queens, and two Kings. Each player starts with a bet of 1 chip. The game consists of two rounds. In the first round, each player is dealt one private card, and in the second round, one public card is revealed. In each round, players have the option to call, raise, or fold. If either player folds, the hand ends, and the opponent takes the pot. If both players call, the game proceeds to the next round. In the second round, it advances to showdown, where the winner is determined by the strength of their cards. Each player can raise 2 chips in the first round and 4 chips in the second round. In each round, each player is allowed to raise only once (a total of 2 raises per round). Therefore, the maximum number of chips that can be bet for each player is  $1 + 2 \times 2 + 4 \times 2 = 13$ .
Observation. The observation consists of a binary vector of size 34 (Table 15).
Table 15: Leduc hold'em features.
Feature	Size
P1 hand	3
Public cards	3
P1 chip	14
P2 chip	14
Total	34
Action. Call (0), Raise (1), or Fold (2).
Game	Othello
Number of players	2
Observation shape	8 × 8 × 2
Number of actions	65
Tag	perfect information
Othello is a two-player perfect information game played on an 8x8 board.
Rules. Players take turns placing discs. The player with more discs at the end wins. Players can place a disc in an empty position where it can sandwich the opponent's discs between their own discs, and the sandwiched discs are flipped to their own color. If a player has no valid move to make, they must pass. The game ends when neither player can make a legal move.
Observation. Each plane represents an  $8 \times 8$  board and encodes the features in Table 16.
Table 16: Othello features.
Feature	# Planes
P1 discs	1
P2 discs	1
Action. The first 64 actions represent placing a disc on each square of the board. The last action represents pass.
Game	Shogi
Number of players	2
Observation shape	9 × 9 × 119
Number of actions	2187
Tag	perfect information
Shogi is a two-player perfect information game on a 9x9 board.
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. Each piece has its own specific movement rules. Captured pieces can be reused by dropping them to the board (instead of moving a piece) on the player's turn. The game ends when a player achieves checkmate or when the game reaches a draw by four-fold repetition. Unlike chess, shogi does not have a draw by stalemate (when the king is not in check but has no legal moves). Not only pawns but also other pieces (except Gold and King) can promote by entering the opponent's territory (1-3 rows).
Observation. The observation design follows the dlshogi observation design [48]. Each plane represents a  $9 \times 9$  board and encodes the position-dependent features in Table 17. Also, the observation
Table 17: Shogi position-dependent features.
# Planes	Description
P1 piece	14
Attacked by P1 piece	14
Attacked by N or more P1 pieces (N = 1, 2, 3)	3
P2 piece	14
Attacked by P2 piece	14
Attacked by N or more P2 pieces (N = 1, 2, 3)	3
Total	62
has the position-independent feature planes (Table 18).
Table 18: Shogi position-independent features.
# Planes	Description
P1 hand	28
P2 hand	28
P1 king is checked	1
Total	57
Action. The action design also follows the dlshogi action design [48]. There are  $81 \times 27 = 2187$  distinct actions. Each action consists of 81 destination to which the piece moves and 27 directions from which the piece moves. The direction is one of 10 moves (8 King moves and 2 Knight moves), 10 moves with promotion, or 7 drops.
Sparrow mahjong
Game	Sparrow mahjong
Number of players	3
Observation shape	11 × 15
Number of actions	11
Tag	imperfect information
Sparrow mahjong is an imperfect information game played with 44 tiles. It is a simplified version of Japanese mahjong. Sparrow mahjong is designed for human players who are not familiar with the rules of full-size Japanese mahjong but requires the same essential skills as Japanese mahjong. It can be played with 2 to 6 players, but Pgx uses a 3-player version.
Rules. Sparrow mahjong has 11 types of tiles, and each tile has 4 copies, for a total of 44 tiles. The 11 types of tiles are bamboo 1-9 (1s-9s) and red dragon (rd), green dragon (gd). Each player starts with 5 tiles in their hand, and the game proceeds by each player drawing a tile from the deck and discarding one tile from their hand to a public river. The objective is to "accomplish" a hand of 6 tiles faster than other players and with a higher score. There are two ways to win the game: ron (winning by using a tile discarded by another player) and tsumo (winning by drawing a tile from the deck). A hand is "accomplished" when the hand consists of either (1) one chow (a sequence of three tiles) and one pung (a set of three identical tiles), (2) two chows, or (3) two pungs. For example, "2s 3s 4s 6s 6s 6s" (one chow and one pung), "1s 2s 3s 7s 8s 9s" (two chows), and "1s 1s 1s rd rd rd" (two
$pungs$  ) are accomplished hands. The accomplished hand is scored according to its difficulty. There is a minimum score required to win, and it is key to infer the opponent's hand. There is a unique rule in Japanese mahjong called furiten: a player cannot win by ron with a tile discarded previously by themselves. The game ends when a player wins or when the deck is empty.
Observation. The observation consists of 15 feature planes, where each plane represents 11 tile types (Table 19). P1 represents the current player, while P2 and P3 represent the opponents.
Table 19: Sparrow mahjong features.
Feature	# Planes
P1 hand	4
Red dora in P1 hand	1
Dora	1
All discarded tiles by P1	1
All discarded tiles by P2	1
All discarded tiles by P3	1
Discarded tiles in the last 3 steps by P2	3
Discarded tiles in the last 3 steps by P3	3
Total	15
Action. Tile to discard.
Reward. Each player gets the game payoff normalized to  $[-1, 1]$ .
Game	Tic-tac-toe
Number of players	2
Observation shape	3 × 3 × 2
Number of actions	9
Tag	perfect information
Tic-tac-toe is a two-player perfect information game played on a  $3 \times 3$  board.
Rules. Players take turns, with one marking X and the other marking O. The objective is for a player to place their mark in a vertical, horizontal, or diagonal line. The player who achieves this first is the winner. If all nine squares are filled and neither player has made a line, the game ends with a draw.
Observation. Each plane represents a  $3 \times 3$  board and encodes the features in Table 20.
Table 20: Tic-tac-toe features.
Feature	# Planes
Marked by P1	1
Marked by P2	1
Action. Square index to place a mark.
D PPO training example
Figure 10: PPO training in MinAtar suite. Each model is trained up to 20M frames on a single A100 GPU. The shaded area represents the standard error of five runs.
In Sec. 5, we conducted learning experiments with AlphaZero for two-player, perfect information games such as  $9 \times 9$  Go. Here, we present the results of RL training using the Proximal Policy Optimization (PPO) algorithm [49] in the MinAtar [20] environments as an example of model-free RL with Pgx.
MinAtar is a miniature version of Atari, specifically used by researchers to perform detailed ablation studies on RL methods using visual inputs (like Atari), without requiring large computational resources [38, 40]. MinAtar includes five games: Asterix, Breakout, Freeway, Seaquest, and Space Invader. It may be worth noting that Freeway and Seaquest are highlighted as possible useful benchmarks for exploration methods [38]. Pgx has re-implemented all five environments in JAX. The versions of  $pgx$  and  $pgx$ -minatar used here were v0.9.0 and v0.2.1, respectively.
PPO implementation details. The implementation of PPO was based on the code available at https://github.com/luchris429/purejaxrl [50], with some modifications. As Pgx focuses on achieving high-speed training through vectorized simulation, we utilized a large batch size of 4096, similar to the training example in Brax [8]. We fine-tuned hyperparameters using the Asterix environment in the preliminary experiments (with a different seed used in the main experiments), and we maintained the same hyperparameters across all five games. The hyperparameters are listed in Table 21. Given the large batch size, the training was conducted up to 20M frames, which is longer compared to previous studies. However, it should be noted that the execution time is remarkably short, as discussed later.
Table 21: Hyperparameters in PPO training.
Hyperparameter	Value
Rollout batch size (i.e., number of vectorized environments)	4096
Rollout length	128
Training minibatch size	4096
Number of epochs	3
Optimizer	Adam
Learning rate	0.0003
Gradient clipping max norm	0.5
Discount factor (γ)	0.99
GAE lambda (λ)	0.95
Clipping parameter (ε)	0.2
Value function coefficient	0.5
Entropy coefficient	0.01
Results. Fig. 10 displays the learning results of PPO in the five MinAtar games. We conducted five learning runs with different seeds for each game. At each point of the learning process, we conducted 100 evaluation runs using the learned (stochastic) policy at that time and plotted the average score. The shaded area illustrates the standard error of the mean scores over the five runs. The runtime includes the RL training time and JIT-ompilation time of JAX code. In all games, PPO with Pgx achieves sufficiently high scores in less than a minute, except for Seaquest, which takes more than 80 seconds to elapse 20M frames. These results demonstrate the effectiveness of vectorized execution in other game types than two-player, perfect information games implemented in Pgx.
E AlphaZero training details
This section describes the details of Gumbel AlphaZero training (Sec. 5). Our study fundamentally follows the original AlphaZero [3] framework except for the following points:
- To enable training on a single GPU, we alternated between self-play and learning, rather than executing them as two independent processes.
- To examine the impact of self-play speed on AlphaZero training speed, we discarded data after a single use, instead of employing a replay buffer.
- We utilized ResNet v2 [51] in place of ResNet v1.
Table 22 shows the hyperparameters used in the training. Note that we used the same hyperparameters for all five games. The GPUs used for training vary depending on the game. Table 23 shows the GPUs and runtime for each game. The version of Pgx used in the training is v0.8.0.
Table 22: Hyperparameters in AlphaZero training.
Hyperparameter	Value
Number of residual blocks	6
Number of channels of conv. layer	128
Self-play batch size (i.e., number of vectorized environments)	1024
Self-play length in each iteration	256
Number of simulations per move	32
Training minibatch size	4096
Optimizer	Adam
Learning rate	0.001
Completed Q-values value scale	0.1
Completed Q-values rescale values	False
Table 23: GPUs used in AlphaZero training.
Environment	GPUs	Runtime (hours)
Animal shogi	A4000 x 1	6.2
Gardner chess	A4000 x 4	14.3
Go 9x9	A100 x 1	8.6
Hex	A4000 x 1	17.6
Othello	A4000 x 1	11.4