iLQR (parallel)#

Solve iLQR with parallel (associative-scan) LQR subproblems

diffilqrax.parallel_ilqr.total_cost(model, params, Xs, Us)[source]#

Total cost of a state-input trajectory, evaluated in parallel over time

Return type:

Array

diffilqrax.parallel_ilqr.dynamics_defects(model, params, Xs, Us)[source]#

Dynamics defects d_t = f(x_t, u_t) - x_{t+1} [Txn], evaluated in parallel over time

Return type:

Array

diffilqrax.parallel_ilqr.value_function_grad(delta_lqr, values)[source]#

Gradient hu_t = r_t + B_tᵀ λ_{t+1} of the cost-to-go wrt u_t at the current iterate, with the costates λ_{t+1} = J_{t+1} a_t - η_{t+1} of the local LQR in deviation coordinates (δx = δu = 0, so δx_{t+1} = a_t). It vanishes at every t if and only if dL/dU does, but unlike the open-loop adjoint it stays bounded when the linearised dynamics are unstable.

Parameters:
Returns:

hu [Txm].

Return type:

Array

diffilqrax.parallel_ilqr.multiple_shooting_step(delta_lqr, gains, values)[source]#

Full step of the local LQR in deviation coordinates (δx_0 = 0, dynamics offsets a = defects), by an associative scan, and its multipliers λ_t = J_t δx_t - η_t from the value functions. Unlike the open-loop adjoint recursion λ_t = A_tᵀ λ_{t+1} + …, this does not overflow when the linearised dynamics are unstable over a long horizon.

Parameters:
Returns:

State step [(T+1)xn], input step [Txm] and multipliers λ [(T+1)xn] of the linearised dynamics (λ_{t+1} belongs to the constraint f(x_t, u_t) - x_{t+1} = 0).

Return type:

Tuple[Array, Array, Array]

diffilqrax.parallel_ilqr.merit_linesearch(model, params, Xs, Us, step, lqr, cost_init, defects, rho, alpha_init, tau=0.5, beta=0.1, alpha_min=1e-08, max_iter_linesearch=None, loop='while', n_parallel=0)[source]#

Armijo backtracking on the augmented Lagrangian merit φ = J + λᵀd + ρ/2 ‖d‖² along (X + α dX, U + α dU), with the multipliers λ of the local LQR. Its slope at α = 0 is -pᵀHp - ρ‖d‖², negative whenever the (regularised) local LQR is strictly convex, for any ρ ≥ 0. Step sizes follow diffilqrax.ilqr.armijo_linesearch(): tau * alpha_init, tau² * alpha_init, …

Parameters:
Returns:

Whether a step was accepted, the (accepted or initial) state and input trajectories, and their cost.

Return type:

Tuple[Array, Tuple[Array, Array], Array]

diffilqrax.parallel_ilqr.pilqr_solver(parallel_model, params, Us_init, max_iter=40, convergence_thresh=None, alpha_init=1.0, verbose=False, use_linesearch=True, regularise=True, kkt_tol=1e-10, mu_max=10000000000.0, max_reg_iter=20, make_psd=None, psd_delta=1e-06, rollout='feedback', defect_tol=None, merit_rho=1.0, loop='while', linesearch='backtrack', return_converged=False, return_info=False, **linesearch_kwargs)[source]#

Solve the iterative Linear Quadratic Regulator (iLQR) problem with parallel LQR subproblems.

The backward pass is an associative Riccati scan of the local LQR problem regularised with R -> R + μI (diffilqrax.plqr.plqr_backward_pass()), which keeps it an exact LQR so the scan applies. μ is adapted as in diffilqrax.ilqr.ilqr_solver(): it is increased until every Huu is positive definite, increased after a rejected step and decreased after an accepted one.

The new trajectory is found in one of two ways (rollout):

  • “feedback”: sequential non-linear closed-loop rollout u = ū + K(x - x̄) + αk, as in diffilqrax.ilqr.ilqr_solver(), with an Armijo line search on the cost. Iterates are feasible.

  • “parallel”: multiple shooting, with no sequential rollout. The states are free variables, the dynamics defects d_t = f(x_t, u_t) - x_{t+1} enter the local LQR as offsets, the step (X + α dX, U + α dU) is found by associative scans and accepted by an Armijo line search on the augmented Lagrangian merit J + λᵀd + ρ/2 ‖d‖² (merit_linesearch()). Iterates are feasible only at convergence.

