Differentiable Model Predictive Control on the GPU
Emre Adabag, Marcus Greiff, John Subosits, Thomas Jonathan Lew
Abstract
Differentiable model predictive control (MPC) offers a powerful framework for combining learning and control. However, its adoption has been limited by the inherently sequential nature of traditional optimization algorithms, which are challenging to parallelize on modern computing hardware like GPUs. In this work, we tackle this bottleneck by introducing a GPU-accelerated differentiable optimization tool for MPC. This solver leverages sequential quadratic programming and a custom preconditioned conjugate gradient (PCG) routine with tridiagonal preconditioning to exploit the problem's structure and enable efficient parallelization. We demonstrate substantial speedups over CPU- and GPU-based baselines, significantly improving upon state-of-the-art training times on benchmark reinforcement learning and imitation learning tasks. Finally, we showcase the method on the challenging task of reinforcement learning for driving at the limits of handling, where it enables robust drifting of a Toyota Supra through water puddles.
DiffMPC provides GPU-accelerated differentiable MPC solver leveraging problem structure for efficient parallelization.
- GPU-accelerated differentiable optimization tool for model predictive control
- Sequential quadratic programming with custom preconditioned conjugate gradient routine
- Exploits time-induced sparsity in optimal control problems for efficient parallelization
- Substantial speedups over CPU and GPU baselines on RL and imitation learning benchmarks
- Model predictive control
- Differentiable optimization
- Sequential quadratic programming
- Preconditioned conjugate gradient
Inequality constraints require penalization in cost or control bounds in dynamics
from the paperDifferentiating through inequality constraints remains challenging with gradient discontinuities
from the paperRuns slower on CPU than GPU due to JAX implementation
from the paperDoes not explicitly support tuning solver hyperparameters
from the paperPoor initial guesses can result in solver divergence and hindered downstream training
from the paper
Handle inequality constraints via augmented Lagrangian or interior-point methods
from the paperRewrite solver in C/C++ for CPU performance improvements
from the paperDevelop robust initialization methods for differentiable optimization pipelines
from the paper
Author keywords
- differentiable optimization
- model predictive control
- optimal control
- gpu-accelerated optimization
- reinforcement learning
- imitation learning
- robotics
Related orals
Mastering Sparse CUDA Generation through Pretrained Models and Deep Reinforcement Learning
SparseRL leverages deep RL and pretrained models to generate high-performance CUDA code for sparse matrix operations.
Overthinking Reduction with Decoupled Rewards and Curriculum Data Scheduling
DECS framework reduces reasoning model overthinking by decoupling necessary from redundant tokens via curriculum scheduling.
MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent
MemAgent uses RL-trained memory modules to enable LLMs to extrapolate from 8K to 3.5M token contexts with minimal performance degradation.
DiffusionNFT: Online Diffusion Reinforcement with Forward Process
DiffusionNFT enables efficient online reinforcement learning for diffusion models via forward process optimization with up to 25x efficiency gains.
Hyperparameter Trajectory Inference with Conditional Lagrangian Optimal Transport
Hyperparameter Trajectory Inference uses conditional Lagrangian optimal transport to reconstruct neural network outputs across hyperparameter spectra without expensive retraining.