iLQR (sequential)#

iterative LQR solver

diffilqrax.ilqr.bounded_loop(cond_fun, body_fun, init, max_steps, loop='while')[source]#

Apply body_fun while cond_fun(carry) holds, at most max_steps times.

Parameters:

  • cond_fun (Callable) – Maps the carry to a boolean scalar; the loop continues while it is True.

  • body_fun (Callable) – Maps the carry to the next carry.

  • init (Any) – Initial carry.

  • max_steps (int) – Maximum number of applications of body_fun.

  • loop (str) – “while” (default) runs a lax.while_loop, which executes only the steps needed (under vmap, until the last batch element stops). It supports forward-mode but not reverse-mode differentiation. “scan” runs a fixed-length lax.scan that skips finished steps with lax.cond; it is reverse-mode differentiable, but stores every step and under vmap executes all max_steps steps.

Returns:

The final carry.

Return type:

Any

Try step sizes α = alpha_first, tau * alpha_first, … until trial accepts one.

The search stops at the first accepted α, after the first rejected α below alpha_min, or after max_steps trials. With n_parallel > 0 the first n_parallel step sizes are tried at once (vmap over α) and the largest accepted one is kept, which gives the same result as trying them in turn; only if none is accepted does the search continue one α at a time.

Parameters:

  • trial (Callable) – Maps α to (accepted, Xs, Us, cost) of the trial step.

  • init (Tuple[Array, Array, Array]) – (Xs, Us, cost) returned if no step is accepted.

  • alpha_first (float) – First step size tried.

  • max_steps (int) – Maximum number of step sizes tried.

  • tau (float) – Reduction factor for α, by default 0.5.

  • alpha_min (float) – Minimum α, by default 1e-8.

  • loop (str) – “while” or “scan”, see bounded_loop(), by default “while”.

  • n_parallel (int) – Number of step sizes tried at once, by default 0 (one at a time).

Returns:

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

Return type:

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

diffilqrax.ilqr.increase_regularisation(delta, mu)[source]#

Increase Levenberg-Marquardt regularisation: δ ← max(2, 2δ), μ ← max(μ_min, μδ).

Parameters:

  • delta (Array) – Regularisation scaling factor.

  • mu (Array) – Regularisation strength.

Returns:

Updated (delta, mu).

Return type:

Tuple[Array, Array]

diffilqrax.ilqr.decrease_regularisation(delta, mu)[source]#

Decrease Levenberg-Marquardt regularisation: δ ← min(1/2, δ/2), μ ← μδ.

Parameters:

  • delta (Array) – Regularisation scaling factor.

  • mu (Array) – Regularisation strength.

Returns:

Updated (delta, mu), with mu set to zero once it falls below MU_MIN.

Return type:

Tuple[Array, Array]

diffilqrax.ilqr.stationary(hu_max, dLdU_max, tol)[source]#

Convergence test of the iLQR solvers: max|dL/dU| and max|hu| below tol, or max|hu| alone where dL/dU is not finite.

Both vanish exactly at a stationary point, but away from it dL/dU = hu + B^T (λ - v) carries hu back through the open-loop dynamics: with unstable dynamics a small max|hu| leaves a much larger max|dL/dU|, so both are required. Over long unstable horizons the open-loop adjoint λ overflows and dL/dU is never finite; there the bounded max|hu| decides alone.

Parameters:
Returns:

Boolean convergence flag.

Return type:

Array

diffilqrax.ilqr.kkt_input_residual(lqr_lin)[source]#

Gradient of the total non-linear cost wrt the inputs, i.e. the dL/dU KKT residual.

The adjoint is propagated through the linearised dynamics so that dL/dX = 0 and dL/dλ = 0 hold by construction along a rollout, leaving dL/dU as the only non-zero KKT condition.

Parameters:

lqr_lin (LQR) – LQR approximation without offsets (from approx_lqr()), i.e. q=Cx, r=Cu, qf=Cfx.

Returns:

dL/dU residual [Txm].

Return type:

Array

diffilqrax.ilqr.sum_cost_to_go(x)[source]#

Sum linear and quadratic components of cost-to-go tuple.

Parameters:

x (CostToGo) – The cost-to-go tuple containing linear and quadratic components.

Returns:

The sum of the linear and quadratic components.

Return type:

Array

diffilqrax.ilqr.compute_lqr_params(model, Xs, Us, params, include_dyn_offset=False, include_cost_offset=False)[source]#

