Parallel Branch Model Predictive Control on GPUs

技术摘要：基于 GPU 的并行分支模型预测控制

问题陈述
分支模型预测控制（BMPC）是一种处理动态环境下不确定性的强大规划框架，例如在自动驾驶中，通过生成其中分支对应于不同不确定性实现的轨迹树。然而，BMPC 的广泛部署受到计算负担的阻碍，特别是在处理长规划时界和众多预测场景时。现有的求解器通常难以高效利用内在的树状结构，或无法实现时间并行性，从而限制了其在实时应用中的适用性。此外，在并行硬件上处理树状结构优化控制框架中的一般阶段约束仍然是一个挑战。

方法论
作者提出了一种基于 GPU 的 BMPC 求解器，该求解器集成了多重射击法（multiple-shooting formulation）和用于约束处理的增广拉格朗日（AL）法。其核心方法依赖于两个专门设计的内层线性二次调节器（LQR）求解器，旨在利用树状稀疏结构：

1. 并行树 LQR 求解器：

   • SLQR（场景级并行）： 该求解器执行从叶节点到根节点的修正 Riccati 递归。它在每个阶段聚合来自子节点的价值函数，从而允许并行解决每个节点处的独立最小化问题。这种方法对 GPU 资源的需求较低，适用于资源受限的场景。
   • STLQR（场景与时间并行）： 该求解器利用并行扫描算法（parallel scan algorithm）在反向（Riccati）和正向（rollout）传递中同时实现场景级和时间级并行。它利用条件价值函数（CVF）和树状组合规则来计算价值函数和仿射控制律，其时间复杂度为 $O(\log N)$。该方法提供了更高的并行度，但对 GPU 资源的需求更高。

2. 通过增广拉格朗日进行约束处理：
   为了处理一般的阶段约束，作者采用了增广拉格朗日（AL）法。内层循环使用迭代 LQR（iLQR）方法，其中通过 Powell-Hestenes-Rockafellar (PHR) 惩罚函数将受约束问题近似为无约束树 LQR 问题。通过线性前向传递（linear rollout）计算最优扰动，从而实现 GPU 上的高效并行化。外层循环根据约束违反情况，遵循 BCL 规则自适应地更新拉格朗日乘子和惩罚权重。

3. 实现：
   该求解器使用 JAX 实现，利用其自动微分和 XLA 编译器进行 GPU 加速。该框架支持单精度（FP32）和双精度（FP64）算术运算。

主要贡献
本文概述了三个主要贡献：

1. 双重并行求解器： 开发了两种并行树 LQR 求解器（SLQR 和 STLQR），它们提供不同程度的并行性，允许用户根据问题规模和可用计算资源选择合适的方法。
2. 约束非线性 BMPC 求解器： 将这些树 LQR 求解器集成到非线性 BMPC 问题的多重射石迭代求解器中，并结合了用于鲁棒约束处理和热启动能力的增广拉格朗日法。
3. 基准测试与开源： 将所提出的求解器与现有的 iLQR 求解器（TRAJAX, MPX）以及高性能 CPU 求解器（HPIPM）进行了全面的基准测试，并发布了开源实现。

数值结果
作者在两个不同的任务上评估了该求解器：无约束树 LQR 问题以及单轮车（unicycle）和四连杆倒立摆（quad-pendulum）的约束轨迹规划。

• Tree LQR 性能： 基于 GPU 的求解器性能高度依赖于问题规模和硬件。在小规模问题（例如 $Z = 2$ 个树路径）上，由于 GPU 内存访问延迟和开销，求解器明显慢于基于 CPU 的 HPIPM 求解器，其中 STLQR 比其慢 5 倍以上，SLQR 比其慢 20 倍以上（在 NVIDIA RTX 5060 Ti 上）。然而，在大规模实例上，性能发生了逆转：在 RTX 5060 Ti 上，当规模较大时（$Z \ge 128$），SLQR 的表现最高可超越 HPIPM 至 2 倍。同样，在 RTX 4090 等高端 GPU 上，对于中大规模树规模（$Z \ge 16$），STLQR 实现了高达 1.9 倍的 HPIPM 加速。
• 约束处理： 在轨迹规划任务中，所提出的求解器（ILQRJAX）展示了与最先进的 CPU 求解器 IPOPT 相当的收敛行为，但显著降低了单次迭代的计算时间（例如，将单轮车的平均迭代时间从 3.80 ms 降至 1.87 ms）。该求解器成功处理了所有测试实例，而其他基于 GPU 的求解器（TRAJAX, MPX）在面对更具挑战性的实例时表现挣扎，由于公式限制或缺乏自适应更新方案，往往无法收敛。

意义与主张
本文声称，通过在 GPU 上利用并行算法充分发挥树状结构的优势，所提方法为大规模问题的实时 BMPC 提供了一条可行路径。作者强调，在能够有效并行化树结构的规模化实例中，其方法实现了优于基于 CPU 的高性能求解器的性能。作者承认，基于并行扫描的求解器对 GPU 资源需求较高，如果资源饱和可能会限制其扩展性；并且对于小规模问题，基于 CPU 的求解器可能仍然更具优势。这项工作旨在通过平衡计算效率与对约束及不确定性的严格处理，为复杂现实应用中的不确定性感知规划提供可行性。未来的工作确定为在 CUDA C++ 中实现该方法，以进一步优化资源利用，并探索混合精度算术以提高在 FP32 优化硬件上的数值稳定性。