Blog

2021.09.07

Research

PFNにおける分散強化学習への取り組み

Keisuke Fukuda

Engineer

はじめに

分散深層学習がなぜ必要か

強化学習(Reinforcement Learning,RL)は,教師あり学習,教師なし学習と並んで,機械学習の3大ジャンルとして知られています.RLでは,エージェント(またはアクターと呼ばれます)が環境(Environment)を探索して,エクスペリエンス(経験)を集め,それを使ってモデルを学習します.教師あり学習と違い,RLはデータに明示的なラベルを必要とせず,動的に探索空間を探索することによって,動的に報酬を得て学習を進めます.

機械学習を使って難しいタスクに取り組むとき,1つの有効な方法として,訓練を高速化するために計算資源をスケールさせるというものがあります.RLの性質上,訓練のための計算資源をスケールさせるには2つの課題をクリアしなければなりません.

1つは,モデルの訓練部分をスケールする方法です.大量のデータを長時間にわたって訓練しなければいけないときに有効です.これは,他のタスク(画像分類の教師あり学習など)と同じです.

2つ目は,環境の探索をスケールさせることです.RLは,複雑な現実環境でのタスクによく使われ,その1つのれはロボットハンドの操作でしょう.このようなタスクでは,環境は物理ベースのシミュレーターとなります.シミュレーターは,一般に,計算が重くシミュレーションに時間がかかることが多いです.この場合,仮に訓練部分を複数のGPUを用いてスケールさせたとしても,経験の生成が間に合わず,訓練部分が必要とするデータを供給できない可能性があります.多くの場合GPUは計算機の中でも最も高価なパーツですから,そのような計算資源の無駄遣いは避ける必要がああります.また,このようなシミュレーターはCPUでしか動かないことがほとんどですから,計算に用いるCPUを増やす必要があります.最終的に,訓練用のGPUと環境用のCPUの両方を最大限位に有効活用できるバランスを見つけて調整する必要があります.

PFNでの分散強化学習のへの取り組みの歴史

PFNは,ロボット操作などの強化学習タスクに取り組んできました.Amazon Picking Challenge 2016や,Autonomous Tidying-up Robot System CEATEC JAPAN 2018 などが良い例です.また,IROS 2020においては,“Distributed Reinforcement Learning of Targeted Grasping with Active Vision for Mobile Manipulators” [1] という論文を発表しました.この論文では,強化学習のアプリケーションを128GPUと1024CPU コアまでスケールさせた結果が掲載されています.

PFNによる強化学習コミュニティへの貢献としては,フレームワークの開発があります.かかつて,PFNは深層学習フレームワークであるChainerファミリーを開発していました. 我々の論文[1]は,ChainerMNとChainerRLを用いたものでした.

2019年,PFNはChainerの開発を停止し,開発環境をPyTorchへ移行することを発表しました.ChainerRLのコンポーネントはPFRLへ移植され,PFRLとして開発が継続されています.このライブラリは,再現性のある研究をサポートし,包括的なアルゴリズムの実装済みコードと,新しいアルゴリズムを開発しやすい柔軟性を念頭において開発されています.

PFRLについての詳細は,PFNのブログポスト Introducing PFRL: A PyTorch-based Deep RL Library | by Prabhat Nagarajan | PyTorch [6] をご覧ください

設計と実装

背景

既に述べたように,既存のファイルからすべてのデータを読み取ることができる従来の機械学習とは異なり、強化学習は次の学習エポックの経験を生成しながら学習しています。このような特性は、学習プロセス以外に、強化学習の新たなボトルネックになりえます。もう一度,強化学習をスケールさせる時の課題を整理しましょう.

遅い環境(アクター):一部の重い環境では、アクターと呼ばれるコンポーネントによるエクスペリエンスの生成が非常に遅くなる可能性があります.
遅い学習:生成されたエクスペリエンスが多くなるか、モデルが大きくなると、学習時間が長くなります。

