Primal-Dual iLQR for GPU-Accelerated Learning and Control in Legged Robots

TL;DR

This work presents a GPU-accelerated, primal-dual iLQR MPC framework for legged robots that uses a parallel associative scan to solve the KKT system, achieving logarithmic scaling in horizon length and square-log scaling in state/control dimensions. The method delivers substantial runtime improvements over state-of-the-art solvers, enables centralized control of multiple robots, and supports large-scale multi-environment learning through a JAX-based implementation. Key contributions include a rigorous parallel solver with exact backward passes, practical handling of inequalities via barrier terms, and open-source code for real-time, learning-enabled locomotion. The results demonstrate real-time performance for up to 16 robots and large-scale learning scenarios, highlighting the approach’s practicality for complex, multi-robot locomotion 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 $\mathcal{O}(\log^2(n)\log{N} + \log^2(m))$ complexity, instead of $\mathcal{O}(N(n + m)^3)$, where $n$, $m$, and $N$ are the dimension of the system 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 environments, 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.

Figures (7)

• Figure 1: Top left: the quadruped robot Aliengo performing a barrel roll. Top right: the humanoid robot Talos walking blindly on uneven terrain. Bottom left: an example of the massive parallelization achievable with our JAX formulation. Bottom right: Aliengo robot walking blindly on uneven terrain in a lab experiment.
• Figure 2: Comparison of the optimal cost, evaluated from the same initial condition with acados acados (in blue) and with our implementation (in orange). Both solvers are evaluated at each call of the and the costs are recorded.
• Figure 3: Path of the two robots tracking their desired trajectories that will lead to a collision. We visualize the distance between the two robots with dashed gray lines of different shades, the darker the smaller distance.
• Figure 4: Breakdown of the solving times for our implementation, acados (SQP) and crocoddyl (FDDP).
• Figure 5: Average solving times when varying the horizon length (number of nodes) for the model (top) and model (bottom). In blue and green, we plot the solution time of two different state-of-the-art solvers (respectively acados acados and crocoddyl crocoddyl), while in orange, we plot the one of our implementation.