pith. sign in

arxiv: 2604.04310 · v2 · pith:PKJ2LZ4Znew · submitted 2026-04-05 · 💻 cs.RO

frax: Fast Robot Kinematics and Dynamics in JAX

Daniel Morton , Marco Pavone This is my paper

Pith reviewed 2026-05-19 17:15 UTC · model grok-4.3

classification 💻 cs.RO
keywords robot dynamicsJAXkinematicsvectorizationautomatic differentiationGPU accelerationreal-time controlrobotics library
0
0 comments X

The pith

A JAX library computes robot kinematics and dynamics in low microseconds on CPU for kilohertz control and up to 100 million evaluations per second on GPU.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper introduces frax, a JAX-based library for robot kinematics and dynamics that uses full vectorization to run efficiently across hardware. This approach delivers computation times in the low microseconds on CPUs, which supports real-time control loops at kilohertz rates, while the identical code scales on GPUs to process thousands of robot instances in parallel. The library remains pure Python and adds automatic differentiation to aid optimization and learning methods. A sympathetic reader would care because it supplies one easy interface that avoids separate code paths for CPUs versus accelerators in robotics applications.

Core claim

frax is a pure-Python JAX library for rigid-body kinematics and dynamics that achieves low-microsecond computation times on CPU suitable for kilohertz control rates, outperforms common Python libraries and approaches optimized C++ speeds, and scales the same code to thousands of instances on GPU reaching upwards of 100 million dynamics evaluations per second, all while supporting automatic differentiation for optimization-based methods.

What carries the argument

The fully-vectorized approach to robot dynamics in JAX, which enables the same pure-Python code to execute efficiently on CPU, GPU, and TPU while providing automatic differentiation.

Load-bearing premise

The fully vectorized JAX implementation preserves numerical accuracy and stability equivalent to established non-vectorized rigid-body dynamics libraries while delivering the stated speedups on the tested hardware configurations.

What would settle it

A side-by-side comparison of computed joint torques, accelerations, and positions from frax against a reference library such as Pinocchio on identical robot models and trajectories, checking for discrepancies beyond floating-point error, or direct timing measurements of single and batched evaluations on standard CPU and GPU hardware.

Figures

Figures reproduced from arXiv: 2604.04310 by Daniel Morton, Marco Pavone.

Figure 1
Figure 1. Figure 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 thr… view at source ↗
Figure 2
Figure 2. Figure 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 from CPU/GPU transfer. loops, for fast evaluation and compilation. Timing comparisons for equivalent loop-based and vectorized implementations of… view at source ↗
Figure 3
Figure 3. Figure 3: Compute timing for controllers written with frax and other common libraries. On CPU, Python-based controllers are up to 2-3x 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. TABLE II COMPUTE TIMES FOR F… view at source ↗
read the original 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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 3 minor

Summary. The paper introduces frax, a JAX-based library for rigid-body kinematics and dynamics computations in robotics. It presents a fully vectorized pure-Python implementation that targets CPU, GPU, and TPU, claims low-microsecond latency on CPU for real-time control and scaling to 100 million dynamics evaluations per second on GPU, supports automatic differentiation, and validates the approach on a Franka Panda manipulator and Unitree G1 humanoid while releasing the code as open source.

Significance. If the performance numbers and numerical fidelity hold under scrutiny, frax would offer a practical, unified framework that lowers the barrier to high-throughput and differentiable dynamics in JAX-based robotics pipelines. The open-source release and emphasis on cross-architecture compatibility without separate code paths are concrete strengths that could accelerate research in parallel simulation, learning, and optimization.

major comments (2)
  1. [§5] §5 (Experiments): The reported CPU and GPU timing results (low-μs per evaluation, 100 M evals/s) are presented without benchmark protocol details such as number of trials, warm-up iterations, timing method (e.g., timeit vs. wall-clock), hardware SKUs, JAX/XLA versions, or error bars; this information is required to substantiate the central performance claims and allow reproduction.
  2. [§4] §4 (Implementation) and §5.2 (Validation): No quantitative numerical equivalence tests are shown comparing frax outputs (mass matrix, bias terms, forward kinematics) against reference implementations such as Pinocchio or MuJoCo for the Franka and G1 models; vectorized reordering and JAX rewrites can alter floating-point accumulation, so explicit residual norms or max-absolute-difference tables are needed to confirm the accuracy assumption underlying all speed claims.
