frax

快速、可微分的機器人運動學、動力學與逆向運動學 JAX 函式庫

授權條款

MIT

概述

frax 是 Stanford 的 Autonomous Systems Lab 以 JAX 打造、純 Python 撰寫的機器人運動學與動力學函式庫。它從 URDF 計算座標轉換、雅可比矩陣、質量矩陣、重力項與科氏力項，以及正向與逆向動力學，全部都經過 JIT 編譯且可微分。典型的控制器在 CPU 上只需幾微秒，而同一份程式碼可在 GPU 或 TPU 上對數千台機器人進行批次運算。0.2.0 版新增 frax.ik，這是建構在 ElastiQP 求解器上的微分逆向運動學與動作重定向工具組。

frax 由 Stanford Autonomous Systems Lab 的 Daniel Morton 與 Marco Pavone 開發，是一個以純 Python 在 JAX 上撰寫的機器人運動學與剛體動力學函式庫。你可以從 URDF 載入機器人，或使用內建模型，例如 Franka Panda 或 Unitree G1，接著向它查詢關節、連桿與座標框架的座標轉換、雅可比矩陣、質量矩陣、重力項與科氏力項，或正向與逆向動力學。由於每一項都是 JAX 函式，因此可以 JIT 編譯、在 CPU、GPU 或 TPU 上執行、一次對多台機器人向量化，並以自動微分求導。

作者鎖定兩類工作。針對單一機器人的控制，CPU 計算典型控制器各項只需幾微秒，論文指出這足以支撐千赫級的控制頻率；README 建議即使有 GPU，這種情況下也使用 CPU。針對批次運算，同一份程式碼在 GPU 上每秒可達超過 1 億次動力學評估。為 Franka Panda/FR3 與 Unitree G1 預先建好的球體化碰撞模型，讓控制器能把自我碰撞與避障當成機器人運動學的一般函式加入。

0.2.0 版（2026 年 11 月 4 日）新增 frax.ik，這是一個精神上類似 mink 與 pink 的微分逆向運動學工具組，但它可微分，並建構在 ElastiQP 之上；ElastiQP 是同一實驗室的 QP 求解器，在約束互相衝突時仍能保持可行。它能以超過 1,000 Hz 處理線上控制與遠端操作，並以每秒超過 100,000 幀進行離線動作重定向：文件示範在單一 CPU 上約四秒就把整個 LAFAN1 人類動作資料集重定向到 Unitree G1。安裝方式為 pip install "frax[ik]"。

為何重要

機器人控制程式碼通常得在快速的 C++ 動力學函式庫（例如 Pinocchio）與 JAX 中的 GPU 模擬器之間二選一。frax 把動力學各項本身提供為可微分的 JAX 函式，在單一 CPU 上很快，也能在加速器上批次運算，因此同一個模型可以同時用於即時控制器、需要梯度的最佳化，以及大型動作重定向工作。有了 frax.ik，它也與 mink 和 PyRoki 並列，成為人形機器人動作重定向的開放選項，而且速度快到讓處理整個動作資料集成為例行工作。

解決的問題

要取得快到足以進行千赫級控制、可微分以供最佳化、又能平行處理大批次的機器人運動學與動力學，通常得組合好幾個函式庫。frax 以單一的 Python 機器人模型同時提供這三項。

重點

• 從 URDF 取得的運動學與動力學各項，皆為經 JIT 編譯、可微分的 JAX 函式
• 在 CPU 上計算控制器各項只需幾微秒；在 GPU 上批次運算時，每秒可達超過 1 億次動力學評估
• 0.2.0 版的 frax.ik：建構在永遠可行的 ElastiQP 求解器上的微分 IK 與動作重定向，離線處理每秒超過 100,000 幀
• 為 Franka Panda/FR3 與 Unitree G1 提供球體化碰撞模型，用於自我碰撞與避障
• MIT 授權，以 frax 名稱發布在 PyPI 上，並有 arXiv 論文（2604.04310）

適用對象

• 以 Python 設計逆向運動學、逆向動力學與安全濾波器控制器的控制研究人員
• 將人類動作資料重定向到 Unitree G1 等人形機器人的機器人學習團隊
• 需要在 GPU 或 TPU 上進行可微分或批次剛體動力學運算的開發者

限制

• 目前只支援 URDF；MJCF 支援列在作者的待辦清單上
• 碰撞幾何必須是球體，因此其他機器人需要先建立球體化碰撞模型
• 不支援封閉運動鏈，且只支援旋轉、滑動、固定與自由浮動關節
• 夾爪與其他不在主鏈上的關節，使用前必須在 URDF 中設為固定
• 效能取決於 JAX 版本與設定：README 建議純 CPU 動力學使用 JAX 0.4.30
• 介面仍在變動：0.2.0 版改變了自由浮動關節的慣例與部分名稱