frax

用于快速、可微的机器人运动学、动力学和逆运动学的 JAX 库

许可证

MIT

概述

frax 是斯坦福大学 Autonomous Systems Lab 基于 JAX 构建的纯 Python 机器人运动学与动力学库。它从 URDF 计算坐标变换、雅可比矩阵、质量矩阵、重力项和科里奥利项以及正向和逆向动力学，全部经过 JIT 编译且可微。典型控制器在 CPU 上的运行时间为几微秒，同一套代码还能在 GPU 或 TPU 上批量处理数千台机器人。0.2.0 版新增了 frax.ik，这是一个基于 ElastiQP 求解器的微分逆运动学与重定向工具包。

frax 由斯坦福大学 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 思路相近的微分逆运动学工具包，但它可微，并基于同一实验室的 QP 求解器 ElastiQP；该求解器在约束相互冲突时仍能保持可行。它能以超过 1,000 Hz 处理在线控制和遥操作，并以每秒超过 100,000 帧的速度完成离线动作重定向：文档中用一个 CPU 在约 4 秒内把整个 LAFAN1 人体动作数据集重定向到 Unitree G1 上。安装命令为 pip install "frax[ik]"。

为何重要

机器人控制代码通常要在 Pinocchio 这样的快速 C++ 动力学库和基于 JAX 的 GPU 仿真器之间做选择。frax 把动力学各项本身作为可微的 JAX 函数提供，在单个 CPU 上运行很快，也能在加速器上批量计算，因此同一个模型既能用于实时控制器，也能用于需要梯度的优化，还能用于大规模重定向任务。借助 frax.ik，它也与 mink 和 PyRoki 一起成为人形机器人重定向的开源选择，其速度让处理整个动作数据集成为日常工作。

解决的问题

要获得足够快以支持千赫兹控制、可微以用于优化、又能并行处理大批量的机器人运动学和动力学，通常需要组合多个库。frax 用一个 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 修改了浮动关节的约定和部分名称