首个支持批量求解的GPU非线性规划工具,让机器人优化提速32倍
Scaling Nonlinear Optimization: Many Problems One GPU
- 用JAX重写IPOPT算法,消除分支控制,实现批处理
- 在四旋翼任务中实测吞吐量提升32.85倍
- 适合需要大批量约束优化的机器人学习系统
众多机器人问题,如轨迹优化、逆运动学和接触密集型运动规划,均可归结为非线性规划(NLP)。成熟的求解器如IPOPT能提供硬约束满足、最优性保证,并具备良好问题维度扩展性。然而这些求解器仍依赖CPU,一次仅能处理一个问题,难以融入基于GPU批处理的学习框架。相比之下,采样方法如强化学习、模型预测路径积分和模仿学习因可利用GPU批处理仿真器而成为现代机器人研究核心,每秒生成的动态回放数量显著增加。若存在一个支持批处理的GPU-NLP求解器,便能带来类似的速度提升。为此,我们提出 exttt{jaxipm},首个基于IPOPT的GPU批处理非线性规划求解器,采用JAX实现。通过引入异构迭代融合以消除控制流,以及迭代级批处理减少GPU空闲时间,实现了高效并行。我们在多种四旋翼非线性模型预测控制基准上进行了评估,包括障碍物环境下的参考轨迹跟踪、无碰撞多四旋翼导航及复杂环境导航。结果表明,与IPOPT相比,吞吐量最高提升32.85倍。完整开源代码已发布于https://github.com/johnviljoen/jaxipm。
原文摘要 · Abstract (English)
Many robotics problems, including trajectory optimization, inverse kinematics, and contact-rich motion planning, reduce to nonlinear programs (NLPs). Mature NLP solvers such as IPOPT can solve these problems, offering hard constraint satisfaction, optimality guarantees, and favorable scaling with problem dimension. These solvers underpin gradient-based methods in robotics, yet remain CPU-bound and solve only one problem at a time, preventing their integration into GPU-batched learning pipelines. On the other hand, sampling-based approaches such as reinforcement learning, model predictive path integral, and imitation learning have become the core of modern robotics research due to their ability to leverage GPU-batched simulators. These simulators can generate orders of magnitude more dynamics rollouts per second than was previously possible. If a GPU-batched NLP solver existed, it would unlock similar speedups in the number of constrained, locally optimal solutions generated per second. This regime of solving many problems concurrently versus solving a single problem at a time is a key requirement for integrating NLP solvers in modern GPU-batched robotics frameworks. To this end, we introduce \texttt{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 \textit{heterogeneous iteration fusion}, and by minimizing GPU idle time with \textit{iteration level batching}. We evaluate \texttt{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 $32.85\times$ increase in throughput over IPOPT. Our complete open-source codebase is available at https://github.com/johnviljoen/jaxipm.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。