💻 computer science

Batched Differentiable Rigid Body Dynamics in PyTorch for GPU-Accelerated Robot Learning

本論文では、Pinocchioのような既存のCPU依存型のソリューションをスループットと学習速度の両面で大幅に上回りつつ、数値的な精度を維持し、効果的な勾配ベースのシステム同定を可能にする、バッチ処理可能な微分可能剛体ダイナミクスのための自己完結型GPU最適化PyTorchライブラリであるBARDを紹介する。

原著者： Yue Wang, Yanran Xu, Wenbo Wu, Chuanhang Qiu, Zhaoxing Li

公開日 2026-06-01

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Yue Wang, Yanran Xu, Wenbo Wu, Chuanhang Qiu, Zhaoxing Li

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

ロボットに歩行、走行、あるいは物体を掴む動作を教えようとしている場面を想像してみてください。これを行うには、ロボットの脳（AI）が、「もし足をこのように動かしたら、バランスはどうなるだろうか？」と常に自問自答する必要があります。これには、**剛体ダイナミクス（rigid-body dynamics）**と呼ばれる複雑な数学が必要です。

長い間、この計算を行うことは、ティースプーン一杯でスイミングプールを満たそうとするようなものでした。科学者が使用してきた標準的なツール（Pinocchioのようなライブラリ）は、一度に一つの計算を行う単一の「脳」（CPU）で動作するように作られていました。研究者が、一度に数千台のロボットを訓練するために、現代の超高速グラフィックスカード（GPU）を使用しようとしたとき、これらの古いツールはボトルネックとなってしまいました。それは、フェラーリをトラクターの渋滞の中に走らせようとするようなものでした。

この論文では、この交通渋滞を解消するために、PyTorch（人気の高いAIフレームワーク）向けに特別に構築された新しいツールである bard (Batched Articulated Rigid-body Dynamics) を紹介しています。

bard の仕組みを、簡単な比喩を用いて説明します：

1. 「怠け者」のシェフ（階層的遅延評価 / Tiered Lazy Evaluation）

4,000人のゲストのために大規模な宴会を準備するシェフを想像してください。

従来の方法： シェフは、ゲストAがスープだけを欲しがり、ゲストBがサラダだけを欲しがる場合でも、すべてのゲストのために「すべての料理」を作ります。これは膨大な時間とエネルギーを無駄にします。
bard の方法： シェフは賢い意味で「怠け者」です。彼らは実際に要求されたものだけを作ります。もしAIがロボットの手の位置（順運動学 / Forward Kinematics）だけを知りたいのであれば、シェフはロボットの筋肉内部の力までは計算しません。もし力が知りたいのであれば、シェフは手の位置の計算をスキップします。これにより、不要な作業をスキップすることで、膨大な時間を節約できます。

2. 「あらかじめ混ぜられた」塗料（行列演算を用いない変換 / Matmul-Free Transforms）

従来のツールでは、関節が動くたびに、コンピュータは新しい位置を特定するために、2つの大きな数字のグリッド（行列）の複雑な掛け算を行わなければなりませんでした。これを何千回も行うのは低速です。

bard の方法： 著者たちは、ロボットの関節の「形」は変わらず、角度だけが変わることに気づきました。そこで、ショーが始まる前に「塗料（定数）」をあらかじめ混ぜておきました。関節が動くたびに毎回塗料を混ぜるのではなく、あらかじめ混ぜられたベースに、少しの「角度（正弦と余弦）」を加えるだけにしました。これにより、重くて遅い「混ぜるプロセス」が、素早く単純な「かき混ぜ」へと変わりました。

3. 「組み立てライン」対「単独の作業員」（レベル並列伝播 / Level-Parallel Propagation）

ロボットは木のような構造（胴体から腕や足が伸びている構造）をしています。従来のツールは、幹から先端に向かって、一つの枝が終わるのを待ってから次のステップを開始するというように、一本ずつ計算していました。

bard の方法： 組み立てラインを想像してください。一人で木全体を担当するのではなく、bard は同じ高さ（深さ）にあるすべての枝をグループ化します。四足歩行ロボットの4本の脚すべてを同時に計算し、一つずつではなく、木全体の「レベル」を一つの巨大なバッチとして処理します。これにより、GPUの膨大なパワーを活用できるのです。

