Reproduce “WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled $LDL^\top$ Factorization” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Henrik Hose, Se Hwan Jeon, Charles Khazoom, Sangbae Kim, Sebastian Trimpe
Researchers developed WarpMPC, a JAX-based toolbox utilizing a Warp backend to optimize Model-Predictive Control (MPC) on GPUs for large batches of problems. This system achieves up to 279,000 QP solves per second by leveraging unrolled sparse LDL⊤ factorization and custom CasADi-to-JAX translation, facilitating rapid training of neural network controllers for robotic systems.
Key insights
- Exploiting the fixed sparsity pattern across large batches of MPC problems allows for specialized, unrolled GPU kernels that eliminate dynamic operations, significantly outperforming generic sparse solvers.
- A symbolic-major memory layout, which makes all batch elements for a specific symbolic operation contiguous, dramatically improves data locality and vectorization on GPUs, boosting throughput for large batches.
- Pre-calculating dependency levels for triangular solves and optimizing segmentation for factorization kernels can enhance GPU utilization, particularly for smaller batch sizes or complex sparse structures, by increasing parallelism.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.