[論文レビュー] Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation

C. Daniel Freeman, Erik Frey+4 more|arXiv (Cornell University)|Jun 24, 2021

Context-Aware Activity Recognition Systems被引用数 61

ひとこと要約

BraxはJAXで構築されたオープンソースの微分可能な剛体物理エンジンで、加速機上で動作し、大規模な並列環境シミュレーションとオンデバイス学習を可能にする。

ABSTRACT

We present Brax, an open source library for rigid body simulation with a\nfocus on performance and parallelism on accelerators, written in JAX. We\npresent results on a suite of tasks inspired by the existing reinforcement\nlearning literature, but remade in our engine. Additionally, we provide\nreimplementations of PPO, SAC, ES, and direct policy optimization in JAX that\ncompile alongside our environments, allowing the learning algorithm and the\nenvironment processing to occur on the same device, and to scale seamlessly on\naccelerators. Finally, we include notebooks that facilitate training of\nperformant policies on common OpenAI Gym MuJoCo-like tasks in minutes.\n

研究の動機と目的

• 加速器上で物理と学習を共置することにより、RL研究をより迅速でアクセスしやすく促進する。
• 勾配ベースの RL 手法を可能にする微分可能でオープンソースのエンジンを提供する。
• ProtoBufベースの環境仕様とスケーラブルなタスクのGym風インターフェイスを提供する。
• Braxの加速器横断のスケーラビリティと性能を実証し、MuJoCoと比較する。
• 微分可能な物理エンジンの改善を導くために、制限と今後の課題を調査する。

提案手法

• 最大座標リジッドボディダイナミクスをQPデータプリミティブとベクトル化変換で実装する。
• 自動ベクトル化、デバイス並列性、JIT、自動微分のためにJAXを使用して、加速器上で数千の環境を実行する。
• 物理ステップを、Brax_systemに二次のシンプレティック・オイラー積分器を用いて、ジョイント、アクチュエータ、コリライダを適用した並列変換として定義する。
• ProtoBufベースのシステム仕様と、物理タスクの構築・解決のためのGym風Envを提供する。
• RLアルゴリズム（PPO、SAC、ES、APG）をJAXで束ね、オンデバイスのロールアウトと勾配ベースの更新を伴い、完全に加速器上で実行する。
• MuJoCo風のAnt、Humanoid、Halfcheetah、Grasp、Fetchを含む環境スイートをデモンストレーションする。

実験結果

リサーチクエスチョン

• RQ1 加速器上で動作する微分可能な物理エンジンは、環境シミュレーションと学習を共置することでRLトレーニングの速度を劇的に高めることができるか。
• RQ2 一般的なRLベンチマークにおけるTPU/GPU上のBraxの性能とスケーラビリティは、MuJoCoのような従来のCPUベースのエンジンとどう比較されるか。
• RQ3 Braxで実装した場合、勾配可能性が学習アルゴリズム（PPO、SAC、ES、APG）に与える影響は何か。
• RQ4 大規模な微分可能リジッドボディシミュレーションの実用的な制限と、必要なチューニングは何か。
• RQ5 ProtoBuf環境仕様とGym風インターフェイスが、多様な物理ベースタスクの構築にどれほど有効か。

主な発見

環境名	観測次元	行動次元	タイプ
Halfcheetah	25	7	continuous
Ant	87	8	continuous
Humanoid	299	17	continuous
Grasp	139	19	continuous
Fetch	101	10	continuous

• Braxは特定の環境で単一の加速器上で毎秒数百万のシミュレーションステップを達成する。
• 4x2 TPUv3では、計算を加速器に分散することで環境ステップを数億回/秒にスケールする。
• PPOとSACの実装をBrax向けにコンパイルすると、デバイス上のロールアウトと勾配更新を最小限のデータ転送オーバーヘッドで実現し、非加速ベースラインよりはるかに高速なトレーニングを達成する。
• SACで評価した場合、同等のハイパーパラメータ下で、Brax環境はMuJoCoの対応物と定性的に一致する。
• MuJoCo Antと比較して、最適化されたPPOを用いたBraxは約10秒で移動可能になる一方、標準PPO実装では約30分かかる。
• Braxは運動量保存とエネルギー・角運動量の挙動が他エンジンと比較して競争力がある。