Define observation/control alignment for discrete-time models

[MatthieuDarcy]

Add observation_control_alignment for discrete-time Simulator (#312).

Mathematical summary

Implements the same_time vs previous time distinction. same_time (default, preserves the existing behavior) implements:

$$ \begin{aligned} x_0 &\sim p_0 \\ y_0 &\sim p(\cdot |x_0, u_0) \end{aligned} $$

and subsequently

$$ \begin{aligned} x_{k+1} &\sim p(\cdot | x_k, u_k)\\ y_k &\sim p(\cdot| x_k, u_k) \end{aligned} $$

previous_time implements

$$ \begin{aligned} x_0 &\sim p_{0}\\ x_{k+1} &\sim p(\cdot | x_k, u_k) \\ y_{k+1} &\sim p(\cdot| x_{k+1}, u_k) \end{aligned} $$

Summary of changes

previous_transition contract

Adds an explicit property of DynamicalModel defined at initialization called observation_control_alignment: Literal[ "same_time", "previous_transition" ] = "same_time"

When defined as previous_transition, this results in the following behavior

1. $y_0$ cannot be sampled. Hence observations $y$ is of size $T-1$ while the states $x$ are of size $T$.
2. I added the controls $u$ in the returnedSimulatedResults. This is very practical when doing MPC. This leads to a similar behavior as observations: when using "same_time" it is of size $T$, when using "previous_time" it is of size $T-1$. Only available when running DiscreteTimeSimulator.

This means that states and times are of different sizes to controls and observations (specifically size $T$ vs $T-1$) when using previous_transition.

1. MPPI is updated to use this convention in Fixed MPPI #348

Updated DiscreteControlLoopSimulator

DiscreteControlLoopSimulator now deliberately always uses the previous-transition convention, independent of dynamics.observation_control_alignment:

$$ \begin{aligned} x_0 &\sim p_0, \qquad \hat p_0 = p_0, \\ (u_k, s_{k+1}) &= \pi(\hat p_k, t_k, t_{k+1}, s_k), \\ x_{k+1} &\sim p(\cdot \mid x_k, u_k, t_k, t_{k+1}), \\ y_{k+1} &\sim p(\cdot \mid x_{k+1}, u_k, t_{k+1}). \end{aligned} $$

Previously, it was assumed that the observation function could support u=None. A closed-loop result therefore has $T$ states, times, and beliefs, and $T-1$ aligned observations and controls. This matches the previous_transition convention and avoids imposing an implicit requirement that the observation model accept u=None or be control-independent, but is change in the behavior.

Future work

1. previous_transition only works for simulation, no conditioning and not filtering. Worth delegating to a seperate PR?
2. Adding controls in SimulatedResults for ODE/SDE/Continuous time. It might also be worthwile to expand support for continuous control beyond zero-order hold.
3. LTI_discrete does not accept observation_control_alignment as a keyword, I did not touch it but worth looking at it once we are satisfied with the contract.

[mattlevine22]

"DiscreteControlLoopSimulator now deliberately always uses the previous-transition convention, independent of dynamics.observation_control_alignment:"

• Hmm, seems better to keep this aligned with reality somehow. I suppose if observation_control_alignment is not specified by a user, then we are free to choose the "right" default; if a user does specify it, we should error if they made an incompatible choice.

[mattlevine22]

"previous_transition only works for simulation, no conditioning and not filtering. Worth delegating to a seperate PR?"

• Yes, that seems like the right way to do it. Could you create an issue and then make it a sub-issue of Enforce universal consistency of time-handling #272?

[mattlevine22]

"Adding controls in SimulatedResults for ODE/SDE/Continuous time. It might also be worthwile to expand support for continuous control beyond zero-order hold."

• Yes, this is basically Add support for flexible control u interpolation schemes #96

[mattlevine22]

Overall, seems solid to me.

Main asks are:

1. Clean up the latex math and make sure it renders in markdown (I didn't read all of it because of that) and is fully specified including all zero cases (e.g. s_0 (even with a note that pi auto-initializes s0 in practice))

2. Try to return useful time-indexes that offer some redundancy to users in SimulatedResult (rather than needing to write things like t[1:])

3. We should always at least issue a warning if we override a user's choice; here w.r.t. dynamics.observation_control_alignment.

4. Look to see if it is easy for the user to simply supply t0 instead of appending it to predict_times in previous-mode; my guess is not, which is fine. May just be better dealt all at once with #272. Would be great to get your thoughts on this though.

5. Add some is_finite checks to the tests (esp hierarchical)...we've gotten fooled in the past and had passing tests that were producing NaNs.

NB: there are a lot of new tests, which we've been resistant of. However, I think they seem reasonable, and in particular, when messing with indexing I feel better about more guardrails and edge case tests, so I'm happy with these.

[mattlevine22]

This isn't rendering on the GitHub or cursor markdown viewers. Mind massaging this somehow to get it to appear? sometimes you need to wrap long latex lines in ..., not really sure though.

[mattlevine22]

Would rather us give an error if the user selects the wrong dynamics.observation_control_alignment. We should allow them to omit it and then default to the appropriate choice at the right time...hopefully not easier said than done?

[mattlevine22]

I don't love that the user has to know to do t[1:]. I've messed stuff like that up so many times over the years.

What do you think about having explicit fields for the observation times, state times, and control times?

@DanWaxman I know this may be a bit wasteful, but (a) maybe it can be optional to return them and (b) I really think ease of use here may be worth it. At least times are scalar so it really just scales with length of the series (not dimensionality of the system).

[mattlevine22]

Should still define all the 0 variables so that this system of relations is fully specified. I'll try to flag that where I see it, but maybe ask uncle AI to find those spots.

[mattlevine22]

After thinking on it, perhaps for now leave this as is with same_time default; When doing control stuff, if it is same_time, issue a warning that you are changing it to previous.

[mattlevine22]

I would assert that things are not NaNs. That has bitten us before.

[MatthieuDarcy]

Added a finiteness check

[mattlevine22]

I'd assert more clearly what these should be, e.g. finite arrays

[mattlevine22]

again, maybe assert these are finite arrays

[mattlevine22]

might even be worth a helper that is assert_finite(array, shape=Optional)

[MatthieuDarcy]

"DiscreteControlLoopSimulator now deliberately always uses the previous-transition convention, independent of dynamics.observation_control_alignment:"

* Hmm, seems better to keep this aligned with reality somehow. I suppose if `observation_control_alignment` is not specified by a user, then we are free to choose the "right" default; if a user does specify it, we should error if they made an incompatible choice.

After thinking about this, I think that we can support both conventions for online control (and keep the default same_time) at the cost of deviating from the existing convention we introduced. Here is my summary:

same_time convention

Notation:

• $\tilde{p}_k \approx p(x_k \mid y_0, \dots, y_{k-1},\ u_0, \dots, u_{k-1})$is the predicted distribution.
• $\hat{p}_k \approx p(x_k \mid y_0, \dots, y_k,\ u_0, \dots, u_k)$ is the filtered distribution.

In this convention, the control $u_k$ drives the transition into the next state $x_{k+1}$ and generates the observation $y_k$ (they are aligned).

$$ \begin{aligned} &x_0 \sim p_0, \quad \tilde{p}_0 = p_0 \\ &u_k = \pi(\tilde{p}_k) \\ &y_k \sim p(y_k \mid x_k, u_k) \\ &\hat{p}_k = \text{FilterUpdate}(y_k, \tilde{p}_k, u_k) \\ &x_{k+1} \sim p(x_{k+1} \mid x_k, u_k) \\ &\tilde{p}_{k+1} = \text{PredictionUpdate}(\hat{p}_k, u_k) \end{aligned} $$

In this case, the number of observations matches the number of states.

Under previous_transition convention

This is the closest to what we decided originally, but doesn't align with Dynestyx default for offline control.

• $\tilde{p}_k \approx p(x_k \mid y_1, \dots, y_{k-1},\ u_0, \dots, u_{k-1})$ is the predicted distribution.
• $\hat{p}_k \approx p(x_k \mid y_1, \dots, y_k,\ u_0, \dots, u_{k-1})$ is the filtered distribution.

In this convention, the control $u_k$ drives the transition into the next state $x_{k+1}$ and generates the observation $y_{k+1}$. Hence, $y_0$ never exists.

$$ \begin{aligned} &x_0 \sim p_0, \quad \tilde{p}_0 = \hat{p}_0 = p_0 \\ &u_k = \pi(\hat{p}_k) \\ &x_{k+1} \sim p(x_{k+1} \mid x_k, u_k) \\ &\tilde{p}_{k+1} = \text{PredictionUpdate}(\hat{p}_k, u_k) \\ &y_{k+1} \sim p(y_{k+1} \mid x_{k+1}, u_k) \\ &\hat{p}_{k+1} = \text{FilterUpdate}(y_{k+1}, \tilde{p}_{k+1}, u_k) \end{aligned} $$

Here, the number of observations does not match the number of states, but it does match the number of controls. The number of controls is exactly the number of transitions (i.e., one less than the number of states).

I made a drawing because this stuff can be confusing.

Thoughts?

[mattlevine22]

What you wrote above looks right to me and sounds consistent with this PR.

Sounds like the only update is that you are proposing to be more accommodating to time-alignment in the closed loop controller?

Seems good to me.

[mattlevine22]

RIP now it gets mad: "The following macros are not allowed: operatorname"

[MatthieuDarcy]

Summary of changes

Open-loop control now supports both the same_time and previous_transition conventions (see below for a definition of both). DiscreteControlLoopSimulator (closed-loop control) now supports the previous_transition convention and uses it by default, warning the user if they do not specify the model's observation_control_alignment and raising an error if they specify same_time. There are some difficulties in supporting same_time.

SimulatedResult now has obs_times and ctrl_times fields, which are populated when relevant (at the suggestion of @mattlevine22). This has a small cost (we have to store these) but makes the results more useful and eases some of the technicalities (see below).

A unified principle is that states and times should align: no giving $N$ times and getting $N-1$ states (but overvations/controls might differ).

I have simplified and reworked some of the tests to properly verify the correctness of the different conventions.

In open-loop control

same_time

$$ \begin{aligned} & \text{Times} = [t_0, \dots, t_{N}]\\ &\text{Controls} = [u_0, \dots, u_N]\\ &x_0 \sim p_0\\ &\text{for $k = 0,\dots N-1$}:\\ &y_{k} \sim p(y_k \mid x_k, u_k, t_k) \quad \text{emit observation}\\ &x_{k+1} \sim p(x_{k+1} \mid x_k, u_k, t_k, t_{k+1}) \quad \text{State transition}\\ &y_{N} \sim p(y_N \mid x_{N}, u_{N}, t_N) \quad \text{Final observation} \end{aligned} $$

Here there are $N+1$ states, observations and controls, all living on the grid $[t_0, \dots, t_N]$. The final control $u_N$ is only ever used for the observation.

previous_transition

$$ \begin{aligned} & \text{Times} = [t_0, \dots, t_{N}]\\ &\text{Controls} = [u_0, \dots, u_{N-1}]\\ &x_0 \sim p_0\\ &\text{for $k = 0,\dots N-1$}:\\ &x_{k+1} \sim p(x_{k+1} \mid x_k, u_k, t_k, t_{k+1}) \quad \text{State transition}\\ &y_{k+1} \sim p(y_{k+1} \mid x_{k+1}, u_k, t_{k+1}) \quad \text{emit observation} \end{aligned} $$

We have $N+1$ states (on the time grid $[t_0, \dots, t_N]$).
We have $N$ controls (on the time grid $[t_0, \dots, t_{N-1}]$).
We have $N$ observations (on the time grid $[t_1, \dots, t_N]$).

In closed-loop control

same_time convention

• $\tilde{p}_k \approx p(x_k \mid y_0, \dots, y_{k-1},\ u_0, \dots, u_{k-1})$ is the predicted distribution.
• $\hat{p}_k \approx p(x_k \mid y_0, \dots, y_k,\ u_0, \dots, u_k)$ is the filtered distribution.

The control $u_k$ drives the transition into the next state $x_{k+1}$ and generates the observation $y_k$ (they are aligned).

$$ \begin{aligned} & \text{Times} = [t_0, \dots, t_{N}]\\ &x_0 \sim p_0, \quad \tilde{p}_0 = p_0 \quad \text{Initialization step} \\ &\text{for $k = 0,\dots N-1$}:\\ &u_{k}, s_{k+1} =\pi(\tilde{p}_{k}, t_{k}, t_{k+1}, s_k) \quad \text{Select the control}\\ &y_{k} \sim p(y_k \mid x_k, u_k, t_k) \quad \text{emit observation}\\ &\hat{p}_k = \text{FilterAnalysis}(y_k, \tilde{p}_k, u_k) \quad \text{Update the filtering distribution using the observation} \\ &x_{k+1} \sim p(x_{k+1} \mid x_k, u_k, t_k, t_{k+1}) \quad \text{State transition} \\ &\tilde{p}_{k+1} = \text{PredictionUpdate}(\hat{p}_k, u_k) \quad \text{Predict the filtering distribution} \end{aligned} $$

Here the policy uses the predicted distribution.
We have $N+1$ states (on the time grid $[t_0, \dots, t_N]$).
We have $N$ controls (on the time grid $[t_0, \dots, t_{N-1}]$); we can never produce a control at time $t_N$.
We have $N$ observations (on the time grid $[t_0, \dots, t_{N-1}]$); we can never produce an observation at time $t_N$, because $u_N$ would need $t_{N+1}$.

This is not implemented because Cuthbert's Filter does not expose separate predict and filter_update steps. This seems possible, but I did not want to overload this PR. I added dummy functions that show how this could be done, but this will raise an error if a user tries to do this.

previous_transition (the default, with a warning)

• $\hat{p}_k \approx p(x_k \mid y_1, \dots, y_k,\ u_0, \dots, u_{k-1})$ is the filtered distribution.

$$ \begin{aligned} & \text{Times} = [t_0, \dots, t_{N}]\\ &x_0 \sim p_0, \quad \hat{p}_0 = p_0 \quad \text{Initialization step} \\ &\text{for $k = 0,\dots N-1$}:\\ &u_{k}, s_{k+1} =\pi(\hat{p}_{k}, t_{k}, t_{k+1}, s_k) \quad \text{Select the control}\\ &x_{k+1} \sim p(x_{k+1} \mid x_k, u_k, t_k, t_{k+1}) \quad \text{State transition} \\ &y_{k+1} \sim p(y_{k+1} \mid x_{k+1}, u_k, t_{k+1}) \quad \text{emit observation}\\ &\hat{p}_{k+1} = \text{FilterUpdate}(y_{k+1}, \hat{p}_k, u_k) \quad \text{Update the filtering distribution using the observation} \end{aligned} $$

Here the policy uses the filtering distribution. In this convention, $y_0$ never exists.
We have $N+1$ states (on the time grid $[t_0, \dots, t_N]$).
We have $N$ controls (on the time grid $[t_0, \dots, t_{N-1}]$); we can never produce a control at time $t_N$.
We have $N$ observations (on the time grid $[t_1, \dots, t_N]$). We do not produce an observation at time $t_0$, because $u_0$ is used for $x_1$ and $y_1$.

Sharp edges

The new previous_transition convention removes the artificial first observation $y_0$ (which used no control). However as a result, the following simple model becomes somewhat tricky to implement

I think one way of doing this could be to add an explicit keyword observe_y0 and then implement the following loop:

@mattlevine22 can you let me know if you think this makes sense? (also I tried fixing all the rendering issues).

[mattlevine22]

Is same_time the default for open-loop?

[MatthieuDarcy]

Yes, to preserve the old behavior.

[mattlevine22]

Overall, my recommendation is to accept this PR as is and move on.

• I don't think it breaks any existing stuff...and if we discover that it does, we will just fix it.
• observe_y0 boolean property of DynamicalModel is a good idea, but better as a new PR. I think this PR is about as big as I can generally handle reviewing haha.

Notes:

• I suppose that IF you set observe_y0=True, then there should be a warning that says: "y0 will be generated with a 0 control value (ignore this warning if your observation model does not depend on control inputs".
• Doing observe_y0 will also require dealing with more branches of the code (e.g., asserting it is True in all the Filters/Smoothers/OpenLoopSimulators unless you want to start supporting it True...which is even more code).
• DynamicalModel has a t0 property...is this being considered at all? It is a bit dicey, and needs to become a more powerful grounding property soon (perhaps not in this PR other than maintaining current internal consistency).

@DanWaxman could you please look at this briefly and decide whether you'd like us to (a) do some more careful reviewing or (b) go for it?

I read the code in detail in my last iteration, but now it is a bit tough for me to follow what has changed (@MatthieuDarcy this is an example where fewer targeted commits helps a lot in the review process)

[MatthieuDarcy]

Yes my apologies for this, I know this is too much but it turned out to be a bit of a rabbit hole.

@mattlevine22 @DanWaxman IF you want a more careful review, I can break it up into smaller PRs, doing one for open loop and one for closed loop (and removing some of the superfluous elements), but this will take me a bit of time.

[mattlevine22]

My gut is that this is good enough but want to get Dan's take

[DanWaxman]

This seems reasonable to me overall. I left comments, mostly about docs. But also:

• does it make sense to use an enum instead of a string?
• I think the belief stuff can be implemented for the EnKF already, it gives its predicted states (as it is necessary for the EnRTS algorithm)

[DanWaxman]

I feel okay with just having this still be in type checking instead of a check

[MatthieuDarcy]

Personally I like it because it gives an explicit error message for when a user gives a bad policy

[DanWaxman]

Okay, I don't feel super strongly about it (but would remove the jnp.asarray call, that should happen way earlier... also lets u be typed better in this function)

[MatthieuDarcy]

yes that part I completely agree with and is removed

[DanWaxman]

Could you put in an array-like type for u then?

[MatthieuDarcy]

Yes, I had to add Distribution to allow for the Distribution check

[DanWaxman]

Can you write what the "predicted state" and "filtered state" mean here? Like $x_{k|k}$ etc.

[MatthieuDarcy]

Let me know if you find this sufficiently clear. I'm not as familiar as you are regarding the accepted terminology/notation

[DanWaxman]

Yup looks good!

[DanWaxman]

I made one more reply, commenting about a missing type for control in the validation function. Otherwise, this looks good to me. Thank you Matthieu!!

[mattlevine22]

Looks great, thank you for the hard work on this one!