pgx

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

Note

If you find this project helpful, we would be grateful for your support through a GitHub star to help us grow the community and motivate further development!

v1 A simplified, children-friendly Mahjong. Tic-tac-toe
"tic_tac_toe" v0 Three in a row wins.

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:

• RobertTLange/gymnax: JAX implementation of popular RL environments (classic control, bsuite, MinAtar, etc) and meta RL tasks
• google/brax: Rigidbody physics simulation in JAX and continuous-space RL tasks (ant, fetch, humanoid, etc)
• instadeepai/jumanji: A suite of diverse and challenging RL environments in JAX (bin-packing, routing problems, etc)
• flairox/jaxmarl: Multi-Agent RL environments in JAX (simplified StarCraft, etc)
• corl-team/xland-minigrid: Meta-RL gridworld environments in JAX inspired by MiniGrid and XLand
• MichaelTMatthews/Craftax: (Crafter + NetHack) in JAX for open-ended RL
• epignatelli/navix: Re-implementation of MiniGrid in JAX

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

• Anakin framework: Highly efficient RL framework that works with JAX-native environments on TPUs
• deepmind/mctx: JAX-native MCTS implementations, including AlphaZero and MuZero
• deepmind/rlax: JAX-native RL components
• google/evojax: Hardware-Accelerated neuroevolution
• RobertTLange/evosax: JAX-native evolution strategy (ES) implementations
• adaptive-intelligent-robotics/QDax: JAX-native Quality-Diversity (QD) algorithms
• luchris429/purejaxrl: Jax-native RL implementations

Limitation

Currently, some environments, including Go and chess, do not perform well on TPUs. Please use GPUs instead.

Citation

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

@inproceedings{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},
  booktitle={Advances in Neural Information Processing Systems},
  pages={45716--45743},
  volume={36},
  year={2023}
}

LICENSE

Apache-2.0

Expand