画像分類などの教師あり学習のスケーリングは遅い学習のほうに重点を置いており、多くの研究は複数のGPUを使用して遅い学習問題を解決しています.たとえば、pytorchのDistributedDataParallelなどです。そして,遅いアクターの課題はほとんど注目されていません。強化学習(特にon-policyのアルゴリズム)は,学習エポックを開始する前に、十分な経験が生成されるのを待つ必要があります.アクターが遅い場合、待ち時間(学習の間隔)が大きくなり、学習エポックと全体的な学習時間を長くする恐れがあります。さらに、ほとんどのアクターはCPUで実行されるため、CPUとGPUの間の計算能力の相違により、このような重い学習はより困難になります。したがって、今回の実験では、並列アクターで遅いアクター問題を解決することに重点をおくことにします。

並列アクター

アクターの並列処理の詳細について説明します。重い学習の場合は、学習ワークロードを複数のGPUに分散するのと同様に、複数のアクターを追加することで、問題に取り組みます。我々のシステムでは、アクターがシステムのボトルネックになった場合、アクターの数またはGPUあたりのアクターの比率を自由に増やせるようになっています.

実は,GPUの追加とは異なり、アクターの追加は簡単ではありません。なぜなら,アクターと学習用GPUの通信を考える必要があるからです.学習用のGPUの速度とアクターようのCPUの速度のバランスを保つためには,物理的な1台の計算機に搭載されているより多いCPUが必要になることがあります.環境のシミュレーターが重い場合は特にそうですし,また学習自体のためにもCPUが必要でもあるからです.

この問題を解決するために、我々はリモートアクター(remote actor)を実装しています。GPUと同じノードに乗っている従来のローカルアクター(local actor)とは対照的に、リモートアクターは異なるノードに起動されるアクターです。リモートアクターは、gRPCを介して、学習用プロセスとの間に,下記の2種類の通信をします:
学習プロセスから定期的に新しいモデルを受け取る
生成されたエクスペリエンスを学習プロセスに送信する
リモートアクターを使用することで、単一ノードのリソースに制限されることなく、アクターの数を自由に増やすことができます。
PFN分散強化学習システムの概要
PFNでの分散強化学習の概要を紹介する。 PFNでは、計算リソースをKubernetesで管理するため、Kubernetesとdocker上に分散型強化学習環境を構築します。以下は、分散型強化学習システムの概要です。

我々のシステムには、下記の3種類のプロセスがあります:

  • アクタープロセス:CPUで実行され、指定された環境とモデルからエクスペリエンスを生成する。アクターには2つのタイプがあり:ローカルアクターは、学習プロセスと同じポッドに配置される;リモートアクターは、別のCPU podに配置され、gRPCを介して学習プロセス/コントローラーに接続している。
  • 学習プロセス:アクターからエクスペリエンスを収集し、そのエクスペリエンスでモデルを更新する。分散データ並列学習の場合、すべての学習プロセスは Pytorch DistributedDataParallelを介してつながっている。
  • コントローラー:リモートアクターが接続するための各GPUポッドのゲートウェイとして機能している。コントローラーは、同じGPUポッドに複数の学習プロセスがいる場合のワークロードのバランスを取るのにも役にもなる。

性能のケーススタディ

アプリケーション例:分子生成タスク

ここでは,分子生成を呼ばれるタスクを取り上げます.分子の構造を生成するタスクは,創薬やマテリアルサイエンスにおいて重要です.そのようなタスクは,探索空間が離散で,事前にラベル付きのデータを準備することもできません.また,You [2] によれば,探索空間は 1023 を大きく超えると想定されています.このようなタスクに機械学習を適用する場合,強化学習は重要な選択肢となります.

先行研究としては, You [2] があり,分子生成のタスクに強化学習のアルゴリズムの1つであるPPOを適用したものです.本項では,このアルゴリズムを社内で独自に実装した強化学習アプリケーションを,gRPCとKubernetesを用いてスケールさせる実験をおこないます.また,それ以外の関連した文献として,[3], [4], [5] なども参考にしています.

実験

アクターの並列処理の有効性を示すための実験として,分散学習をサポートするように分子生成アプリケーションの社内実装を改造し,すべての実験は、元の実装と同じパラメーターを使用して行っています.分散学習の結果となる精度については、元の実装と比較し、許容可能なエラー範囲内で良好な精度を達成していることを確認しています.

計算は,PFNの内部クラスターmn2(学習ジョブ)とmn3a(リモートアクタージョブ)上で実行します.

