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)

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:

  • JAX-native. All step functions are JIT-able
  • Super fast in parallel execution on accelerators
  • Various game support including Backgammon, Chess, Shogi, and Go
  • Beautiful visualization in SVG format

Install

pip install pgx

Usage

Open In Colab

import jax
import pgx

env = pgx.make("go-19x19")
init = jax.jit(jax.vmap(env.init))  # vectorize and JIT-compile
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.all():
    action = model(state.current_player, state.observation, state.legal_action_mask)
    state = step(state, action)  # state.reward (2,)

Supported games and road map

:warning: Pgx is currently in the beta version. Therefore, API is subject to change without notice. We aim to release v1.0.0 in April 2023. Opinions and comments are more than welcome!

Use pgx.available_games() to see the list of currently available games.

Game Environment Visualization
2048 :white_check_mark: :white_check_mark:
Animal Shogi :white_check_mark: :white_check_mark:
Backgammon :white_check_mark: :white_check_mark:
Bridge Bidding :construction: :white_check_mark:
Chess :white_check_mark: :white_check_mark:
Connect Four :white_check_mark: :white_check_mark:
Go :white_check_mark: :white_check_mark:
Hex :white_check_mark: :white_check_mark:
Kuhn Poker :white_check_mark: :white_check_mark:
Leduc hold'em :white_check_mark: :white_check_mark:
Mahjong :construction: :construction:
MinAtar/Asterix :white_check_mark: :white_check_mark:
MinAtar/Breakout :white_check_mark: :white_check_mark:
MinAtar/Freeway :white_check_mark: :white_check_mark:
MinAtar/Seaquest :white_check_mark: :white_check_mark:
MinAtar/SpaceInvaders :white_check_mark: :white_check_mark:
Othello :white_check_mark: :white_check_mark:
Shogi :white_check_mark: :white_check_mark:
Sparrow Mahjong :white_check_mark: :white_check_mark:
Tic-tac-toe :white_check_mark: :white_check_mark:

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

@article{koyamada2023pgx,
  title={Pgx: Hardware-accelerated parallel game simulation 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


Release history Release notifications | RSS feed

Download files

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

Source Distribution

pgx-0.5.3.tar.gz (208.0 kB view hashes)

Uploaded Source

Built Distribution

pgx-0.5.3-py3-none-any.whl (289.8 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