An efficient algorithm for solving linear equality-constrained LQR problems

Abstract

We consider finite-horizon linear-quadratic regulator (LQR) problems with arbitrary stagewise linear equality constraints. We present a two-part reduction to an unconstrained LQR problem. First, every stage is represented by an affine relation between its endpoint states. Composition of adjacent relations eliminates their shared state and is associative. An associative suffix scan therefore computes, for every node, an affine parameterization $x_i=T_i z_i+t_i$ of the states from which the remaining horizon is feasible. Second, once these parameterizations are known, every stage can be transformed independently. The original mixed constraints and the requirement that the successor lie in its feasible domain yield an affine control map $u_i=Y_i z_i+Z_i v_i+y_i$. Substitution produces an ordinary unconstrained LQR problem in the reduced state and control variables $(z_i,v_i)$, with smaller or equal dimensions, possibly varying by stage. We prove equivalence of the two problems and preservation of the standard convexity assumptions. The reduction has linear work and logarithmic depth in the horizon. Finally, we show that multipliers can be recovered by applying the same affine-relation contraction and expansion to the original KKT stationarity equations. We validate work-efficient sequential and logarithmic-depth parallel implementations against dense KKT solutions and original KKT residuals. An optimized C++/CUDA implementation exploits the varying reduced dimensions and provides a substantial end-to-end GPU speedup over sequential C++ in our experiments. We release the JAX and C++/CUDA packages under the MIT License.