Reproduce “frax: Fast Robot Kinematics and Dynamics in JAX” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Daniel Morton, Marco Pavone
Researchers from Stanford University developed `frax`, a JAX-based library for robot kinematics and dynamics. It delivers low-microsecond computation times on CPU for real-time control and scales to over 100 million dynamics evaluations per second on GPU for parallel workloads, all while natively supporting automatic differentiation.
Key insights
- The vectorized dynamics approach, while theoretically increasing algorithmic complexity (O(n) to O(n^2)), translates to practical performance gains on accelerators due to XLA's ability to exploit fine-grained parallelism.
- A pure-Python JAX-native library can offer competitive real-time performance for single-robot CPU control tasks compared to optimized C++ libraries, while providing superior scalability and automatic differentiation support for parallel GPU workloads.
- Efficient JIT compilation and automatic differentiation within JAX simplify the development and deployment of complex, gradient-reliant control and planning algorithms in robotics.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.