[论文解读] Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning

Sotetsu Koyamada, Shinri Okano+5 more|arXiv (Cornell University)|Mar 29, 2023

Artificial Intelligence in Games被引用 5

一句话总结

Pgx 是一个基于 JAX 的硬件加速、并行游戏模拟器套件，专为强化学习设计，针对 GPU/TPU 加速进行了优化。通过利用 JAX 的自动向量化和 JIT 编译功能，Pgx 的模拟吞吐量比现有的基于 Python 的库（如 OpenSpiel 和 PettingZoo）快 10–100 倍，从而实现了对 Go、国际象棋和将棋等复杂棋盘游戏的 Gumbel AlphaZero 和 PPO 等强化学习算法的快速训练。

ABSTRACT

We propose Pgx, a suite of board game reinforcement learning (RL) environments written in JAX and optimized for GPU/TPU accelerators. By leveraging JAX's auto-vectorization and parallelization over accelerators, Pgx can efficiently scale to thousands of simultaneous simulations over accelerators. In our experiments on a DGX-A100 workstation, we discovered that Pgx can simulate RL environments 10-100x faster than existing implementations available in Python. Pgx includes RL environments commonly used as benchmarks in RL research, such as backgammon, chess, shogi, and Go. Additionally, Pgx offers miniature game sets and baseline models to facilitate rapid research cycles. We demonstrate the efficient training of the Gumbel AlphaZero algorithm with Pgx environments. Overall, Pgx provides high-performance environment simulators for researchers to accelerate their RL experiments. Pgx is available at http://github.com/sotetsuk/pgx.

研究动机与目标

• 解决由缓慢的、CPU 限制的游戏模拟器导致的强化学习研究性能瓶颈，这些模拟器阻碍了大规模自对弈和数据收集。
• 实现在 GPU 和 TPU 等加速器上对复杂棋盘游戏（尤其是具有大离散状态空间的游戏）的高吞吐量、并行模拟。
• 通过将模拟和学习任务部署在相同硬件上，消除 CPU 与加速器之间的数据传输开销。
• 提供统一、易用的 Python API，支持高性能游戏环境，以促进强化学习研究中的快速原型设计和基准测试。
• 提供标准化、开源的游戏环境，包含内置基线模型和微型变体，以加速研究周期并实现可复现的评估。

提出的方法

• 在 JAX 中实现游戏模拟器，以利用自动向量化和针对 GPU/TPU 的硬件特定 JIT 编译。
• 设计环境以支持使用 JAX 的函数式编程范式和自动微分，对数千个游戏实例进行批量并行模拟。
• 优化状态和动作表示，以在加速器上实现高效的内存访问和计算，最小化延迟并最大化吞吐量。
• 通过与常见算法（如 PPO 和 AlphaZero）的兼容性，将模拟器集成到标准强化学习训练流水线中。
• 提供游戏的微型版本（如 5x5 国际象棋）和预训练的基线模型，以减少训练时间并支持消融研究。
• 在 https://github.com/sotetsuk/pgx 开源完整套件，以确保可复现性并促进社区采纳。

Figure 1: Example games included in Pgx.

实验结果

研究问题

• RQ1 基于 JAX 的硬件加速模拟器是否能显著提升模拟吞吐量，相比传统的基于 Python 的游戏环境？
• RQ2 在相同加速器（GPU/TPU）上集成模拟与学习，对强化学习中的训练效率和数据传输成本有何影响？
• RQ3 Pgx 在复杂棋盘游戏上对样本密集型强化学习算法（如 AlphaZero 和 PPO）的训练加速程度如何？
• RQ4 微型游戏变体和内置基线模型在多大程度上能有效减少强化学习研究中快速实验和评估所需的时间与精力？
• RQ5 在高端硬件（如 DGX-A100）上运行时，Pgx 相较于 OpenSpiel 和 PettingZoo 等现有库的性能提升是多少？

主要发现

• 在 DGX-A100 工作站上，Pgx 相较于现有的基于 Python 的库（如 OpenSpiel 和 PettingZoo）实现了 10–100 倍的模拟吞吐量提升。
• 在 MinAtar 环境中，使用批量大小为 4096 的 PPO 训练，每款游戏的训练时间均少于一分钟，仅 Seaquest 超过 80 秒（训练 2000 万帧）。
• 在单张 A4000 GPU 上，9x9 围棋、 Othello 等游戏的 Gumbel AlphaZero 训练在 18 小时内完成，表明其在硬件上具有良好的可扩展性。
• JAX 的自动向量化和 JIT 编译功能实现了对数千个并行游戏环境的高效并行执行，无需手动线程管理或数据传输瓶颈。
• 微型游戏环境和预训练基线模型的引入，显著缩短了强化学习实验中超参数调优和评估所需的时间。
• Pgx 与加速器的集成消除了 CPU-GPU 间的数据传输开销，从而实现了更快、更具可扩展性的端到端训练流水线。

Figure 2: Basic usage of Pgx. The init function generates the initial state object. The state object has an attribute current player that indicates the agent which acts next. In this case, since we are using a batch size of 1024 for a 2-player game, current player is a binary vector whose size is 10