LQR (parallel)#
LQR solver using associative parallel scan [Särkkä & García-Fernández, 2023, arXiv:2104.03186].
The value functions are found with an associative scan over conditional value function elements, and the optimal trajectory with an associative scan over the closed-loop affine dynamics. The cross term S and the linear input cost r are removed first by the change of variables u = ũ - G x - g, with G = R⁻¹Sᵀ and g = R⁻¹r, which leaves an equivalent problem with S = 0, r = 0.
- diffilqrax.plqr.remove_cross_term(lqr)[source]#
Change of input variables u = ũ - G x - g, with G = R⁻¹Sᵀ and g = R⁻¹r, which removes the cross term and the linear input cost. The value functions and the adjoints are unchanged.
- diffilqrax.plqr.build_associative_riccati_elements(model)[source]#
Conditional value function elements (A, b, C, η, J) for the associative Riccati scan, with the value function V_t(x) = ½ xᵀ J_t x - η_tᵀ x.
- diffilqrax.plqr.associative_riccati_scan(model)[source]#
Obtain value function through associative riccati scan.
- diffilqrax.plqr.optimal_gains(model, etas, Js)[source]#
Optimal feedback u_t = K_t x_t + k_t given the value functions, in the original input variables.
- diffilqrax.plqr.get_dcosts(model, etas, Js, alpha=1.0)[source]#
Expected change in cost terms, as returned by the sequential backward pass: V = Σ ½ kᵀ Huu k and v = Σ kᵀ hu, with hu = -Huu k. Scale with alpha for a backtracking step.
- diffilqrax.plqr.regularise_lqr(lqr, mu)[source]#
Levenberg-Marquardt regularisation in input space, R -> R + μI, i.e. a proximal cost ½μ‖u‖² on the input step. Unlike a regulariser applied to Huu during the Riccati recursion, it does not depend on the value function, so the regularised problem is still an LQR and the associative scan solves it exactly. In the Riccati recursion it adds μI to every Huu.
- diffilqrax.plqr.project_psd(mat, delta=0.0)[source]#
Project symmetric matrices […xnxn] onto the PSD cone, clipping eigenvalues at
delta.- Return type:
- diffilqrax.plqr.project_psd_lqr(lqr, delta=1e-06, which='full')[source]#
Convexify the local LQR stage by stage. R is projected to eigenvalues ≥ δ and, if
which == "full", the Schur complement Q - S R⁻¹Sᵀ and Qf too. The fully projected problem is a strictly convex LQR, so every element of the associative scan is valid and every Huu is positive definite without retries. Linear terms are untouched, so the gradient of the local problem is unchanged, but it is no longer a Newton model where the projection is active. Withwhich == "R"only negative input curvature is removed: the projection is inactive wherever R ≻ δI, where the Newton step is kept.
- diffilqrax.plqr.plqr_backward_pass(model, mu=0.0, pd_thresh=1e-08, return_values=False)[source]#
Parallel counterpart of
diffilqrax.lqr.regularised_backward_pass(): associative Riccati scan of the problem regularised withregularise_lqr(), and its optimal gains.- Parameters:
model (
LQRParams) – LQR model parameters.mu (
float) – Regularisation strength, by default 0.0 (exact LQR solution).pd_thresh (
float) – Minimum eigenvalue of the regularised Huu for it to count as positive definite, by default 1e-8.return_values (
bool) – Also return the value functions (η, J) of the regularised problem, by default False.
- Returns:
Expected change in cost terms (V: Σ ½ kᵀ Huu k with the unregularised Huu, v: Σ kᵀ hu), the gains (u = Kx + k), and a flag that is True if the regularised Huu is positive definite at every time step and the gains are finite. With
return_values, also (η, J) [(T+1)xn], [(T+1)xnxn], whose costates λ_t = J_t x_t - η_t are the multipliers of the regularised problem.- Return type:
- diffilqrax.plqr.build_associative_lin_dyn_elements(model, gains, alpha=1.0)[source]#
Closed-loop affine dynamics elements x_{t+1} = F_t x_t + c_t under u_t = K_t x_t + α k_t. The first element absorbs the initial state.
- diffilqrax.plqr.closed_loop_scan(model, gains, alpha=1.0)[source]#
States x_1..x_T of the linear dynamics under u_t = K_t x_t + α k_t, by an associative scan.
- diffilqrax.plqr.associative_opt_traj_scan(model, etas, Js, alpha=1.0)[source]#
Obtain the optimal state trajectory through an associative scan of the closed-loop dynamics.
- Parameters:
- Returns:
States x_1..x_T, and the gains with u_t = K_t x_t + α k_t.
- Return type:
- diffilqrax.plqr.solve_plqr_swap_x0(model)[source]#
Run backward forward sweep to find optimal control freeze x0 to zero.
- Parameters:
model (
LQRParams) – LQR model parameters.- Returns:
State trajectory, control inputs, and adjoint variables.
- Return type:
Tuple[Array, Array, Array]
- diffilqrax.plqr.build_fwd_lin_dyn_elements(lqr_params, Us_init, a_term=None)[source]#
Generate sequence of elements {c} for forward integration.
- diffilqrax.plqr.parallel_forward_lin_integration(lqr_params, Us_init, a_term)[source]#
Associative scan for forward linear dynamics.
- diffilqrax.plqr.build_rev_lin_dyn_elements(lqr_params, xs_traj, us_traj)[source]#
Generate sequence of elements {c} for reverse integration of adjoints.