frax用JAX实现机器人动力学的高性能跨平台计算，兼顾实时控制与大规模并行。
frax: Fast Robot Kinematics and Dynamics in JAX
- 基于JAX的全向量化设计，统一支持CPU/GPU/TPU上的机器人运动学与动力学计算。
- 在CPU上实现微秒级计算，支持千赫兹控制；在GPU上每秒可处理超1亿次动力学评估。
- 开源提供纯Python接口，适用于强化学习、优化控制等需自动微分的场景。
在机器人控制、规划与学习中，亟需兼具高性能、易用性且兼容多架构（CPU/加速器）的刚体动力学库。现有库往往仅擅长低延迟CPU执行或高吞吐GPU运算，难以兼顾。为此，我们提出frax——一个基于JAX的机器人运动学与动力学库，提供跨CPU、GPU、TPU的高性能纯Python接口。通过全向量化方法，frax支持高效实时控制与大规模并行，同时支持自动微分以用于优化方法。在CPU上，frax实现微秒级计算时间，满足千赫兹控制需求，性能优于常见Python库并接近优化后的C++实现；在GPU上，相同代码可扩展至数千实例，每秒可达1亿次以上动力学评估。我们在Franka Panda机械臂和Unitree G1人形机器人上验证了性能，并开源了frax。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。