Skip to content

PolicyCallable

Any callable matching PolicyCallable can be used as a policy. We provide some baseline implementations in Provided Policies.

Bases: Protocol

Structural protocol for a control policy \(\pi\).

\[u_k, s_{k+1} = \pi(\tilde x_k, t_k, t_{k+1}, s_k)\]

\(\tilde x_k\) is the loop's current state or state estimate. Which one it is depends on use_true_state and on the convention.

When using use_true_state=True, \(\tilde x_k\) is a Delta distribution centered on the true state \(x_k\) (including at \(k=0\)).

When using use_true_state=False (default), this is a distribution depending on the filter configuration:

  • under same_time, \(\tilde x_k\) is the predicted state \(\hat{x}_{k|k-1}\)
  • under previous_transition, \(\tilde x_k\) is the filtered state \(\hat{x}_{k|k}\)

where \(\hat{x}_{k|j}\) is the state estimate at time \(t_k\) given observations up to time \(t_j\). At \(k=0\), this is the model's initial-state distribution \(p_0\).

With use_true_state=True there is no filtering: \(\tilde x_k\) is a Delta centered on the true state \(x_k\).

It is always a NumPyro Distribution. Use x_hat.mean for a family-agnostic point estimate. t_now/t_next are the current and next time points. Any plain callable matching this signature works, including an equinox.Module with a matching __call__ or a plain Python function. control_policy must return a concrete value.

s is an internal state for the policy (such as a PRNG key), which is passed back to it on the next call. A stateless policy can ignore it and return None for the next state.

Policy helpers

Helpers for turning the filter state handed to a policy into a point estimate or a distribution.

filter_state_dist

Full-belief NumPyro distribution for a filter state.

Only filter_source="cuthbert" is supported today, matching DiscreteControlLoopSimulator's own restriction. Converts the filter state to a NumPyro Distribution in the same way that ConditionedResult.dists does for the recorded filtered states.

The shared conversion is time-indexed, so the state is given a leading axis of length one and the single distribution unwrapped.

Parameters:

Name Type Description Default
state Any

A single, unbatched filter state produced by filter_config.

required
filter_config BaseFilterConfig

The config that produced state. Selects the backend and carries recorded_filtered_states_cov_jitter for ensemble states.

required

Returns:

Type Description
Distribution

The belief as a NumPyro Distribution.

Raises:

Type Description
ValueError

If filter_config.filter_source is not "cuthbert".