frax: Fast Robot Kinematics and Dynamics in JAX

2026-04-05

AI Summary

Existing robotics dynamics libraries struggle to simultaneously achieve high performance, ease of use, and a unified interface across CPUs and accelerators such as GPUs or TPUs. This work proposes the first purely Python-based robot kinematics and dynamics library built on JAX, leveraging pervasive vectorization and automatic differentiation to enable efficient computation across heterogeneous hardware. The framework uniquely supports, within a single codebase, both microsecond-latency CPU-based control suitable for kHz-rate real-time applications and GPU-accelerated parallel dynamics computations exceeding hundreds of millions of operations per second. Experimental validation on the Franka Panda and Unitree G1 platforms demonstrates its efficacy, and the implementation is publicly released as open-source software.

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.

Problem

• robot kinematics
• rigid-body dynamics
• multi-architecture compatibility
• high-performance computing
• robot control

Innovation

• JAX
• robot dynamics
• vectorized computation
• automatic differentiation
• multi-architecture acceleration