💻 computer science

frax: Fast Robot Kinematics and Dynamics in JAX

This paper introduces frax, a high-performance, pure-Python JAX library for robot kinematics and dynamics that delivers low-latency CPU execution and massive GPU parallelization while supporting automatic differentiation for optimization-based control and learning.

Original authors: Daniel Morton, Marco Pavone

Published 2026-05-18

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Daniel Morton, Marco Pavone

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to move its arm, walk, or catch a ball. To do this, you need a "brain" that can instantly calculate the physics: how heavy the arm is, how fast it's swinging, and exactly how much force is needed to stop it or change its direction. This is called robot dynamics.

For a long time, the tools available to engineers for this calculation were like having two different types of calculators:

The "Speedy" Calculator: It was incredibly fast but only worked on a specific type of computer chip (CPU). If you tried to use it on a modern graphics card (GPU), it would break.
The "Mass-Production" Calculator: It was great for doing thousands of calculations at once (like simulating 1,000 robots at the same time), but it was slow and clunky when you just wanted to control a single robot.

The paper introduces frax, a new tool that acts like a universal translator for robot physics. It is built using a technology called JAX, which allows it to be written in simple Python code but run with the raw speed of high-performance C++ code.

Here is how frax works, using some everyday analogies:

1. The "Assembly Line" vs. The "Team Huddle"

Most old robot libraries calculate physics like a single person doing a huddle. They ask: "Okay, joint 1 moves, so joint 2 moves, which makes joint 3 move..." They do this step-by-step, one after another. This is fast for one robot, but if you try to do it for 1,000 robots, the line gets clogged.

frax changes the game. Instead of a huddle, it acts like a massive, synchronized assembly line. It calculates the movement of every joint for every robot all at the exact same time.

The Trick: It uses a "mask" (a digital filter) to figure out which joints affect which others, allowing the computer to do all the math in parallel.
The Result: On a single computer, it's fast enough to control a robot in real-time (thousands of times a second). On a powerful graphics card, it can simulate 100 million robot movements in a single second.

2. The "Magic Mirror" (Automatic Differentiation)

One of the hardest parts of robot learning is figuring out how to improve the robot's movement. Usually, you have to manually write complex math formulas to see how changing one tiny thing affects the whole system.

frax comes with a "Magic Mirror" built-in. Because it is built on JAX, it can automatically look at the robot's movement and instantly tell you exactly how to tweak it to get better results. You don't need to be a math genius to derive these formulas; the library does the heavy lifting for you. This makes it much easier to design robots that can learn and adapt.

3. Speed and Flexibility

The authors tested frax on two real robots: a Franka Panda (a robotic arm) and a Unitree G1 (a humanoid robot).

On a standard computer (CPU): frax is so fast that it beats other popular Python tools by 2 to 5 times. It gets close to the speed of the very best tools written in complex C++ code, but without the headache of writing C++.
On a graphics card (GPU): It performs just as well as the best tools designed specifically for graphics cards, but it starts up and compiles much faster, saving engineers time when they are testing new ideas.

The Bottom Line

frax is a new, open-source library that lets engineers write robot control software in simple Python. It is fast enough to control a single robot in real-time on a standard laptop, but powerful enough to simulate thousands of robots at once on a supercomputer. It bridges the gap between "easy to use" and "extremely fast," making it easier to build robots that can move, learn, and interact with the real world safely.

Technical Summary: frax: Fast Robot Kinematics and Dynamics in JAX

Problem Statement

In the fields of robot control, planning, and learning, there is a critical need for rigid-body dynamics libraries that are simultaneously high-performance, easy to use, and compatible with diverse hardware architectures (CPUs, GPUs, and TPUs). Existing solutions face significant trade-offs:

CPU-Optimized Libraries: Libraries like Pinocchio and MuJoCo excel at low-latency execution on CPUs but cannot be directly executed on accelerators. Their Python bindings often introduce overhead and complicate integration with automatic differentiation (AD) workflows.
GPU-Optimized Libraries: Libraries such as MJX and BRAX leverage GPU parallelism for high-throughput simulation but often suffer from high overhead when used for single-robot control on CPUs, making them unsuitable for real-time, kilohertz-rate control loops.
AD Compatibility: Integrating C++-based dynamics libraries into Python-based machine learning or planning workflows often creates friction due to the difficulty of applying automatic differentiation across language boundaries.

There is a lack of a unified framework that targets multiple architectures without compromising performance or ease of use.

Methodology