Compute LQR parameters with optional offset for dynamics and cost.

Parameters:

  • model (Any) – The system model.

  • Xs (Array) – The state trajectory.

  • Us (Array) – The control trajectory.

  • params (iLQRParams) – The iLQR parameters.

  • include_dyn_offset (bool) – Whether to include dynamic offset, by default False.

  • include_cost_offset (bool) – Whether to include cost offset, by default False.

Returns:

The LQR parameters.

Return type:

LQR

diffilqrax.ilqr.approx_lqr_offset(model, Xs, Us, params)[source]#

Approximate non-linear model as LQR with offset.

Parameters:
Returns:

The LQR parameters.

Return type:

LQR

diffilqrax.ilqr.approx_lqr(model, Xs, Us, params)[source]#

Approximate non-linear model as LQR without offset.

Parameters:
Returns:

The LQR parameters.

Return type:

LQR

diffilqrax.ilqr.approx_lqr_dyn(model, Xs, Us, params)[source]#

Approximate non-linear model as LQR with dynamic offset.

Parameters:
Returns:

The LQR parameters.

Return type:

LQR

diffilqrax.ilqr.ilqr_simulate(model, Us, params)[source]#

Simulate forward trajectory and cost with nonlinear parameters.

Parameters:

  • model (System) – The system model.

  • Us (Array) – The control trajectory.

  • params (iLQRParams) – The iLQR parameters.

Returns:

A tuple containing the updated state and control trajectory, and the total cost of the trajectory.

Return type:

Tuple[Tuple[Array, Array], float]

diffilqrax.ilqr.ilqr_forward_pass(model, params, Ks, Xs, Us, alpha=1.0)[source]#

Perform a forward pass of the iterative Linear Quadratic Regulator (iLQR) algorithm.

Parameters:

  • model (System) – The nonlinear system model.

  • params (iLQRParams) – The parameters of the system.

  • Ks (Gains) – The gains obtained from the LQR controller.

  • Xs (Array) – The target state trajectory.

  • Us (Array) – The control trajectory.

  • alpha (float) – The linesearch parameter, by default 1.0.

Returns:

A tuple containing the updated state trajectory and control trajectory, and the total cost of the trajectory.

Return type:

Tuple[Tuple[Array, Array], float]

diffilqrax.ilqr.ilqr_solver(model, params, Us_init, max_iter=100, convergence_thresh=None, alpha_init=1.0, verbose=False, use_linesearch=True, kkt_tol=1e-10, mu_max=10000000000.0, max_reg_iter=20, loop='while', linesearch='backtrack', return_converged=False, return_info=False, **linesearch_kwargs)[source]#

Solve the iterative Linear Quadratic Regulator (iLQR) problem.

The non-linear system is approximated to be a local LQR problem. The optimal control gains are iteratively added to the x and u perturbations until the KKT conditions are satisfied. The input Hessian is regularised with an adaptive Levenberg-Marquardt term until positive definite, and the step size is chosen by an Armijo backtracking line search (see armijo_linesearch()). A rejected step increases the regularisation and retries; an accepted step decreases it.

Parameters:

  • model (System) – The system model.

  • params (iLQRParams) – The parameters of the system.

  • Us_init (Array) – The initial control trajectory.

  • max_iter (int) – The maximum number of iterations, by default 100. The default step schedule (first trial α = tau * alpha_init = 0.5) converges linearly at first, so harder problems may need more.

  • convergence_thresh (Optional[float]) – Deprecated and ignored; 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 line search for the forward pass, by default True. If False, every step is taken with α = 1 and a step that increases the cost is rejected, increasing the regularisation instead.

  • kkt_tol (float) – Converged when max|dL/dU| (kkt_input_residual()) and max|hu| both fall below this tolerance, by default 1e-10. Use ~1e-4 for float32. hu is the gradient of the cost-to-go wrt the inputs from the unregularised backward pass (diffilqrax.lqr.regularised_backward_pass()); it vanishes if and only if dL/dU does, and stays finite when the open-loop adjoint of dL/dU overflows (unstable dynamics, long horizons), where the test falls back to max|hu| alone (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.

  • loop (str) – How the iterations, regularisation retries and line search are looped (see bounded_loop()). “while” (default) runs only the iterations needed, also under vmap, but is not reverse-mode differentiable; forward mode and the implicit gradients of diffilqrax.diff_ilqr.dilqr() work. “scan” runs max_iter iterations that skip finished work with lax.cond, so reverse mode differentiates the unrolled iterations and stores every one of them.

  • linesearch (str) – “backtrack” (default) tries the step sizes one at a time. “parallel” rolls out n_parallel (default 8) step sizes at once and keeps the largest accepted one, which takes the same step but needs fewer sequential rollouts (useful on GPU). Ignored if use_linesearch is False.

  • 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 armijo_linesearch() (tau, beta, alpha_min, max_iter_linesearch, n_parallel).

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 + 1], 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]

