Pgx: Hardware-accelerated parallel game simulation for reinforcement learning
Graduate School of Informatics, Kyoto University, Kyoto, Japan
arXiv:2303.17503 [cs.AI], (29 Mar 2023)
@misc{koyamada2023pgx,
title={Pgx: Hardware-accelerated parallel game simulation for reinforcement learning},
author={Sotetsu Koyamada and Shinri Okano and Soichiro Nishimori and Yu Murata and Keigo Habara and Haruka Kita and Shin Ishii},
year={2023},
eprint={2303.17503},
archivePrefix={arXiv},
primaryClass={cs.AI}
}
We propose Pgx, a collection of board game simulators written in JAX. Thanks to auto-vectorization and Just-In-Time compilation of JAX, Pgx scales easily to thousands of parallel execution on GPU/TPU accelerators. We found that the simulation of Pgx on a single A100 GPU is 10x faster than that of existing reinforcement learning libraries. Pgx implements games considered vital benchmarks in artificial intelligence research, such as Backgammon, Shogi, and Go. Pgx is available.
April 2, 2023 by hgpu
