[論文レビュー] Podracer architectures for scalable Reinforcement Learning

Matteo Hessel, Matthias Kroiß+6 more|arXiv (Cornell University)|Apr 13, 2021

Reinforcement Learning in Robotics参考文献 28被引用数 9

ひとこと要約

この論文では、JAXを用いてTPUポッド上で強化学習を効率的にスケーリングするためのPodracerアーキテクチャ—AnakinとSebulba—を紹介している。大規模な並列処理、大容量バッチ、最適化されたTPU利用により、フルTPUポッドで最大4300万フレーム/秒を達成し、1分未満でエージェントを訓練する。従来のシステムと比較してコストを削減し、データ効率を向上させた。

ABSTRACT

Supporting state-of-the-art AI research requires balancing rapid prototyping, ease of use, and quick iteration, with the ability to deploy experiments at a scale traditionally associated with production systems.Deep learning frameworks such as TensorFlow, PyTorch and JAX allow users to transparently make use of accelerators, such as TPUs and GPUs, to offload the more computationally intensive parts of training and inference in modern deep learning systems. Popular training pipelines that use these frameworks for deep learning typically focus on (un-)supervised learning. How to best train reinforcement learning (RL) agents at scale is still an active research area. In this report we argue that TPUs are particularly well suited for training RL agents in a scalable, efficient and reproducible way. Specifically we describe two architectures designed to make the best use of the resources available on a TPU Pod (a special configuration in a Google data center that features multiple TPU devices connected to each other by extremely low latency communication channels).

研究の動機と目的

• スケーラブルなディープ強化学習（RL）学習における増大する計算リソース要件に対処すること。
• TPUポッドを活用して、RL研究における迅速なプロトタイピングと高スループット学習を可能にすること。
• 使いやすさとプロダクションレベルのスケーラビリティ・再現可能性の両立を図るシステム設計をすること。
• 最適化されたTPU利用とアーキテクチャ選択を通じて、データ効率の向上と学習コストの低減を実現すること。
• MuZeroのようなモデルフリーおよび探索ベースのエージェントを含む多様なRLワークロードをサポートすること。

提案手法

• TPUポッド上でオンラインおよびエージェント・ラーナーRLエージェントを学習するため、PodracerアーキテクチャとしてAnakinとSebulbaを設計すること。
• JAXを用いて合成可能なプログラム変換、自動微分、TPU上のハードウェア加速を実現すること。
• アクターおよびラーナーのバッチサイズの増加、より長い軌道、ネットワークの深さを活用してTPUの利用を最大化することで、学習をスケーリングすること。
• アクティングと学習のバッチサイズを分離することで、特に探索ベースのエージェントに有益なスループットの線形スケーリングを実現すること。
• 検索用に純粋なJAXベースのMCTSを実装し、検索部分にカスタムC++コードを一切使用しないこと。
• 最大2048コアのTPUコアを横断して実験を再現・スケーリングし、高スループットと短いウォルクロック時間の両立を実現すること。

実験結果

リサーチクエスチョン

• RQ1 研究の柔軟性を保ちながら、TPUポッドをどのように効果的に活用してディープ強化学習の学習をスケーリングできるか？
• RQ2 大規模なTPU上のRL学習において、スループットとデータ効率を最大化するためのアーキテクチャ的選択とは何か？
• RQ3 複雑なRLエージェント（例：MuZero）に対して、純粋なJAX実装はハイブリッドC++/Pythonシステムを同等または上回る性能を達成できるか？
• RQ4 バッチサイズおよびネットワーク容量の増加が、TPUベースのRL学習におけるデータ効率とコストに与える影響は何か？
• RQ5 アクター・ラーナーRLフレームワークにおいて、スケーラビリティとデータ効率をどの程度分離できるか？

主な発見

• Sebulbaは2048コアのTPUポッドで最大4300万フレーム/秒を達成し、Pongエージェントを1分未満で訓練した。
• 8コアのTPU上でV-traceエージェントを学習したところ、アクターのバッチサイズを32から128に増加させることで、20万フレーム/秒に到達した。
• より大きなネットワークを採用することで、バッチサイズの増加なしにデータ効率を向上させ、TPU時間やコストを増加させずに済ませた。
• MuZeroエージェントはSebulba上で16コアのTPUで9時間で2億フレームのAtariフレームを達成し、プリエンプティブルインスタンスでコストは約40ドルであった。
• MuZeroではスループットがTPUコア数に比例して線形にスケーリングされ、効果的な水平スケーリングが実証された。
• アクティングと学習のバッチサイズを分離することで、データ効率を維持しながら、レプリケーションによる高速化を実現した。