[논문 리뷰] Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning

Sotetsu Koyamada, Shinri Okano외 5명|arXiv (Cornell University)|2023. 03. 29.

Artificial Intelligence in Games인용 수 5

한 줄 요약

Pgx는 GPU/TPU 가속을 최적화한 JAX 기반 하드웨어 가속, 병렬 게임 시뮬레이터 세트로, 강화학습을 위한 고속 시뮬레이션을 가능하게 한다. JAX의 자동 벡터화 및 JIT 컴파일 기능을 활용하여 기존의 파이썬 기반 라이브러리인 OpenSpiel과 PettingZoo보다 10–100배 빠른 시뮬레이션 처리량을 달성하여, 체스, 고, 쇼기와 같은 복잡한 보드 게임에서 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와 가속기 간의 데이터 전송 오버헤드를 제거하기 위해.
• 고성능 게임 환경을 위한 통합적이고 접근성 있는 파이썬 API를 제공하여 강화학습 연구에서의 빠른 프로토타이핑과 벤치마킹을 가능하게 하기 위해.
• 연구 주기를 가속화하고 재현 가능성을 확보하기 위해 표준화된 오픈소스 환경과 함께 내장된 베이스라인 모델 및 미니어처 변형을 제공하기 위해.

제안 방법

• JAX를 사용하여 게임 시뮬레이터를 구현하여 GPU/TPU 실행을 위한 자동 벡터화 및 하드웨어 전용 JIT 컴파일 기능을 활용하기 위해.
• JAX의 기능형 프로그래밍 패러다임과 자동 미분 기능을 활용하여 수천 개의 게임 인스턴스를 병행하여 배치 처리할 수 있도록 환경을 설계하기 위해.
• 가속기에서의 메모리 접근 및 계산 효율성을 높이기 위해 상태 및 행동 표현 방식을 최적화하여 지연 시간을 최소화하고 처리량을 극대화하기 위해.
• PPO 및 AlphaZero와 같은 일반적인 알고리즘과의 호환성을 확보하여 표준 강화학습 훈련 파이프라인에 시뮬레이터를 통합하기 위해.
• 훈련 시간을 단축하고 분석 연구를 용이하게 하기 위해 게임의 미니어처 버전(예: 5x5 체스)과 사전 훈련된 베이스라인 모델을 제공하기 위해.
• 전체 세트를 https://github.com/sotetsuk/pgx 에서 오픈소스로 공개하여 재현 가능성과 커뮤니티 수용을 보장하기 위해.

Figure 1: Example games included in Pgx.

실험 결과

연구 질문

• RQ1 JAX 기반 하드웨어 가속 시뮬레이터는 기존의 파이썬 기반 게임 환경보다 뚜렷하게 높은 시뮬레이션 처리량을 달성할 수 있는가?
• RQ2 동일한 가속기(GPU/TPU)에서 시뮬레이션과 학습을 통합함으로써 강화학습에서 훈련 효율성과 데이터 전송 비용에 어떤 영향을 미치는가?
• RQ3 Pgx는 체스, 고, 기타 복잡한 보드 게임에서 샘플 집약적인 강화학습 알고리즘인 AlphaZero 및 PPO의 훈련을 얼마나 빠르게 가속화할 수 있는가?
• RQ4 미니어처 게임 변형과 내장된 베이스라인 모델은 강화학습 연구에서의 빠른 실험 및 평가에 필요한 시간과 노력을 얼마나 줄일 수 있는가?
• RQ5 고성능 하드웨어인 DGX-A100에서 실행할 경우, Pgx는 기존의 OpenSpiel 및 PettingZoo 라이브러리 대비 어떤 성능 향상을 보이는가?

주요 결과

• Pgx는 DGX-A100 워크스테이션에서 기존의 파이썬 기반 라이브러리인 OpenSpiel 및 PettingZoo보다 10–100배 빠른 시뮬레이션 처리량을 달성했다.
• 배치 크기가 4096인 MinAtar 환경에서 PPO 훈련은 각 게임당 1분 이내에 완료되었으며, Seaquest를 제외한 모든 게임은 1분 이내로 완료되었다. Seaquest는 2000만 프레임 훈련에 80초 이상 소요되었다.
• 단일 A4000 GPU에서 9x9 고, 올리고, 기타 게임에서 Gumbel AlphaZero 훈련이 18시간 이내에 완료되어 하드웨어 성능에 따라 효율적으로 확장됨을 보여주었다.
• JAX의 자동 벡터화 및 JIT 컴파일 기능 덕분에, 수천 개의 동시에 실행되는 게임 환경에서 수동 스레딩이나 데이터 전송 병목 현상 없이 효율적인 병렬 실행이 가능했다.
• 미니어처 게임 환경과 사전 훈련된 베이스라인 모델의 통합으로 강화학습 실험에서 하이퍼파라미터 튜닝 및 평가에 소요되는 시간이 크게 감소했다.
• 가속기와의 통합으로 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