frax: Fast Robot Kinematics and Dynamics in JAX

frax is a fast kinematics and dynamics library in pure Python, using JAX for JIT-compilation and automatic differentiation.

With frax, you can design high-performance inverse-kinematics and inverse-dynamics controllers at the speed of Pinocchio, ease of use of Python, and differentiation and parallelization-compatibility of MJX.

On CPU, you can expect compute times for typical controllers in the low microseconds range (~25-100 kHz), and on GPU or TPU, frax can compute dynamics terms at upwards of 100 million computations per second, depending on your batch size.

New: Inverse kinematics and retargeting with frax.ik (at >100,000 FPS)

frax 0.2.0 now includes new utilities for fast, flexible IK pipelines! For full details, check out the the frax.ik guide and the examples in examples/ik (humanoid retargeting, as well as interactive manipulator and quadruped IK demos).

Installation

From PyPI

pip install "frax[ik]"

If you're getting started, frax[ik] will pull in all frax functionality, including the new IK logic. But if you just want the core kinematics/dynamics terms, or compatibility with older JAX (before 0.9.1), you can install frax directly.

From source

git clone https://github.com/StanfordASL/frax
cd frax
pip install -e ".[ik]"

If you're installing from source, I highly recommend using a uv-managed virtual environment. In this case, pip commands can be replaced with uv pip

For more fine-grained control over the JAX install for GPU/TPU, you can also use the [cuda12], [cuda13], or [tpu] tags.

Examples

Note

Remember to install from source with the [examples] dependencies: pip install -e ".[examples]"

frax's "hello world": Compute your robot's mass matrix (AKA joint-space inertia matrix)

import frax
import jax

jax.config.update("jax_enable_x64", True)  # Recommmended for high accuracy

robot = frax.Robot("path/to/your/robot.urdf")  # Or use a built-in model, e.g. frax.load_panda()
q = robot.default_configuration
M = robot.mass_matrix(q)
print(M)

For any operations with frax and jax more broadly, you should wrap your code in a jitted region, for instance,

@jax.jit
def jit_mass_matrix(q_):
    return robot.mass_matrix(q_)

M = jit_mass_matrix(q)
print(M)

See the Performance Tips section below for more advice on making your code fast.

Many more kinematics and dynamics terms are available (joint/link/frame transforms and Jacobians, gravity vector, centrifugal/coriolis forces, and many other values relevant to robot control). We also provide Manipulator, Humanoid, and Quadruped classes for useful helper functions based on your robot's form-factor, and frax comes pre-loaded with the Franka Panda, Kuka iiwa, Unitree G1, and Unitree A2.

More advanced and interactive demos are included in the examples directory, as seen below

Examples of using frax for typical robot controllers:

• Differential inverse kinematics: examples/control/diff_ik_demo.py
• Operational space control: examples/control/osc_demo.py

When these scripts are launched in trajectory mode, you'll see the following sinusoidal tracking demo:

panda_ik_demo.mp4

This can also be launched in manual mode, for interactive mouse control of the target.

However, these controllers are fairly simple, and Pinocchio or MuJoCo could have been used to give the same Jacobians and inertial values. The real benefit of frax comes from the automatic differentiation through the kinematics and dynamics for flexible controller design with minimal manual Jacobian derivations. In the below demo (examples/control/oscbf_demo.py), I've reimplemented OSCBF, which uses JAX's autodiff under the hood to form the CBF constraints.

oscbf_demo.mp4

Here, we're enforcing

• Singularity avoidance
• Joint limit avoidance
• Collision avoidance with the floor
• Collision avoidance with an external obstacle
• Self-collision avoidance

... all with easy prototyping of the constraint design, in simple functions of the robot kinematics and dynamics (such as, robot.ee_manipulability_index(q) >= eps).

Usage tips

• If you would like to use frax's collision methods, you must first define a spherized collision model of your robot (or, use our pre-built collision models for the Franka Panda/FR3, Unitree G1, and more). Check out this page for more info!
• For now, if you have joints in your URDF that are not part of the primary kinematic chain/tree being controlled (for instance, gripper joints), please set these as fixed so that they can be ignored, and so their child links' inertias can be fused into the parent. In the future, we will allow for fixing joints programmatically.
• Free-floating bases with a quaternion rotation representation follow MuJoCo's convention: q = [position, WXYZ quaternion, joints] and v = [world-frame linear velocity, body-frame angular velocity, joint velocities]. The original implementation of frax used 6 "virtual joints" (3 prismatic, 3 revolute) to represent a floating base, in which case nq = nv and q_dot = v. However, the intrinsic Euler XYZ rotation representation can suffer from singularities / gimbal lock.
• If working with quaternions, use robot.integrate(q, v, dt) rather than q + v * dt, and if you differentiate a function f(q) with autodiff, map the result back to velocity space with jax.jacobian(f)(q) @ robot.velocity_to_qdot_map(q).

