frax: Fast Robot Kinematics and Dynamics in JAX

Daniel Morton; Marco Pavone

cs.ROApr 5, 2026

frax: Fast Robot Kinematics and Dynamics in JAX

Authors: Daniel Morton, Marco Pavone

Organizations: Daniel Morton and Marco Pavone are with the Departments of Mechanical Engineering and Aeronautics & Astronautics, Stanford University, Stanford, CA 94305.

Abstract

In robot control, planning, and learning, there is a need for rigid-body dynamics libraries that are highly performant, easy to use, and compatible with CPUs and accelerators. While existing libraries often excel at either low-latency CPU execution or high-throughput GPU workloads, few provide a unified framework that targets multiple architectures without compromising performance or ease-of-use. To address this, we introduce frax, a JAX-based library for robot kinematics and dynamics, providing a high-performance, pure-Python interface across CPU, GPU, and TPU. Via a fully-vectorized approach to robot dynamics, frax enables efficient real-time control and parallelization, while supporting automatic differentiation for optimization-based methods. On CPU, frax achieves low-microsecond computation times suitable for kilohertz control rates, outperforming common libraries in Python and approaching optimized C++ implementations. On GPU, the same code scales to thousands of instances, reaching upwards of 100 million dynamics evaluations per second. We validate performance on a Franka Panda manipulator and a Unitree G1 humanoid, and release frax as an open-source library.

Figures & tables

Fig. 1: frax is a high-performance library for robot kinematics and dynamics on CPU, GPU, and TPU, supporting manipulators, humanoids, and more . Via JAX, frax enables fast robot control and planning with automatic differentiation through arbitrary functions of the kinematics and dynamics. Shown above (Franka Panda): collision and singularity avoidance in an optimization-based inverse dynamics controller, fully through jax.jvp .

	CPU	GPU/TPU	Automatic Differentiation	Interface	Compilation	Best For
Pinocchio (C++)	✓	✗	✗ †	C++	✓	Fast dynamics and analytical derivatives on CPU
MuJoCo (C++)	✓	✗	✗	C++	✓	Fast simulation and dynamics on CPU
Pinocchio (Python)	✓	✗	✗	Python	✗ ∗	Controller prototyping on CPU
MuJoCo (Python)	✓	✗	✗	Python	✗ ∗	Simulation and controller prototyping on CPU
BRAX	✓	✓	✓	Python	✓	GPU/TPU-parallelized simulation and RL
MJX	✓	✓	✓	Python	✓	GPU/TPU-parallelized simulation and RL

TABLE I: Feature and Compatibility Comparison for Common Dynamics Libraries

Fig. 2: Rigid-body dynamics performance across batch sizes . frax ’s vectorized dynamics methods are designed for high performance on CPU and on GPU, including calls to automatic differentiation ( jax.jvp ) on these methods. All timing values for GPU include any I/O overhead.

Fig. 3: Compute timing for controllers written with frax and other common libraries . On CPU, Python-based controllers are up to 2-5x faster with frax than Pinocchio or MuJoCo’s bindings, coming close to matching their underlying performance on C++. And, on GPU, frax is on par with MJX or BRAX; enabling high-speed control for both single-robot deployment and massively-parallel training.

Controller	Franka Panda	Unitree G1
Inverse kinematics [Diff. IK]	4.09	6.71
Inverse dynamics [OSC]	7.03	24.6

TABLE II: Compute Times for FRAX Controllers on CPU (

\mu

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Notation	Definition	Shape
$\mathbf{U}$	Ancestor mask: $\mathbf{U}_{ij}=1$ if $j\in\text{Anc}(i)$ else $0$	$n\times n$
$\mathbf{S}$	Spatial axes	$n\times 6$
$\mathbf{I}$	Spatial inertias	$n\times 6\times 6$
$\mathbf{V}$	Spatial velocities	$n\times 6$
$\mathbf{A}$	Spatial accelerations	$n\times 6$
$\mathbf{F}$	Spatial forces	$n\times 6$

Appendix

TABLE III: Notation for Algorithms 1 and 2

Robot Kinematics

See all 30 papers

Batched Differentiable Rigid Body Dynamics in PyTorch for GPU-Accelerated Robot Learning
May 29, 2026Yue Wang, Yanran Xu, Wenbo Wu +2Multibody Dynamics Robot Dynamics Learning
University of Southampton, Southampton, SO17 1BJ, UK
Fast Asymptotically Optimal Kinodynamic Planning via Vectorization
Jul 4, 2026Yitian Gao, Andrew Lu, Zachary KingstonConstrained Motion Planning Kinodynamic Motion Planning
Department of Computer Science, Purdue University
FlashNeRD: Performance-First Contact-Rich Neural Robot Dynamics
Oct 6, 2026Mohammadmehdi Ataei, Mohammad Amin Nabian, Jie Xu +9Robot Dynamics Learning Robotics Simulation
NVIDIA

May 29, 2026J/K move · Enter open · S save

Batched Differentiable Rigid Body Dynamics in PyTorch for GPU-Accelerated Robot Learning

Yue Wang, Yanran Xu, Wenbo Wu, Chuanhang Qiu+1

University of Southampton, Southampton, SO17 1BJ, UK

Multibody Dynamics Robot Dynamics Learning Differentiable Programming Differentiable Physics Robot Learning

Open paper AskPDF arXiv

As robot control shifts toward large-scale reinforcement learning with in-loop dynamics computation, the community's reliance on CPU-bound libraries such as Pinocchio creates a throughput bottleneck in GPU-based training pipelines. We present BARD (Batched Articulated Rigid-body Dynamics), a self-contained PyTorch implementation of Featherstone's rigid-body dynamics algorithms, optimized for batched GPU evaluation and automatic differentiation. Three design choices make this efficient: a tiered lazy-evaluation cache that avoids redundant tree traversals, matmul-free joint transforms via pre-computed Rodrigues constants, and level-parallel propagation that reduces sequential operations to tree-depth batched steps. On five robot models (7-23 DOFs), BARD matches Pinocchio numerically while reaching up to 64x higher throughput for Forward Kinematics and 63x for Jacobians at batch size 4096 on an NVIDIA H200. We validate differentiability through gradient-based system identification on a 7-DOF manipulator, recovering link masses to 1.24% mean error under 5% torque noise, and integrate BARD into an Isaac Lab AMP training pipeline for an 11-DOF spined quadruped with 4096 parallel environments, where it is 8.5x faster than Pinocchio and 2.0x faster than ADAM for in-loop dynamics. BARD is open-sourced at: https://github.com/YueWang996/bard-pytorch-dynamics.

Abstract

Robot Kinematics

Explore similar work

Batched Differentiable Rigid Body Dynamics in PyTorch for GPU-Accelerated Robot Learning

Fast Asymptotically Optimal Kinodynamic Planning via Vectorization

FlashNeRD: Performance-First Contact-Rich Neural Robot Dynamics