Pgx：面向强化学习的硬件加速并行博弈模拟器
人工智能
2024-01-17 v4 机器学习
摘要
我们提出 Pgx，一套用 JAX 编写并针对 GPU/TPU 加速器优化的棋类博弈强化学习（RL）环境。通过利用 JAX 在加速器上的自动向量化与并行化，Pgx 能够高效地扩展到加速器上数千个同步模拟。在 DGX-A100 工作站上的实验中，我们发现 Pgx 模拟 RL 环境的速度比 Python 中现有实现快 10–100 倍。Pgx 包含 RL 研究中常用作基准的 RL 环境，例如西洋双陆棋、国际象棋、将棋和围棋。此外，Pgx 提供迷你博弈集合与基线模型以促进快速研究循环。我们展示了利用 Pgx 环境对 Gumbel AlphaZero 算法的高效训练。总体而言，Pgx 为研究者提供高性能环境模拟器以加速其 RL 实验。Pgx 可在 http://github.com/sotetsuk/pgx 获取。
引用
@article{arxiv.2303.17503,
title = {Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning},
author = {Sotetsu Koyamada and Shinri Okano and Soichiro Nishimori and Yu Murata and Keigo Habara and Haruka Kita and Shin Ishii},
journal= {arXiv preprint arXiv:2303.17503},
year = {2024}
}