Using diffrax to integrate system dynamics for Model Predictive Control

[Hs293Go]

Diffrax docs' Kalman Filter Example encouraged me to apply diffrax to other linear-quadratic control problems, such as the following

\min_{u(t)} x_N^TQx_N + \sum_{k=0}^N x_k^T Q x_k + u_k^T R u_k\quad \text{subject to}\ \dot{x}(t) = f(x(t), u(t)),\ x(0) = x_{init}

essentially finding the optimal control $u$ over a prediction horizon $T$ discretized into $N$ segments.

Diffrax seems like a good fit for solving the IVP $\dot{x}(t) = f(x(t), u(t)),\ x(0) = x{init}$ over the prediction horizon to compute $x{0,\ldots,N}$. In this case...

• I would evaluate $f$ at fixed time points so I don't want to pay any overhead of error estimation. (It appears as low-order as Heun already has an embedded error estimator)
• The control terms $u_{0, \ldots, N}$, being my optimized variables, are known beforehand at the discretization points; I don't want pay the overhead of interpolation per interpolate_us in the KF example
• My $f$ is nonlinear in $u$ and I'm unsure if ControlTerm is useful here

Is diffrax the right tool for this job? Is yes, what is the best approach; Are there any idioms?

[patrick-kidger]

Yup, this sounds exactly like what Diffrax is for. You might also find this example a useful reference for how to insert the control input u.

Touching on your other points:

• If you don't want to pay the computational cost of using an error estimator, then feel free to use any solver + a fixed step size. Their built-in error estimators will actually be compiled out if they're not used.
• The overhead of interpolation is probably best handled by writing your own function

  def u(t):
      # lookup values of `u_i` here given `t`. Implement arbitrary behaviour for those `t` that aren't in the discretisation points.

• You don't want ControlTerm; this is used for SDEs etc. (The word "control" here is in the sense of a "controlled differential equation", e.g. the SDE-like objects dy(t) = f(y(t)) dx(t).) You just want an ODE like in the example above.

I hope that helps!