GPU加速大批量MPC求解,速度提升3到25倍。
WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled $LDL^\top$ Factorization

- 通过展开稀疏分解和求解,优化ADMM求解器性能。
- 单次迭代达8000至25万次/秒,实测提升3-25倍。
- 适合需要高速在线控制的机器人系统应用。
本文针对在相同结构下求解大规模批量(10,000至超过100,000个)序列二次规划(SQP)迭代的问题,提出一系列数值优化方法以最大化GPU上的吞吐量。基于批量中所有MPC实例具有相同时间、代价与约束稀疏结构的洞察,提出展开稀疏线性因子分解与求解过程,该部分主导了ADMM求解器的运行时间。通过优化内存布局、减少填充的分段处理及依赖层级调度回代,有效避免内存访问瓶颈并减少冗余计算,进一步加速灵敏度计算。在非线性小车、四旋翼和人形机器人基准测试中,实现每秒8,000至250,000次SQP迭代,相较基线提升3×至25×。通过合成数据集,在4分钟内训练出神经网络近似控制器,成功在硬件实验中稳定控制纳米四旋翼飞行器。
原文摘要 · Abstract (English)
This paper introduces numerical optimizations for maximizing throughput on GPU when solving large batches (10,000 to over 100,000) of sequential quadratic programming (SQP) iterations, where all problems have the same structure. The optimizations are implemented in a toolbox WarpMPC for model-predictive control (MPC) in JAX and Warp. Based on the insight that all MPC problem instances in a batch share the same sparsity in time, cost, and constraints, we propose unrolling sparse linear factorizations and solves, which dominate alternating direction method of multipliers (ADMM) solver runtime. We avoid memory access bottlenecks and wasting computations via optimized memory layout, padding-reducing segmentation of the unrolled factorization, and dependency level scheduled backsolves, additionally accelerating sensitivity computation. We achieve throughputs of 8,000 to 250,000 SQP iterations per second on nonlinear cartpole, quadrotor, and humanoid robot benchmarks, outperforming baselines by 3$\times$ to 25$\times$. We illustrate practical usefulness by synthesizing a dataset and training a neural network approximation of an MPC in under 4 minutes that stabilizes a nano quadrotor in hardware experiments.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。