Differentiable LQR#
Module contains functions for solving the differential linear quadratic regulator (DLQR) problem.
- diffilqrax.diff_lqr.offset_lqr(lqr, x_stars, u_stars)[source]#
Adjust linear terms of LQR cost along nominal trajectory.
- diffilqrax.diff_lqr.build_ajoint_lqr(dims, params, tau_star, lambs, tau_bar, solve_swap_x0=<function solve_lqr_swap_x0>)[source]#
Helper function to build LQR problem with reverse gradients.
- Parameters:
dims (
ModelDims) – The dimensions of the model.params (
LQRParams) – The parameters of the model.tau_star (
Array) – The optimal state-control trajectory.lambs (
Array) – The adjoint variables.tau_bar (
Array) – The gradients with respect to tau.solve_swap_x0 (
Callable[[LQRParams],Tuple[Array,Array,Array]]) – Solver of the reverse LQR problem (see_get_qra_bar()), by default sequential.
- Returns:
The LQR parameters with reverse gradients.
- Return type:
- diffilqrax.diff_lqr.fwd_dlqr(dims, params, tau_guess)[source]#
Solves the forward differential linear quadratic regulator (DLQR) problem.
- Parameters:
- Returns:
A tuple containing the optimal state-control trajectory and the updated parameters and solution.
- Return type:
- diffilqrax.diff_lqr.rev_dlqr(dims, res, tau_bar)[source]#
Reverse mode for DLQR. Find the reverse gradient by solving LQR problem on the optimal trajectory and the gradients wrt to tau_star.
Notes
\(\overline{q}\), \(\overline{r}\), \(\overline{a}\) from solving the reverse LQR problem where \(q_{\text{rev}} = \overline{x}\), \(r_{\text{rev}} = \overline{u}\), \(a_{\text{rev}} = \overline{\lambda}\) which is set to 0.
Use
build_ajoint_lqr()to build the reverse LQR problem with gradients. Where \(\overline{c} = [\overline{q}, \overline{r}]\), \(\overline{F}\), where \(F = [A, B]\), and \(\overline{C}\) (where \(C = [Q, R]\)) as,\[\overline{C} = 0.5 \left( \overline{c} \tau_{\star}^{T} + \tau_{\star} \overline{c}^T \right)\]\[\overline{F}_t = \lambda_{\star_{t+1}} \overline{c}^T + f_{t+1} \tau_{\star_t}^T\]- Parameters:
- Returns:
The LQR parameters with reverse gradients.
- Return type:
- diffilqrax.diff_lqr.fwd_dllqr(dims, params, tau_star)[source]#
Forward pass of
dllqr(): returnstau_starunchanged, as the primal does, and saves the residuals for the reverse pass. The adjoints are evaluated attau_starwithdiffilqrax.lqr.lqr_adjoint_pass()rather than by re-solving the LQR, so the output does not depend on whether gradients are taken, even iftau_staris not exactly optimal.- Parameters:
- Returns:
The state-control trajectory, and the offset LQR parameters with the trajectory and adjoints for the reverse pass.
- Return type: