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

In plain words

A legged robot has to re-plan its movements many times a second. The planning step is an optimisation over a short look-ahead window, and the usual solvers work through that window one step after another, so a longer window costs more time. This paper rewrites the core calculation so that its pieces are merged in pairs, in a tree-like pattern that a graphics card can run all at once. The same idea is used across the window and across the robot's state variables, so longer windows and bigger problems add much less time.

The solver is written in JAX, so it can also run thousands of simulated robots in parallel on the GPU, which matters when a learning method needs the planner inside its training loop. The gain is not free: moving data to and from the GPU has a fixed cost, so for small problems a CPU solver can still be faster.

TL;DR

The paper presents a GPU-parallel model predictive controller for legged robots. It replaces the sequential backward pass of an LQR-style solver with an associative scan, and wraps it in a sequential quadratic programming (SQP) loop with multiple shooting, a barrier for inequality constraints and a parallel filter line search. The authors report faster solve times than acados and crocoddyl, enough speed for 16 centralised robots, and a JAX implementation usable inside simulated training. All speed-ups are author-reported and depend on the benchmark setup.

Problem

Model predictive control (MPC) turns a task into an optimal control problem (OCP) over a horizon of N nodes, with state x, input u, dynamics h and constraints §1 eq. 1. The paper states that for solvers such as OCS2, acados and crocoddyl the cost grows linearly with the horizon length and cubically with the state dimension, which limits the horizon, the update rate and the number of robots that can be controlled together §1.

The standard LQR-type solution has complexity O(N(n+m)^3) with a sequential Riccati recursion, where n and m are the state and control dimensions [abstract] §2.3 eq. 5. Earlier GPU attempts were partial: they parallelise only the line search or gradients, batch independent problems without parallelism over time or state, or need offline precomputation (Relu-QP) §1. Single-shooting DDP additionally needs a feasible initial guess and can have numerical problems, which the multiple-shooting formulation used here avoids §1 §2.2.

The paper's target is legged locomotion, both single-rigid-body (SRBD) and whole-body (WB) models §3. Its control problem also has barrier-type constraints (friction cones, torque and joint limits) that the solver must handle without while-loops, so JAX can differentiate through it §2.6. The paper names friction cones and joint limits as examples §2.6.

Method

Section numbers §N follow the order of the \section commands in the LaTeX source (1 Introduction, 2 Primal-Dual LQR, 3 Model Predictive Control, 4 Results, 5 Conclusion). The printed paper uses Roman numerals for sections, Arabic numbers for equations and figures, and Roman numerals for tables.

Problem and SQP loop. The OCP is a finite-horizon problem with a stage cost l_i, dynamics x_{i+1} = h_i(x_i, u_i), state and input sets, and an initial condition §2.1 eq. 1. Each control call performs one SQP iteration on the quadratic model of the Lagrangian with linearised dynamics, giving the QP in eq. 3 for the search directions (δx, δu) §2.2 eq. 2 eq. 4. The paper calls the state and control directions δx and δu and the multiplier direction δλ §2.2.

Primal and dual solution by Riccati recursion. The QP is solved backward over the horizon with the Riccati recursion of eq. 5, which yields feedback gains K_i and feed-forward terms k_i; a forward pass then gives δu_i = K_i δx_i

• k_i eq. 6. Once P_i and p_i are known for all nodes, the multiplier step is δλ_i = P_i δx_i + p_i, computed for all i in parallel eq. 7 §2.3.

Parallel associative scan. The paper observes that the cost-to-go between any two nodes (a "conditional value function") is again a quadratic of the same form, and that combining two such pieces with the min-plus-style operator of eq. 9 is associative §2.4. Each adjacent pair (i, i+1) gives a starting element eq. 12 and the last node has its own initialisation eq. 13; the combination rule eq. 11 then lets a scan produce P_i and p_i for every node in O(log N) depth §2.4 eq. 8. Matrix inversions inside the scan are cited to a parallel inversion method §2.4.

Forward rollout by scan. Substituting K and k into the dynamics gives an affine closed-loop recursion with Ā_i = A_i + B_i K_i and b̄_i = B_i k_i + b_i eq. 14. Conditional trajectories are composed by the rule in eq. 16 and initialised as in eq. 17; δx_i then comes from a scan over the initial element and eq. 18. The paper gives this as O(log n log N + log m) §2.4.

Line search. The step is x^{k+1} = x^k + α δx, λ^{k+1} = λ^k + α δλ, u^{k+1} = u^k + α δu eq. 19. The paper uses a filter-type line search after Ipopt and Perceptive MPC §2.5: the constraint violation θ eq. 20 is compared with a threshold, and a step is accepted by an Armijo test or by a decrease in cost or θ, depending on whether the direction is a descent direction §2.5. All ten candidate step sizes 2^0 to 2^-9 are evaluated in parallel and the largest accepted one is taken §2.5.

Gauss-Newton and barriers. The Hessian blocks use a Gauss-Newton approximation of the least-squares cost with weight W §2.6 eq. 21. Each inequality ξ < 0 is handled with the relaxed log-barrier of eq. 22, and its curvature is added to the Hessian blocks as in eq. 23 to keep the QP convex §2.6.

Warm start and control loop. One SQP iteration is run per control loop, warm-started with the previous solution shifted by one time step, following the real-time iteration scheme §2.6. The first control is applied; the paper says only the initial state is sent to the GPU and torques, joint targets and velocities are returned, adding about 1 ms §4.1.

# One MPC call: one warm-started SQP iteration  [§2.6]
# Input:  x0 (n,), previous guess (X, U, lam)
# Output: u0 (m,), guess shifted for the next call

def mpc_step(x0, X, U, lam):
    # Shift the previous solution by one time step  [§2.6]
    X, U, lam = shift_by_one_step(X, U, lam)
    # Linearise dynamics and build cost and barrier blocks for every
    # node  [§2.2] [§2.6] [eq. 3] [eq. 21] [eq. 22] [eq. 23]
    stage_qp = build_stage_qp(x0, X, U, lam)
    # Backward scan for the value function P_i, p_i  [§2.4] [eq. 8]-[eq. 13]
    P, p = reverse_scan_value_function(stage_qp)
    # Feedback gains and feed-forward terms per node  [§2.3] [eq. 5]
    K, k = policy_gains(stage_qp, P, p)
    # Forward scan for state, then control and multiplier directions
    # [§2.3] [§2.4] [eq. 6] [eq. 7] [eq. 14]-[eq. 18]
    dX, dU, dlam = forward_scan_directions(stage_qp, K, k)
    # Choose the step size on a fixed grid  [§2.5] [eq. 19] [eq. 20]
    alpha = filter_line_search(X, U, lam, dX, dU, dlam)
    X = X + alpha * dX
    U = U + alpha * dU
    lam = lam + alpha * dlam
    return U[0], (X, U, lam)

# Filter line search on a fixed grid, evaluated in one batch  [§2.5]
# Output: step size alpha (0 if none passes)

def filter_line_search(X, U, lam, dX, dU, dlam):
    # Ten candidate steps, from 2^0 down to 2^-9  [§2.5]
    alphas = 0.5 ** arange(10)
    # Roll out all candidates at once, without while-loops  [§2.5] [§2.6]
    trials = rollout_batch(X, U, lam, dX, dU, dlam, alphas)
    # Keep the largest candidate that passes a filter test  [§2.5]
    # Threshold and Armijo constant are not stated in the paper.
    return first_accepted(trials, alphas)

Equations as printed. The LaTeX source is quoted here in the paper's own notation. Symbols that look inconsistent are listed in Pitfalls.

Riccati recursion, one node eq. 5:

\begin{aligned} G_i &= R_i + B_i^\top P_{i+1} B_i, \\ H_i &= S_i + B_i^\top P_{i+1} A_i, \\ h_i &= B_i^\top (p_{i+1} + P_{i+1} b_i) + r_i, \\ K_i &= -G_i^{-1} H_i, \qquad k_i = -G_i^{-1} h_i, \\ P_i &= Q_i + A_i^\top P_{i+1} A_i + K_i^\top H_i, \\ p_i &= q_i + A_i^\top (p_{i+1} + P_{i+1} b_i) + K_i^\top h_i. \end{aligned} eq. 5

Associative operator on conditional value functions eq. 9:

V_{i \to k} \otimes V_{k \to j} = \min_{x_k} \{ V_{i \to k} + V_{k \to j} \} eq. 9

Combination rule for the conditional value function (as printed) eq. 11:

\begin{aligned} \tilde{P}_{i,j} &= \tilde{A}_{i,k}^\top (I + \tilde{P}_{k,j} \tilde{C}_{i,k})^{-1} \tilde{P}_{k,j} \tilde{A}_{i,k} + \tilde{P}_{i,k}, \\ \tilde{p}_{i,j} &= \tilde{A}_{k,j}^\top (I + \tilde{P}_{k,j} \tilde{C}_{i,k})^{-1} (\tilde{p}_{k,j} - \tilde{P}_{k,j} b_{i,j}) + \tilde{p}_{i,k}, \\ \tilde{A}_{i,j} &= \tilde{A}_{k,j} (I + \tilde{C}_{i,k} \tilde{P}_{k,j})^{-1} \tilde{A}_{i,k}, \\ \tilde{C}_{i,j} &= \tilde{A}_{k,j} (I + \tilde{C}_{i,k} \tilde{P}_{k,j})^{-1} \tilde{C}_{i,k} \tilde{A}_{k,j}^\top + \tilde{C}_{k,j}, \\ \tilde{b}_{i,j} &= \tilde{A}_{k,j} (I + \tilde{C}_{i,k} \tilde{P}_{k,j})^{-1} (\tilde{b}_{i,k} - \tilde{C}_{i,k} p_{k,j}) + \tilde{b}_{j,k}. \end{aligned} eq. 11

Initial elements for adjacent nodes eq. 12:

\begin{aligned} \tilde{A}_{i,i+1} &= A_i - B_i R_i^{-1} S_i, & \tilde{P}_{i,i+1} &= Q_i - S_i^\top R_i^{-1} S_i, \\ \tilde{C}_{i,i+1} &= B_i R_i^{-1} B_i^\top, & \tilde{p}_{i,i+1} &= q_i - S_i^\top R_i^{-1} r_i, \\ \tilde{b}_{i,i+1} &= b_i - B_i R_i^{-1} r_i. \end{aligned} eq. 12

Last node eq. 13: P̃_{N,N+1} = Q_{N+1}, p̃_{N,N+1} = q_{N+1}, and Ã_{N,N+1}, b̃_{N,N+1}, C̃_{N,N+1} are zero eq. 13.

Closed-loop matrices and affine term eq. 14: Ā_i = A_i + B_i K_i and b̄_i = B_i k_i + b_i.

Composition of conditional trajectories, as printed eq. 16: Ā_{i,j} = Ā_{i,k} Ā_{k,j} and b̄_{i,j} = Ā_{k,j} b̄_{i,k} + b̄_{k,j}.

Gauss-Newton blocks for the cost ε with weight W eq. 21:

\begin{aligned} Q_{N+1} &= \nabla_x \epsilon_{N+1}^\top W_{N+1} \nabla_x \epsilon_{N+1}, \\ Q_i &= \nabla_x \epsilon_i^\top W_i \nabla_x \epsilon_i, \\ S_i &= \nabla_x \epsilon_i^\top W_i \nabla_u \epsilon_i, \\ R_i &= \nabla_u \epsilon_i^\top W_i \nabla_u \epsilon_i. \end{aligned} eq. 21

Relaxed log-barrier for a constraint ξ(x, u) eq. 22. The paper writes the constraint as ξ < 0 and says μ and δ are tuned per constraint §2.6:

\mathcal{B}(\xi, \mu, \delta) = \begin{cases} -\mu \ln(\xi), & \xi \geq \delta, \\ \frac{\mu}{2}\left(\left(\frac{\xi - 2\delta}{\delta}\right)^2 - 1\right) - \mu \ln(\delta), & \xi < \delta. \end{cases} eq. 22

Constraint curvature added to the blocks eq. 23 (as printed):

\begin{aligned} Q_{N+1} &= \nabla_x \epsilon_{N+1}^\top W_{N+1} \nabla_x \epsilon_{N+1} + \nabla_x \xi_{N+1}^\top \nabla_{\xi\xi} \mathcal{B}_{N+1} \nabla_x \xi_{N+1}, \\ Q_i &= \nabla_x \epsilon_i^\top W_i \nabla_x \epsilon_i + \nabla_x \xi_i^\top \nabla_{\xi\xi} \mathcal{B}_i \nabla_x \xi_i, \\ S_i &= \nabla_x \epsilon_i^\top W_i \nabla_x \epsilon_i + \nabla_x \xi_i^\top \nabla_{\xi\xi} \mathcal{B}_i \nabla_u \xi_i, \\ R_i &= \nabla_u \epsilon_i^\top W_i \nabla_u \epsilon_i + \nabla_u \xi_i^\top \nabla_{\xi\xi} \mathcal{B}_i \nabla_u \xi_i. \end{aligned} eq. 23