4. 結果：速度と精度

研究者たちは、非常に強力なコンピュータチップである NVIDIA H200 GPU を使用して、4,096台のロボット・シミュレーションを同時に実行するテストを行いました。

速度： ロボットの手の位置を特定する場合、bard は従来の標準よりも 64倍速く なりました。ロボットの内部力を計算する場合も、依然として大幅に高速であり（実際のトレーニングシナリオでは最大8.5倍）、非常に高速でした。
精度： これほど高速でありながら、bard は従来の低速なツールと同等の精度を保っています。数学的な誤差は極めて小さく、実質的にゼロでした（砂粒と山の違いのようなものです）。
実世界でのテスト： 彼らは、強化学習と呼ばれる手法を用いて、11本の脚を持つロボット（脊椎を持つ四足歩行ロボット）に動き方を教えるために bard を使用しました。bard が非常に高速であったため、トレーニングプロセスは従来のツールを使用した場合よりも 8.5倍速く 完了しました。

5. ななぜこれが重要なのか

この論文は、bard によって、かつては大規模なコンピュータ・クラスターを必要とした複雑なロボット訓練を、単一のGPUで実行できることを示しています。これは「ドロップイン・リプレイスメント（そのまま置き換え可能なもの）」として機能するため、科学者はコード全体を書き直すことなく、古いツールを bard に入れ替えることができます。

要約すると、bard はロボット物理学の重くて遅い数学を、現代のAIのスピードで動作するように最適化し、ロボットがより速く、より効率的に学習できるようにするものです。

技術要約：PyTorchにおけるバッチ化された微分可能な剛体動力学（bard）

問題提起
ロボット制御が、ループ内での動力学計算を伴う大規模な強化学習（RL）へと移行するにつれ、分野は深刻なスループットのボトルネックに直面している。Isaac LabやMuJoCo Playgroundに見られる現在のトレーニングパイプラインは、バッチ化されたGPU実行のためにPyTorchに依存している。しかし、標準的な動力学ライブラリであるPinocchioはCPUにバインドされており、逐次的（sequential）である。このアーキテクチャは、評価のたびに線形なCPU–GPUデータ転送オーバーヘッドを強制し、GPUコアを逐次的な木構造の走査中にアイドル状態にしてしまう。BraxやMuJoCo MJXのような既存のGPUシミュレータは高いスループットを実現しているが、これらは非PyTorchエコシステム（JAX）に紐付けられているか、あるいはカスタムの損失関数や物理情報に基づいた報酬に必要なスタンドアロンの微分可能演算子（ヤコビ行列や質量行列など）を公開しないモノリシックなシミュレータとして機能している。対照的に、ADAMのような既存のPyTorch互換の微分可能ライブラリは、バックエンドの汎用性のためにGPUスループットを犠牲にしており、カーネル融合を活用したり、PyTorchコンパイラのグラフ・トレーシングを破壊するデータ依存の制御フローを回避したりすることに失敗している。

手法
著者らは、Featherstoneの剛体動力学アルゴリズム（FK、RNEA、CRBA、ABA、およびヤコビ行列）を実装した、純粋なPyTorchライブラリであるbard（Batched Articulated Rigid-body Dynamics）を提案する。本ライブラリは、バッチ化されたGPU評価とtorch.compileへの完全な互換性を目的として設計されている。性能向上は、以下の4つの具体的な設計戦略によって達成される：

