StanfordASL/frax

What it solves

frax is designed to provide high-performance robot kinematics and dynamics computations. It addresses the need for a library that combines the speed of C++ libraries like Pinocchio, the ease of use of Python, and the native support for automatic differentiation and parallelization provided by JAX.

How it works

Written in pure Python using JAX, the library allows for JIT-compilation and automatic differentiation. It can compute dynamics terms (such as mass matrices, Jacobians, and gravity vectors) for robots defined via URDF files. It supports various joint types including revolute, prismatic, fixed, and free-floating.

To achieve maximum performance, users are encouraged to wrap their controller logic in @jax.jit decorators and use double precision where necessary. On CPU, it can achieve compute times in the low microseconds range, while on GPU/TPU, it can handle upwards of 100 million computations per second.

Who it’s for

Robotics researchers and engineers who need to design and prototype flexible, high-performance controllers (such as inverse kinematics or operational space control) without manually deriving Jacobians.

Highlights

• JAX-powered: Native support for automatic differentiation and parallelization.
• Extreme Speed: High-frequency control loops (25-100 kHz) on CPU and massive throughput on GPU/TPU.
• Robot-specific Helpers: Includes specialized Manipulator and Humanoid classes.
• Collision Avoidance: Supports spherized collision models for enforcing safety constraints.
• Pre-loaded Models: Comes with built-in support for Franka Panda and Unitree G1 robots.