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

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

・Braxは「シミュレーション時に分岐が発生しない」ことを保証する事で効率化を実現
・計算が厳密に同じであるため複雑さを軽減でき、通信時の余分な負荷を削減できる
・BraxはTPUだけでなくハイエンドGPUでも良好に機能して規模拡大が可能

２．Braxの概要

以下、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のデザイン

前述の観察の結果、Braxは、「シミュレーション時に分岐が発生しない」ことを保証することにより、その物理計算が数千の並列環境のそれぞれでまったく同じになるように設計されています。(つまり、「if」ロジックはシミュレーション環境の状態の結果を受けて分岐します)

物理エンジンの分岐の例は、ボールと壁の間に接触力を適用することです。ボールが壁に接触しているかどうかに応じて、異なるコードが実行されます。つまり、ボールが壁に接触すると、壁からのボールの跳ね返りをシミュレートするための別のコードが実行されます。Braxは、分岐を回避するために、次の3つの戦略を組み合わせて採用しています。

(1)個別の分岐ロジックを連続関数に置き換えます。例えば、符号付き距離関数を使用してボールと壁の接触力を概算するなどです。このアプローチにより、最も効率が向上します。

(2)JAXのジャストインタイムコンパイル中に分岐を評価します。2つの物体が衝突する可能性があるかどうかなど、環境の静的属性に基づく多くの分岐は、シミュレーションを実行する前に評価できる場合があります。

(3)シミュレーション中に分岐の両側を実行し、必要な結果のみを選択します。これは最終的に使用されないコードを実行するため、上記に比べて操作が無駄になります。

計算が厳密に同じであることが保証されると、トレーニングアーキテクチャ全体の複雑さが軽減され、単一のTPUまたはGPUで実行できます。そうすることで、マシン間の通信の計算オーバーヘッドと遅延がなくなります。実際には、これらの変更により、同等のワークロードでトレーニングのコストが100倍から1000倍低くなります。

Brax環境

環境(Environments)とは、RLエージェントが学習するタスクを定義する小さなパッケージ化された世界です。環境には、世界をシミュレートする手段だけでなく、世界を観察する方法やその世界の目標の定義などの機能も含まれています。

近年、新しいRLアルゴリズムをテストし、研究科学者が一般的に理解している指標を使用してそれらのアルゴリズムの影響を評価するための、いくつかの標準的なベンチマーク環境が登場しました。Braxには、人気のあるOpenAIジムに付属する、Ant、HalfCheetah、Humanoid、Reacherの4つがすぐに使用できる環境が含まれています。

左から右へ：Ant、HalfCheetah、Humanoid、およびReacherは、強化学習研究で一般的に基本となる環境です。

Braxには、3つの新しい環境も含まれています。

物体を器用に操作(ロボット工学における一般的な課題)、汎用的な移動(その周囲のどこかに配置されたターゲットにエージェントを移動する)、および産業用ロボットアームのシミュレーションです。

左：Grasp、クロウハンドを器用に操作する事を学ぶ
中：Fetch、箱状の犬型おもちゃで汎用的な目標に向かう移動ポリシーを学習
右：UR5e、産業用ロボットアームのシミュレーション

パフォーマンスベンチマーク

Braxのパフォーマンスを分析するための最初のステップは、大規模バッチ環境をシミュレートできるだけの速度があるかを測定することです。学習者が迅速に経験を消費して学習出来なければボトルネックが発生してしまうため、十分な速度で処理できる必要があります。

以下の2つのグラフは、より多くの環境を並行してシミュレートするタスクを実行するときに、Braxが生成できる物理ステップ(環境状態の更新)の数を示しています。

左側のグラフは、Braxが並行環境の数に比例して1秒あたりのステップ数を規模拡大出来ている事を示しています。10,000回の環境でメモリ帯域幅のボトルネックが発生するだけです。これは、単一のエージェントをトレーニングするのに十分であるだけでなく、エージェントの母集団全体をトレーニングするのにも適しています。

右側のグラフは、2つのことを示しています。1つは、BraxがTPUだけでなくハイエンドGPU(V100およびP100曲線を参照)でも良好に機能すること。もう1つは、JAXのデバイス並列処理プリミティブを活用することでBraxがスムーズに規模拡大出来ることです。

複数のデバイス間で、1秒あたり数億の物理ステップ(高速相互接続を介して互いに直接接続された64個のTPUv3チップであるTPUv3 8×8の曲線を参照)に達します。