Skip to content

Fixed MPPI - #348

Merged
mattlevine22 merged 3 commits into
md-control-alignmentfrom
md-mppi-fix
Sep 24, 2026
Merged

Fixed MPPI#348
mattlevine22 merged 3 commits into
md-control-alignmentfrom
md-mppi-fix

Conversation

@MatthieuDarcy

@MatthieuDarcy MatthieuDarcy commented Aug 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Solves part of #337 .

MPPI now uses the previous_transition convention from dsx.simulate directly.

The implementation is now mathematically correct: it does not duplicate controls.

The initial state $x_0$ are not present in the SimulatedResult fieldsstates, which will avoid a user using them in their loss accidentally. The initial condition $x_0$ is available through res.x_0. This is intuitive from a user perspective (you basically never want to use the initial condition in your loss), but differs from the usual contract. The times field reflects this: you don't have $t_0$.

@MatthieuDarcy MatthieuDarcy mentioned this pull request Aug 31, 2026
Merged
@MatthieuDarcy

MatthieuDarcy commented Sep 1, 2026

Copy link
Copy Markdown
Contributor Author

Now uses vmap directly over SimulatedResult thanks to #350

@MatthieuDarcy @claude
Plan MPPI rollouts under previous_transition (#312)
baf7f25 
MPPI's rollouts now use the observation/control convention the closed
loop runs: each candidate's u_k drives x_{k+1} and is paired with
y_{k+1}, with one control per transition and no padding. loss_fn gets
the rollout with t_0 and x_0 dropped, so times, states, observations and
controls share one index; x_0 stays available as result.x_0.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
mattlevine22
mattlevine22 requested changes Sep 24, 2026
View reviewed changes
Comment thread dynestyx/control/mppi.py
lambda m: (m.initial_condition, m.observation_control_alignment),
self.dynamics,
(dist.Delta(x0, event_dim=1), "same_time"),
(dist.Delta(x0, event_dim=1), "previous_transition"),

@mattlevine22 mattlevine22 Sep 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This silent override might confuse a user---is there a check + warning of the override that happens earlier on?

@MatthieuDarcy MatthieuDarcy Sep 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The override is in the description of MPPI itself (as far as I know its the most natural definition of the MPPI algorithm), but that's about it. The override is also purely internal to MPPI, it doesn't modify the dynamics otherwise.

@MatthieuDarcy MatthieuDarcy Sep 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we could add it in the class definition with default previous transition (matching the closed loop control default).

@mattlevine22 mattlevine22 Sep 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

well the point is that you can hand it dynamics with "same_time" and not realize it is going to overwrite that.

Maybe add a __post_init__ check that warns of the override if a user has specified "same_time"?

@MatthieuDarcy MatthieuDarcy Sep 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes that's easy, done

Comment thread tests/test_discrete_control.py Outdated
(nominal, jr.PRNGKey(0)),
)

# Plain SimulatedResult now carries the controls; no ControlledSimulatedResult.

@mattlevine22 mattlevine22 Sep 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

unnecessary

@MatthieuDarcy MatthieuDarcy Sep 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removed

@mattlevine22
mattlevine22 self-requested a review September 24, 2026 21:47
mattlevine22
mattlevine22 approved these changes Sep 24, 2026
View reviewed changes

@mattlevine22 mattlevine22 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good!

@mattlevine22
mattlevine22 merged commit 0039496 into main Sep 24, 2026
3 checks passed
mattlevine22 pushed a commit that referenced this pull request Sep 24, 2026
@MatthieuDarcy @claude
Define observation/control alignment for discrete-time models (#331)
e6dbc85 
Add observation_control_alignment for discrete-time Simulator (#312).

# Mathematical summary
Implements the `same_time` vs `previous time distinction`. `same_time`
(default, preserves the existing behavior) implements:

$$
\begin{aligned}
x_0 &\sim p_0 \\
y_0 &\sim p(\cdot |x_0, u_0)
\end{aligned}
$$

and subsequently

$$
\begin{aligned}
x_{k+1} &\sim p(\cdot | x_k, u_k)\\
y_k &\sim p(\cdot| x_k, u_k)
\end{aligned}
$$

`previous_time` implements

$$
\begin{aligned}
x_0 &\sim p_{0}\\
x_{k+1} &\sim p(\cdot | x_k, u_k) \\
y_{k+1} &\sim p(\cdot| x_{k+1}, u_k)
\end{aligned}
$$

# Summary of changes


## `previous_transition` contract

Adds an explicit property of `DynamicalModel` defined at initialization
called `observation_control_alignment: Literal[
            "same_time", "previous_transition"
        ] = "same_time"`

When defined as `previous_transition`, this results in the following
behavior

1. $y_0$ cannot be sampled. Hence observations $y$ is of size $T-1$
while the states $x$ are of size $T$.
2. I added the controls $u$ in the returned`SimulatedResults`. This is
very practical when doing MPC. This leads to a similar behavior as
observations: when using "same_time" it is of size $T$, when using
"previous_time" it is of size $T-1$. Only available when running
`DiscreteTimeSimulator`.

This means that `states` and `times` are of different sizes to
`controls` and `observations` (specifically size $T$ vs $T-1$) when
using `previous_transition`.

1. MPPI is updated to use this convention in #348

## Updated `DiscreteControlLoopSimulator` 

`DiscreteControlLoopSimulator` now deliberately always uses the
previous-transition convention, independent of
`dynamics.observation_control_alignment`:

$$
\begin{aligned}
x_0 &\sim p_0, \qquad \hat p_0 = p_0, \\
(u_k, s_{k+1}) &= \pi(\hat p_k, t_k, t_{k+1}, s_k), \\
x_{k+1} &\sim p(\cdot \mid x_k, u_k, t_k, t_{k+1}), \\
y_{k+1} &\sim p(\cdot \mid x_{k+1}, u_k, t_{k+1}).
\end{aligned}
$$

Previously, it was assumed that the observation function could support
`u=None`. A closed-loop result therefore has $T$ states, times, and
beliefs, and $T-1$ aligned observations and controls. This matches the
`previous_transition` convention and avoids imposing an implicit
requirement that the observation model accept `u=None` or be
control-independent, but is change in the behavior.

# Future work

1. `previous_transition` only works for simulation, no conditioning and
not filtering. Worth delegating to a seperate PR?
2. Adding controls in `SimulatedResults` for ODE/SDE/Continuous time. It
might also be worthwile to expand support for continuous control beyond
zero-order hold.
3. `LTI_discrete` does not accept `observation_control_alignment` as a
keyword, I did not touch it but worth looking at it once we are
satisfied with the contract.

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mattlevine22 mattlevine22 mattlevine22 approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@MatthieuDarcy @mattlevine22