frax：JAXによる高速ロボットの運動学・ダイナミクス

要旨: ロボット制御、計画、および学習において、高性能で、使いやすく、CPUとアクセラレータに対応した剛体ダイナミクス用ライブラリが必要とされています。既存のライブラリはしばしば、低遅延のCPU実行、または高スループットのGPUワークロードのいずれかにおいて優れている一方で、複数のアーキテクチャを性能を損なうことなく、また使いやすさを犠牲にすることなく同時に対象とする統一的な枠組みを提供するものは多くありません。そこで本研究では、ロボットの運動学およびダイナミクスのためのJAXベースのライブラリである frax を導入します。frax は、CPU、GPU、TPUにわたって純粋なPythonのインターフェースで高性能を実現します。ロボットのダイナミクスに対して完全にベクトル化されたアプローチを用いることで、frax は効率的なリアルタイム制御と並列化を可能にし、最適化ベースの手法のための自動微分にも対応します。CPUでは、frax はキロヘルツ級の制御レートに適した低マイクロ秒の計算時間を達成し、Pythonにおける一般的なライブラリを上回り、最適化されたC++実装に近づきます。GPUでは、同じコードが数千のインスタンスへとスケールし、1秒あたり1億回を超えるダイナミクス評価に到達します。性能は、Franka Panda マニピュレータおよび Unitree G1 ヒューマノイドで検証し、オープンソースのライブラリとして frax を公開します。