Pgx: Hardware-accelerated parallel game simulation for reinforcement learning

03/29/2023
by   Sotetsu Koyamada, et al.
0

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 at https://github.com/sotetsuk/pgx.

READ FULL TEXT

Please sign up or login with your details

Forgot password? Click here to reset