Search papers, labs, and topics across Lattice.
The paper introduces frax, a JAX-based library for robot kinematics and dynamics that achieves high performance across CPU, GPU, and TPU. Frax leverages a fully-vectorized approach to robot dynamics, enabling efficient real-time control, parallelization, and automatic differentiation. Benchmarks on Franka Panda and Unitree G1 robots show frax achieving low-microsecond CPU computation times and scaling to millions of dynamics evaluations per second on GPU, outperforming existing Python libraries and approaching optimized C++ implementations.
Finally, a single robotics library delivers both low-latency CPU performance rivaling C++ and massive GPU throughput, all in pure Python.
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.