minor comments (3)
  1. [Abstract] Abstract: The phrase 'outperforming common libraries in Python' should name the specific libraries and versions used in the comparison.
  2. [Figures] Figure 3 (or equivalent scaling plot): Axis labels and legends are too small for readability; consider increasing font size and adding a table of raw throughput numbers.
  3. [§3] §3 (Related Work): The discussion of existing JAX robotics libraries is brief; adding a short table contrasting API style, supported features, and reported throughputs would improve clarity.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for their constructive comments, which have helped us improve the clarity and reproducibility of our work. We address each major comment in turn below.

read point-by-point responses

  1. Referee: [§5] §5 (Experiments): The reported CPU and GPU timing results (low-μs per evaluation, 100 M evals/s) are presented without benchmark protocol details such as number of trials, warm-up iterations, timing method (e.g., timeit vs. wall-clock), hardware SKUs, JAX/XLA versions, or error bars; this information is required to substantiate the central performance claims and allow reproduction.

    Authors: We agree that the benchmark protocol details were insufficiently described. In the revised manuscript we will expand §5 with a dedicated experimental setup subsection that specifies the number of trials, warm-up iterations, timing method (using jax.block_until_ready to ensure XLA execution), hardware SKUs, JAX and XLA versions, and error bars (mean ± one standard deviation over repeated runs). These additions will fully substantiate the reported CPU and GPU performance figures and enable reproduction. revision: yes

  2. Referee: [§4] §4 (Implementation) and §5.2 (Validation): No quantitative numerical equivalence tests are shown comparing frax outputs (mass matrix, bias terms, forward kinematics) against reference implementations such as Pinocchio or MuJoCo for the Franka and G1 models; vectorized reordering and JAX rewrites can alter floating-point accumulation, so explicit residual norms or max-absolute-difference tables are needed to confirm the accuracy assumption underlying all speed claims.

    Authors: We acknowledge the value of explicit quantitative checks. Although the manuscript already performs validation on the Franka Panda and Unitree G1, it does not report residual norms or max-absolute-difference tables. In the revised §5.2 we will add such tables, reporting maximum absolute differences and L2 norms for the mass matrix, bias terms, and forward-kinematics outputs against Pinocchio (and MuJoCo where applicable) over a representative set of joint configurations. This will directly address potential floating-point discrepancies introduced by vectorization and JAX rewrites. revision: yes

Circularity Check

0 steps flagged

No circularity: implementation and empirical benchmarks only

full rationale

The paper describes an open-source JAX library implementing standard rigid-body dynamics (RNEA/ABA) via vectorization for CPU/GPU/TPU. No mathematical derivation chain, first-principles predictions, or fitted parameters are claimed. Performance numbers are presented as direct empirical measurements on Franka and G1 hardware rather than outputs derived from the library's own results. Any self-citations are incidental and non-load-bearing; the central contribution is engineering and benchmarking, which remains self-contained against external reference libraries.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

The central claims rest on standard rigid-body dynamics formulations from prior literature and on the correctness of the JAX runtime; no new free parameters, ad-hoc axioms, or invented physical entities are introduced.

axioms (1)
  • standard math Standard rigid-body kinematics and dynamics equations hold and can be implemented via vectorized operations without loss of correctness
    Invoked implicitly when claiming that the vectorized JAX code computes the same quantities as conventional libraries.

pith-pipeline@v0.9.0 · 5717 in / 1358 out tokens · 78666 ms · 2026-05-19T17:15:16.619507+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Reference graph

Works this paper leans on

