Differentiable iLQR#

Module solves the differential iterative linear quadratic regulator (DiLQR) problem.

diffilqrax.diff_ilqr.add_dynamics_curvature(model, local_lqr, Xs, Us, Lambs, params)[source]#

Replace the cost Hessians of the local LQR with the Hessians of the Lagrangian.

Adds the dynamics curvature term λ_{t+1}^T ∇²f(x_t, u_t) to Q, S and R, and shifts the linear terms q and r so the LQR’s KKT residual at (Xs, Us) and its derivative wrt the parameters are unchanged. The resulting LQR has the exact KKT Jacobian of the non-linear problem, so the implicit gradients through diffilqrax.diff_lqr.dllqr() are exact.

Parameters:

  • model (System) – The system model.

  • local_lqr (LQR) – LQR approximation with offsets around (Xs, Us), from approx_lqr_offset().

  • Xs (Array) – Optimal state trajectory [(T+1)xn].

  • Us (Array) – Optimal input trajectory [Txm].

  • Lambs (Array) – Optimal adjoint trajectory [(T+1)xn].

  • params (iLQRParams) – The iLQR parameters.

Returns:

Local LQR with Lagrangian Hessians.

Return type:

LQR

diffilqrax.diff_ilqr.dilqr(model, params, Us_init, exact_hessian=True, **kwargs)[source]#

Solves the differential iLQR problem.

Parameters:

  • model (System) – The system model.

  • params (iLQRParams) – The iLQR parameters.

  • Us_init (Array) – The initial control sequence.

  • exact_hessian (bool) – If True (default), the implicit gradients use the Hessian of the Lagrangian, including the dynamics curvature λ^T ∇²f, and match direct differentiation. If False, only the cost Hessians are used (Gauss-Newton); this is cheaper but exact only for linear dynamics.

  • **kwargs – Additional keyword arguments for the iLQR solver.

Returns:

The optimized state-control trajectory [(T+1)x(n+m)].

Return type:

Array

Notes

The gradients are implicit: they are exact only if the solver has converged (see kkt_tol and max_iter of diffilqrax.ilqr.ilqr_solver()). Us_init receives no gradient, and the solver is never differentiated, so memory does not grow with the number of iterations even if Us_init depends on the parameters (e.g. a warm start). This is why the solver can run as a lax.while_loop (loop="while", the default), which is not reverse-mode differentiable itself. The parameters must enter the problem through params, not through arrays the model closes over.

diffilqrax.diff_ilqr.implicit_dilqr(model, params, solution, exact_hessian=True)[source]#

The differentiable part of dilqr(): the optimal trajectory, with implicit gradients wrt params.

Parameters:

  • model (System) – The system model.

  • params (iLQRParams) – The iLQR parameters.

  • solution (Tuple[Array, Array, Array]) – Optimal (Xs, Us, Lambs) of the solver for params; they receive no gradient.

  • exact_hessian (bool) – See dilqr(), by default True.

Returns:

The optimized state-control trajectory [(T+1)x(n+m)].

Return type:

Array