CusADi2: Latency-Optimized Batch Computation of Robot Kinematics, Jacobians, and Hessians on the GPU
Organizations: Department of Mechanical Engineering, University of California, Berkeley, CA, USA.
Abstract
Trajectory optimization methods, from sampling-based zero-order to gradient-based first- and second-order approaches, require many evaluations of dynamical system functions, i.e., functions derived from a robot's dynamical model, such as forward kinematics and their Jacobians and Hessians. These evaluations are largely independent, e.g., across sampled trajectories or time steps, and can thus be evaluated in parallel as a batch. When used to make decisions online, these methods require low latency, i.e., each batch must complete in as little wall-clock time as possible. Existing frameworks utilize modern GPU hardware primarily to parallelize across batch instances, thereby optimizing throughput but not latency. We present CusADi2, a software library that reduces the batch evaluation latency by additionally parallelizing computation within each batch instance. We benchmark forward kinematics, Jacobians, and Hessians for all 67~robots in the MuJoCo Menagerie, including arms, hands, quadrupeds, humanoids, and biomechanical models, against MuJoCo XLA, MuJoCo Warp, frax, CusADi, and CasADi. CusADi2 achieves lower latency than all baselines at batch sizes of 64 to 4096, while matching them at larger batch sizes. We also demonstrate the performance advantage of CusADi2 in a sampling-based control application. All code is open-source at: https://github.com/johnviljoen/cusadi2
Figures & tables
GPU Acceleration
GPU acceleration uses parallel processors, specialized kernels, and memory-aware systems to speed computation across AI and scientific workloads.
See all 167 papersHJCD-IK: GPU-Accelerated Inverse Kinematics through Batched Hybrid Jacobian Coordinate Descent
Scaling Nonlinear Optimization: Many Problems One GPU
jaxipm, the first GPU-batched NLP solver, based on IPOPT, and implemented in JAX. We accomplish this by redesigning IPOPT's algorithm to eliminate control flow with heterogeneous iteration fusion, and by minimizing GPU idle time with iteration level batching. We evaluate jaxipm on a variety of quadrotor nonlinear model predictive control benchmarks, including reference tracking in the presence of obstacles, multi-quadrotor navigation without collision, and navigation in a cluttered environment. We demonstrate up to a increase in throughput over IPOPT. Our complete open-source codebase is available at https://github.com/johnviljoen/jaxipm.