27 extracted references · 27 canonical work pages · 1 internal anchor

  1. [1]

    The pinocchio c++ library – a fast and flexible implementation of rigid body dynamics algorithms and their analytical derivatives

    Justin Carpentier, Guilhem Saurel, Gabriele Buondonno, Joseph Mirabel, Florent Lamiraux, Olivier Stasse, and Nicolas Mansard. The pinocchio c++ library – a fast and flexible implementation of rigid body dynamics algorithms and their analytical derivatives. InIEEE International Symposium on System Integrations (SII), 2019

    work page 2019

  2. [2]

    Martin L. Felis. Rbdl: an efficient rigid-body dynam- ics library using recursive algorithms.Autonomous Robots, pages 1–17, 2016. ISSN 1573-7527. doi: 10.1007/s10514-016-9574-0. URL http://dx.doi.org/10. 1007/s10514-016-9574-0

    work page doi:10.1007/s10514-016-9574-0 2016

  3. [3]

    Marco Frigerio, Jonas Buchli, Darwin G Caldwell, and Claudio Semini. Robcogen: a code generator for efficient kinematics and dynamics of articulated robots, based on domain specific languages.Journal of Software Engineering for Robotics (JOSER), 7(1):36–54, 2016

    work page 2016

  4. [4]

    Rigidbodydy- namics.jl, 2016

    Twan Koolen and contributors. Rigidbodydy- namics.jl, 2016. URL https://github.com/tkoolen/ RigidBodyDynamics.jl

    work page 2016

  5. [5]

    ALTRO: A Fast Solver for Constrained Trajectory Optimization,

    Sabrina M. Neuman, Twan Koolen, Jules Drean, Ja- son E. Miller, and Srinivas Devadas. Benchmarking and workload analysis of robot dynamics algorithms. In 2019 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 5235–5242, 2019. doi: 10.1109/IROS40897.2019.8967694

    work page doi:10.1109/iros40897.2019.8967694 2019

  6. [6]

    Todorov, T

    Emanuel Todorov, Tom Erez, and Yuval Tassa. Mu- joco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026–5033. IEEE, 2012. doi: 10.1109/IROS.2012.6386109

    work page doi:10.1109/iros.2012.6386109 2012

  7. [7]

    Drake: Model-based design and verification for robotics, 2019

    Russ Tedrake and the Drake Development Team. Drake: Model-based design and verification for robotics, 2019. URL https://drake.mit.edu

    work page 2019

  8. [8]

    Pybullet, a python module for physics simulation for games, robotics and machine learning

    Erwin Coumans and Yunfei Bai. Pybullet, a python module for physics simulation for games, robotics and machine learning. http://pybullet.org, 2016–2021

    work page 2016

  9. [9]

    Isaac Lab: A GPU-Accelerated Simulation Framework for Multi-Modal Robot Learning

    Mayank Mittal, Pascal Roth, James Tigue, Antoine Richard, Octi Zhang, Peter Du, Antonio Serrano-Mu ˜noz, Xinjie Yao, Ren ´e Zurbr ¨ugg, Nikita Rudin, Lukasz Wawrzyniak, Milad Rakhsha, Alain Denzler, Eric Hei- den, Ales Borovicka, Ossama Ahmed, Iretiayo Akinola, Abrar Anwar, Mark T. Carlson, Ji Yuan Feng, Ani- mesh Garg, Renato Gasoto, Lionel Gulich, Yijie...

    work page internal anchor Pith review Pith/arXiv arXiv 2025

  10. [10]

    Newton: Gpu-accelerated physics simulation for robotics and simulation research,

    The Newton Contributors. Newton: Gpu-accelerated physics simulation for robotics and simulation research,

    work page

  11. [11]

    URL https://github.com/newton-physics/newton

    work page

  12. [12]

    Genesis: A generative and universal physics engine for robotics and beyond, December 2024

    Genesis Authors. Genesis: A generative and universal physics engine for robotics and beyond, December 2024. URL https://github.com/Genesis-Embodied-AI/Genesis

    work page 2024

  13. [13]

    Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem

    C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax - a differentiable physics engine for large scale rigid body simulation, 2021. URL http://github.com/google/brax

    work page 2021

  14. [14]

    curobo: Parallelized collision-free minimum-jerk robot motion generation, 2023

    Balakumar Sundaralingam, Siva Kumar Sastry Hari, Adam Fishman, Caelan Garrett, Karl Van Wyk, Valts Blukis, Alexander Millane, Helen Oleynikova, Ankur Handa, Fabio Ramos, Nathan Ratliff, and Dieter Fox. curobo: Parallelized collision-free minimum-jerk robot motion generation, 2023

    work page 2023

  15. [15]

    Wil Thomason, Zachary Kingston, and Lydia E. Kavraki. Motions in microseconds via vectorized sampling-based planning. InIEEE International Conference on Robotics and Automation, pages 8749–8756, 2024. doi: 10.1109/ ICRA57147.2024.10611190. URL http://arxiv.org/abs/ 2309.14545

    work page arXiv 2024

  16. [16]

    Pyroki: A modular toolkit for robot kinematic optimization

    Chung Min Kim*, Brent Yi*, Hongsuk Choi, Yi Ma, Ken Goldberg, and Angjoo Kanazawa. Pyroki: A mod- ular toolkit for robot kinematic optimization. In2025 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), 2025. URL https://arxiv.org/abs/ 2505.03728

    work page arXiv 2025

  17. [17]

    Pink: Python inverse kinematics based on Pinocchio, 2026

    St ´ephane Caron, Yann De Mont-Marin, Rohan Budhiraja, Seung Hyeon Bang, Ivan Domrachev, Simeon Nedelchev, Peter Du, Adrien Escande, Joris Vaillant, Bruce Wingo, Santosh Patapati, and Daniel San Jos ´e Pro. Pink: Python inverse kinematics based on Pinocchio, 2026. URL https: //github.com/stephane-caron/pink

    work page 2026

  18. [18]

    Mink: Python inverse kinematics based on MuJoCo, February 2026

    Kevin Zakka. Mink: Python inverse kinematics based on MuJoCo, February 2026. URL https://github.com/ kevinzakka/mink

    work page 2026

  19. [19]

    Analytical derivatives of rigid body dynamics algorithms

    Justin Carpentier and Nicolas Mansard. Analytical derivatives of rigid body dynamics algorithms. In Robotics: Science and systems (RSS 2018), 2018

    work page 2018

  20. [20]

    Neuman, Thomas Bourgeat, Scott Kuindersma, Srinivas Devadas, and Vijay Janapa Reddi

    Brian Plancher, Sabrina M. Neuman, Thomas Bourgeat, Scott Kuindersma, Srinivas Devadas, and Vijay Janapa Reddi. Accelerating robot dynamics gradients on a cpu, gpu, and fpga.IEEE Robotics and Automation Letters, 6 (2):2335–2342, 2021. doi: 10.1109/LRA.2021.3057845

    work page doi:10.1109/lra.2021.3057845 2021

  21. [21]

    Varadarajan, A

    Brian Plancher, Sabrina M. Neuman, Radhika Ghosal, Scott Kuindersma, and Vijay Janapa Reddi. Grid: Gpu- accelerated rigid body dynamics with analytical gradi- ents. In2022 International Conference on Robotics and Automation (ICRA), pages 6253–6260, 2022. doi: 10.1109/ICRA46639.2022.9812384

    work page doi:10.1109/icra46639.2022.9812384 2022

  22. [22]

    Crocoddyl: An Efficient and Versatile Framework for Multi-Contact Optimal Control

    Carlos Mastalli, Rohan Budhiraja, Wolfgang Merkt, Guilhem Saurel, Bilal Hammoud, Maximilien Naveau, Justin Carpentier, Ludovic Righetti, Sethu Vijayakumar, and Nicolas Mansard. Crocoddyl: An Efficient and Versatile Framework for Multi-Contact Optimal Control. InIEEE International Conference on Robotics and Au- tomation (ICRA), 2020

    work page 2020

  23. [23]

    Proxddp: Proximal constrained trajectory optimization

    Wilson Jallet, Antoine Bambade, Etienne Arlaud, Sarah El-Kazdadi, Nicolas Mansard, and Justin Carpentier. Proxddp: Proximal constrained trajectory optimization. IEEE Transactions on Robotics, 41:2605–2624, 2025. doi: 10.1109/TRO.2025.3554437

    work page doi:10.1109/tro.2025.3554437 2025

  24. [24]

    Implementing torque control with high-ratio gear boxes and without joint- torque sensors

    Andrea Del Prete, Nicolas Mansard, Oscar E Ramos, Olivier Stasse, and Francesco Nori. Implementing torque control with high-ratio gear boxes and without joint- torque sensors. InInt. Journal of Humanoid Robotics, page 1550044, 2016. URL https://hal.archives-ouvertes. fr/hal-01136936/document

    work page 2016

  25. [25]

    Safe, task-consistent manipulation with operational space control barrier func- tions

    Daniel Morton and Marco Pavone. Safe, task-consistent manipulation with operational space control barrier func- tions. In2025 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 187–194,

    work page

  26. [26]

    doi: 10.1109/IROS60139.2025.11246389

    work page doi:10.1109/iros60139.2025.11246389 2025

  27. [27]

    Springer, 2008

    Roy Featherstone.Rigid body dynamics algorithms. Springer, 2008. APPENDIX A. Timing All timing values were recorded on a desktop PC with an i9-14900KF CPU, 64GB RAM, and an RTX 4090 GPU. Python methods were tested withtime.perf_counter over 10,000 iterations, withjax.block_until_ready() called on any JAX functions to avoid timing the asynchronous dispatch...

    work page 2008