StanfordASL/frax

What it solves

frax is designed to provide high-performance robot kinematics and dynamics calculations in pure Python. It addresses the need for a library that combines the speed of C++ based tools like Pinocchio with the ease of use of Python and the the ability to perform automatic differentiation and parallelization via JAX.

How it works

The library is built using JAX, allowing it to leverage 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 robot form-factors through specialized classes like Manipulator, Humanoid, and Quadruped and handles free-floating bases using MuJoCo's quaternion convention.

Who it’s for

Robotics researchers and engineers who need to design fast, differentiable robot controllers (e.g., inverse kinematics, operational space control) and those who want to avoid manual Jacobian derivations through the use of autodiff.

Highlights

• High Performance: Achieves compute times in the low microseconds range on CPU and upwards of 100 million computations per second on GPU/TPU.
• JAX-Powered: Full compatibility with JAX's JIT-compilation and automatic differentiation.
• Built-in Robot Models: Comes pre-loaded with models for Franka Panda, Kuka iiwa, Unitree G1, and Unitree A2.
• Spherized Collision Model: Includes methods for collision avoidance, provided a spherized collision model is defined.
• Flexible Controller Design: Enables easy prototyping of constraints (e.g., singularity and self-collision avoidance) using simple functions of kinematics and dynamics.