Performance tips

• If you're not planning on using frax.ik and only need fast dynamics on CPU, I recommend JAX version 0.4.30. This is due to a degradation in the XLA compiler performance in recent JAX versions, but the JAX/XLA team is currently working on addressing this: see jax-ml/jax#26021
• If you're running on CPU, take advantage of double precision with jax.config.update("jax_enable_x64", True) or by setting the environment variable JAX_ENABLE_X64=True. On GPU, this is a more nuanced decision. Generally, running in double precision on GPU will lead to about a 2-6x slowdown, depending on your batch size, because GPUs are so well-suited for single-precision operations. This choice will be application-dependent -- you might need high precision, or maybe you can get away with being a little less precise.
• I've also found that for typical CPU workflows with frax, it's better to configure JAX/XLA to run on a single thread. Generally, for robot controller design, the matrices involved are not quite large enough to require multithreading (unless you are doing something more parallelizable like MPPI). In this case, restricting it to a single thread can improve performance and also make it easier to share CPU resources across processes (multiple ROS nodes, for instance). You can use source scripts/setup_cpu_env.sh to load my recommended config, or feel free to modify as needed.
• If you only need to solve for the dynamics of one robot instance, and you have a GPU on your computer, it will be significantly faster to use CPU. To force CPU usage, use jax.config.update("jax_platforms", "cpu") or set the environment variable JAX_PLATFORMS=cpu.
• For certain batched operations, it often is fastest to parallelize across multiple CPU devices, rather than running on GPU. Check out how we handle batched dataset retargeting across multiple CPU devices in frax.ik!
• frax by default does not wrap every method with a @jax.jit decorator. When calling any JAX code, add a JIT decorator to the top-most function call to ensure the best performance.
• If you have any code that is outside of a jitted region, use numpy operations and arrays. Inside a jitted region, use jax.numpy.

For general advice on JAX, check out the quickstart guide and the sharp bits.

New, breaking changes:

• I recently improved the handling of free-floating joints and added a position + WXYZ quaternion representation (following MuJoCo) rather than just position + intrinsic Euler XYZ. This should lead to fewer issues with gimbal lock, but I needed to update some terminology now that nq != nv. For instance:
  • Joint limits are specified now only in terms of the actuated joints
  • Terms like num_joints are deprecated in favor of nq and nv
• I've updated the Unitree G1's collision model: there are now more collision bodies, more self-collision pairs, and better contact geometry at the feet. This new model is better for the retargeting in frax.ik.
• The convention for fext has changed slightly in forward_dynamics, for the better. Previously, this was poorly documented and untested, and had an uncommon frame convention. Now, it uses a more intuitive convention: wrenches are applied at the link CoM frame, with axes aligned with the world frame.

I'll update this list if more changes come up.

TODOs / upcoming features

frax is (by design) more minimal than other libraries -- resulting in high performance on the restricted setting of interest. But, there are a few more features I'd like to add in the future:

• Add MJCF support
• Add support for collision primitives other than just spheres
• Analytical Jacobians of forward/inverse dynamics (see: Pinocchio)

The following features are unplanned:

• Closed kinematic chains
• Joint types other than revolute, prismatic, fixed, or free-floating

Other recommended resources

frax might not serve your needs exactly -- that's fine! Here are some other useful repositories to look at

• stack-of-tasks/pinocchio -- Robot kinematics + dynamics (C++/Python)
• google/mujoco/mjx -- Parallelized simulation (JAX)
• google/brax -- Parallelized simulation (JAX)
• chungmin99/pyroki -- Global IK and kinematic optimization (JAX)
• stephane-caron/pink -- Differential IK with Pinocchio (Python)
• kevinzakka/mink -- Differential IK with MuJoCo (Python)
• StanfordASL/cbfpy -- Control barrier functions (JAX)
• StanfordASL/oscbf -- Safe manipulator control (JAX)

Citation

@article{morton2026frax,
  author={Morton, Daniel and Pavone, Marco},
  title={frax: Fast Robot Kinematics and Dynamics in JAX},
  journal={arXiv preprint arXiv:2604.04310},
  year={2026},
  note={ICRA 2026 Workshop on Frontiers of Optimization for Robotics},
}

If you use frax.ik, please also cite ElastiQP, the solver that makes this possible:

@article{morton2026elastiqp,
  author={Morton, Daniel and Arrizabalaga, Jon and Manchester, Zachary and Pavone, Marco},
  title={Elasti{QP}: An Always-Feasible QP Solver for Constrained Robot Control},
  journal={arXiv preprint arXiv:2609.19080},
  year={2026},
}