面向 GPU 加速的 Primal-Dual iLQR 用于四足机器人的学习与控制

本文引入一种 novel 的 Model Predictive Control (MPC) 实现, 用于四足机器人 locomotion, 利用 GPU 并行化。我们的方法通过纳入并行关联扫描来求解 primal-dual Karush-Kuhn-Tucker (KKT) 系统, 从而实现时序和状态空间的并行化。如此, 最优控制问题的复杂度从 $\mathcal{O}(N(n + m)^3)$ 降低到 $\mathcal{O}(\log^2(n)\log{N} + \log^2(m))$, 其中 $n$, $m$, 和 $N$ 分别是系统状态维数、控制向量维数和预测时域长度。我们展示了该实现相对于两种最先进求解器 (acados 和 crocoddyl) 的优势, 在 Whole Body Dynamics (WB)-MPC 和 Single Rigid Body Dynamics (SRBD)-MPC 中分别实现了最高可达 60% 和 700% 的运行时改进, 随着预测时域长度的变化。所提出的表述在问题状态维数方面也高效比例扩展, 使我们能够定义用于最多 16 个四足机器人的集中式控制器, 在低于 25 ms 的计算时间内即可计算完成。此外, 由于采用 JAX 实现, 求解器支持跨多个环境的大规模并行化, 允许在 GPU 中直接将学习与 MPC 一起进行。与本工作相关的 code 可在 https://github.com/iit-DLSLab/mpx 查找。