💻 computer science

frax: Fast Robot Kinematics and Dynamics in JAX

本論文は、最適化ベースの制御や学習のための自動微分をサポートしつつ、低遅延の CPU 実行と大規模な GPU 並列化を実現する、高性能かつ純粋な Python 製のロボット運動学および力学用 JAX ライブラリである frax を紹介する。

原著者： Daniel Morton, Marco Pavone

公開日 2026-05-18

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

CC BY 4.0

原著者： Daniel Morton, Marco Pavone

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

ロボットに腕の動きや歩行、ボールのキャッチを教えることを想像してください。そのためには、腕の重さ、振り子の速さ、それを止めるか方向を変えるために必要な力の量を瞬時に計算できる「脳」が必要です。これをロボットダイナミクスと呼びます。

長らく、エンジニアがこの計算に利用できたツールは、2 種類の異なる電卓を持っているようなものでした：

「速い」電卓：非常に高速でしたが、特定の種類のコンピュータチップ（CPU）でのみ動作しました。これを現代のグラフィックカード（GPU）で使おうとすると、機能しませんでした。
「大量生産」電卓：一度に数千の計算を行うのに優れていました（1,000 体のロボットを同時にシミュレーションするなど）。しかし、単一のロボットを制御したいだけの場合、遅く、扱いにくいものでした。

この論文は、ロボット物理学のための**「万能翻訳機」として機能する新しいツールfraxを紹介しています。これはJAX**という技術で構築されており、シンプルな Python コードで記述しながらも、高性能な C++ コードの生々しい速度で実行できるようにしています。

frax がどのように機能するかを、日常の比喩を用いて説明します：

1. 「組立ライン」対「チームの集まり」

ほとんどの古いロボットライブラリは、一人の人が集まって話し合うように物理学を計算します。「関節 1 が動けば、関節 2 が動き、それが関節 3 を動かす…」と尋ねます。これらは一歩ずつ、順番に実行されます。これは 1 体のロボットにとっては速いですが、1,000 体のロボットで行おうとすると、ラインが詰まってしまいます。

fraxはゲームを変えます。集まりではなく、大規模で同期された組立ラインのように機能します。すべてのロボットのすべての関節の動きを、完全に同時に計算します。

トリック：どの関節が他のどの関節に影響を与えるかを特定する「マスク（デジタルフィルター）」を使用し、コンピュータがすべての計算を並列で行えるようにします。
結果：単一のコンピュータ上では、リアルタイム（1 秒間に数千回）でロボットを制御するのに十分な速度です。強力なグラフィックカード上では、1 秒間に 1 億回のロボット運動をシミュレートできます。

2. 「魔法の鏡」（自動微分）

ロボット学習の最も難しい部分の一つは、ロボットの動きを改善する方法を特定することです。通常、小さな変化がシステム全体にどのような影響を与えるかを見るために、複雑な数式を手動で記述する必要があります。

frax には、**「魔法の鏡」**が内蔵されています。JAX を基盤としているため、ロボットの動きを自動的に分析し、より良い結果を得るためにどのように微調整すべきかを瞬時に教えてくれます。これらの数式を導き出すために数学の天才である必要はありません。ライブラリが重労働を担ってくれます。これにより、学習し適応できるロボットを設計することがはるかに容易になります。

3. 速度と柔軟性

著者は frax を、Franka Panda（ロボットアーム）とUnitree G1（ヒューマノイドロボット）という 2 つの実際のロボットでテストしました。

標準的なコンピュータ（CPU）上：frax は非常に高速で、他の人気のある Python ツールを 2 倍から 5 倍の速度で凌駕します。複雑な C++ コードで書かれた最高級のツールの速度に近づきつつも、C++ を書くという頭痛を伴いません。
グラフィックカード（GPU）上：グラフィックカード専用に設計された最高級のツールと同等のパフォーマンスを発揮しますが、起動とコンパイルがはるかに速く、エンジニアが新しいアイデアをテストする際に時間を節約できます。

結論

fraxは、エンジニアがシンプルな Python でロボット制御ソフトウェアを記述できる新しいオープンソースライブラリです。標準的なラップトップ上でリアルタイムに単一のロボットを制御するのに十分な速度でありながら、スーパーコンピュータ上で一度に数千のロボットをシミュレートするのに十分な強力さを持っています。これは「使いやすさ」と「極めて高速」という間のギャップを埋め、安全に現実世界で動き、学習し、相互作用するロボットを構築しやすくします。

技術概要：frax：JAX における高速ロボット運動学および動力学

問題提起

ロボット制御、計画、学習の分野では、高性能でありながら使いやすく、かつ多様なハードウェアアーキテクチャ（CPU、GPU、TPU）と互換性のある剛体動力学ライブラリが不可欠です。既存のソリューションは重大なトレードオフに直面しています：

CPU 最適化ライブラリ: Pinocchio や MuJoCo などのライブラリは CPU 上での低遅延実行に優れていますが、アクセラレータ上で直接実行することはできません。それらの Python バインディングはオーバーヘッドを導入し、自動微分（AD）ワークフローとの統合を複雑化することが多いです。
GPU 最適化ライブラリ: MJX や BRAX などのライブラリは、GPU の並列性を活用して高スループットなシミュレーションを実現しますが、CPU 上での単一ロボット制御に使用される場合、オーバーヘッドが高くなる傾向があり、リアルタイムのキロヘルツ制御ループには不適切です。
AD 互換性: C++ ベースの動力学ライブラリを Python ベースの機械学習や計画ワークフローに統合する際、言語境界を越えて自動微分を適用する難しさにより、摩擦が生じることがよくあります。

