LQR (sequential)

Solve LQR problem via dynamic programming

diffilqrax.lqr.bmm(arr1, arr2)

Batch matrix multiplication

Return type:

Array

diffilqrax.lqr.simulate_trajectory(dynamics, Us, params, dims)

Simulate forward pass with LQR params.

Parameters:

• dynamics (Callable[[Array, Array, Array, Any], Array]) – Function of dynamics with args t, x, u, params.

• Us (Array | ndarray | bool | number | bool | int | float | complex) – Input timeseries shape [Txm].

• params (LQRParams) – Parameters containing x_init, horizon and theta.

• dims (ModelDims) – Parameters containing shape of system n, m, horizon and dt.

Returns:

State trajectory [(T+1)xn].

Return type:

Array

diffilqrax.lqr.lin_dyn_step(t, x, u, lqr)

State space linear step

Return type:

Array

diffilqrax.lqr.lqr_adjoint_pass(Xs, Us, params)

Adjoint backward pass with LQR params.

Parameters:

• Xs (Array | ndarray | bool | number | bool | int | float | complex) – State timeseries shape [(T+1)xn].

• Us (Array | ndarray | bool | number | bool | int | float | complex) – Input timeseries shape [Txm].

• params (LQRParams) – LQR state and cost matrices.

Returns:

Adjoint λs [(T+1)xn].

Return type:

Array

diffilqrax.lqr.lqr_forward_pass(gains, params)

LQR forward pass using gain state feedback.

Parameters:

• gains (Gains) – K matrices.

• params (LQRParams) – LQR state and cost matrices.

Returns:

Updated state [(T+1)xn] and inputs [Txm].

Return type:

Tuple[Array, Array]

diffilqrax.lqr.calc_expected_change(dJ, alpha=0.5)

Expected change in cost [Tassa, 2020].

Parameters:

• dJ (CostToGo) – Cost to go.

• alpha (float) – Scaling factor, by default 0.5.

Returns:

Expected change in cost.

Return type:

float

diffilqrax.lqr.regularised_backward_pass(lqr, mu=0.0, pd_thresh=1e-08, return_grad=False)

LQR backward pass with Levenberg-Marquardt regularisation of the input Hessian.

Only Huu is regularised, Huu_reg = Huu + mu B^T B. The value function update uses the full form so it stays consistent with the regularised gains, and the expected change in cost uses the unregularised Huu (Tassa et al., 2012). With mu=0 this is the exact LQR solution.

Parameters:

• lqr (LQR) – LQR parameters.

• mu (float) – Regularisation strength, by default 0.0.

• pd_thresh (float) – Minimum eigenvalue of Huu_reg for it to count as positive definite, by default 1e-8.

• return_grad (bool) – Also return max_t |hu_t|, by default False. hu_t = r_t + B_t^T (v_{t+1} + V_{t+1} a_t) is the gradient of the cost-to-go wrt u_t; it vanishes at every t if and only if dL/dU does, but v is propagated through the closed loop A + BK, so it stays bounded for unstable open-loop dynamics.

Returns:

Expected change in cost terms (V: 0.5 k^T Huu k, v: k^T hu summed over time), the optimal feedback gains, and a boolean flag that is True if Huu_reg is positive definite at every time step. With return_grad, also max_t |hu_t|.

Return type:

Tuple

diffilqrax.lqr.lqr_backward_pass(lqr)

LQR backward pass learn optimal Gains given LQR cost constraints and dynamics.

Parameters:

lqr (LQR) – LQR parameters.

Returns:

Expected change in cost terms and optimal feedback gains.

Return type:

Tuple[CostToGo, Gains]

diffilqrax.lqr.kkt(params, Xs, Us, Lambs)

Define KKT conditions for LQR problem.

Parameters:

• params (LQRParams) – LQR parameters.

• Xs (Array) – State trajectory.

• Us (Array) – Control inputs.

• Lambs (Array) – Adjoint variables.

Returns:

Gradients with respect to states, inputs, and adjoint variables.

Return type:

Tuple[Array, Array, Array]

diffilqrax.lqr.solve_lqr(params)

Run backward forward sweep to find optimal control.

Parameters:

params (LQRParams) – LQR parameters.

Returns:

State trajectory, control inputs, and adjoint variables.

Return type:

Tuple[Array, Array, Array]

diffilqrax.lqr.solve_lqr_swap_x0(params)

Run backward forward sweep to find optimal control freeze x0 to zero.

Parameters:

params (LQRParams) – LQR parameters.

Returns:

State trajectory, control inputs, and adjoint variables.

Return type:

Tuple[Array, Array, Array]