Notes

To differentiate the solution with loop="while", use diffilqrax.diff_ilqr.dilqr(). The parameters must then enter the problem through params (e.g. params.theta), not through arrays the model closes over, so that no tangents reach the loop.

diffilqrax.ilqr.armijo_linesearch(update, Ks, Xs_init, Us_init, alpha_init, cost_init, expected_dJ, tau=0.5, beta=0.1, alpha_min=1e-08, max_iter_linesearch=None, loop='while', n_parallel=0)[source]#

Armijo backtracking line search.

Step sizes α = tau * alpha_init, tau^2 * alpha_init, … are tried until the cost reduction exceeds beta times the expected reduction, or α drops below alpha_min. The expected reduction is the textbook quadratic model ΔJ(α) = -(α k^T hu + ½ α² k^T Huu k) (diffilqrax.lqr.calc_expected_change()). When the expected reduction is below the floating point resolution of the cost the comparison is meaningless, so the step is accepted provided the cost does not increase beyond that resolution.

References

Tassa, Y., Erez, T., & Todorov, E. (2012). Synthesis and stabilization of complex behaviors through online trajectory optimization. IEEE. https://doi.org/10.1109/iros.2012.6386025

Parameters:

  • update (Callable) – Rollout function which returns new Xs, Us and cost.

  • Ks (Gains) – Gains obtained from the LQR controller.

  • Xs_init (Array) – State trajectory.

  • Us_init (Array) – Input trajectory.

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

  • cost_init (float) – Cost of initial trajectory.

  • expected_dJ (CostToGo) – Expected change in cost terms from the LQR backward pass.

  • tau (float) – Reduction factor for alpha, by default 0.5.

  • beta (float) – Armijo sufficient decrease constant, by default 0.1.

  • alpha_min (float) – Minimum alpha value, by default 1e-8.

  • max_iter_linesearch (Optional[int]) – Maximum number of rollouts. Defaults to the number needed to reach alpha_min.

  • loop (str) – “while” (default) or “scan”, see bounded_loop().

  • n_parallel (int) – Number of step sizes rolled out at once, by default 0 (one at a time). The accepted step is the same either way (see backtracking_search()); parallel rollouts suit GPUs, at the price of up to n_parallel rollouts per call.

Returns:

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

Return type:

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

diffilqrax.ilqr.linesearch(update, Ks, Xs_init, Us_init, alpha_init, cost_init, expected_dJ, beta=0.5, max_iter_linesearch=12, tol=0.1, alpha_min=1e-06, verbose=False)[source]#

Implementation of the line search backtracking algorithm for iLQR algorithm.

Each iteration of the line search algorithm performs a forward pass with the new control inputs defined by the gains and an alpha value. The change in cost is compared to the expected change in cost and the alpha is selected based on the ratio of the two values above a given tolerance. Otherwise, the alpha value is reduced by a factor of beta.

References

Tassa, Y., Erez, T., & Todorov, E. (2012). Synthesis and stabilization of complex behaviors through online trajectory optimization. IEEE. https://doi.org/10.1109/iros.2012.6386025

Parameters:

  • update (Callable) – Rollout function which returns new Xs, Us and cost.

  • Ks (Gains) – Gains obtained from the LQR controller.

  • Xs_init (Array) – State trajectory.

  • Us_init (Array) – Input trajectory.

  • alpha_init (float) – Initial alpha value.

  • cost_init (float) – Cost of initial trajectory.

  • expected_dJ (CostToGo) – Expected change in cost from LQR controller.

  • beta (float) – Reduction factor for alpha, by default 0.5.

  • max_iter_linesearch (int) – Maximum iterations of linesearch, by default 12.

  • tol (float) – Tolerance of ratio of actual to expected cost change to accept alpha value, by default 0.1.

  • alpha_min (float) – Minimum alpha value, by default 1e-6.

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

Returns:

Returns the updated state and control trajectories, the alpha value, the total cost, and the cost history.

Return type:

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