階層的遅延評価（Tiered Lazy Evaluation）: すべての運動学的量を毎回の呼び出しで計算するのではなく、bardは運動学的状態を、有効性フラグによって保護された3つのティアに分割する：(1) ローカルな親子間の変換、(2) 空間随伴（動力学アルゴリズムにのみ必要）、(3) ワールドフレームのポーズ（FK/ヤコビ行列のクエリにのみ必要）。これにより、RNEAのみを呼び出すRLループが、ワールドフレームの伝播に伴う計算コストを回避できることを保証する。
行列乗算を用いない関節変換（Matrix-Multiplication-Free Joint Transforms）: GPU上での数千の小さな $4 \times 4$ 行列乗算のローンチ・オーバーヘッドを排除するため、本ライブラリはモデル構築時にロドリゲスの回転定数を事前計算する。実行時には、親子の変換を $\sin(q_i)$ および $\cos(q_i)$ を用いたスカラー・行列乗算のみを用いて復元することで、PyTorchコンパイラがすべての関節およびバッチ要素にわたって操作を単一のカーネルへと融合することを可能にする。
レベル並列ツリー伝播（Level-Parallel Tree Propagation）: Featherstoneのアルゴリズムは本質的に逐次的であるが、bardは木構造の走査をレベル並列なバッチ操作へと再構成する。同じツリーの深さにあるノードをグループ化することで、本ライブラリは $n_b$ （ボディ数）の逐次的ステップを、 $d$ （ツリーの深さ）のバッチ化されたステップへと削減する。逆方向パス（力の累積）はデータ依存性により逐次的であり続けるが、すべての $B$ 個の環境は各ノードにおいて同時に処理される。
直接的なヤコビ列の構築（Direct Jacobian Column Construction）: 単一自由度（single-DOF）関節の場合、本ライブラリは完全な $6 \times 6$ の随伴行列の構築をバイパスする。代わりに、事前計算されたオフセットと利用可能なワールドフレームのポーズを使用して、効率的な外積と行列・ベクトル乗算に置き換えることで、ヤコビ列を直接計算する。

本ライブラリは、バッチファーストのメモリレイアウトを利用し、キャッシュラインの無駄を避けるために転置された空間変換を連続的に格納する。また、エージモード（最大スループットのために勾配追跡をスキップする）とファンクショナルモード（完全な自動微分）の両方をサポートしており、静的な制御フローと定数インデックステンソルを使用することでtorch.compileと完全に互換性を持つ。

主な貢献

bard: Featherstoneのアルゴリズムのバッチ化されたGPU評価を可能にする、オープンソースの純粋なPyTorch動力学ライブラリ。
システムレベルの技術: キネマティクス・キャッシング、事前割り当てされたワークスペース、およびレベル並列伝播を含む、GPU上での再帰的なツリー計算を実行するための新しい手法。
検証: 勾配ベースのシステム同定および4,096個の並列環境を用いた大規模な強化学習における包括的なテスト。

実験結果
著者らは、NVIDIA H200 GPU上で、5つのロボットモデル（7〜23自由度）に対してPinocchioおよびADAMと比較してbardを評価した：

数値的精度: bardは、浮動小数点丸め誤差の範囲内でPinocchioと一致した。最大誤差は、FKおよびヤコビ行列においてマシンイプシロン（ $\approx 10^{-16}$ ）付近であり、RNEAおよびCRBAでは $10^{-13}$ 以下であった。ABAは、その3パス構造による誤差蓄積のため、わずかに大きな不一致（最大 $10^{-10}$ ）を示したが、実用的な範囲内に留まった。
スループット: バッチサイズ4,096において、bardはForward KinematicsでPinocchioに対して最大64倍、ヤコビ行列で63倍の高いスループットを達成した。動力学アルゴリズムについては、3倍から7倍の高速化を実現した。
システム同定: 5%のトルクノイズを伴う7自由度マニピュレータの勾配ベースの同定タスクにおいて、bardはリンク質量を平均誤差**1.24%**で正常に復元した。
RL統合: 4,096個の並列環境を持つ11自由度の四足歩行ロボットのIsaac Lab AMPトレーニングパイプラインに統合した結果、bardはPinocchioよりも8.5倍、ADAMよりも2.0倍高いスループットを達成した。

意義と主張
本論文は、bardがPyTorchのためのネイティブで高性能かつモジュール化された動力学ライブラリを提供することで、ロボティクス・エコシステムの決定的なギャップを埋めるものであると主張している。これは、これまでCPUのボトルネックやエコシステムの不適合によって困難であった、単一のGPU上でのループ内動力学を伴う大規模な強化学習を可能にする。著者らは、スループットの優位性は、大規模なバッチサイズにおけるキネマティクス主体のワークロードにおいて最も顕著になることを指摘している。また、低価格帯のGPU（例：NVIDIA L4）ではレジスタ・スピリングによりtorch.compileが動力学主体のワークロードの性能を低下させる可能性があるため、そのような場合にはエージモードの使用を推奨している。本研究は、逐次的なアルゴリズムを並列なGPU実行用に再構成することで、数値的な精度や微分可能性を損なうことなく、大幅な性能向上を実現できることを示している。