Dynamics model (SRBD and WB). The SRBD model keeps the centre-of-mass position, velocity, orientation (quaternion, so the state is in R^7) and angular velocity, and its control is the ground reaction forces, so m = 3 n_c. For the quadruped the paper gives n = 13 and m = 12 §3. The WB model adds joint positions, velocities and torques: n = 13 + 2 n_joint and m = 3 n_c + n_joint, which the paper gives as n = 37 and m = 24 for the quadruped §3. The paper's own equations for both models are commented out in the source and are not printed here. The implementation's WB model is in Paper vs code.

Hyperparameters

Results

Author-reported throughout. Baselines: acados (SQP) for SRBD and multi-robot comparisons, crocoddyl (DDP) for WB horizon comparisons §4. The numbers below are from the text and abstract; the bar and line values in the figures were not read (figures are images here, see Pitfalls).

Paper vs code

Official code: https://github.com/iit-DLSLab/mpx (linked in the abstract), BSD-3-Clause, pinned at code_commit. The LICENSE is BSD 3-Clause by DLS Lab, IIT. The solver lives in a git submodule, mpx/jax_ocp_solvers, pinned at commit 2b7cbac2bc34750903bc83722affff2237a12b15 (iit-DLSLab/jax_ocp_solvers). That submodule is not in the shallow clone used for this card, so its scan, line-search and Hessian code were not verified line by line. Its licence was not checked. Discrepancies below come from the main repository only.

• Horizon. The paper uses N = 50 for the Go2 quadruped §4. The Go2 config sets N = 25 with dt = 0.02 config_go2.py. The barrel-roll config uses N = 100 with dt = 0.01 config_barrel_roll.py. The paper does not state dt.
• State and control sizes. The paper's WB quadruped has n = 37 and m = 24 §3. The Go2 config defines n = 13 + 2·12 + 6·4 = 61 and m = 12, with the ground reaction forces as part of the state config_go2.py. The code's state therefore carries an extra 6·n_c block; the paper does not describe this block.
• Hessian approximation. The paper uses Gauss-Newton blocks §2.6. The Go2 config passes hessian_approx = None config_go2.py, and mpc_wrapper forwards it to the solver mpc_wrapper.py. What None selects is inside the unverified submodule.
• Barrier defaults. The paper leaves μ and δ unspecified §2.6. The code's penalty defaults are α = 0.1 and σ = 5 objectives.py, and the friction-cone constraint uses them with a clip to [1e-6, 1e6] objectives.py. The clip is not in the paper.
• Penalty sign. The code's penalty takes a positive margin, using log of max(constraint, σ) objectives.py. The paper says ξ < 0 is the constraint §2.6 and then writes −μ ln ξ, which needs ξ > 0 (see Pitfalls).
• Warm start. The paper shifts the previous solution by one step §2.6. The code computes the shift as 1 / (dt · mpc_frequency) mpc_wrapper.py and fills the tail by repeating the last element mpc_wrapper.py. With dt = 0.02 and 50 Hz the shift is 1 step. The paper does not state the tail rule.
• Solver modes. The paper describes only the primal-dual SQP solver. The wrapper also has an FDDP mode and defaults to primal_dual mpc_wrapper.py mpc_wrapper.py.
• Gait and cost weights. The paper gives no gait or cost weights. The Go2 config sets duty factor 0.65, step frequency 1.35 Hz, step height 0.065 m config_go2.py and cost weights such as Qp = diag(0, 0, 1e4) config_go2.py.
• Torque limits. The Go2 config uses ±25 config_go2.py. The paper does not state torque limits for its experiments.

Pitfalls

