Update to filter update and

Make filter_obj/filter_config mutually exclusive; add filter_state_dist

compute_cuthbert_filter_update now takes filter_config: BaseFilterConfig |
None and requires exactly one of filter_config/filter_obj, matching how
DiscreteControlLoopSimulator.simulate() already calls it with a pre-built
filter_obj.

Also add filter_state_dist(state) -> provides a numpyro distribution from the states. This will try to infer the right distribution given the state (mean, cov -> Gaussian, particles -> weighted/uniform particles.

Modified the control loop to take this into account as well as the MPPI definition and the example notebooks

Author: MatthieuDarcy

docs/tutorials/control/controller_demo.ipynb

@@ -4,6 +4,7 @@
     ControlledSimulatedResult,
     DiscreteControlLoopSimulator,
     PolicyCallable,
+    filter_state_dist,
     filter_state_mean,
 )
 from dynestyx.control.mppi import MPPI
@@ -13,5 +14,6 @@
     "DiscreteControlLoopSimulator",
     "MPPI",
     "PolicyCallable",
+    "filter_state_dist",
     "filter_state_mean",
 ]

docs/tutorials/control/mpc_demo.ipynb

@@ -21,10 +21,14 @@
 cuthbert's `Filter.filter_prepare`/`filter_combine` primitives one step at a
 time instead of over a whole pre-supplied trajectory. This works for any
 filter family exposed there (`KFConfig`, `EKFConfig`, `EnKFConfig`,
-`PFConfig`); the belief `x_hat` returned each step is therefore whatever
-state type that family produces (e.g. a Kalman-family state with a `.mean`
-property, or a `ParticleFilterState` with `.particles`/`.log_weights`) --
-see `filter_state_mean` below for a family-agnostic point estimate.
+`PFConfig`); the raw belief each step produces is family-specific (e.g. a
+Kalman-family state with a `.mean`/`.chol_cov`, or a `ParticleFilterState`
+with `.particles`/`.log_weights`), so before it's handed to `control_policy`
+it's converted via `filter_state_dist` into a family-agnostic NumPyro
+`Distribution` (`MultivariateNormal` for the Gaussian families,
+`WeightedParticles` for `PFConfig`) -- a policy can call `.mean` for a point
+estimate, or use the full distribution (e.g. `.sample`) for risk-aware
+planning.
 
 Important: the control passed to `dynamics.observation_model` for
 `y_{k+1}` is `u_k` (the control that drove the transition into `x_{k+1}`),