性能や使いやすさを犠牲にすることなく複数のアーキテクチャを対象とした統合フレームワークが存在していません。

手法

著者は、ロボット運動学および動力学のための高性能インターフェースを提供する、JAX 上で構築された純粋な Python ライブラリであるfraxを導入します。手法上の核心的な革新は、ほとんどの C++ ライブラリで用いられている従来の再帰的木走査手法（例：再帰的ニュートン・オイラー法）とは異なる、剛体動力学に対する完全ベクトル化アプローチです。

主要な技術的要素：

ベクトル化定式化: 再帰的ループの代わりに、frax は祖先マスク（ $U \in \{0, 1\}^{n \times n}$ $U \in {0, 1}^{n \times n}$ ）を用いて運動学的木構造を符号化します。これにより、動力学計算を行列乗算および配列ブロードキャストを通じて実行することが可能になります。
- 逆動力学（RNEA）: マスク $U$ を用いて、木を上下に力と加速度を合計することで計算され、逐次的な依存関係を回避します。
- 質量行列（CRBA）: 合成慣性の蓄積をベクトル化することで計算されます。
計算量トレードオフ: このアプローチは、RNEA などのアルゴリズムの理論的計算量を $O(n)$ から $O(n^2)$ に、CRBA を $O(nd) $から$ O(n^2) $に意図的に増加させます（ここで$ n$ は自由度）。しかし、著者は、現代のハードウェア（XLA コンパイル経由）における微細な並列性による性能向上と、ループオーバーヘッドの排除が、冗長な計算を大幅に上回ると主張しています。
JAX 統合: JAX を活用することで、frax は以下をネイティブにサポートします：
- Just-In-Time (JIT) コンパイル: CPU、GPU、TPU 上での最適化された実行のため。
- 自動微分（AD）: 長い展開ループをトレースするコストを伴わずに、最適化ベースの制御のための効率的な勾配計算を可能にします。
- ヤコビアン・ベクトル積（JVPs）: 完全なヤコビアンの計算なしに、制御制約（例：制御バリア関数）に対するリー微分の効率的な評価を可能にします。

主要な貢献

統合アーキテクチャ: frax は、CPU、GPU、TPU 上で効率的に実行される単一のコードベースを提供し、低遅延制御と高スループットシミュレーションの間のギャップを埋めます。
性能:
- CPU 上: キロヘルツ制御レートに適した低マイクロ秒の計算時間を達成し、Pinocchio および MuJoCo の Python バインディングを 2～5 倍上回り、最適化された C++ 実装の生性能に近づきます。
- GPU/TPU 上: 数千の並列インスタンスにスケールし、1 秒あたり 1 億回を超える動力学評価に達し、MJX や BRAX などの専用 GPU シミュレータの性能と同等かそれ以上です。
使いやすさ: 純粋な Python ライブラリとして、プロトタイピングを簡素化し、Python エコシステムとシームレスに統合します。高性能を維持しながら C++ バインディングの必要性を排除します。
オープンソース: 本ライブラリはオープンソースソフトウェアとしてリリースされ、Franka Panda や Unitree G1 などの特定のロボットへのサポート、調整済み衝突モデルを含みます。

結果

著者は、Franka PandaマニピュレータとUnitree G1ヒューマノイド上で frax を検証しました。

遅延: CPU 上では、frax は逆運動学（微分 IK）の時間で Franka が4.09 µs、G1 が6.71 µs、逆動力学（作業空間制御）の時間で Franka が7.03 µs、G1 が24.6 µsを達成しました。
比較:
- Pinocchio および MuJoCo の Python API と比較して、frax は JIT コンパイルとコンパイラ融合により2～5 倍高速でした。
- Pinocchio/MuJoCo の C++ インターフェースと比較して、frax は生性能に近づきつつも、Python の柔軟性と AD 機能を保持します。
- MJX および BRAX と比較して、frax は GPU 上で同等かやや高速な性能を示しましたが、JIT コンパイル時間を大幅に短縮しました（6～12 秒対 1～2 秒）。これにより、より迅速な反復開発が可能になりました。
スループット: GPU 上では、大規模なバッチサイズ（4096）において 1 秒あたり 1 億回を超える評価のスループット率を示し、大規模並列トレーニング環境への適合性を示しています。

意義と主張

本論文は、fraxがロボット制御と学習ワークフローの統合に向けた重要な一歩であると主張しています。その主な意義は、同じコードベース内で、CPU 上でのリアルタイム単一ロボット制御を可能にしながら、同時にアクセラレータ上での大規模並列シミュレーションをサポートする点にあります。

著者は、frax が研究者に以下を可能にすると強調しています：

ハードウェア展開に必要な性能を犠牲にすることなく、自動微分を用いて Python でコントローラをプロトタイピングすること（例：非線形軌道最適化や MPC 向け）。
安全上重要な制約に対して、完全なヤコビアン計算を効率的な JVP に置き換えること。
アルゴリズムを書き換えることなく、異なるハードウェアバックエンド（実行用 CPU、トレーニング用 GPU/TPU）全体に同じロジックを配備すること。

本論文は、frax が Pinocchio のような包括的なライブラリ（例：グローバル IK ソルバや完全なシミュレーション環境の欠如）よりも現在では最小限であり、C++ 実装が依然としてわずかに高速な制御レートを提供する可能性があることを認め、控えめな立場を維持しています。しかし、frax を、現代のロボット制御と機械学習の交差点のために特別に設計された、柔軟で高性能な基盤として位置づけています。