Skip to content

Closed-loop control

Dynestyx can interleave simulation, observation, filtering, and control for a single discrete-time trajectory. At each step it performs

\[ \begin{aligned} x_0 &\sim p(x_0), \\ y_0 \mid x_0 &\sim p(y_0 \mid x_0, t_0), \\ \hat{x}_{0\mid0} &= \operatorname{FilterUpdate}(y_0, t_0), \\ (u_k, s_{k+1}) &= \pi(\hat{x}_{k\mid k}, t_k, t_{k+1}, s_k), \\ x_{k+1} \mid x_k,u_k &\sim p(x_{k+1}\mid x_k,u_k,t_k,t_{k+1}), \\ y_{k+1} \mid x_{k+1},u_k &\sim p(y_{k+1}\mid x_{k+1},u_k,t_{k+1}), \\ \hat{x}_{k+1\mid k+1} &= \operatorname{FilterUpdate} (\hat{x}_{k\mid k},u_k,y_{k+1},t_k,t_{k+1}). \end{aligned} \]

The observation at t[k + 1] receives u[k], the control that produced its state. This differs from the same-index convention used for a precomputed open-loop control trajectory. This is a temporary difference: [Issue

312](https://github.com/BasisResearch/dynestyx/issues/312) tracks aligning

closed-loop control with the simulator convention and requiring controlled DynamicalModel observation models to follow that convention.

Controlled simulation currently supports one trajectory at a time. Its online filter update is implemented with Cuthbert and supports KFConfig, EKFConfig, EnKFConfig, and PFConfig. dsx.plate and n_simulations > 1 are rejected explicitly. Controlled simulation requires filter_source="cuthbert" and rejects configurations that request another backend.

Simulator and policy protocol

dynestyx.control.discrete_controller_simulators.DiscreteControlLoopSimulator

Bases: BaseSimulator

Closed-loop simulator: simulate, observe, filter, and decide controls online.

Unlike DiscreteTimeSimulator, which requires the entire control trajectory as a pre-supplied ctrl_values array, DiscreteControlLoopSimulator computes each \(u_k\) online from the filtered belief \(\hat x_{k|k}\) via control_policy. See the closed-loop control API page for the full loop equations and the control-index convention used by dynamics.observation_model.

The online loop uses \(u_k\) for both the transition into \(x_{k+1}\) and the observation \(y_{k+1}\). This control-observation alignment is temporary; Issue #312 tracks aligning it with the regular simulator convention and requiring controlled DynamicalModel observation models to follow that convention. The one-step filter update currently uses Cuthbert and supports KFConfig, EKFConfig, EnKFConfig, and PFConfig. Plated controlled simulation is not yet supported; see Issue #318.

Attributes:

Name Type Description
control_policy

Control policy \(\pi\); see PolicyCallable. Its initial state \(s_0\) is exactly simulate's initial_policy_state argument (default None, for a stateless policy) -- control_policy is never introspected for an initial_state() method; a stateful policy's initial state must always be passed explicitly.
filter_config

Selects the filtering algorithm (KFConfig/EKFConfig/EnKFConfig/PFConfig). Defaults to _default_filter_config(dynamics) when None. The online one-step update currently requires filter_source="cuthbert". Its record_filtered_states_mean/record_max_elems fields gate whether the filtered_states_mean output is recorded, exactly as they do for Filter (see dynestyx.utils._should_record_field).
n_simulations int

Currently only 1 is supported.

simulate(dynamics: DynamicalModel, *, rng_key: PRNGKeyArray, ctrl_times: Real[Array, ' ctrl_time'] | None = None, ctrl_values: Real[Array, 'ctrl_time control_dim'] | Real[Array, ' ctrl_time'] | None = None, predict_times: Real[Array, ' predict_time'] | None = None, initial_policy_state: PyTree | None = None, **kwargs: Any) -> ControlledSimulatedResult

Simulate one online controlled trajectory.

Parameters:

Name Type Description Default
dynamics DynamicalModel

Discrete-time dynamical model.

 required
rng_key PRNGKeyArray

Root key for environment and fallback filter randomness.

 required
ctrl_times Real[Array, ' ctrl_time'] | None

Unsupported because controls are selected online.

 None
ctrl_values Real[Array, 'ctrl_time control_dim'] | Real[Array, ' ctrl_time'] | None

Unsupported because controls are selected online.

 None
predict_times Real[Array, ' predict_time'] | None

Strictly increasing simulation times.

 None
initial_policy_state PyTree | None

Initial state passed to control_policy.

 None
**kwargs Any

Additional shared simulator-handler metadata, ignored here.

 {}

Returns:

Type Description
ControlledSimulatedResult

States, observations, controls, filter means, and policy states.

Raises:

Type Description
ValueError

If inputs are incompatible with online discrete control.
NotImplementedError

If the requested simulation mode is unsupported.

dynestyx.control.discrete_controller_simulators.PolicyCallable

Bases: Protocol

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 a NumPyro Distribution -- MultivariateNormal for KFConfig/EKFConfig/EnKFConfig, WeightedParticles for PFConfig (see 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 callable matching this signature works, including an equinox.Module with a matching __call__ (e.g. a learned neural policy) or a plain Python function (e.g. an LQR gain lookup).

control_policy never receives a PRNG key and must return a concrete value, not a NumPyro Distribution (returning one raises a ValueError -- not yet supported). A stochastic policy instead carries any randomness it needs (e.g. MPPI's exploration noise) inside s, splitting/advancing it internally on every call and sampling from its own distributions itself before returning a value. The policy owns and seeds this randomness entirely by itself -- see dynestyx.control.mppi.MPPI's seed attribute for the pattern.

dynestyx.control.discrete_controller_simulators.ControlledSimulatedResult dataclass

Bases: SimulatedResult

SimulatedResult extended with the control loop's extra outputs.

Registered as deterministic sites the same generic way as SimulatedResult's own fields (dynestyx.simulation.utils. _register_simulated_result_sites iterates every dataclass field and skips None values) -- so the existing recording-gating logic just means passing None for a field instead of conditionally omitting a dict key, as the old (pre-refactor) version of this class did.

Policy helpers

dynestyx.control.discrete_controller_simulators.filter_state_mean(state: Any) -> Real[Array, ...]

Point-estimate summary of a cuthbert filter state, any family.

Kalman-family states (KFConfig, EKFConfig, EnKFConfig) expose a .mean property directly. PFConfig states (ParticleFilterState) have no such property -- they represent the belief as a weighted particle cloud (.particles, .log_weights), so the point estimate is the weighted mean instead. Broadcasts over any leading batch/time axis, so it works on both a single belief and a whole scanned-out sequence of them.

dynestyx.control.discrete_controller_simulators.filter_state_dist(state: Any) -> 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.

MPPI-inspired policy

dynestyx.control.mppi.MPPI

Bases: Module

Model Predictive Path Integral (MPPI) controller.

At each call: sample n_samples candidate control sequences of length horizon as Gaussian perturbations around a nominal sequence (the policy state s, warm-started from the previous call), roll each one forward horizon steps through dynamics.state_evolution (built internally -- the caller only ever supplies the one-step dynamics, never a hand-written rollout), score the resulting trajectories with loss_fn, and combine them via the standard MPPI weighting

\[w_i \\propto \\exp(-\\mathrm{loss}_i / \\lambda), \\qquad u_{0:H-1} = \\sum_i w_i\\, u^{(i)}_{0:H-1}\]

i.e. a softmax over the (negated, temperature-scaled) per-sample losses. Only the first control of that weighted-mean sequence is applied this step (receding horizon); the remainder becomes next step's nominal sequence, shifted left by one with the last entry repeated.

Attributes:

Name Type Description
dynamics DynamicalModel

a DynamicalModel (the same model used for the real simulation or some approximate). Each candidate rollout is computed by calling dsx.simulate. If dynamics holds trainable parameters you're also fitting via the outer simulation, they remain in the differentiable pytree so gradients through planning are tracked too.
loss_fn MPPILossFn

MPPILossFn, i.e. (result: ControlledSimulatedResult) -> scalar, called once per sample (vmapped) on that candidate's full rollout. Every field carries a leading n_simulations=1 axis -- e.g. result.states.shape == (1, horizon + 1, state_dim) -- matching how dsx.simulate never drops that axis, even for one trajectory; jnp.sum(result.states**2)-style reductions don't need to care, but explicit indexing does (result.controls[0, 0] is the whole first control vector, not a scalar). times/states/observations have length horizon + 1 (including the starting state) and controls has length horizon, matching ControlledSimulatedResult's own control_time = time - 1 convention.
horizon int

Planning horizon length H -- the number of internal one-step dynamics calls per rollout. Defaults to 10.
noise_std Real[Array, ''] | Real[Array, ' control_dim']

Standard deviation of the Gaussian perturbations added to the nominal sequence, scalar or shape (control_dim,). Defaults to 1.0.
n_samples int

Number of sampled control sequences per call. Defaults to 20.
dt float

Fixed planning step size. Defaults to 1.0.
temperature float

MPPI's \(\\lambda\); higher values flatten the weights toward a uniform average, lower values concentrate weight on the lowest-loss samples.
batched bool

Whether the n_samples candidate rollouts are computed with jax.vmap (default, fast, requires dynamics.state_evolution to be vmap-compatible) or jax.lax.map (a sequential loop -- slower, but works for a dynamics.state_evolution that isn't vmap-compatible, e.g. wraps an external simulator via jax.pure_callback).
seed int

Seeds MPPI's own PRNG key, carried inside the policy state s (as (nominal_sequence, key)) and split internally on every call.

initial_state() -> tuple[Real[Array, 'horizon control_dim'], PRNGKeyArray]

Zero nominal control sequence plus MPPI's own seeded PRNG key. Pass this call's result as initial_policy_state to simulate/ dsx.simulate -- DiscreteControlLoopSimulator never calls this automatically, so it must be supplied explicitly.

plan_step(x_hat: Distribution, t_now: Real[Array, ''], s: tuple[Real[Array, 'horizon control_dim'], PRNGKeyArray]) -> tuple[Real[Array, ' control_dim'], tuple[Real[Array, 'horizon control_dim'], PRNGKeyArray], ControlledSimulatedResult]

Do MPPI's full planning step and also return the batch of every candidate rollout considered (n_samples-wide ControlledSimulatedResult) -- useful for debugging/plotting what MPPI weighed, or diagnosing a loss_fn. __call__ (used by DiscreteControlLoopSimulator) is a thin wrapper around this that drops the rollout batch, since PolicyCallable's return signature can't carry a third value.

Note: the returned result's leading axis indexes candidates, not independent draws from the true generative process -- it's not a real simulated trajectory. filtered_states_mean/policy_states/ predicted_* are always None (not meaningful for a planning rollout).