frax

JAX library for fast, differentiable robot kinematics, dynamics and inverse kinematics

License

MIT

Overview

frax is a robot kinematics and dynamics library in pure Python, built on JAX by Stanford's Autonomous Systems Lab. From a URDF it computes transforms, Jacobians, the mass matrix, gravity and Coriolis terms and forward and inverse dynamics, all JIT-compiled and differentiable. Typical controllers run in the low microseconds on a CPU, and the same code batches across thousands of robots on a GPU or TPU. Version 0.2.0 adds frax.ik, a differential inverse kinematics and retargeting toolkit built on the ElastiQP solver.

frax, from Daniel Morton and Marco Pavone at Stanford's Autonomous Systems Lab, is a library for robot kinematics and rigid-body dynamics written in pure Python on JAX. You load a robot from a URDF, or one of the built-in models such as the Franka Panda or the Unitree G1, and ask it for joint, link and frame transforms, Jacobians, the mass matrix, gravity and Coriolis terms, or forward and inverse dynamics. Because every term is a JAX function, it can be JIT-compiled, run on CPU, GPU or TPU, vectorised across many robots at once and differentiated with automatic differentiation.

The authors aim at two kinds of work. For control on one robot, a CPU computes the terms of a typical controller in the low microseconds, which the paper puts at kilohertz control rates; the README recommends the CPU here even when a GPU is present. For batches, the same code reaches upwards of 100 million dynamics evaluations per second on a GPU. Spherized collision models, prebuilt for the Franka Panda/FR3 and the Unitree G1, let a controller add self-collision and obstacle avoidance as plain functions of the robot's kinematics.

Version 0.2.0 (4 Nov 2026) adds frax.ik, a differential inverse kinematics toolkit in the spirit of mink and pink, but differentiable and built on ElastiQP, a QP solver from the same lab that stays feasible when constraints conflict. It handles online control and teleoperation at over 1,000 Hz and offline motion retargeting at over 100,000 frames per second: the docs retarget the whole LAFAN1 human-motion dataset to a Unitree G1 in about four seconds on one CPU. Installation is pip install "frax[ik]".

Why it matters

Robot control code usually chooses between a fast C++ dynamics library such as Pinocchio and a GPU simulator in JAX. frax gives the dynamics terms themselves as differentiable JAX functions that are fast on a single CPU and also batch on accelerators, so the same model serves a real-time controller, an optimisation that needs gradients and a large retargeting job. With frax.ik it also joins mink and PyRoki as an open option for humanoid retargeting, at a speed that makes processing whole motion datasets routine.

Problem it solves

Getting robot kinematics and dynamics that are fast enough for kilohertz control, differentiable for optimisation and parallel for large batches usually means combining several libraries. frax provides all three from one Python model of the robot.

Key takeaways

• Kinematics and dynamics terms from a URDF as JIT-compiled, differentiable JAX functions
• Low-microsecond controller terms on a CPU; upwards of 100 million dynamics evaluations per second batched on a GPU
• frax.ik in 0.2.0: differential IK and retargeting on the always-feasible ElastiQP solver, over 100,000 frames per second offline
• Spherized collision models for the Franka Panda/FR3 and Unitree G1, for self-collision and obstacle avoidance
• MIT licence, on PyPI as frax, with a paper on arXiv (2604.04310)

Who it is for

• Control researchers designing inverse-kinematics, inverse-dynamics and safety-filter controllers in Python
• Robot-learning groups retargeting human motion data to humanoids such as the Unitree G1
• Developers who need differentiable or batched rigid-body dynamics on GPU or TPU

Limitations

• URDF only for now; MJCF support is on the authors' to-do list
• Collision geometry must be spheres, so other robots need a spherized collision model first
• No closed kinematic chains, and only revolute, prismatic, fixed and free-floating joints
• Gripper and other off-chain joints must be set to fixed in the URDF before use
• Performance depends on JAX version and settings: the README recommends JAX 0.4.30 for CPU-only dynamics
• The interface is still changing: 0.2.0 changed the free-floating joint convention and some names