Parallel Branch Model Predictive Control on GPUs

Technical Summary: Parallel Branch Model Predictive Control on GPUs

Problem Statement
Branch Model Predictive Control (BMPC) is a powerful planning framework for handling uncertainty in dynamic environments, such as automated driving, by generating trajectory trees where branches correspond to different uncertainty realizations. However, the widespread deployment of BMPC is hindered by the significant computational burden required to solve these problems, particularly when dealing with long planning horizons and numerous predicted scenarios. Existing solvers often struggle to exploit the inherent tree structure efficiently or fail to achieve temporal parallelism, limiting their suitability for real-time applications. Furthermore, handling general stage-wise constraints within a tree-structured optimal control framework on parallel hardware remains a challenge.

Methodology
The authors propose a GPU-based solver for BMPC that integrates a multiple-shooting formulation with an augmented Lagrangian (AL) method for constraint handling. The core of the approach relies on two tailored inner Linear Quadratic Regulator (LQR) solvers designed to exploit the tree-sparse structure:

1. Parallel Tree LQR Solvers:

   • SLQR (Scenario-Level Parallelization): This solver performs a modified Riccati recursion from leaf nodes to the root. It aggregates value functions from child nodes at each stage, allowing independent minimization problems at each node to be solved in parallel. This approach requires fewer GPU resources and is suitable for scenarios where resources are limited.
   • STLQR (Scenario & Temporal Parallelization): This solver leverages the parallel scan algorithm to achieve both scenario-level and temporal parallelism in both backward (Riccati) and forward (rollout) passes. It utilizes Conditional Value Functions (CVFs) and a tree-structured combination rule to compute value functions and affine control laws in $O(\log N)$ time complexity. This method offers higher parallelism but demands more GPU resources.

2. Constraint Handling via Augmented Lagrangian:
   To address general stage-wise constraints, the authors employ an Augmented Lagrangian (AL) method. The inner loop uses an iterative LQR (iLQR) approach where the constrained problem is approximated as an unconstrained tree LQR problem using the Powell-Hestenes-Rockafellar (PHR) penalty function. A linear rollout is used to compute optimal perturbations, enabling efficient parallelization on GPUs. The outer loop adaptively updates Lagrange multipliers and penalty weights based on constraint violations, following the BCL rule.

3. Implementation:
   The solver is implemented in JAX, utilizing its automatic differentiation and XLA compiler for GPU acceleration. The framework supports both single-precision (FP32) and double-precision (FP64) arithmetic.

Key Contributions
The paper outlines three primary contributions:

1. Dual Parallel Solvers: The development of two parallel tree LQR solvers (SLQR and STLQR) that offer different levels of parallelism, allowing users to select the appropriate method based on problem size and available computational resources.
2. Constrained Nonlinear BMPC Solver: The integration of these tree LQR solvers into a multiple-shooting iterative solver for nonlinear BMPC problems, incorporating an augmented Lagrangian method for robust constraint handling and warm-starting capabilities.
3. Benchmarking and Open Source: A comprehensive benchmarking of the proposed solver against existing iLQR solvers (TRAJAX, MPX) and a high-performance CPU-based solver (HPIPM), along with the release of an open-source implementation.

Numerical Results
The authors evaluated the solver on two distinct tasks: unconstrained tree LQR problems and constrained trajectory planning for a unicycle and a quad-pendulum.

• Performance on Tree LQR: The performance of the GPU-based solvers is highly dependent on problem size and hardware. On small problem sizes (e.g., $Z = 2$ tree paths), the solvers are significantly slower than the CPU-based HPIPM solver, with STLQR being over 5$\times$ slower and SLQR over 20$\times$ slower on an NVIDIA RTX 5060 Ti due to GPU memory-access latency and overhead. However, on large-scale instances, the performance reverses: SLQR can outperform HPIPM by up to 2$\times$ on large-scale instances ($Z \ge 128$) on the RTX 5060 Ti. Similarly, on high-end GPUs like the RTX 4090, STLQR achieves up to a 1.9$\times$ speedup over HPIPM for moderate-to-large tree sizes ($Z \ge 16$).
• Constraint Handling: In trajectory planning tasks, the proposed solver (ILQRJAX) demonstrated convergence behavior comparable to the state-of-the-art CPU solver IPOPT but with significantly reduced per-iteration computation time (e.g., reducing average iteration time from 3.80 ms to 1.87 ms for the unicycle). The solver successfully handled all test instances, whereas other GPU-based solvers (TRAJAX, MPX) struggled with more challenging instances, often failing to converge due to formulation limitations or lack of adaptive update schemes.

Significance and Claims
The paper claims that the proposed approach offers a viable path toward real-time BMPC for large-scale problems by fully exploiting the tree structure through parallel algorithms on GPUs. The authors emphasize that their method achieves superior performance compared to high-performance CPU-based solvers specifically on large-scale instances where the tree structure can be effectively parallelized. They acknowledge, however, that the parallel scan-based solver has high GPU resource demands, which can limit scalability if resources are saturated, and that for small problem sizes, CPU-based solvers may still outperform. The work positions itself as a step toward making uncertainty-aware planning feasible for complex, real-world applications by balancing computational efficiency with the rigorous handling of constraints and uncertainty. Future work is identified as implementing the method in CUDA C++ to further optimize resource utilization and potentially exploring mixed-precision arithmetic to improve numerical stability on FP32-optimized hardware.