@@ -44,20 +48,23 @@
 """
 
 import dataclasses
-from typing import Any, Protocol, runtime_checkable
+from typing import Protocol, runtime_checkable
 
 import jax
 import jax.numpy as jnp
 import jax.random as jr
+import numpyro.distributions as dist
 from jax import Array
 from jaxtyping import PRNGKeyArray, PyTree, Real
 from numpyro.distributions import Distribution
 
 from dynestyx.inference.configs.filter import BaseFilterConfig
 from dynestyx.inference.filters import _default_filter_config
 from dynestyx.inference.integrations.cuthbert.discrete_filter import (
+    build_cuthbert_filter,
     compute_cuthbert_filter_update,
 )
+from dynestyx.inference.integrations.utils import WeightedParticles
 from dynestyx.models import DynamicalModel
 from dynestyx.simulation.base import BaseSimulator
 from dynestyx.simulation.utils import _ensure_trailing_dim, _tile_times
@@ -83,15 +90,38 @@ def filter_state_mean(state) -> Array:
     raise TypeError(f"Cannot summarize filter state of type {type(state).__name__}")
 
 
+def filter_state_dist(state) -> Distribution:
+    """Full-belief NumPyro distribution for a cuthbert filter state, any family.
+
+    Kalman-family states (`KFConfig`, `EKFConfig`, `EnKFConfig`) expose
+    `.mean`/`.chol_cov`, giving an exact `MultivariateNormal`. `PFConfig`
+    states have no such property -- their belief is a weighted particle
+    cloud (`.particles`, `.log_weights`), represented via `WeightedParticles`
+    (dynestyx's own `Distribution`; NumPyro has no built-in equivalent).
+    Unlike `filter_state_mean`, this does not broadcast over a leading
+    time/batch axis -- call it once per (unbatched) state.
+    """
+    if hasattr(state, "chol_cov"):
+        return dist.MultivariateNormal(state.mean, scale_tril=state.chol_cov)
+    if hasattr(state, "particles") and hasattr(state, "log_weights"):
+        log_weights = jax.nn.log_softmax(state.log_weights, axis=-1)
+        return WeightedParticles(state.particles, log_weights)
+    raise TypeError(
+        f"Cannot build a distribution for filter state of type {type(state).__name__}"
+    )
+
+
 @runtime_checkable
 class PolicyCallable(Protocol):
     r"""Structural protocol for a control policy $\pi$.
 
     $$u_k, s_{k+1} = \pi(\hat x_{k|k}, t_k, t_{k+1}, s_k)$$
 
-    `x_hat` is whatever belief state the chosen `filter_config` family
-    produces (see module docstring); use `filter_state_mean` for a
-    family-agnostic point estimate. `t_now`/`t_next` are the current and next
+    `x_hat` is a NumPyro `Distribution` -- `MultivariateNormal` for
+    `KFConfig`/`EKFConfig`/`EnKFConfig`, `WeightedParticles` for `PFConfig`
+    (see module docstring and `filter_state_dist`); use `x_hat.mean` for a
+    family-agnostic point estimate, or the distribution itself for
+    uncertainty-aware planning. `t_now`/`t_next` are the current and next
     times -- always passed, even to a policy that ignores them, so that a
     policy needing genuine time-dependence (e.g. `dynestyx.control.mppi.MPPI`,
     which plans forward from `t_now`) doesn't need special-casing. Any plain
@@ -110,7 +140,11 @@ class PolicyCallable(Protocol):
     """
 
     def __call__(
-        self, x_hat: Any, t_now: Real[Array, ""], t_next: Real[Array, ""], s: PyTree
+        self,
+        x_hat: Distribution,
+        t_now: Real[Array, ""],
+        t_next: Real[Array, ""],
+        s: PyTree,
     ) -> tuple[Real[Array, " control_dim"], PyTree]:
         raise NotImplementedError()
 
@@ -209,6 +243,7 @@ def simulate(
         )
 
         key, k_x0, k_y0, k_filt0 = jr.split(rng_key, 4)
+        filter_obj = build_cuthbert_filter(dynamics, filter_config, key=rng_key)
 
         x_0 = dynamics.initial_condition.sample(k_x0)
         y_0 = dynamics.observation_model(x_0, None, times[0]).sample(k_y0)
@@ -226,13 +261,14 @@ def simulate(
         dt0 = times[1] - times[0] if T > 1 else jnp.asarray(1.0, dtype=times.dtype)
         x_hat_0 = compute_cuthbert_filter_update(
             dynamics,
-            filter_config,
+            None,
             None,
             k_filt0,
             y=y_0,
             u=None,
             t=times[0],
             t_prev=times[0] - dt0,
+            filter_obj=filter_obj,
         )
         initial_state_fn = getattr(self.control_policy, "initial_state", None)
         s_0 = initial_state_fn() if callable(initial_state_fn) else None
@@ -243,7 +279,9 @@ def _step(carry, t_idx):
             t_now = times[t_idx]
             t_next = times[t_idx + 1]
 
-            u_k, s_next = self.control_policy(x_hat_prev, t_now, t_next, s_prev)
+            u_k, s_next = self.control_policy(
+                filter_state_dist(x_hat_prev), t_now, t_next, s_prev
+            )
             if isinstance(u_k, Distribution):
                 raise ValueError(
                     "Returning a distribution is not yet supported, instead "
@@ -258,13 +296,14 @@ def _step(carry, t_idx):
 
             x_hat_next = compute_cuthbert_filter_update(
                 dynamics,
-                filter_config,
+                None,
                 x_hat_prev,
                 k_filt,
                 y=y_next,
                 u=u_k,
                 t=t_next,
                 t_prev=t_now,
+                filter_obj=filter_obj,
             )
 
             new_carry = (x_next, x_hat_next, s_next, step_key)
@@ -320,5 +359,6 @@ def _step(carry, t_idx):
     "ControlledSimulatedResult",
     "DiscreteControlLoopSimulator",
     "PolicyCallable",
+    "filter_state_dist",
     "filter_state_mean",
 ]

dynestyx/control/__init__.py

@@ -16,9 +16,9 @@
 import jax.numpy as jnp
 import jax.random as jr
 from jax import Array
-from jaxtyping import PRNGKeyArray, PyTree, Real
+from jaxtyping import PRNGKeyArray, Real
+from numpyro.distributions import Distribution
 
-from dynestyx.control.discrete_controller_simulators import filter_state_mean
 from dynestyx.models import DynamicalModel
 
 
@@ -146,7 +146,7 @@ def step(carry, u_and_idx):
 
     def __call__(
         self,
-        x_hat: PyTree,
+        x_hat: Distribution,
         t_now: Real[Array, ""],
         t_next: Real[Array, ""],
         s: tuple[Real[Array, "horizon control_dim"], PRNGKeyArray],
@@ -157,7 +157,7 @@ def __call__(
         # t_next (the real simulation's next observation time) is unused --
         # MPPI plans its own horizon-step lookahead from t_now using its own dt.
         del t_next
-        x0 = filter_state_mean(x_hat)
+        x0 = x_hat.mean
         nominal, key = s
         key, noise_key, rollout_key = jr.split(key, 3)
         control_dim = nominal.shape[-1]

dynestyx/control/discrete_controller_simulators.py

@@ -204,52 +204,60 @@ def _build_cuthbert_filter_obj(
     return filter_obj, parallel
 
 
-def compute_cuthbert_filter_update(
+def build_cuthbert_filter(
     dynamics: DynamicalModel,
     filter_config: BaseFilterConfig,
+    key: jax.Array | None = None,
+):
+    """Build the cuthbert Filter object for `(dynamics, filter_config)`.
+
+    Depends only on `dynamics`/`filter_config`, never on trajectory data, so
+    it can be built once and reused across many `compute_cuthbert_filter_update`
+    calls via that function's `filter_obj=` argument -- avoids rebuilding it
+    on every step. `key` is only required (non-`None`) for `PFConfig`/`EnKFConfig`.
+    """
+    filter_kwargs = _config_to_filter_kwargs(filter_config)
+    filter_obj, _ = _build_cuthbert_filter_obj(
+        dynamics, filter_config, filter_kwargs, key, want_parallel=False
+    )
+    return filter_obj
+
+
+def compute_cuthbert_filter_update(
+    dynamics: DynamicalModel,
+    filter_config: BaseFilterConfig | None,
     prev_state,
     key: jax.Array,
     *,
     y: jax.Array,
     u: jax.Array | None,
     t: jax.Array,
     t_prev: jax.Array | None = None,
+    filter_obj=None,
 ):
-    r"""One-step FilterUpdate: state_k + u_k + y_{k+1} -> state_{k+1}.
-
-    Unlike `compute_cuthbert_filter` (whole-trajectory), this performs exactly
-    one predict+update step using cuthbert's `Filter.filter_prepare`/
-    `filter_combine` primitives directly, without requiring future
-    observations. This is what makes online control possible: the state
-    returned here can be consumed by a policy to choose the next control
-    before the next observation exists.
-
-    Pass `prev_state=None` for the bootstrap call (computing the filtering
-    state after only the first observation, with no control history yet);
-    this internally calls the cuthbert filter's `init_prepare` first.
-
-    Control convention (important, and different from `compute_cuthbert_filter`
-    / `DiscreteTimeSimulator`): `u` is the control that drove the transition
-    *into* the state being filtered, i.e. u_k when producing state_{k+1} from
-    state_k and y_{k+1} -- matching `FilterUpdate(x_hat_k, u_k, y_{k+1}, ...)`
-    in the control-loop equations. `compute_cuthbert_filter`/
-    `DiscreteTimeSimulator` instead pair `ctrl_values[t]` with *both* the
-    observation and the outgoing transition at the same index t, which is
-    only valid when the whole control trajectory is already known in advance.
-    For online control this is impossible: u_{k+1} cannot exist before
-    y_{k+1} is observed, since it is computed by the policy from the filtered
-    state that itself depends on y_{k+1}. So `u` here is used for both
-    `CuthbertInputs.u` and `CuthbertInputs.u_prev` in the single-row input
-    built for this step. Pass `u=None` for the bootstrap call, matching y_0's
-    lack of a control argument in the control-loop equations (numerically
-    equivalent to zeros for models with a control-input matrix, since D=None
-    or u=None are both treated as "no control contribution").
+    r"""One-step FilterUpdate: state_k + u_k + y_{k+1} -> state_{k+1}. Used for online filtering.
+
+    Unlike `compute_cuthbert_filter` (whole-trajectory), this performs one
+    predict+update step directly via cuthbert's `Filter.filter_prepare`/
+    `filter_combine`
+    Pass `prev_state=None` for the bootstrap call
+    (first observation, no control history yet); this runs `init_prepare`
+    first.
+
+    `u` is the control that drove the transition *into* the state being
+    filtered (u_k, producing state_{k+1} from y_{k+1}).
+
+    Provide exactly one of `filter_config` (builds the filter internally) or
+    `filter_obj` (an already-built filter from `build_cuthbert_filter`, to
+    reuse across repeated calls instead of rebuilding it here).
     """
-    filter_kwargs = _config_to_filter_kwargs(filter_config)
+    if (filter_config is None) == (filter_obj is None):
+        raise ValueError("Provide exactly one of filter_config or filter_obj.")
+
     key_state, key_prep = jr.split(key)
-    filter_obj, _ = _build_cuthbert_filter_obj(
-        dynamics, filter_config, filter_kwargs, key_state, want_parallel=False
-    )
+    if filter_obj is None:
+        assert filter_config is not None
+        filter_obj = build_cuthbert_filter(dynamics, filter_config, key_state)
 
     control_dim = dynamics.control_dim
     u_arr = jnp.zeros((control_dim,)) if u is None else jnp.asarray(u)

dynestyx/control/mppi.py

@@ -114,6 +114,12 @@ def sample(self, key, sample_shape=()):
     def log_prob(self, value):
         raise NotImplementedError("log_prob is not implemented for WeightedParticles.")
 
+    @property
+    def mean(self) -> jax.Array:
+        """Weighted mean of the particles."""
+        weights = jax.nn.softmax(self.log_weights, axis=-1)
+        return jnp.sum(weights[..., None] * self.particles, axis=-2)
+
 
 def particles_to_delta_mixtures(
     particles: Real[Array, "*plate time n_particles state_dim"]