Parameters:

  • parallel_model (ParallelSystem | System) – The system model. A ParallelSystem is accepted for backward compatibility; its integrators are not used.

  • params (iLQRParams) – The iLQR parameters.

  • Us_init (Array) – The initial control trajectory. The initial states are its non-linear rollout.

  • max_iter (int) – The maximum number of iterations, by default 40.

  • convergence_thresh (Optional[float]) – If given, also stop once an accepted step reduces the cost by less than this relative amount. By default None, so convergence is determined by kkt_tol.

  • alpha_init (float) – Line search scale; the first step tried is tau * alpha_init, by default 1.0.

  • verbose (bool) – Whether to print debug information, by default False.

  • use_linesearch (bool) – Whether to use the line search, by default True. If False, every step is taken with α = 1 and a step that increases the cost (or merit) beyond round-off is rejected.

  • regularise (bool) – Whether to adapt the regularisation μ, by default True. If False, μ = 0 and the solver stops at the first step that does not decrease the cost.

  • kkt_tol (float) – Converged when max|dL/dU| and max|hu| both fall below this tolerance (and the defects below defect_tol), by default 1e-10. hu is the gradient of the cost-to-go wrt the inputs of the unregularised local LQR (value_function_grad()); where the open-loop adjoint of dL/dU overflows, max|hu| decides alone (diffilqrax.ilqr.stationary()).

  • mu_max (float) – Stop when the regularisation exceeds this value without an accepted step, by default 1e10.

  • max_reg_iter (int) – Maximum number of regularisation increases per backward pass, by default 20.

  • make_psd (Optional[str]) – Convexify the local LQR before the backward pass by projection onto the PSD cone (diffilqrax.plqr.project_psd_lqr()), by default None (no projection). “R” projects the input cost only: it removes negative input curvature without retries, which helps far from a solution, and is inactive where R ≻ δI. If R is indefinite at the solution (with Huu ≻ 0 through BᵀJB) it stays active and final convergence is slower. “full” also projects Q - S R⁻¹Sᵀ and Qf, which makes every Huu positive definite but turns the step into a Gauss-Newton-like step where the state cost is locally non-convex, which can converge only linearly. Composes with regularise.

  • psd_delta (float) – Minimum eigenvalue after the PSD projection, by default 1e-6.

  • rollout (str) – “feedback” (default) or “parallel”, see above.

  • defect_tol (Optional[float]) – Converged only when max|d_t| is below this tolerance, by default kkt_tol. Only binding for rollout="parallel".

  • merit_rho (float) – Penalty ρ on the squared defects in the merit function of rollout="parallel", by default 1.0.

  • loop (str) – “while” (default) or “scan”, as in diffilqrax.ilqr.ilqr_solver(). “while” is not reverse-mode differentiable; use diffilqrax.parallel_dilqr.pdilqr() for gradients.

  • linesearch (str) – “backtrack” (default) or “parallel”, as in diffilqrax.ilqr.ilqr_solver().

  • return_converged (bool) – If True, also return whether the solution passes the convergence test, by default False.

  • return_info (bool) – If True, also return the convergence flag (as return_converged) and the number of iterations run, by default False. Takes precedence over return_converged.

  • **linesearch_kwargs – Additional keyword arguments for the line search (tau, beta, alpha_min, max_iter_linesearch, n_parallel), see diffilqrax.ilqr.armijo_linesearch().

Returns:

A tuple containing the final state trajectory, control trajectory, and the adjoint variables. Also returns the total cost of the trajectory and the cost history [max_iter], which repeats the final cost after the last iteration. With return_converged=True a fourth element, the boolean convergence flag, is appended; with return_info=True the flag and the integer number of iterations are appended.

Return type:

Tuple[Tuple[Array, Array, Array], float, Array]