Skip to main content

GPU/TPU-accelerated parallel game simulators for reinforcement learning (RL)

Project description

ci

A collection of GPU/TPU-accelerated parallel game simulators for reinforcement learning (RL)

🎉 v1.0.0 is released! (2023.6.20)

Why Pgx?

Brax, a JAX-native physics engine, provides extremely high-speed parallel simulation for RL in continuous state space. Then, what about RL in discrete state spaces like Chess, Shogi, and Go? Pgx provides a wide variety of JAX-native game simulators! Highlighted features include:

  • Super fast in parallel execution on accelerators
  • 🎲 Various game support including Backgammon, Chess, Shogi, and Go
  • 🖼️ Beautiful visualization in SVG format

Quick start

Usage

The following code snippet shows a simple example of using Pgx. You can try it out in this Colab. Note that all step functions in Pgx environments are JAX-native., i.e., they are all JIT-able. Please refer to the documentation for more details.

import jax
import pgx

env = pgx.make("go_19x19")
init = jax.jit(jax.vmap(env.init))
step = jax.jit(jax.vmap(env.step))

batch_size = 1024
keys = jax.random.split(jax.random.PRNGKey(42), batch_size)
state = init(keys)  # vectorized states
while not (state.terminated | state.truncated).all():
    action = model(state.current_player, state.observation, state.legal_action_mask)
    state = step(state, action)  # state.reward (2,)

Pgx is a library that focuses on faster implementations rather than just the API itself. However, the API itself is also sufficiently general. For example, all environments in Pgx can be converted to the AEC API of PettingZoo, and you can run Pgx environments through the PettingZoo API. You can see the demonstration in this Colab.

Installation

pip install pgx

Note that the MinAtar suite is provided as a separate extension for Pgx (pgx-minatar). Therefore, please run the following command additionaly to use the MinAtar suite in Pgx:

pip install pgx-minatar

Pgx is provided under the Apache 2.0 License, but the original MinAtar suite follows the GPL 3.0 License. Therefore, please note that the separated MinAtar extension for Pgx also adheres to the GPL 3.0 License.

Supported games

Backgammon Chess Shogi Go

Use pgx.available_envs() -> Tuple[EnvId] to see the list of currently available games. Given an <EnvId>, you can create the environment via

>>> env = pgx.make(<EnvId>)
Game/EnvId Visualization Version Five-word description
2048
"2048"
v0 Merge tiles to create 2048.
Animal Shogi
"animal_shogi"
v0 Animal-themed child-friendly shogi.
Backgammon
"backgammon"
v0 Luck aids bearing off checkers.
Bridge bidding
"bridge_bidding"
v0 Partners exchange information via bids.
Chess
"chess"
v1 Checkmate opponent's king to win.
Connect Four
"connect_four"
v0 Connect discs, win with four.
Gardner Chess
"gardner_chess"
v0 5x5 chess variant, excluding castling.
Go
"go_9x9" "go_19x19"
v0 Strategically place stones, claim territory.
Hex
"hex"
v0 Connect opposite sides, block opponent.
Kuhn Poker
"kuhn_poker"
v0 Three-card betting and bluffing game.
Leduc hold'em
"leduc_holdem"
v0 Two-suit, limited deck poker.
MinAtar/Asterix
"minatar-asterix"
v0 Avoid enemies, collect treasure, survive.
MinAtar/Breakout
"minatar-breakout"
v0 Paddle, ball, bricks, bounce, clear.
MinAtar/Freeway
"minatar-freeway"
v0 Dodging cars, climbing up freeway.
MinAtar/Seaquest
"minatar-seaquest"
v0 Underwater submarine rescue and combat.
MinAtar/SpaceInvaders
"minatar-space_invaders"
v0 Alien shooter game, dodge bullets.
Othello
"othello"
v0 Flip and conquer opponent's pieces.
Shogi
"shogi"
v0 Japanese chess with captured pieces.
Sparrow Mahjong
"sparrow_mahjong"
v0 A simplified, children-friendly Mahjong.
Tic-tac-toe
"tic_tac_toe"
v0 Three in a row wins.
  • Mahjong environments are under development 🚧 If you have any requests for new environments, please let us know by opening an issue
  • Five-word descriptions were generated by ChatGPT 🤖

Versioning policy

Each environment is versioned, and the version is incremented when there are changes that affect the performance of agents or when there are changes that are not backward compatible with the API. If you want to pursue complete reproducibility, we recommend that you check the version of Pgx and each environment as follows:

>>> pgx.__version__
'1.0.0'
>>> env.version
'v0'

See also

Pgx is intended to complement these JAX-native environments with (classic) board game suits:

Combining Pgx with these JAX-native algorithms/implementations might be an interesting direction:

Citation

If you use Pgx in your work, please cite the following paper:

@article{koyamada2023pgx,
  title={Pgx: Hardware-accelerated Parallel Game Simulators for Reinforcement Learning},
  author={Koyamada, Sotetsu and Okano, Shinri and Nishimori, Soichiro and Murata, Yu and Habara, Keigo and Kita, Haruka and Ishii, Shin},
  journal={arXiv preprint arXiv:2303.17503},
  year={2023}
}

LICENSE

Apache-2.0

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

pgx-1.1.0.tar.gz (333.6 kB view hashes)

Uploaded Source

Built Distribution

pgx-1.1.0-py3-none-any.whl (410.7 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page