frax: Fast Robot Kinematics and Dynamics in JAX

• URL: http://arxiv.org/abs/2604.04310v1
• Date: Sun, 05 Apr 2026 23:08:10 GMT
• Title: frax: Fast Robot Kinematics and Dynamics in JAX
• Abstract summary: frax is a JAX-based library for robot kinematics and dynamics.<n>It provides a high-performance, pure-Python interface across CPU, GPU, and TPU.<n>On CPU, frax achieves low-microsecond computation times suitable for kilohertz control rates.<n>On GPU, the same code scales to thousands of instances, reaching upwards of 100 million dynamics evaluations per second.
• Score: 22.00362268347345
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• mjorbit: A Simulation Framework for Space Robotics [9.58318935962458]
  This paper presents a general framework for simulating multi-body space robots with contact.<n>We bring efficient, large-scale robot simulation to in-space servicing, assembly, and manufacturing applications.<n>We present mjorbit, a general, flexible, and performant framework built on the MuJoCo engine widely used in robotics.
  arXiv  Detail & Related papers  (2026-09-07T21:45:19Z)
• UniLab: A Heterogeneous Architecture for Robot RL Beyond GPU-Dominant Paradigms [68.40808444348735]
  In simulation-dominated robot control, the essential question is not which processor runs physics, but whether simulation throughput, policy learning, and runtime synchronization form an efficient end-to-end loop.<n>We present UniLab, a heterogeneous CPU-resident / GPU-learning architecture that decouples CPU-parallel simulation from GPU policy updates through a unified runtime for data movement, buffering, and synchronization.
  arXiv  Detail & Related papers  (2026-05-28T17:53:50Z)
• ManiSkill3: GPU Parallelized Robotics Simulation and Rendering for Generalizable Embodied AI [26.070510680936486]
  ManiSkill3 is the fastest state-visual GPU parallelized robotics simulator with contact-rich physics targeting generalizable manipulation.<n>ManiSkill3 supports GPU parallelization of many aspects including simulation+rendering, heterogeneous simulation, pointclouds/voxels visual input, and more.
  arXiv  Detail & Related papers  (2024-10-01T06:10:39Z)
• JaxMARL: Multi-Agent RL Environments and Algorithms in JAX [105.343918678781]
  We present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge.
  arXiv  Detail & Related papers  (2023-11-16T18:58:43Z)
• Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures [67.47328776279204]
  This work introduces a framework to develop efficient, portable Deep Learning and High Performance Computing kernels. We decompose the kernel development in two steps: 1) Expressing the computational core using Processing Primitives (TPPs) and 2) Expressing the logical loops around TPPs in a high-level, declarative fashion. We demonstrate the efficacy of our approach using standalone kernels and end-to-end workloads that outperform state-of-the-art implementations on diverse CPU platforms.
  arXiv  Detail & Related papers  (2023-04-25T05:04:44Z)
• VRKitchen2.0-IndoorKit: A Tutorial for Augmented Indoor Scene Building in Omniverse [77.52012928882928]
  INDOORKIT is a built-in toolkit for NVIDIA OMNIVERSE. It provides flexible pipelines for indoor scene building, scene randomizing, and animation controls.
  arXiv  Detail & Related papers  (2022-06-23T17:53:33Z)
• WarpDrive: Extremely Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPU [15.337470862838794]
  We present WarpDrive, a flexible, lightweight, and easy-to-use open-source RL framework that implements end-to-end multi-agent RL on a single GPU. Our design runs simulations and the agents in each simulation in parallel. It also uses a single simulation data store on the GPU that is safely updated in-place. WarpDrive yields 2.9 million environment steps/second with 2000 environments and 1000 agents (at least 100x higher throughput compared to a CPU implementation) in a benchmark Tag simulation.
  arXiv  Detail & Related papers  (2021-08-31T16:59:27Z)
• Sample Factory: Egocentric 3D Control from Pixels at 100000 FPS with Asynchronous Reinforcement Learning [68.2099740607854]
  "Sample Factory" is a high- throughput training system optimized for a single-machine setting. Our architecture combines a highly efficient, asynchronous, GPU-based sampler with off-policy correction techniques. We extend Sample Factory to support self-play and population-based training and apply these techniques to train highly capable agents for a multiplayer first-person shooter game.
  arXiv  Detail & Related papers  (2020-06-21T10:00:23Z)
• PolyDL: Polyhedral Optimizations for Creation of High Performance DL primitives [55.79741270235602]
  We present compiler algorithms to automatically generate high performance implementations of Deep Learning primitives. We develop novel data reuse analysis algorithms using the polyhedral model. We also show that such a hybrid compiler plus a minimal library-use approach results in state-of-the-art performance.
  arXiv  Detail & Related papers  (2020-06-02T06:44:09Z)
• Lyceum: An efficient and scalable ecosystem for robot learning [11.859894139914754]
  Lyceum is a high-performance computational ecosystem for robot learning. It is built on top of the Julia programming language and the MuJoCo physics simulator. It is 5-30x faster than other popular abstractions like OpenAI's Gym and DeepMind's dm-control.
  arXiv  Detail & Related papers  (2020-01-21T05:03:04Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.