Differentiable LQR (parallel)#

Differentiable parallel (associative-scan) LQR solver.

The forward problem is solved with diffilqrax.plqr.solve_plqr(). Gradients are implicit, as in diffilqrax.diff_lqr.dlqr(): the reverse pass solves an adjoint LQR problem, either with the associative scan (adjoint="parallel") or with the sequential Riccati recursion (adjoint="sequential"). Both give the same gradients up to round-off.

diffilqrax.parallel_dlqr.pdlqr(dims, params, tau_guess=None, adjoint='parallel')[source]#

Solves the LQR problem with the associative scan, with implicit gradients wrt the LQR parameters.

Parameters:

  • dims (ModelDims) – The dimensions of the model.

  • params (LQRParams) – The parameters of the model.

  • tau_guess (Optional[Array]) – Unused; kept for API compatibility with diffilqrax.diff_lqr.dlqr().

  • adjoint (str) – Solver of the adjoint LQR problem in the reverse pass: “parallel” (associative scan, default) or “sequential” (Riccati recursion).

Returns:

Concatenated optimal state and control sequence along axis=1 [(T+1)x(n+m)].

Return type:

Array

diffilqrax.parallel_dlqr.pdllqr(dims, params, tau_star, adjoint='parallel')[source]#

Identity on a precomputed optimal trajectory, with implicit gradients wrt the LQR parameters. Parallel counterpart of diffilqrax.diff_lqr.dllqr(): the adjoints at tau_star and the adjoint LQR problem of the reverse pass are computed with associative scans (adjoint="parallel") or sequentially (adjoint="sequential").

Parameters:

  • dims (ModelDims) – The dimensions of the model.

  • params (LQRParams) – LQR whose optimum is tau_star (e.g. the local LQR of an iLQR solution).

  • tau_star (Array) – The optimal state-control trajectory. Returned unchanged; it receives no gradient.

  • adjoint (str) – “parallel” (default) or “sequential”.

Returns:

Concatenated optimal state and control sequence along axis=1 [(T+1)x(n+m)].

Return type:

Array