frax: Fast Robot Kinematics and Dynamics in JAX

TL;DR

Frax is a high-performance, versatile robot kinematics and dynamics library built in JAX, enabling real-time control and large-scale parallelization across CPU, GPU, and TPU architectures.

Contribution

Frax introduces a unified, fully-vectorized JAX-based library for robot dynamics that outperforms existing Python libraries and scales efficiently across multiple hardware platforms.

Findings

Frax achieves microsecond computation times on CPU for kilohertz control.

On GPU, it scales to over 100 million dynamics evaluations per second.

Validated on Franka Panda and Unitree G1 robots.

Abstract

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.