[논문 리뷰] TorchBeast: A PyTorch Platform for Distributed RL

Heinrich Küttler, Nantas Nardelli외 5명|arXiv (Cornell University)|2019. 10. 08.

Model-Driven Software Engineering Techniques참고 문헌 13인용 수 29

한 줄 요약

TorchBeast는 분산 강화학습을 위한 PyTorch 기반 플랫폼로, 단순성, 확장성, 성능에 중점을 두고 IMPALA 알고리즘을 구현합니다. 순수 파이썬 버전(MonoBeast)과 고성능 다중 머신 버전(PolyBeast)을 제공하며, 시스템 수준 컴포넌트에 C++를 사용하여 Atari 환경에서 텐서플로우 기반 IMPALA와 동등한 성능을 달성합니다.

ABSTRACT

TorchBeast is a platform for reinforcement learning (RL) research in PyTorch. It implements a version of the popular IMPALA algorithm for fast, asynchronous, parallel training of RL agents. Additionally, TorchBeast has simplicity as an explicit design goal: We provide both a pure-Python implementation ("MonoBeast") as well as a multi-machine high-performance version ("PolyBeast"). In the latter, parts of the implementation are written in C++, but all parts pertaining to machine learning are kept in simple Python using PyTorch, with the environments provided using the OpenAI Gym interface. This enables researchers to conduct scalable RL research using TorchBeast without any programming knowledge beyond Python and PyTorch. In this paper, we describe the TorchBeast design principles and implementation and demonstrate that it performs on-par with IMPALA on Atari. TorchBeast is released as an open-source package under the Apache 2.0 license and is available at \url{https://github.com/facebookresearch/torchbeast}.

연구 동기 및 목표

• 분산 강화학습을 위한 접근성 있고 고성능이며 확장 가능한 PyTorch 기반 구현의 부족을 보완한다.
• 모든 머신러닝 코드를 파이썬과 PyTorch로 유지함으로써 연구자가 시스템 지식 없이도 RL 알고리즘을 프로토타ип링하고 확장할 수 있도록 한다.
• 단일 머신 프로토타이핑과 대규모 다중 머신 학습을 모두 지원하는 최소한의 수정이 가능한 코드베이스를 제공한다.
• OpenAI Gym 환경과의 호환성을 확보하고, 기존 텐서플로우 기반 IMPALA 구현과 동등한 성능을 유지를 한다.

제안 방법

• V-trace를 통한 오프-폴리시 보정을 사용하는 액터-크리틱 학습 방식을 활용해 PyTorch 기반 IMPALA 알고리즘을 구현한다.
• 롤아웃을 구조화된 텐서로 저장하며, 데이터 수집을 담당하는 액터, 모델 업데이트를 담당하는 런너, 배치 처리 메커니즘으로 시스템을 분리한다.
• GIL을 우회하기 위해 공유 메모리와 소켓을 사용하는 다중 프로세싱 아키텍처(MonoBeast) 또는 RPC와 C++ 스레드를 사용하는 아키텍처(PolyBeast)를 도입해 고처리량 학습을 가능하게 한다.
• 각 컴포넌트별로 단일 메인 파일(e.g., polybeast.py)을 기반으로 코드베이스를 구성해 가독성과 수정 용이성을 극대화한다.
• 시스템 수준 컴포넌트(e.g., 배치 처리, 큐잉)를 학습 로직에서 분리해 연구자가 학습 및 환경 코드만 수정하면 되도록 한다.
• 동적 배치 처리와 기울기 업데이트를 지원하는 파이썬 기반 추론 및 학습 루프를 사용해 저수준 병렬 처리를 추상화한다.

실험 결과

연구 질문

• RQ1 PyTorch 기반 IMPALA 구현이 표준 RL 벤치마크에서 원본 텐서플로우 기반 IMPALA와 동등한 성능을 달성할 수 있는가?
• RQ2 저수준 시스템 프로그래밍 지식 없이도 순수 파이썬 구현이 성능을 희생시키지 않고도 확장성 있고 고처리량 분산 RL 학습을 지원할 수 있는 정도는 어느 정도인가?
• RQ3 빠른 프로토타이핑을 가능하게 하면서도 대규모 다중 머신 학습을 지원할 수 있는 연구 우호적이고 최소한의 코드베이스 프레임워크를 어떻게 설계할 수 있는가?
• RQ4 높은 성능을 달성하면서도 단순성과 가독성을 유지하기 위해 분산 RL 시스템에서 어떤 엔지니어링 트레이드오프가 필요한가?
• RQ5 시스템 컴포넌트(e.g., 배치 처리)에 C++를 사용할 경우 순수 파이썬 또는 다른 시스템 대비 성능 및 유지보수성에 어떤 영향을 미치는가?

주요 결과

• TorchBeast는 처리량, 데이터 효율성, 안정성, 최종 에이전트 성능 측면에서 Atari 베드지에서 원본 텐서플로우 기반 IMPALA와 동등한 성능을 달성한다.
• 순수 파이썬 기반의 MonoBeast 버전은 저수준 시스템 프로그래밍 지식 없이도 빠른 프로토타이핑과 실험을 가능하게 한다.
• PolyBeast 버전은 시스템 수준 컴포넌트에 C++를 사용해 다중 머신 간 효과적인 확장성을 확보했으며, 모든 머신러닝 컴포넌트는 PyTorch로 유지하면서 고처리량을 달성한다.
• 설계가 성공적으로 시스템 수준의 복잡성을 격리해, 연구자가 저수준 인fra스트럭처를 수정하지 않고도 에이전트와 환경 코드만 수정할 수 있도록 한다.
• 공유 메모리와 RPC를 사용한 다중 프로세싱 아키텍처를 통해 파이썬의 글로벌 인터프리터 락(GIL)의 영향을 효과적으로 완화해 진정된 병렬 처리를 가능하게 한다.
• Apache 2.0 라이선스 하에 오픈소스로 배포되어 연구 커뮤니티의 광범위한 접근성과 확장성이 보장된다.