1つのGPUで強化学習のトレーニングを数分で完了できます。Googleの新しいエンジンはディープラーニングを1,000倍高速化します

[[413545]]

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

ロボットはどのようにしてこのような動作を実行するのでしょうか?

一般的に、強化学習に基づくシミュレーション環境でシミュレーショントレーニングを実施します。

このとき、単一のマシン上のCPU環境でシミュレーショントレーニングを行うと、数時間から数日かかります。

しかし今では、たった1 つのTPU/GPU で、数千のCPU または GPU のコンピューティング クラスターと同じ速度を実現でき、必要な時間がわずか数分にまで直接短縮されます。

これは強化学習の速度を1,000 倍に増加させることに相当します。

これは、Google の科学者によって開発された物理シミュレーション エンジン、Braxです。

論理分岐を避ける3つの戦略

最近の物理シミュレーション エンジンのほとんどはどのように設計されているのでしょうか?

重力、モーター駆動、関節拘束、物体衝突などのタスクを 1 つのシミュレータに統合し、複数のシミュレーションを並行して実行して、実際の動作システムを近似します。

△シミュレーションの各時間ステップで、力とモーメントが統合される

この場合、各シミュレータでの計算は異なり、データはデータセンター内のネットワークを介して転送される必要があります。

この並列レイアウトにより、待ち時間も長くなります。つまり、学習者はシミュレータから経験を得るまでに10,000 ナノ秒以上待つ必要がある場合があります。

では、この遅延を短縮するにはどうすればよいでしょうか?

Brax は、シミュレーションでの分岐を回避することで、何千もの並列環境での計算の完全な均一性を確保し、トレーニング アーキテクチャ全体の複雑さを軽減することを選択しました。

複雑さが単一の TPU または GPU で実行できるレベルまで低減されると、マシン間通信の計算オーバーヘッドが削減され、遅延を効果的に排除できます。

主な方法は3つあります。

• 連続関数が離散分岐ロジックに取って代わる

たとえば、ボールと壁の間の接触力を計算する場合、次のブランチが生成されます。

ボールが壁に触れると、壁から跳ね返るボールをシミュレートする別のコードが実行されます。

それ以外の場合は、他のコードを実行します。

ここで、符号付き距離関数を使用すると、if/else の離散分岐ロジックの生成を回避できます。

• JAX JIT コンパイルにおけるブランチの評価

シミュレーションの前に、2 つのオブジェクトが衝突する可能性があるかどうかなど、環境の静的プロパティに基づいてブランチを評価します。

• シミュレーションで必要なブランチ結果のみを選択します

これら 3 つの戦略を使用すると、剛体、ジョイント、アクチュエータで構成される環境をシミュレートする物理エンジンが得られます。

また、この環境におけるさまざまな操作（進化戦略、直接軌道最適化など）を実装する学習アルゴリズムでもあります。

それで、Brax のパフォーマンスはどうでしょうか?

最大1000倍のスピード

Brax テストで使用されるベンチマークは、OpenAI Gym の Ant、HalfCheetah、Humanoid、および Reacher です。

物理法則の巧みな操作、一般的な移動（オブジェクトが配置されている周囲の任意の場所に移動するなど）、産業用ロボット アームのシミュレーションなど、3 つの新しい環境も追加されました。

研究者たちはまず、Brax がより多くの環境を並行してシミュレートする際に、どれだけの物理ステップ (または環境の状態の更新) を生成できるかをテストしました。

テスト結果の TPUv3 8x8 曲線は、Brax が複数のデバイス間でシームレスにスケーリングして、1 秒あたり数億の物理ステップに到達できることを示しています。

TPU だけでなく、V100 および P100 曲線からも、Brax はハイエンドGPUでも優れたパフォーマンスを発揮することがわかります。

さらに、Brax が単一のワークステーションで強化学習実験を実行するのにかかる時間もあります。

ここで研究者らは、Ant ベンチマーク環境でトレーニングされた Brax エンジンと MuJoCo 物理エンジンを比較しました。

ご覧のとおり、MuJoCo (青線) に必要な約 3 時間と比較すると、Brax のアクセラレータ ハードウェアを使用した最速時間はわずか10 秒です。

Brax を使用すると、シングルコアトレーニングの効率が向上するだけでなく、大規模な並列シミュレーショントレーニングにも拡張できます。

論文の宛先:
https://arxiv.org/abs/2106.13281

ダウンロード：
https://github.com/google/brax