Scaling Nonlinear Optimization: Many Problems One GPU

arXiv:2606.26341v2 Announce Type: replace Abstract: Modern GPUs and software frameworks have made sampling-based methods, such as reinforcement learning and model predictive path integral control, increasingly popular in robotics. However, the logic these methods can execute per sample is limited to what has been implemented on the GPU. Nonlinear programs (NLPs) are one of the most expressive tools in robotics, reliably minimizing a cost while exactly enforcing general constraints. NLP optimizers were among the field's main workhorses, but were not designed for modern GPU hardware. Embedding NLPs such as trajectory optimization and inverse kinematics directly within each sample would be highly valuable, yet the only current option is to call a CPU solver per sample, possibly in parallel on a multi-core CPU. This bottlenecks the CPU, requires GPU-to-CPU-to-GPU transfers when the rest of the pipeline resides on the GPU, and leaves the GPU intermittently idle. To address this issue, in this work we present jaxipm, the first GPU-batched solver for generally constrained nonlinear optimization problems. Our approach is a JAX implementation of the battle-tested IPOPT algorithm, facilitated by eliminating control flow through heterogeneous iteration fusion, and minimizing GPU idle time through iteration-level batching. Our implementation enables solving hundreds of NLPs in parallel on a single GPU. We validate jaxipm on a range of quadrotor nonlinear model predictive control benchmarks, where on a single GPU we achieve between 30.7% and 69.9% of the throughput of 64 IPOPT instances running in parallel on a 64-core CPU. open-source codebase: https://github.com/johnviljoen/jaxipm