• Table I and text disagree. Table I gives 370× at 50 Hz and 570× at 25 Hz for SRBD tab. I. The text says the 50 Hz SRBD run gives 570 s of simulated time per second §4.4. Check the source before quoting either figure.
• "700% improvement" is undefined. A runtime reduction cannot exceed 100%, so this is likely a speed-up ratio. The paper does not define it [abstract].
• Figures were not read. The horizon sweeps fig. 5 and the timing breakdown fig. 4 are images. Figure values are not in the text, so this card reports only the text and abstract numbers.
• Asymptotic claims are not wall-clock claims. The O(log² n log N + log² m) bound assumes a parallel scan with parallel matrix inversion §2.4. The paper reports a GPU host-device overhead of about 1 ms §4.1 and says the GPU loses to crocoddyl below 80 nodes for WB §4.2. Small problems are a bad fit for this approach.
• Dimension of the WB model. Use the paper's n = 37 and m = 24 §3 only with the paper's model. The code's 61 / 12 split changes the sizes of the matrices that the scan carries.
• Max versus min. The text writes V_i(x_i) = max over x_{N+1} §2.4, while the operator uses min over x_k eq. 9, and the conditional value function is a max over η eq. 10. The sign convention for costs is not stated; check against a rollout before relying on either.
• Eq. 11 symbols look inconsistent. The p̃ rule uses b_{i,j} and p_{k,j} without tildes, and the b̃ rule ends with b̃_{j,k}, where b̃_{k,j} looks intended eq. 11. The paper does not resolve these. Implement from the code or derive the rule and test it on a small problem.
• Eq. 16 matrix order. The rule eq. 16 composes Ā_{i,j} = Ā_{i,k} Ā_{k,j}. For x_j = Ā_{i,j} x_i + b̄_{i,j} to hold with the recursion of eq. 14, the matrix product must be Ā_{k,j} Ā_{i,k}. Check the order against a rollout.
• Index conventions. eq. 17 gives the initial elements for i = 1…N while the scan elements run over 0…N §2.4. eq. 18 then uses a_0 … a_{i-1}. Choose one 0-based convention and check the first and last stages.
• Terminal weight name. The QP calls the terminal Hessian P_{N+1} eq. 3, while the scan initialisation uses Q_{N+1} eq. 13 and the Gauss-Newton block is called Q_{N+1} eq. 21. Assume they are the same matrix; the paper does not say so.
• Symbol h. h is used for the dynamics §2.1 and for the linear term of the Riccati recursion eq. 5. Keep them apart in code.
• Probable typo in eq. 23. The S_i block is printed with ∇_x ε_i^T W_i ∇_x ε_i, while eq. 21 defines S_i with ∇_x ε_i^T W_i ∇_u ε_i. The correct term is likely the eq. 21 form, but the paper does not confirm it eq. 21 eq. 23.
• Constraint violation index. θ eq. 20 sums i from 0 to N+1, but h_i is only defined for i up to N. Sum over the defined nodes.
• Barrier sign. The paper writes ξ < 0 for a satisfied constraint §2.6, but eq. 22 has −μ ln ξ, which needs ξ > 0. The code uses positive margins objectives.py. Use the code's convention.
• Line-search constants. The Armijo constant, the violation threshold and the filter margin are not stated §2.5. Copy them from the code or the solver submodule, once checked.
• Clipping not in the paper. The code clips friction-cone penalties and their Hessians to [1e-6, 1e6] objectives.py objectives.py. The paper does not mention this clip, and its effect was not tested here.
• No while-loops. The paper avoids while-loops so that the solver stays differentiable through JAX §2.6. It uses a fixed count of ten candidate steps §2.5.
• Table I timings include the simulator. Table I runs the MPC inside MuJoCo XLA simulations §4.4. The paper does not separate simulator time from MPC time, so do not read Table I as a pure solver speed.

Lineage

• Builds on: the associative-scan LQR of Särkkä et al. (cited as the source of the scan in §2.4); the multiple-shooting iLQR cited as GNSM §1; feasibility-driven multiple-shooting DDP in crocoddyl §1; the filter line search of Ipopt and Perceptive MPC §2.5; the real-time iteration scheme §2.6.
• Compared against: acados (SQP) and crocoddyl (DDP) §4; batched acados with OpenMP §4.4. Also discussed: Relu-QP, the GPU ILQR variants of Lee et al. and Jeon et al., and the parallel ILQR of Frasch et al. §1.
• Superseded by: not known from this card's sources.