Online and control with cd-dynamax filter - #314
Draft
MatthieuDarcy wants to merge 47 commits into
Draft
Conversation
Implements the discrete-time control loop from the design issue: a new DiscreteControlLoopSimulator (alongside DiscreteTimeSimulator) interleaves simulation, observation, filtering, and control online, deciding each u_k from the filtered belief via a user-supplied policy rather than requiring the whole control trajectory up front. FilterUpdate is powered by a new compute_cuthbert_filter_update, which drives cuthbert's existing Filter.filter_prepare/filter_combine primitives one step at a time instead of over a whole pre-supplied trajectory -- generic across KFConfig/EKFConfig/EnKFConfig/PFConfig. Verified PFConfig and EnKFConfig support genuinely black-box state transitions (no log_prob required), matching the design issue's black-box dynamics requirement. Includes a demo notebook (docs/tutorials/control/controller_demo.ipynb) showing a linear feedback policy driving a 1D linear-Gaussian system to 0. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…transitions Adds a tutorial section demonstrating a genuinely continuous-time SDE (ContinuousTimeStateEvolution + Discretizer(discretize=euler_maruyama)) composed with DiscreteControlLoopSimulator, showing the loop only ever sees discrete times while the underlying dynamics are a true SDE solved between them. Building this surfaced a real bug: the bootstrap FilterUpdate call left t_prev defaulting to t (dt=0). For fixed-covariance transitions this is harmless, but for a dt-scaled transition (like the Euler-Maruyama discretization here), it constructs a zero-covariance distribution whose NaN log-density leaks through the gradient in EKF's Taylor linearization (jnp.where evaluates both branches, unlike jax.lax.cond) -- corrupting the filtered state and, downstream, the policy's control and the simulated trajectory itself. Fixed by giving the bootstrap call a non-degenerate t_prev, borrowing the width of the first real interval (matching compute_cuthbert_filter's own dummy-row convention). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ccuracy Documents how the continuous-time integration is actually defined: Discretizer takes exactly one Euler-Maruyama step over the whole gap between requested times (needed so filtering gets an explicit one-step transition density), unlike SDESimulator/solve_sde's substepped solvers, which only support pure forward simulation without filtering. Adds a section 6 demonstrating the same recipe on a genuinely nonlinear 2D system with drift = A x^2 + u: x=0 is an unstable equilibrium (x^2 >= 0 always pushes away from the origin), so the uncontrolled run diverges while the same linear feedback policy as before stabilizes it locally. Uses a finer time grid (dt=0.1 vs 1.0) since the one-step EM approximation needs a small enough gap to stay accurate for a nonlinear drift. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…rt_filter_update 35 tests across 7 groups: filter_state_mean unit tests, compute_cuthbert_filter_update correctness (including regression tests for the dt=0 EKF NaN bug), validation/error paths, end-to-end shape/output-key tests (including the stateless-policy regression), behavioral/control correctness (closed-loop stabilization, the control-index convention, determinism, eqx.Module policies), continuous-time/Discretizer composition, and black-box transition compatibility. Surfaced one more real gap while writing these: a genuinely nested-pytree policy_state_init (e.g. a dict of arrays) can't actually round-trip today -- BaseSimulator's shared _run_single_member_simulation enforces a dict[str, Array] | None return type via jaxtyping at runtime, so policy_states must stay a flat Array. Documented in the corresponding test rather than widening the shared base class's type contract, which is out of scope here. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ulator there Consolidates control-loop + control-policy code under dynestyx/control/: moves discrete_controller_simulators.py there, and adds a basic MPPI (Model Predictive Path Integral) controller (dynestyx/control/mppi.py) that plugs into the same control_policy= slot. Deliberately simple: samples candidate control sequences as Gaussian perturbations around a nominal sequence, rolls each through a user-supplied dynamics_model, and returns the softmax-weighted mean control (standard MPPI weighting). dynamics_model can be batched (one call handles all samples) or not (driven via jax.lax.map, so it never needs to support batching itself). MPPI needs fresh randomness every step, which PolicyCallable's signature didn't provide -- added a key parameter to __call__(x_hat, s, key), passed in by DiscreteControlLoopSimulator's per-step scan. This also keeps MPPI's own policy state a flat array (just the nominal control sequence), avoiding a real limitation found while testing last time: a nested-pytree policy_state_init can't actually be recorded as output today, since BaseSimulator's shared return-type check requires flat Array values. Verified: existing linear-feedback policies (tests, tutorial notebook) updated for the new signature, full regression suite still passes (117 tests), and a smoke test confirms MPPI stabilizes a marginally-unstable linear system identically under both batched and non-batched dynamics_model. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
# Conflicts: # dynestyx/inference/integrations/cuthbert/discrete_filter.py
…lation) main/control_v2 landed a major refactor (commit 3fb1cd7, "Numpyro-Free Usage") that deleted dynestyx/simulators.py entirely, replacing it with a new dynestyx/simulation/ package: BaseSimulator subclasses now implement a public `simulate(dynamics, *, rng_key, ...) -> SimulatedResult` (not `_simulate(name, ..., obs_times=..., ...) -> dict`), take rng_key directly instead of calling numpyro.prng_key() internally, and numpyro site registration happens via a generic deferred callback (dataclasses.fields() iteration over the returned SimulatedResult) rather than base-class magic over a free-form dict. Ported DiscreteControlLoopSimulator to the new contract: - Plain class with explicit __init__ (matching DiscreteTimeSimulator's new shape) instead of a bare @dataclasses.dataclass. - New ControlledSimulatedResult(SimulatedResult) subclass carrying the extra controls/filtered_states_mean/policy_states fields -- confirmed by reading base.py directly that _run_single_member_simulation's common path is a bare `return self.simulate(...)` with no reconstruction, so subclass fields flow through untouched and get registered generically the same way as SimulatedResult's own fields (None-valued fields are auto-skipped). - Dropped the now-redundant manual `numpyro.prng_key()`/obs_values checks: the base class already validates/rejects these before simulate() is ever called. compute_cuthbert_filter_update/_build_cuthbert_filter_obj (added to discrete_filter.py on the `control` branch) merged in via `git merge control` with a single trivial conflict (re-adding the `jax.random` import) -- confirmed via git merge-tree dry run and direct diffing that every symbol they depend on exists byte-identical post-refactor, just re-imported from dynestyx.inference.configs.filter instead of dynestyx.inference.filter_configs. Discretizer required no changes at all (already ported upstream). Verified: full regression suite passes on control_v2 (118 tests: this file's 35 plus test_filters.py/test_filter_simulator.py/test_discretizers.py/ test_predictive_filter_simulator_shapes.py), the tutorial notebook re-executes cleanly end-to-end, and MPPI produces numerically identical results to its control-branch run under both batched and non-batched dynamics_model. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Section 7 demonstrates dynestyx.control.MPPI plugged into the same control_policy= slot as the earlier linear-feedback policy, on the same 1D system from section 1. Builds a planning rollout directly from dynamics.state_evolution's deterministic mean (a standard MPPI simplification -- the planner doesn't need a faithful stochastic simulation, just a reasonable prediction of where a candidate control sequence leads) and a simple quadratic loss, then compares against the K=0 no-control baseline already computed in section 3. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
DynamicalModel.state_evolution is declared as the broader ContinuousTimeStateEvolution | DiscreteStateTransition union, so `ty` couldn't statically narrow continuous_dynamics.state_evolution to StochasticContinuousTimeStateEvolution before passing it to euler_maruyama() (which requires the narrower type) -- even though this is exactly how Discretizer._sample_ds itself resolves it at runtime. Added the same isinstance narrowing check used there. Also includes an incidental ruff-format tweak to mppi.py (collapsing a one-line call that had been wrapped unnecessarily). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…into control_v2
…book Section 8 demonstrates that DiscreteControlLoopSimulator is differentiable end-to-end: rolls out a short horizon with the linear policy from section 2, defines a loss as the norm of the final true state, and takes jax.grad of it with respect to the gain K -- no special machinery needed beyond plain JAX autodiff. Verified independently against a finite-difference check before writing the notebook cells (gradient matched to ~1e-4), and shows one gradient-descent step meaningfully reducing the loss (0.62 -> 0.18). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Section 3 (differentiating through the closed loop) computed grad_K via jax.grad but never independently verified it. Added a cell checking a central finite difference on K[0, 0] against the autodiff value -- confirms the gradient through the whole closed loop (policy -> transition -> observation -> filter update) is actually correct, not just plausible. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Section 3's rollout_final_state_norm now calls sim.simulate(dynamics, rng_key=..., predict_times=...) directly instead of going through dsx.sample/numpyro.handlers.seed -- unneeded here since there's no conditioning/MCMC, and it makes the PRNG key an explicit argument instead of an implicit fixed rng_seed=0. More importantly, this fixes a real issue in section 4's training loop: every epoch previously reused the same fixed seed, so gradient descent could overfit K to one specific noise realization rather than the underlying dynamics. Now each epoch draws a fresh key via jax.random.split, so K_opt has to generalize across realizations. Verified the loop still converges to a sensible stabilizing gain with this fresh-key-per-epoch setup before editing the notebook. Section 4's final comparison plot (run()/dsx.sample) is left unchanged -- there, reusing the same seed for both the unoptimized and optimized rollouts is what makes the before/after comparison fair. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
run() (updated separately) now returns a ControlledSimulatedResult from sim.simulate() directly instead of a numpyro trace dict, so the plotting cell needed the equivalent attribute-access update: trace["loop_X"]["value"] -> result.X. Also removed a leftover duplicate plotting cell still using the old trace-dict access, which was left after the updated one and caused the notebook to fail on execution. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Replaces the with sim: dsx.sample(...) + numpyro.handlers.seed/trace pattern across every section (discrete demo, SDE, nonlinear 2D, MPPI) with direct sim.simulate(dynamics, rng_key=key, predict_times=...) calls returning ControlledSimulatedResult objects -- no NumPyro handler or model function needed, since every section here is a pure generative rollout with nothing to condition on. Renamed trace_* variables to result_* throughout, since they're no longer NumPyro trace dicts, and updated all downstream plotting cells from trace["name_field"]["value"] to result.field attribute access. For the continuous-time sections, this also removes the Discretizer handler composition: instead of wrapping dsx.sample in `with Discretizer(...)`, the continuous-time DynamicalModel is discretized once up front via euler_maruyama(state_evolution) into a plain discrete-time DynamicalModel, which is then simulated exactly like any other discrete-time model. Updated the surrounding markdown to match (no more handler-chain explanation needed). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Unified non-linear dynamics, black box and partial observations example
Unify simulate API with control_policy/filter_config; add distribution-returning policies; expand controller tutorial - dsx.simulate() and Simulator now accept control_policy/filter_config, routing to DiscreteControlLoopSimulator the same way dynamics.state_evolution auto-routes to SDESimulator/ODESimulator/DiscreteTimeSimulator. - A policy may return a NumPyro Distribution instead of a value, sampled automatically; MPPI now owns and advances its own exploration randomness via a `seed` field, - Policies self-initialize via an optional initial_state() (replacing the removed policy_state_init argument), mirroring optax's optimizer.init(). - controller_demo.ipynb: sections now use dsx.simulate() directly; added a new random-controller example (Section 5) demonstrating a distribution-returning policy. unchanged on the nonlinear 2D black-box SDE (Section 6). Removed the MPPI example from the basic tutorial (now lives in a separate tutorial.
Rework MPPI to take one-step dynamics; complete and stabilize mpc_demo - MPPI now takes dynamics: DynamicalModel (one-step transition) instead of a hand-rolled full-rollout function; it builds the horizon-unrolling and n_samples batching internally via jax.lax.scan/vmap, using .mean when available and falling back to .sample() (own internal key) for black-box transitions. - All fields but dynamics/loss_fn now have defaults (horizon=10, noise_std=1.0, n_samples=20, dt=1.0, temperature=1.0, batched=True, seed=0). - stability: Non-finite (nan/inf) rollout losses are clamped before the softmax weighting, fixing NaN control - mpc_demo.ipynb: completed (previously referenced undefined names, never ran), updated to the new MPPI API, added MPC/MPPI citations.
Update the policy to now take in state, t_now, t_next, s. Changed how stochastic policies operate: sample within the policy, do not return a distribution. If a distribution is returned, raise an error to indicate to the user that they should sample from it instead.
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
Removed the ability to give a filter config to the one step filter update. Now always requires a filter object
Updated the control optimization notebook to use the new syntax. Also added a check to see that gradients flow correctly through both the dynamics and the MPC planning loop.
Removed the need/requirement to have an initial state attribute to a policy. Instead this must be passed explicitly to the simulator. This will avoid hidden operations and is more in line with other coding practices.
…into control_v2
MatthieuDarcy
changed the title
Merged
MatthieuDarcy
force-pushed
the
control_v2
branch
from
August 13, 2026 22:27
8c31c0b to
a70c572
Compare
mattlevine22
added a commit
that referenced
this pull request
Aug 14, 2026
* Add online control loop support via DiscreteControlLoopSimulator Implements the discrete-time control loop from the design issue: a new DiscreteControlLoopSimulator (alongside DiscreteTimeSimulator) interleaves simulation, observation, filtering, and control online, deciding each u_k from the filtered belief via a user-supplied policy rather than requiring the whole control trajectory up front. FilterUpdate is powered by a new compute_cuthbert_filter_update, which drives cuthbert's existing Filter.filter_prepare/filter_combine primitives one step at a time instead of over a whole pre-supplied trajectory -- generic across KFConfig/EKFConfig/EnKFConfig/PFConfig. Verified PFConfig and EnKFConfig support genuinely black-box state transitions (no log_prob required), matching the design issue's black-box dynamics requirement. Includes a demo notebook (docs/tutorials/control/controller_demo.ipynb) showing a linear feedback policy driving a 1D linear-Gaussian system to 0. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Add continuous-time SDE demo and fix EKF bootstrap NaN for dt-scaled transitions Adds a tutorial section demonstrating a genuinely continuous-time SDE (ContinuousTimeStateEvolution + Discretizer(discretize=euler_maruyama)) composed with DiscreteControlLoopSimulator, showing the loop only ever sees discrete times while the underlying dynamics are a true SDE solved between them. Building this surfaced a real bug: the bootstrap FilterUpdate call left t_prev defaulting to t (dt=0). For fixed-covariance transitions this is harmless, but for a dt-scaled transition (like the Euler-Maruyama discretization here), it constructs a zero-covariance distribution whose NaN log-density leaks through the gradient in EKF's Taylor linearization (jnp.where evaluates both branches, unlike jax.lax.cond) -- corrupting the filtered state and, downstream, the policy's control and the simulated trajectory itself. Fixed by giving the bootstrap call a non-degenerate t_prev, borrowing the width of the first real interval (matching compute_cuthbert_filter's own dummy-row convention). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Add nonlinear 2D SDE example and note on Euler-Maruyama integration accuracy Documents how the continuous-time integration is actually defined: Discretizer takes exactly one Euler-Maruyama step over the whole gap between requested times (needed so filtering gets an explicit one-step transition density), unlike SDESimulator/solve_sde's substepped solvers, which only support pure forward simulation without filtering. Adds a section 6 demonstrating the same recipe on a genuinely nonlinear 2D system with drift = A x^2 + u: x=0 is an unstable equilibrium (x^2 >= 0 always pushes away from the origin), so the uncontrolled run diverges while the same linear feedback policy as before stabilizes it locally. Uses a finer time grid (dt=0.1 vs 1.0) since the one-step EM approximation needs a small enough gap to stay accurate for a nonlinear drift. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Add test coverage for DiscreteControlLoopSimulator and compute_cuthbert_filter_update 35 tests across 7 groups: filter_state_mean unit tests, compute_cuthbert_filter_update correctness (including regression tests for the dt=0 EKF NaN bug), validation/error paths, end-to-end shape/output-key tests (including the stateless-policy regression), behavioral/control correctness (closed-loop stabilization, the control-index convention, determinism, eqx.Module policies), continuous-time/Discretizer composition, and black-box transition compatibility. Surfaced one more real gap while writing these: a genuinely nested-pytree policy_state_init (e.g. a dict of arrays) can't actually round-trip today -- BaseSimulator's shared _run_single_member_simulation enforces a dict[str, Array] | None return type via jaxtyping at runtime, so policy_states must stay a flat Array. Documented in the corresponding test rather than widening the shared base class's type contract, which is out of scope here. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Add MPPI controller in dynestyx/control/, move DiscreteControlLoopSimulator there Consolidates control-loop + control-policy code under dynestyx/control/: moves discrete_controller_simulators.py there, and adds a basic MPPI (Model Predictive Path Integral) controller (dynestyx/control/mppi.py) that plugs into the same control_policy= slot. Deliberately simple: samples candidate control sequences as Gaussian perturbations around a nominal sequence, rolls each through a user-supplied dynamics_model, and returns the softmax-weighted mean control (standard MPPI weighting). dynamics_model can be batched (one call handles all samples) or not (driven via jax.lax.map, so it never needs to support batching itself). MPPI needs fresh randomness every step, which PolicyCallable's signature didn't provide -- added a key parameter to __call__(x_hat, s, key), passed in by DiscreteControlLoopSimulator's per-step scan. This also keeps MPPI's own policy state a flat array (just the nominal control sequence), avoiding a real limitation found while testing last time: a nested-pytree policy_state_init can't actually be recorded as output today, since BaseSimulator's shared return-type check requires flat Array values. Verified: existing linear-feedback policies (tests, tutorial notebook) updated for the new signature, full regression suite still passes (117 tests), and a smoke test confirms MPPI stabilizes a marginally-unstable linear system identically under both batched and non-batched dynamics_model. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Port control-loop work onto post-refactor architecture (dynestyx.simulation) main/control_v2 landed a major refactor (commit 3fb1cd7, "Numpyro-Free Usage") that deleted dynestyx/simulators.py entirely, replacing it with a new dynestyx/simulation/ package: BaseSimulator subclasses now implement a public `simulate(dynamics, *, rng_key, ...) -> SimulatedResult` (not `_simulate(name, ..., obs_times=..., ...) -> dict`), take rng_key directly instead of calling numpyro.prng_key() internally, and numpyro site registration happens via a generic deferred callback (dataclasses.fields() iteration over the returned SimulatedResult) rather than base-class magic over a free-form dict. Ported DiscreteControlLoopSimulator to the new contract: - Plain class with explicit __init__ (matching DiscreteTimeSimulator's new shape) instead of a bare @dataclasses.dataclass. - New ControlledSimulatedResult(SimulatedResult) subclass carrying the extra controls/filtered_states_mean/policy_states fields -- confirmed by reading base.py directly that _run_single_member_simulation's common path is a bare `return self.simulate(...)` with no reconstruction, so subclass fields flow through untouched and get registered generically the same way as SimulatedResult's own fields (None-valued fields are auto-skipped). - Dropped the now-redundant manual `numpyro.prng_key()`/obs_values checks: the base class already validates/rejects these before simulate() is ever called. compute_cuthbert_filter_update/_build_cuthbert_filter_obj (added to discrete_filter.py on the `control` branch) merged in via `git merge control` with a single trivial conflict (re-adding the `jax.random` import) -- confirmed via git merge-tree dry run and direct diffing that every symbol they depend on exists byte-identical post-refactor, just re-imported from dynestyx.inference.configs.filter instead of dynestyx.inference.filter_configs. Discretizer required no changes at all (already ported upstream). Verified: full regression suite passes on control_v2 (118 tests: this file's 35 plus test_filters.py/test_filter_simulator.py/test_discretizers.py/ test_predictive_filter_simulator_shapes.py), the tutorial notebook re-executes cleanly end-to-end, and MPPI produces numerically identical results to its control-branch run under both batched and non-batched dynamics_model. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Add MPPI section to the control-loop tutorial notebook Section 7 demonstrates dynestyx.control.MPPI plugged into the same control_policy= slot as the earlier linear-feedback policy, on the same 1D system from section 1. Builds a planning rollout directly from dynamics.state_evolution's deterministic mean (a standard MPPI simplification -- the planner doesn't need a faithful stochastic simulation, just a reasonable prediction of where a candidate control sequence leads) and a simple quadratic loss, then compares against the K=0 no-control baseline already computed in section 3. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Fix ty type-check error in test_discrete_control.py DynamicalModel.state_evolution is declared as the broader ContinuousTimeStateEvolution | DiscreteStateTransition union, so `ty` couldn't statically narrow continuous_dynamics.state_evolution to StochasticContinuousTimeStateEvolution before passing it to euler_maruyama() (which requires the narrower type) -- even though this is exactly how Discretizer._sample_ds itself resolves it at runtime. Added the same isinstance narrowing check used there. Also includes an incidental ruff-format tweak to mppi.py (collapsing a one-line call that had been wrapped unnecessarily). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Add gradient-through-the-rollout section to the control tutorial notebook Section 8 demonstrates that DiscreteControlLoopSimulator is differentiable end-to-end: rolls out a short horizon with the linear policy from section 2, defines a loss as the norm of the final true state, and takes jax.grad of it with respect to the gain K -- no special machinery needed beyond plain JAX autodiff. Verified independently against a finite-difference check before writing the notebook cells (gradient matched to ~1e-4), and shows one gradient-descent step meaningfully reducing the loss (0.62 -> 0.18). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * added notebook on optimizing the policy using gradients * Add finite-difference gradient check to control_optimization.ipynb Section 3 (differentiating through the closed loop) computed grad_K via jax.grad but never independently verified it. Added a cell checking a central finite difference on K[0, 0] against the autodiff value -- confirms the gradient through the whole closed loop (policy -> transition -> observation -> filter update) is actually correct, not just plausible. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Switch control_optimization.ipynb to explicit rng_key control Section 3's rollout_final_state_norm now calls sim.simulate(dynamics, rng_key=..., predict_times=...) directly instead of going through dsx.sample/numpyro.handlers.seed -- unneeded here since there's no conditioning/MCMC, and it makes the PRNG key an explicit argument instead of an implicit fixed rng_seed=0. More importantly, this fixes a real issue in section 4's training loop: every epoch previously reused the same fixed seed, so gradient descent could overfit K to one specific noise realization rather than the underlying dynamics. Now each epoch draws a fresh key via jax.random.split, so K_opt has to generalize across realizations. Verified the loop still converges to a sensible stabilizing gain with this fresh-key-per-epoch setup before editing the notebook. Section 4's final comparison plot (run()/dsx.sample) is left unchanged -- there, reusing the same seed for both the unoptimized and optimized rollouts is what makes the before/after comparison fair. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Update final comparison plot to match run()'s new sim.simulate() syntax run() (updated separately) now returns a ControlledSimulatedResult from sim.simulate() directly instead of a numpyro trace dict, so the plotting cell needed the equivalent attribute-access update: trace["loop_X"]["value"] -> result.X. Also removed a leftover duplicate plotting cell still using the old trace-dict access, which was left after the updated one and caused the notebook to fail on execution. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Switch controller_demo.ipynb to direct sim.simulate() syntax throughout Replaces the with sim: dsx.sample(...) + numpyro.handlers.seed/trace pattern across every section (discrete demo, SDE, nonlinear 2D, MPPI) with direct sim.simulate(dynamics, rng_key=key, predict_times=...) calls returning ControlledSimulatedResult objects -- no NumPyro handler or model function needed, since every section here is a pure generative rollout with nothing to condition on. Renamed trace_* variables to result_* throughout, since they're no longer NumPyro trace dicts, and updated all downstream plotting cells from trace["name_field"]["value"] to result.field attribute access. For the continuous-time sections, this also removes the Discretizer handler composition: instead of wrapping dsx.sample in `with Discretizer(...)`, the continuous-time DynamicalModel is discretized once up front via euler_maruyama(state_evolution) into a plain discrete-time DynamicalModel, which is then simulated exactly like any other discrete-time model. Updated the surrounding markdown to match (no more handler-chain explanation needed). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Improved the control_optimization tutorial * Improved controller demo Unified non-linear dynamics, black box and partial observations example * Enabled controllersimulator so use dsx.simulate and inproved tutorial Unify simulate API with control_policy/filter_config; add distribution-returning policies; expand controller tutorial - dsx.simulate() and Simulator now accept control_policy/filter_config, routing to DiscreteControlLoopSimulator the same way dynamics.state_evolution auto-routes to SDESimulator/ODESimulator/DiscreteTimeSimulator. - A policy may return a NumPyro Distribution instead of a value, sampled automatically; MPPI now owns and advances its own exploration randomness via a `seed` field, - Policies self-initialize via an optional initial_state() (replacing the removed policy_state_init argument), mirroring optax's optimizer.init(). - controller_demo.ipynb: sections now use dsx.simulate() directly; added a new random-controller example (Section 5) demonstrating a distribution-returning policy. unchanged on the nonlinear 2D black-box SDE (Section 6). Removed the MPPI example from the basic tutorial (now lives in a separate tutorial. * Rework MPPI and notebook Rework MPPI to take one-step dynamics; complete and stabilize mpc_demo - MPPI now takes dynamics: DynamicalModel (one-step transition) instead of a hand-rolled full-rollout function; it builds the horizon-unrolling and n_samples batching internally via jax.lax.scan/vmap, using .mean when available and falling back to .sample() (own internal key) for black-box transitions. - All fields but dynamics/loss_fn now have defaults (horizon=10, noise_std=1.0, n_samples=20, dt=1.0, temperature=1.0, batched=True, seed=0). - stability: Non-finite (nan/inf) rollout losses are clamped before the softmax weighting, fixing NaN control - mpc_demo.ipynb: completed (previously referenced undefined names, never ran), updated to the new MPPI API, added MPC/MPPI citations. * Updated policy: time and stochasticity Update the policy to now take in state, t_now, t_next, s. Changed how stochastic policies operate: sample within the policy, do not return a distribution. If a distribution is returned, raise an error to indicate to the user that they should sample from it instead. * Typos fixed * Update controller_demo.ipynb * Update controller_demo.ipynb * 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 * updated notebooks * One step filter updates requires filter object Removed the ability to give a filter config to the one step filter update. Now always requires a filter object * Updated the discrete filters constructor for cuthbert filters * Update test_discrete_control.py * improved tutorials * Update controller_demo.ipynb * Updated control optimization: added gradients through MPC Updated the control optimization notebook to use the new syntax. Also added a check to see that gradients flow correctly through both the dynamics and the MPC planning loop. * Added initial policy state to simulator Removed the need/requirement to have an initial state attribute to a policy. Instead this must be passed explicitly to the simulator. This will avoid hidden operations and is more in line with other coding practices. * Update controller_demo.ipynb * Updated cuthbert filters building * MPPI loss now uses the result from dsx.sim * improved MPPI doc * Update discrete_filter.py * revised typing and small fixes * require cuthbert backend for filter in control simulator * reduce intrusiveness; rollback some changes * add back a long comment * make cuthbert filter source explicit * Update mppi.py * Linked to #314 * add links to outstanding issues * bootrap -> initial filter setup with no prev state --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Co-authored-by: Matthew Levine <mattlevine22@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is a draft pull request to implement support for online control + filtering with CD-Dynamax filters.
The main difficulty is that CD-Dynamax does not support one step updates nor does it provide public functions that enable this. Hence a one step update needs to import hidden utilities in the code which might be changed in the future.
I think that a more elegant solution would be to modify CD-Dynamax's codebase to provide public facing function (might be useful for other purposes as well). Hence in the meantime, this PR is a draft mostly to illustrate the ideas and issues.