PyRoFFI: A Python Library for Accelerated Rigid Body Motion Planning using the Foreign Function Interface
By Weihang Guo, Sai Coumar
PyRoFFI is a toolkit optimized with Jax JIT tracing for accelerated motion planning research with NVIDIA GPUs. This repository is expands on pyroki and is backward-compatible. PyRoFFI replaces core kinematics/dynamics kernels in PyRoKi with raw CUDA implementations through the Jax FFI interface to retain the benefits of Jax's compiler optimizations while adding the flexibility of CUDA for low-level GPU optimization of primitive operations. Additionally, we add an extended suite of optimization solvers for inverse kinematics and trajectory optimization. We provide accelerated kernels for the following kinematics functions:
- Forward Kinematics
- Inverse Kinematics: HJCD-IK, LS, SQP, MPPI, IKFlow
- Trajectory Optimization: SCO, CHOMP, L-BFGS, LS
Additional Features:
- SRDF parsing
- Runtime Neural SDFs
- VAMP Compatible Primitive Sets
- Improved Jax performance
git clone https://github.com/commalab/pyroffi.git
cd pyroffi
pip install -e .