Smoothed ADMM: Enhancing Differentiable MPC Policies on GPU

Paper: Efficient Learning of Differentiable MPC Policies on the GPU via Smoothed ADMM and First-Order Reinforcement Learning

Authors: Jianghan Zhang, Roland Schwan, Gabriel Bravo-Palacios, John Zhang, Thomas Lew, Brian Plancher

Executive Summary

Recent research introduces an innovative approach to improving differentiable model predictive control (MPC) policies, which are crucial for optimizing decisions in complex systems while adhering to specific constraints. The study addresses two significant challenges that have hindered the effectiveness of learning MPC policies: the difficulties posed by nonsmooth sensitivities during active-set transitions and the high computational cost associated with policy evaluations, both of which typically require extensive optimization solutions.

The research presents a novel solution called SmoothedTurboMPC, a GPU-accelerated differentiable MPC solver. This new tool combines sequential quadratic programming (SQP) with a log-barrier-smoothed alternating direction method of multipliers (ADMM) for its inner solver. The design of SmoothedTurboMPC allows for smoother gradients and precise sensitivities at convergence. Moreover, its execution on GPUs significantly enhances computational efficiency, making the process faster and more effective.

To further streamline the training of MPC policies, the researchers integrated their solver with a first-order reinforcement learning method known as Short-Horizon Actor-Critic (SHAC). This approach allows for the reuse of previously computed MPC solutions, reducing the need for continual re-evaluation, which is often a time-consuming part of policy training. In comparative tests against leading solvers, SmoothedTurboMPC demonstrated up to 4.4 times greater training throughput.

The practical implications of this research are particularly evident in a vehicle racing task where the new method not only outperformed the existing MPC approach (specifically MPC Proximal Policy Optimization, or PPO) but also achieved an impressive 46 times speedup. This efficiency allowed the training process to complete on 3.3 million transitions in just 7 minutes. The researchers also illustrated the potential for developing a more reliable MPC racing policy through reinforcement learning combined with domain randomization, validating the approach on a full-scale race car.

The advancements outlined in this research are integrated into the open-source TurboMPC solver, making these innovations accessible for further development and application within the field. This work highlights the potential for significant improvements in the efficiency and effectiveness of decision-making systems across various applications, though real-world implementations and industry-specific deployments remain to be fully explored.

Academic Abstract

Differentiable model predictive control (MPC) enables structured, constraint-aware policies to be optimized end-to-end. However, learning MPC policies faces two bottlenecks: (1) nonsmooth sensitivities at active-set transitions and (2) expensive policy evaluation, since each requires an optimization solve. We first introduce SmoothedTurboMPC, a GPU-parallel differentiable MPC solver combining sequential quadratic programming (SQP) with a log-barrier-smoothed alternating direction method of multipliers (ADMM) inner solver. Its design enables smooth gradients, exact sensitivities at convergence, and efficient execution on GPUs. Second, to efficiently train MPC policies, we pair our solver with first-order reinforcement learning (FORL) using Short-Horizon Actor-Critic (SHAC), which reuses rollout MPC solutions instead of repeatedly re-evaluating them. Comparisons with state of the art solvers show that the approach achieves up to $4.4\times$ higher end-to-end training throughput. On a vehicle racing task, our method achieves higher returns than MPC PPO with a $46\times$ speedup, training on 3.3 million transitions in only 7 minutes. Finally, we use the framework to train a more reliable MPC racing policy via RL with domain randomization, and validate it on a full-scale race car. We integrate our advances into the open-source TurboMPC solver.