[論文レビュー] 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 は、GPU/TPU での加速に最適化された、JAX をベースにしたハードウェアアクセラレート型で並列処理可能な強化学習向けゲームシミュレータースイートです。JAX の自動ベクトル化および JIT コンパイルを活用することで、OpenSpiel や PettingZoo などの既存の Python ベースのライブラリと比較して、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 は OpenSpiel や PettingZoo などの既存の Python ベースのライブラリと比較して、10–100倍の高速なシミュレーションスループットを達成した。
• MinAtar 環境における PPO トレーニングでは、バッチサイズ 4096 で 1 ゲームあたり 1 分未塔で完了したが、Seaquest を除き、2000万フレームの学習に 80 秒以上を要した。
• 1 枚の A4000 GPU を用いた 9x9 ゴーやオセロなどのゲームで、Gumbel AlphaZero のトレーニングが 18 時間未塔で完了し、ハードウェアスケーリングの効率性を示した。
• JAX の自動ベクトル化と JIT コンパイルの活用により、手動でのスレーディングやデータ転送のボトル neck が生じないまま、数千個の同時ゲーム環境で効率的な並列実行が可能になった。
• ミニチュアゲーム環境と事前学習済みベースラインモデルの導入により、強化学習実験におけるハイパーパramータチューニングと評価に要する時間が顕著に短縮された。
• アクセラレータとの統合により、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