2025_NIPS_Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning

文章主要内容和创新点总结

一、主要内容

本文提出了Pgx，一套基于JAX编写、针对GPU/TPU加速器优化的棋盘游戏强化学习（RL）环境套件。该套件旨在解决现有Python RL环境库在复杂离散状态游戏模拟中存在的并行化不足、CPU与加速器间数据传输成本高、速度慢等问题。

Pgx包含20余种游戏，涵盖完美信息游戏（如国际象棋、围棋、将棋）、含随机事件的游戏（如双陆棋、2048）、不完全信息游戏（如桥牌叫牌、库恩扑克）以及类Atari游戏（来自MinAtar套件），还提供迷你版游戏环境（如迷你国际象棋）以适配快速研究周期。其核心优势在于借助JAX的自动向量化、加速器并行化和即时编译（JIT）特性，实现了极高的模拟吞吐量。

实验验证显示，在NVIDIA DGX-A100工作站上，Pgx的模拟速度比OpenSpiel、PettingZoo等现有Python库快10-100倍，且支持多加速器扩展，8块A100 GPU的吞吐量较单块GPU平均提升7.4倍。此外，Pgx提供基线模型，已成功支持Gumbel AlphaZero算法在多种游戏环境中的高效训练，且在多加速器场景下能显著缩短RL训练时间（如9x9围棋训练中，8块GPU较单块GPU提速约4倍）。

二、创新点

1. 硬件加速的离散状态游戏环境：首次基于JAX构建了覆盖多种类型棋盘游戏的硬件加速RL环境套件，填补了JAX生态中缺乏综合棋盘游戏环境库的空白，同时解决了传统Python库无法高效利用GPU/TPU并行计算的问题。
2. 极致的模拟性能