cluster name mn2
CPU Intel(R) Xeon(R) Gold 6254 CPU @ 3.10GHz 36 cores
Memory <376G/td>
GPU Nvidia Tesla V100 SXM2 x 8
Network 100GbE x 4

MN2仕様

cluster name mn3a
CPU Intel(R) Xeon(R) Platinum 8260M CPU @ 2.40GHz 48 cores
Memory 376G
Network 100GbE x 4

MN3-a仕様

スケーラビリティ実験

この実験では、並列アクターの有効性を検証し、アクターを追加することによって、強化学習の実行時間がどのように短縮されるかを示します。分子生成の社内実装では,on-policyのアルゴリズムであるPPOを利用しているため、アクターは学習プロセスと並行して実行できません(on-policy,off-policyの違いについてはここでは述べません)。しかし、アクター用の計算資源を追加すると、モデルの更新後にエクスペリエンスを蓄積する速度が上昇するため、更新の間隔が短くなり、全体の実行時間が短縮されます。この実験では、単一のGPUで学習を実行し、アクターの数をスケールアウトしています。各アクターが一つのCPUコアを割り当て,物理ノードのCPU数以上のCPUは,ネットワークで接続されたCPU計算ノード上に起動したアクター用PodがGPUプロセスを通信を行います.

上のグラフは、この実装をスケールさせた結果を示しています。Y軸は、”Cumulative Stpes per second” で、これは1秒間あたりに学習できた経験の数を示しています(higher is better)。X軸は、アクターのために使用したCPUの合計を示しています。今回の実験と実装では,GPUを使う学習プロセスが3 CPUコアを使ってプロセス内でアクターを実行しており,それに追加してリモートアクターのCPUコア数を増やしていきました.例えば, “16 (+3)” のような表示は,「リモートアクター用に16CPUコアを利用,それに加えて学習プロセスが3コア使用」という意味となります.なお,強化学習は、乱数シードによって結果や性能が大きく異なる可能性があるので、今回は4個のシードの平均をとっています。

また、並列化効率については、”0 (+3)”〜”16 (+3)” 程度までは良好なスケーリングとなっており、その後は徐々に低下し始めることが確認されました。さらに、”64 (+3)”からは,1秒間あたりに学習できた経験の数が下がり始めていることがわかります。今回の実験で並列化可能な部分はエクスペリエンスを生成している部分です。PPOはOn-policyアルゴリズムのため、学習部分は今回の実験では並列化していません。よって、アムダールの法則から、ある一定以上は原理的に並列化不可能であり、かつ並列化のオーバーヘッドによって性能が劣化していくと考えられます。また、この「最適なアクター用CPUの数」は、アプリケーションの性質に大きく依存します。

Conclusion

このブログ記事では,PFNで行われている,分散強化学習をスケールさせる手法と実装についてご紹介しました.環境となるシミュレーターの計算が遅いような時,PFRL(旧ChainerRL)とgRPC通信を組み合わせ,アクター(環境)をスケールさせる方法を紹介しました.Off-policyなアプリケーションであったDQNをスケールさせた結果を掲載したIROS 2020の論文を紹介し,また,On-policyなアルゴリズムであるPPOをスケールさせる実験も紹介しました.

参考文献

  • [1] “Distributed Reinforcement Learning of Targeted Grasping with Active Vision for Mobile Manipulators” https://arxiv.org/abs/2007.08082
  • [2] “Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation” https://arxiv.org/abs/1806.02473
  • [3] “Decentralized Distributed PPO” http://learningsys.org/neurips19/assets/papers/25_CameraReadySubmission_ddppo_neurips_sysml_cr.pdf
  • [4] “DD-PPO: LEARNING NEAR-PERFECT POINTGOAL
    NAVIGATORS FROM 2.5 BILLION FRAMES” https://openreview.net/pdf?id=H1gX8C4YPr

  • [5] “Massively Large-Scale Distributed Reinforcement Learning with Menger ”https://ai.googleblog.com/2020/10/massively-large-scale-distributed.html
  • [6] Introducing PFRL: A PyTorch-based Deep RL Library | by Prabhat Nagarajan | PyTorch
  • Twitter
  • Facebook