[論文レビュー] TorchBeast: A PyTorch Platform for Distributed RL

Heinrich Küttler, Nantas Nardelli+5 more|arXiv (Cornell University)|Oct 8, 2019

Model-Driven Software Engineering Techniques参考文献 13被引用数 29

ひとこと要約

TorchBeast は、分散強化学習のための PyTorch ベースのプラットフォームであり、IMPALA アルゴリズムをシンプルさ、スケーラビリティ、パフォーマンスに重点を置いて実装している。純粋な Python バージョン（MonoBeast）と、システムレベルのコンponent に C++ を使用する高性能でマルチマシン対応のバージョン（PolyBeast）を提供しており、Atari 環境における TensorFlow ベースの 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 を使用してすべての機械学習コードを Python に保つことで、研究者がシステム知識に依存せずに RL アルゴリズムのプロトタイピングとスケーリングを可能にすること。
• 単一マシンでのプロトタイピングと大規模なマルチマシン学習を両立できる最小限のフォーク可能なコードベースを提供すること。
• OpenAI Gym 環境との互換性を維持するとともに、既存の TensorFlow ベースの IMPALA 実装と同等のパフォーマンスを確保すること。

提案手法

• V-trace を用いたオフポリシー補正を伴うアクター・クリティック学習に基づく、PyTorch ベースの IMPALA アルゴリズムの実装。
• データ収集を行うアクター、モデル更新を行うランナ、バッチ処理メカニズムに分離し、ロールアウトを構造化されたテンソルとして保存。
• GIL を回避し、高スループット学習を実現するため、共有メモリとソケット（MonoBeast）または RPC と C++ スレッド（PolyBeast）を用いたマルチプロセスアーキテクチャを採用。
• 各コンponent に1つのメインファイル（例: polybeast.py）を設けることで、コードの可読性と保守性を最大化。
• システムレベルのコンponent（例：バッチ処理、キューイング）を学習ロジックから分離し、研究者が学習と環境のコードのみを変更できるようにする。
• 動的バッチ処理と勾配更新を備えた Python ベースの推論とトレーニングループを採用し、低レベルの並列処理を抽象化する。

実験結果

リサーチクエスチョン

• RQ1 PyTorch ベースの IMPALA 実装は、標準的な RL ベンチマークにおいて、元の TensorFlow ベースの IMPALA 実装と同等のパフォーマンスを達成できるか？
• RQ2 低レベルのシステムプログラミングの知識が不要な純粋な Python 実装は、パフォーマンスを犠牲にせずに、スケーラブルで高スループットの分散 RL 学習をどの程度サポートできるか？
• RQ3 研究者フレンドリーで最小限のコードベースのフレームワークは、迅速なプロトタイピングと大規模なマルチマシン学習を両立できるようにどのように設計できるか？
• RQ4 高パフォーマンスな分散 RL システムを実現するにあたり、シンプルさと可読性を保ちながら必要なエンジニアリング上の妥当な妥協は何か？
• RQ5 バッチ処理などのシステムコンponent に C++ を使用することのパフォーマンスと保守性への影響は、純粋な Python や他のシステムと比較してどうか？

主な発見

• TorchBeast は、スループット、データ効率、安定性、最終的なエージェントパフォーマンスの観点で、Atari サイズのベンチマークにおいて、元の TensorFlow ベースの IMPALA 実装と同等のパフォーマンスを達成している。
• 純粋な Python の MonoBeast バージョンは、低レベルのシステムプログラミングの知識が不要なため、迅速なプロトタイピングと実験が可能である。
• PolyBeast バージョンは、システムコンponent に C++ を使用することで、複数マシンに効果的にスケーリングでき、高スループットを実現しながらも、すべての機械学習コンponent を PyTorch に保っている。
• この設計により、システムレベルの複雑さが適切に隔離され、研究者が低レベルインfraストラクチャを触らずにエージェントと環境のコードのみを変更できる。
• 共有メモリと RPC を用いたマルチプロセスアーキテクチャにより、Python のグローバルインタプリタロック（GIL）の影響が効果的に軽減され、真の並列処理が可能になっている。
• Apache 2.0 ライセンスでのオープンソースリリースにより、研究コミュニティにおける広範なアクセス性と拡張性が保証されている。