frax: Fast Robot Kinematics and Dynamics in JAX

📄 arXiv: 2604.04310 📥 PDF

作者: Daniel Morton, Marco Pavone

分类: cs.RO

发布日期: 2026-04-07

💡 一句话要点

Frax：基于JAX的高性能机器人动力学库，加速机器人控制与学习。

🎯 匹配领域: 支柱一：机器人控制 (Robot Control)

关键词: 机器人动力学 JAX 自动微分 高性能计算 机器人控制 运动规划 GPU加速

📋 核心要点

  1. 现有机器人动力学库难以兼顾CPU低延迟和GPU高吞吐，缺乏统一框架。
  2. Frax基于JAX，采用全向量化方法，支持CPU、GPU和TPU，并提供自动微分。
  3. Frax在CPU上接近C++性能，GPU上可达每秒1亿次动力学评估，已开源。

📝 摘要（中文）

本文介绍了一个名为frax的基于JAX的机器人运动学和动力学库，旨在为机器人控制、规划和学习提供高性能、易于使用且兼容CPU和加速器的解决方案。现有库通常在低延迟CPU执行或高吞吐量GPU工作负载方面表现出色，但很少有能够提供统一框架，在不牺牲性能或易用性的前提下，同时支持多种架构。Frax通过完全向量化的机器人动力学方法，实现了高效的实时控制和并行化，并支持基于优化的方法的自动微分。在CPU上，frax实现了微秒级的计算时间，适用于千赫兹控制速率，优于常见的Python库，并接近优化的C++实现。在GPU上，相同的代码可以扩展到数千个实例，达到每秒超过1亿次动力学评估。在Franka Panda机械臂和Unitree G1人形机器人上验证了性能，并将frax作为开源库发布。

🔬 方法详解

问题定义：机器人控制、规划和学习需要高性能的刚体动力学库。现有的库通常针对CPU或GPU进行了优化，但难以在不同架构上保持高性能和易用性，缺乏一个统一的框架。这限制了算法的部署和开发效率。

核心思路：Frax的核心思路是利用JAX的自动微分和即时编译（JIT）能力，以及其对CPU、GPU和TPU的广泛支持，构建一个高性能、可移植的机器人动力学库。通过全向量化的方法，实现高效的并行计算和自动微分，从而加速机器人控制、规划和学习算法的开发和部署。

技术框架：Frax的整体架构包括：1) 机器人模型的表示，例如URDF格式的解析；2) 运动学计算，包括正向运动学、逆向运动学和雅可比矩阵的计算；3) 动力学计算，包括正向动力学（计算给定关节力矩下的关节加速度）和逆向动力学（计算给定关节加速度下的关节力矩）；4) 基于JAX的自动微分，用于优化算法。整个框架使用纯Python编写，并利用JAX进行编译和优化。

关键创新：Frax的关键创新在于其基于JAX的全向量化实现，使得相同的代码可以在CPU、GPU和TPU上高效运行，无需针对不同架构进行专门优化。此外，Frax还提供了自动微分功能，方便用户进行基于优化的机器人控制和规划算法的开发。与现有方法相比，Frax在性能、易用性和可移植性方面都具有优势。

关键设计：Frax的关键设计包括：1) 使用JAX的vmap函数实现全向量化，从而充分利用硬件的并行计算能力；2) 使用JAX的jit函数进行即时编译，从而提高代码的执行效率；3) 提供清晰的API，方便用户进行机器人模型的加载、运动学和动力学计算，以及自动微分；4) 针对不同的机器人模型，提供预编译的函数，进一步提高性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Frax在CPU上实现了微秒级的动力学计算时间，性能优于常见的Python库，并接近优化的C++实现。在GPU上，Frax可以扩展到数千个实例，达到每秒超过1亿次动力学评估。在Franka Panda机械臂和Unitree G1人形机器人上的实验验证了Frax的性能和可扩展性。

🎯 应用场景

Frax可应用于机器人控制、运动规划、强化学习等领域。例如，可用于高速高精度的机器人实时控制，也可用于大规模的机器人仿真和训练。该库的开源发布，将促进机器人算法的快速开发和部署，加速机器人技术的进步，并可能应用于工业自动化、服务机器人、自动驾驶等领域。

📄 摘要（原文）

In robot control, planning, and learning, there is a need for rigid-body dynamics libraries that are highly performant, easy to use, and compatible with CPUs and accelerators. While existing libraries often excel at either low-latency CPU execution or high-throughput GPU workloads, few provide a unified framework that targets multiple architectures without compromising performance or ease-of-use. To address this, we introduce frax, a JAX-based library for robot kinematics and dynamics, providing a high-performance, pure-Python interface across CPU, GPU, and TPU. Via a fully-vectorized approach to robot dynamics, frax enables efficient real-time control and parallelization, while supporting automatic differentiation for optimization-based methods. On CPU, frax achieves low-microsecond computation times suitable for kilohertz control rates, outperforming common libraries in Python and approaching optimized C++ implementations. On GPU, the same code scales to thousands of instances, reaching upwards of 100 million dynamics evaluations per second. We validate performance on a Franka Panda manipulator and a Unitree G1 humanoid, and release frax as an open-source library.