Interest in differentiable classic infinite horizon LQ control (generating LGSSM types)?

[jlperla]

In economics we heavily rely on the sorts of features in this amazing package, and I am looking forward to integrating it into teaching and research. For example, I teach a computational macroeconomics course at UBC where I would love to port everything to JAX.

However, one thing that is missing are some (rudimentary) features of classic LQ control. All discounted infinite horizon. One way to think about this is that it takes in a cost function and your LGSSM primitives for the "law of motion" and it returns a control, which you can then use to generate the complete LGSSM setup for all of your downstream features. These are well-established methods and algorithms. While in principle a separate control-theory jax library could be built which extends dynamax, no such package exists - and the minimal set of useful features is actually pretty small.

Would there be interest in me preparing a full PR adding a focused, JAX-native infinite-horizon LQR transformation for time-invariant ParamsLGSSMs?

---

As a strawman: I have in mind an interface approximately like the following:

Consider the centered LGSSM in your existing canonical form

$$ x_{t+1} = F x_t + B u_t + w_t, \qquad y_t = H x_t + D u_t + v_t, $$

and the stationary infinite-horizon objective

$$ \mathbb{E}\sum_{t=0}^{\infty}\beta^t \left(x_t^\top Q x_t + 2u_t^\top N x_t+ u_t^\top R u_t\right) $$

A possible interface would be

closed_loop_params, solution = lgssm_lqr(params, cost)

where cost could be a pytree such as

class StationaryLQCost(NamedTuple):
    Q: Array
    R: Array
    N: Array | None = None
    beta: Scalar

Both params and the returned closed_loop_params would use the existing public type dynamax.linear_gaussian_ssm.ParamsLGSSM

rather than introducing another state-space model type.

For the optimal policy $u_t = -Kx_t$, closed_loop_params would update params to contain

$$ F_{\mathrm{cl}} = F - BK, \qquad H_{\mathrm{cl}} = H - DK, $$

with the control input absorbed and the original initial distribution and noise covariances preserved. The result could therefore be passed directly to existing functions such as lgssm_filter, lgssm_smoother, and lgssm_joint_sample.

solution would contain the Riccati value matrix $P$, feedback gain $K$, and residual diagnostics. Where the P and the K would support co-tangents. The tangents would support all of the things in the StationaryLQCost and ParamsLGSSM.

Alternatively, the lgssm_lqr could just return the K and the user could manually just recreate their own LGSSM parameters where they put in the new control.

The DARE solve would support jax.jit and use an implicit custom JVP/VJP based on the converged Riccati equation rather than differentiating through the solver iterations. Only jax algorithms, and no wrapping of scipy/slicot/etc. This would allow downstream Dynamax computations to differentiate through the LQGR problem with respect to the original dynamics and cost parameters. This would allow you, for example, to do sensitivity analysis or bayesian sampling on the law-of-motion parameters, the control cost parameters, etc. while using all of your downstream kalman filters for likelihoods on observables or differentiable simulations.

Would this fit Dynamax's scope and preferred functional API? If so, I will prepare a complete PR for review with unit tests, implicit AD-rules, and documentation.