frax: Fast Robot Kinematics and Dynamics in JAX

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.

Figures (3)

• Figure 1: frax is a high-performance library for robot kinematics and dynamics on CPU, GPU, and TPU, supporting manipulators, humanoids, and more. Via JAX, frax enables fast robot control and planning with automatic differentiation through arbitrary functions of the kinematics and dynamics. Shown above (Franka Panda): collision and singularity avoidance in an optimization-based inverse dynamics controller, fully through jax.jvp.
• Figure 2: Rigid-body dynamics performance across batch sizes. frax's vectorized dynamics methods are designed for high performance on CPU and on GPU, including calls to automatic differentiation (jax.jvp) on these methods. All timing values for GPU include any I/O overhead from CPU/GPU transfer.
• Figure 3: Compute timing for controllers written with frax and other common libraries. On CPU, Python-based controllers are up to 2-3x faster with frax than Pinocchio or MuJoCo's bindings, coming close to matching their underlying performance on C++. And, on GPU, frax is on par with MJX or BRAX; enabling high-speed control for both single-robot deployment and massively-parallel training.