Control v2

[MatthieuDarcy]

This PR adds at first attempt at support for closed loop control in dynastyx.

Added in /dynestyx/control

1. DiscreteControlLoopSimulator

API:

    def model():
        with DiscreteControlLoopSimulator(
            control_policy=policy,
            policy_state_init=None,
            filter_config=FilterConfig(record_filtered_states_mean=True),
        ):
            return dsx.sample("loop", dynamics, predict_times=predict_times)

2. MPPI

Basic MPPI controller

API

mppi = MPPI(
    dynamics_model=dynamic_rollout
    loss_fn=quadratic_loss,
    horizon=horizon,
    n_samples=n_samples,
    noise_std=1.0,
    temperature=1.0,
)

Added in tutorials
3. Notebook controller_demo.ipynb

Additional elements I will work on

1. Verifying compatibility with different filters.
2. Verifying gradient computation.

Closes #288

[MatthieuDarcy]

The biggest addition to other parts of the code is in cuthbert/discrete_filter.py

compute_cuthbert_filter_update now computes a one step filtering update: takes in the previous estimated state, a control, and observation and provides a new estimated state.

[mattlevine22]

This looks great, thank you @MatthieuDarcy !

1. Can we have @cooijmanstim look at the Jupyter notebook demo, and get his take?

2. My preliminary comments/questions (haven't looked carefully through the code yet):

• I think DiscreteControlLoopSimulator should be made available directly through the Simulator() handler via simulation/auto.py (and, hence, numpyro-free dsx.simulate) interface.
• Should check that we can autodiff through a ClosedLoopSimulation.
• Should probably support n_simulations > 1.
  -DiscreteControlLoopSimulator should another argument for an "approximate model" (assuming the "true" model is in the dsx.sample statement)?
• I like that one can feed very generic controllers (including MPPI); however, I think the current MPPI setup should be better curated so that the user doesn't have to do a bunch of the extra manual work that appears in the demo. I think the following should be possible:

sim_mppi = DiscreteControlLoopSimulator(
    control_policy_config=MPPIConfig(
                                                      loss_fn="quadratic",  # or a callable
                                                      horizon=horizon,
                                                      n_samples=n_samples,
                                                      noise_std=1.0, # is this needed?
                                                      temperature=1.0,
                                                ),
    policy_state_init=mppi_initial_state(horizon, control_dim),
    filter_config=KFConfig(record_filtered_states_mean=True),
)

# MPPI should have make_mppi_rollout internally (don't expose user to it)

[cooijmanstim]

Looks great, this interface works for us. For the notebook it would be good besides no control vs control to see the effectiveness of the controller as the noise level on the observations changes. Just three variants in total will do (no control, control with lots of noise, control with little noise).