The authors introduce frax, a pure-Python library built on JAX that provides a high-performance interface for robot kinematics and dynamics. The core methodological innovation is a fully-vectorized approach to rigid-body dynamics, diverging from the traditional recursive tree-traversal methods (e.g., Recursive Newton-Euler Algorithm) used in most C++ libraries.

Key Technical Components:

Vectorized Formulation: Instead of recursive loops, frax encodes the kinematic tree structure using an ancestor mask ( $U \in \{0, 1\}^{n \times n}$ $U \in {0, 1}^{n \times n}$ ). This allows dynamics calculations to be performed via matrix multiplication and array broadcasting.
- Inverse Dynamics (RNEA): Computed by summing forces and accelerations up and down the tree using the mask $U$ , avoiding sequential dependencies.
- Mass Matrix (CRBA): Computed by vectorizing the accumulation of composite inertias.
Complexity Trade-off: This approach intentionally increases the theoretical computational complexity of algorithms like RNEA from $O(n)$ to $O(n^2)$ and CRBA from $O(nd)$ to $O(n^2)$ (where $n$ is degrees of freedom). However, the authors argue that the performance gains from fine-grained parallelism on modern hardware (via XLA compilation) and the elimination of loop overhead far outweigh the redundant computations.
JAX Integration: By leveraging JAX, frax natively supports:
- Just-In-Time (JIT) Compilation: For optimized execution on CPU, GPU, and TPU.
- Automatic Differentiation (AD): Enabling efficient gradient computation for optimization-based control without the cost of tracing through long unrolled loops.
- Jacobian-Vector Products (JVPs): Allowing efficient evaluation of Lie derivatives for control constraints (e.g., Control Barrier Functions) without computing full Jacobians.

Key Contributions

Unified Architecture: frax provides a single codebase that runs efficiently on CPUs, GPUs, and TPUs, bridging the gap between low-latency control and high-throughput simulation.
Performance:
- On CPU: Achieves low-microsecond computation times suitable for kilohertz control rates, outperforming Python bindings of Pinocchio and MuJoCo by 2–5x and approaching the raw performance of optimized C++ implementations.
- On GPU/TPU: Scales to thousands of parallel instances, reaching upwards of 100 million dynamics evaluations per second, matching or exceeding the performance of dedicated GPU simulators like MJX and BRAX.
Ease of Use: As a pure-Python library, it simplifies prototyping and integrates seamlessly with the Python ecosystem, removing the need for C++ bindings while maintaining high performance.
Open Source: The library is released as open-source software, with support for specific robots like the Franka Panda and Unitree G1, including tuned collision models.

Results

The authors validated frax on a Franka Panda manipulator and a Unitree G1 humanoid.

Latency: On CPU, frax achieved inverse kinematics (differential IK) times of 4.09 µs (Franka) and 6.71 µs (G1), and inverse dynamics (Operational Space Control) times of 7.03 µs (Franka) and 24.6 µs (G1).
Comparison:
- Compared to Python APIs of Pinocchio and MuJoCo, frax was 2–5x faster due to JIT compilation and compiler fusions.
- Compared to C++ interfaces of Pinocchio/MuJoCo, frax approaches their raw performance but retains Python flexibility and AD capabilities.
- Compared to MJX and BRAX, frax showed comparable or slightly faster performance on GPU but significantly reduced JIT compilation times (1–2 seconds vs. 6–12 seconds), facilitating faster iterative development.
Throughput: On GPU, the library demonstrated throughput rates exceeding 100 million evaluations per second at large batch sizes (4096), indicating suitability for massively parallel training environments.

Significance and Claims

The paper claims that frax represents a significant step toward unifying robot control and learning workflows. Its primary significance lies in enabling real-time, single-robot control on CPUs while simultaneously supporting massively parallel simulation on accelerators within the same codebase.

The authors emphasize that frax allows researchers to:

Prototype controllers in Python with automatic differentiation (e.g., for nonlinear trajectory optimization or MPC) without sacrificing the performance required for hardware deployment.
Replace full Jacobian computations with efficient JVPs for safety-critical constraints.
Deploy the same logic across different hardware backends (CPU for execution, GPU/TPU for training) without rewriting algorithms.

The paper maintains a modest stance, acknowledging that frax is currently more minimal than comprehensive libraries like Pinocchio (e.g., lacking global IK solvers and full simulation environments) and that C++ implementations may still offer marginally faster control rates. However, it positions frax as a flexible, high-performance foundation specifically designed for the intersection of modern robotics control and machine learning.