pgx

repository·main·Indexed 20 days ago

https://github.com/sotetsuk/pgx

A collection of GPU-accelerated, JAX-native parallel game simulators for reinforcement learning in discrete state spaces. It supports a wide range of environments including Chess, Go, Shogi, Backgammon, and the MinAtar suite. The library features JIT-able and vectorizable step functions via jax.vmap for high-speed execution on accelerators, provides conversion to the PettingZoo AEC API, and includes utilities for loading baseline models and integrating with mctx for Monte Carlo Tree Search (MCTS).

Tokens
25.3K
Snippets
85
Records
177
Agent score
70%

What's inside pgx

  1. Understand the Kuhn Poker environment

    main

    Kuhn poker is a simplified two-player poker game using three cards: Jack (J), Queen (Q), and King (K). Each player is dealt one card, and the remaining card is unused. The game involves betting and passing, with payoffs determined by the sequence of actions and the final showdown.

    Key Specifications:

    • Version: v1
    • Players: 2
    • Actions: 2 (Bet or Pass)
    • Observation Shape: (7,) (Boolean type)
    • Rewards: {-2, -1, +1, +2} (Zero-sum game)
  2. Go rules and Superko implementation

    main

    Pgx implements Go following the Tromp-Taylor Rules.

    Komi

    By default, the komi is set to 7.5. You can specify a different value when constructing the Go class.

    Superko Rule

    Pgx uses a compromise for the Superko rule to maintain efficiency:

    • Legal Moves: Uses SSK (Situational Superko) to determine legal moves.
    • PSK (Positional Superko): If a move results in a PSK state, it is treated as an immediate loss for the player making the move.
  3. Understand Backgammon observations

    main

    The observation is a vector of size 34. The first 28 dimensions follow the representation used in [Antonoglou+22], where positive numbers represent the current player's chips and negative numbers represent the opponent's chips.

    IndexDescription
    [:24]Number of checkers on each of the 24 positions
    [24:26]Number of checkers on the bar (hit chips)
    [26:28]Number of checkers borne off
    [28:34]Number of available moves for each die number (1-6)
  4. Understand Backgammon actions

    main

    Actions are designed based on micro-actions. An action consists of up to 4 micro-actions (representing the maximum number of dice a player can play per turn).

    Each micro-action is encoded as a single integer using the formula: micro-action = src * 6 + die

    Where:

    • die is the value of the die used.
    • src is the source position (0 to 25):
      • 0: No-op
      • 1: Retrieving a chip from the hit pile (bar)
      • 2-25: Selecting a chip from one of the 24 possible points

    Note: Unlike standard micro-action implementations, a state transition occurs after every micro-action. A player's turn can continue for up to 4 micro-actions.

  5. Understand Gardner Chess observations

    main

    Observations follow the AlphaZero [Silver+18] design. The observation shape is (5, 5, 115). P1 denotes the current player and P2 denotes the opponent.

    IndexDescription
    [:, :, 0:6]P1 board @ 0-steps before
    [:, :, 6:12]P2 board @ 0-steps before
    [:, :, 12:14]Repetitions @ 0-steps before
    ...(@ 1-7 steps before)
    [:, :, 112]Color
    [:, :, 113]Total move count
    [:, :, 114]No progress count
  6. Understand Hex rewards and termination

    main

    Hex does not allow draws. Termination occurs when one player successfully connects their opposite sides of the board.

    Rewards are only provided at terminal states:

    • Win: +1
    • Lose: -1
  7. Understand Tic-tac-toe observations and rewards

    main

    Observation Structure

    The observation is a boolean tensor of shape (3, 3, 2):

    • [:, :, 0]: Represents the (3, 3) squares filled by the current player.
    • [:, :, 1]: Represents the (3, 3) squares filled by the opponent of the current player.

    Rewards

    Non-zero rewards are only issued at terminal states:

    • Win: +1
    • Lose: -1
    • Draw: 0

    Termination

    The game terminates when:

    1. A player places three symbols in a row (horizontally, vertically, or diagonally).
    2. All nine squares are filled.
  8. Understand the 2048 observation format

    main

    The observation is a (4, 4, 31) boolean tensor. This design is based on a binary representation of the board tiles.

    An index [i, j, b] represents that the square at position (i, j) contains a tile with the value 2^b (for b > 0).

  9. Kuhn Poker game rules and payoffs

    main

    The game outcome and rewards depend on the sequence of actions taken by the two players:

    1. bet (1st) - bet (2nd): Showdown occurs; winner takes +2, loser takes -2.
    2. bet (1st) - pass (2nd): 1st player takes +1, 2nd player takes -1.
    3. pass (1st) - pass (2nd): Showdown occurs; winner takes +1, loser takes -1.
    4. pass (1st) - bet (2nd) - bet (1st): Showdown occurs; winner takes +2, loser takes -2.
    5. pass (1st) - bet (2nd) - pass (1st): 2nd player takes +1, 1st player takes -1.