TransformerMPC: Accelerating Model Predictive Control via Transformers
Vrushabh Zinage, Ahmed Khalil, Efstathios Bakolas
TL;DR
TransformerMPC tackles the real-time computation bottleneck of MPC by using transformer-based attention to identify active constraints and to provide warm starts, enabling a reduced MPC problem that preserves feasibility via offline verification. It introduces two dedicated transformers: one for active-constraint prediction and one for warm-start initialization, and provides an analytical solution pathway for QP-based MPC as well as an NMPC acceleration strategy using a single smooth log-sum-exp constraint with GPU parallelization. The approach is solver-agnostic and includes an offline verifier to ensure constraint satisfaction after constraint removal. Empirical results on Upkie, Crazyflie, and Atlas demonstrate up to 35x runtime improvements with significant reductions in inactive constraints and robust constraint satisfaction, highlighting strong potential for real-time robotic control. Overall, TransformerMPC offers a practical, scalable path to real-time nonlinear MPC by combining learning-based constraint screening, warm-starting, and parallelizable acceleration techniques.
Abstract
In this paper, we address the problem of reducing the computational burden of Model Predictive Control (MPC) for real-time robotic applications. We propose TransformerMPC, a method that enhances the computational efficiency of MPC algorithms by leveraging the attention mechanism in transformers for both online constraint removal and better warm start initialization. Specifically, TransformerMPC accelerates the computation of optimal control inputs by selecting only the active constraints to be included in the MPC problem, while simultaneously providing a warm start to the optimization process. This approach ensures that the original constraints are satisfied at optimality. TransformerMPC is designed to be seamlessly integrated with any MPC solver, irrespective of its implementation. To guarantee constraint satisfaction after removing inactive constraints, we perform an offline verification to ensure that the optimal control inputs generated by the MPC solver meet all constraints. The effectiveness of TransformerMPC is demonstrated through extensive numerical simulations on complex robotic systems, achieving up to 35x improvement in runtime without any loss in performance.
