frax

高速で微分可能なロボットの運動学、動力学、逆運動学のための JAX ライブラリ

ライセンス

MIT

概要

frax は、スタンフォード大学の Autonomous Systems Lab が JAX 上に構築した、純粋な Python によるロボットの運動学・動力学ライブラリです。URDF から座標変換、ヤコビアン、質量行列、重力項とコリオリ項、順動力学と逆動力学を計算し、すべて JIT コンパイルされ微分可能です。一般的な制御器は CPU で数マイクロ秒程度で動き、同じコードを GPU や TPU で数千台のロボットにまとめてバッチ処理できます。バージョン 0.2.0 では、ElastiQP ソルバーに基づく微分逆運動学とリターゲティングのツールキット frax.ik が追加されました。

frax は、スタンフォード大学 Autonomous Systems Lab の Daniel Morton と Marco Pavone によるライブラリで、ロボットの運動学と剛体動力学を JAX 上の純粋な Python で実装しています。URDF、または Franka Panda や Unitree G1 などの組み込みモデルからロボットを読み込み、関節・リンク・フレームの座標変換、ヤコビアン、質量行列、重力項とコリオリ項、順動力学と逆動力学を求めます。各項は JAX の関数なので、JIT コンパイルでき、CPU、GPU、TPU で実行でき、多数のロボットにまとめてベクトル化でき、自動微分で微分できます。

作者は 2 種類の用途を想定しています。1 台のロボットの制御では、一般的な制御器の各項を CPU が数マイクロ秒程度で計算し、論文はこれをキロヘルツの制御周期に相当するとしています。README は、GPU がある場合でもこの用途では CPU を推奨しています。バッチ処理では、同じコードが GPU で毎秒 1 億回を超える動力学計算に達します。Franka Panda/FR3 と Unitree G1 向けに用意された球ベースの衝突モデルを使えば、制御器に自己衝突回避と障害物回避をロボットの運動学の単純な関数として追加できます。

バージョン 0.2.0（2026年11月4日）では frax.ik が追加されました。これは mink や pink と同じ発想の微分逆運動学ツールキットですが、微分可能で、同じ研究室の QP ソルバー ElastiQP に基づいています。ElastiQP は制約が矛盾しても実行可能解を保ちます。オンラインの制御と遠隔操作を 1,000 Hz 超で、オフラインのモーションリターゲティングを毎秒 100,000 フレーム超で処理します。ドキュメントでは、人の動作データセット LAFAN1 全体を 1 つの CPU で約 4 秒で Unitree G1 にリターゲットしています。インストールは pip install "frax[ik]" です。

重要な理由

ロボットの制御コードは通常、Pinocchio のような高速な C++ 動力学ライブラリか、JAX による GPU シミュレータのどちらかを選ぶことになります。frax は動力学の各項そのものを微分可能な JAX 関数として提供し、1 つの CPU で高速に動くうえ、アクセラレータでのバッチ処理にも対応します。そのため同じモデルを、リアルタイム制御器にも、勾配が必要な最適化にも、大規模なリターゲティング処理にも使えます。frax.ik によって、mink や PyRoki と並ぶヒューマノイドのリターゲティング用のオープンな選択肢にもなり、動作データセット全体の処理を日常的な作業にできる速度を備えています。

解決する課題

キロヘルツの制御に十分な速さ、最適化のための微分可能性、大規模バッチのための並列性を兼ね備えたロボットの運動学と動力学を得るには、通常いくつものライブラリを組み合わせる必要があります。frax はこの 3 つすべてを 1 つの Python のロボットモデルから提供します。

主なポイント

URDF から得た運動学と動力学の各項を、JIT コンパイルされた微分可能な JAX 関数として提供
CPU では制御器の各項を数マイクロ秒で計算し、GPU のバッチ処理では毎秒 1 億回を超える動力学計算
0.2.0 の frax.ik：常に実行可能な ElastiQP ソルバー上の微分逆運動学とリターゲティング、オフラインで毎秒 100,000 フレーム超
自己衝突回避と障害物回避のための、Franka Panda/FR3 と Unitree G1 向けの球ベースの衝突モデル
MIT ライセンス、PyPI では frax、論文は arXiv（2604.04310）

対象

Python で逆運動学、逆動力学、安全フィルタの制御器を設計する制御の研究者
人の動作データを Unitree G1 などのヒューマノイドにリターゲットするロボット学習グループ
GPU や TPU で微分可能な剛体動力学やバッチ処理が必要な開発者

制限事項

現時点では URDF のみ対応で、MJCF 対応は作者の今後の予定にあります
衝突形状は球で表す必要があり、ほかのロボットではまず球ベースの衝突モデルが必要です
閉リンク機構には対応せず、関節は回転、直動、固定、浮遊ベースのみです
グリッパーなど主な運動連鎖から外れた関節は、使用前に URDF で fixed にする必要があります
性能は JAX のバージョンと設定に左右され、README は CPU のみの動力学計算には JAX 0.4.30 を推奨しています
インターフェースはまだ変化しており、0.2.0 では浮遊関節の表現規約と一部の名前が変更されました