Primal-Dual iLQR for GPU-Accelerated Learning and Control in Legged Robots
Lorenzo Amatucci, Joao Moreira de Sousa Pinto, Giulio Turrisi, Dominique Orban, Victor Barasuol, Claudio Semini
AI summary
Problem
Traditional MPC solvers for legged robots struggle to exploit modern GPU hardware accelerators, leading to high computational complexity that limits prediction horizon length, state dimension scaling, and real-time learning integration.
Approach
The authors develop a multiple-shooting SQP method that solves the primal-dual KKT system using parallel associative scans on GPUs, enabling simultaneous temporal, state, and control parallelization while maintaining an exact backward pass.
Key results
- Achieves O(log²(n) log N + log²(m)) computational complexity instead of O(N(n+m)³)
- Up to 60% runtime improvement for WB-MPC and 700% for SRBD-MPC compared to acados and crocoddyl
- Enables a centralized controller for 16 legged robots computable in under 25 ms
- JAX implementation supports large-scale parallelization for GPU-accelerated learning-in-the-loop
Why it matters
This breakthrough removes computational bottlenecks in legged robot control, enabling real-time, high-fidelity MPC and seamless integration with GPU-based reinforcement learning for complex, multi-robot collaborative tasks.
Abstract
This paper introduces a novel Model Predictive Control (MPC) implementation for legged robot locomotion that leverages GPU parallelization. Our approach enables both temporal and state-space parallelization by incorporating a parallel associative scan to solve the primal-dual Karush-Kuhn- Tucker (KKT) system. In this way, the optimal control problem is solved in O(log2(n) log N + log2(m)) complexity, instead of O(N(n+m)3), where n, m, and N are the dimension of the sys- tem state, control vector, and the length of the prediction horizon. We demonstrate the advantages of this implementation over two state-of-the-art solvers (acados and crocoddyl), achieving up to a 60% improvement in runtime for Whole Body Dynamics (WB)- MPC and a 700% improvement for Single Rigid Body Dynamics (SRBD)-MPC when varying the prediction horizon length. The presented formulation scales efficiently with the problem state dimensions as well, enabling the definition of a centralized controller for up to 16 legged robots that can be computed in less than 25 ms. Furthermore, thanks to the JAX implementation, the solver supports large-scale parallelization across multiple envi- ronments, allowing the possibility of performing learning with the MPC in the loop directly in GPU. The code associated with this work can be found at https://github.com/iit-DLSLab/mpx.