Brax：Colabで大規模分散システムを凌駕する強化学習用物理シミュレーションエンジン(3/3)

１．Brax：Colabで大規模分散システムを凌駕する強化学習用物理シミュレーションエンジン(3/3)まとめ

・BraxはCPUだけで実行しても従来の一般的な環境に比べて1桁以上高速に動作する
・Braxは高速で規模拡大可能な強化学習とロボット工学の研究をはるかに利用しやすくする
・従来は大規模なコンピューティングシステムのみ可能であったものがColabで動かせる

２．BraxとMuJoCoの比較

以下、ai.googleblog.comより「Speeding Up Reinforcement Learning with a New Physics Simulation Engine」の意訳です。元記事は2021年7月15日、C. Daniel FreemanさんとErik Freyさんによる投稿です。

アイキャッチ画像のクレジットはPhoto by Samuel-Elias Nadler on Unsplash

Braxのパフォーマンスを分析する別の方法は、単一のワークステーションで強化学習実験を実行する際にかかる時間を測定することです。ここでは、人気のあるAntベンチマークでBraxをOpenAIのMuJoCo物理エンジン搭載環境と比較します。

下のグラフでは、青い線は標準的なワークステーション上で実行したセットアップを表しています。このセットアップでは学習者はGPUで実行され、シミュレーターはCPUで実行されます。

Antを適度な習熟度(y軸で4,000のスコア)で実行するようにトレーニングするのにかかる時間は、青い線では約3時間ですが、アクセラレータハードウェアでBraxを使用した場合は約10秒に短縮されることがわかります。興味深いことに、CPUだけでも(灰色の線)、Braxは1桁以上高速に動作します。これは学習者とシミュレーターの両方が同じプロセスにいることで恩恵を受けるためです。

Braxに最適化されたPPO(Proximal Policy Optimization)とMuJoCo-Ant-v2環境で学習する標準のGPUでサポートされたPPOの比較
1,000万ステップの評価です。x軸はlog-wallclock-time(秒単位)であることに注意してください。影付きの領域は、5つのレプリカで最もパフォーマンスの低いシードと最もパフォーマンスの高いシードを示し、実線は平均を示しています。

物理学的な忠実度

実世界の動作に一致するシミュレーターを設計することは、本研究では対処できない既知の難しい問題です。それでも、Braxをリファレンスシミュレーターと比較して、少なくとも同じくらい有効な出力を生成していることを確認すると有益です。

この場合も、シミュレーションの品質で高く評価されているMuJoCoとBraxを比較します。MuJoCoでトレーニングされていても、Braxでトレーニングされていても、他のすべてが同じであれば、ポリシーには同様な報酬の軌跡があると予想されます。

MuJoCo-Ant-v2とBrax Antの比較
環境で達成された平均エピソードスコアに対する環境ステップの数をグラフ化します。
両方の環境は、SACの同じ標準実装でトレーニングされました。影付きの領域は、5回の実行で最もパフォーマンスの低いシードと最もパフォーマンスの高いシードを示し、実線は平均を示します。

これらの曲線は、両方のシミュレーターで報酬がほぼ同じ割合で上昇するため、両方のエンジンが同等のレベルの複雑さまたは解決の難しさで物理を計算することを示しています。また、両方の曲線がほぼ同じ報酬で最高になるため、どちらのシミュレーションでも最大限の能力を発揮するエージェントには、同じ一般的な物理的限界が適用されると確信しています。

線形運動量、角運動量、およびエネルギーを保存するBraxの能力を測定することもできます。

線形運動量(左)、角運動量(中央)、エネルギーの非保存(右)
Braxおよび他のいくつかの物理エンジンでの規模拡張性。y軸は予想される計算からの漂流を示し(Y軸が大きいほど漂流が小さく性能が良いことを意味します)、x軸はシミュレーション時間を示します。

物理シミュレーションの品質を評価するこの尺度は、MuJoCoの作成者によって最初に提案されました。これは、シミュレーションがますます長い時間ステップと計算量でタスクを実行するときに、シミュレーションがコースからどのように漂流するかを理解する方法です。ここで、Braxはその同胞と同じように機能します。

結論

Brax Training Colabで独自のポリシーをトレーニングすることにより、研究者にBraxの物理的忠実度について定性的な測定を実行することをお勧めします。学習した軌跡は、OpenAI Gymで見られるものと似ている事が認識できます。

私たちの研究は、高速で規模拡大可能な強化学習とロボット工学の研究をはるかに利用しやすくします。以前は大規模なコンピューティングクラスターを介してのみ可能であったものを、ワークステーションで実行することも、ホストされているGoogle Colaboratoryを介して無料で実行することもできます。

Githubリポジトリには、Braxシミュレーションエンジンだけでなく、高速トレーニング用の多数のリファレンスRLアルゴリズムも含まれています。Braxがどのような新しい研究を可能にするのか楽しみにしています。

謝辞

論文の共著者であるAnton Raichuk, Sertan Girgin, Igor Mordatch, Olivier Bachemに感謝します。また、物理エンジンの構築に関するアドバイスを提供してくれたErwin Coumans、JAXとXLAの最適化ヘルプを提供してくれたHechtmanとJames Bradbury 、アドバイスを提供してくれたLuke MetzとShane Guにも感謝します。また、このプロジェクトのインキュベーションを支援してくれたVijay Sundaram, Wright Bagwell, Matt Leffler, Gavin Dodd, Brad Mckee, そしてLogan Olsonにも感謝します。