Skip to content

MPPI

Bases: Module

Model Predictive Path Integral (MPPI) controller.

At each call: sample n_samples candidate control sequences of length horizon_length as perturbations around a nominal sequence:

\[ u = \bar u + \varepsilon \]

where \(\varepsilon\) is drawn from a noise distribution (by default a Gaussian AR(1)). The nominal sequence \(\bar{u}\) is the control sequence from the previous step (initially zero), shifted by one time step. It is carried in the policy state s. Each control sequence is rolled out over the horizon grid through the dynamics. The resulting trajectories are scored with loss_fn and combined 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 temperature-scaled negative losses. Only the first control of the resulting sequence is applied at each step. The remainder becomes the next step's nominal sequence.

By default all candidates are rolled out with the same PRNG key (common random numbers), so they face the same process and observation noise. The n_simulations rollouts of one candidate draw different noise. Set common_rollout_randomness=False to give each candidate its own noise instead (increases variance).

Each rollout is run under the "previous_transition" observation/ control convention, so a candidate's \(u_k\) influences \(x_{k+1}\) and \(y_{k+1}\).

Customizing. Subclass MPPI and override any of the hooks below; plan_step calls them in this order, and the defaults implement the standard MPPI above. Each hook takes its own inputs followed by the same context (x_hat, t_now, s): the belief handed to the policy, the current time, and the policy state s, an MPPIState (see there for carrying more or different state).

  1. rollout_initial_condition(x_hat, t_now, s): the distribution each rollout's \(x_0\) is drawn from. Default: a Delta at x_hat.mean.
  2. sample_controls(key, x_hat, t_now, s): the candidate control sequences (the proposal). Default: s.nominal_sequence plus draws from noise (scaled by noise_std).
  3. Rollouts and loss_fn (not overridable). Non-finite losses are clamped to the largest finite value.
  4. combine_sequences(losses, candidates, x_hat, t_now, s): chooses the weights and returns the plan, the combined sequence. Default: the softmax weighting above.
  5. update_state(plan, info, s): returns the next policy state. info is an MPPIStepInfo with the rest of the step: x_hat, t_now, candidates, losses and results (the rollouts). Default: s with nominal_sequence set to the plan shifted left by one (last entry repeated), other fields unchanged.

plan_step then applies plan[0] and advances the next state's key.

Hooks run inside jax.lax.scan (and possibly jax.grad), so they must be pure and JAX-traceable and must not draw randomness from s.key (sample_controls gets its own key). The state update_state returns must keep the structure, shapes and dtypes of the one it received. New fields on an MPPI subclass need a default (or eqx.field(kw_only=True)).

Attributes:

Name Type Description
dynamics DynamicalModel

a DynamicalModel (the same model used for the real simulation or an approximation). Each candidate rollout is computed by calling dsx.simulate.
loss_fn MPPILossFn

MPPILossFn, i.e. (result: SimulatedResult) -> scalar, called once per sample (vmapped). Every field carries a leading n_simulations axis -- e.g. result.states.shape == (n_simulations, horizon, state_dim), so (1, horizon, state_dim) by default. The starting state \(x_0\) is not in states; it is available separately as result.x_0, shape (n_simulations, state_dim). states, controls and observations correspond to times, ctrl_times and obs_times respectively.
horizon Real[Array, ' horizon_plus_one']

The planning grid relative to t_now, [0, t_1, ..., t_H] (finite, starting at 0, strictly increasing). The rollout runs on [t_now, t_now + t_1, ..., t_now + t_H]. H is the number of internal one-step dynamics calls per rollout (see horizon_length). Required.
noise_std Real[Array, ''] | Real[Array, ' control_dim']

Scale applied to the perturbations drawn from noise, scalar or shape (control_dim,). The noises in dynestyx.control.utils.distribution_utils have unit marginal variance per step, so for them this is the per-step standard deviation. Defaults to 1.0.
noise Distribution | None

Any NumPyro distribution whose samples have shape (horizon_length, control_dim), drawn n_samples times per call as the perturbations around the nominal sequence. None (default) is replaced by AR1Noise(horizon, dynamics.control_dim, rho=0.5) at construction. The built-in noises take the planning grid as times, e.g. AR1Noise(horizon, control_dim, rho=...).
n_samples int

Number of sampled control sequences per call. Defaults to 10.
n_simulations int

Number of rollouts per candidate control sequence (forwarded to dsx.simulate), each with its own noise draw. Defaults to 1.
common_rollout_randomness bool

If True (default), the j-th rollout uses the same noise for every candidate, so candidates are compared under identical noise and their losses differ only through their controls. If False, every candidate draws its own noise.
temperature float

MPPI's \(\lambda \ge 0\); higher values flatten the weights toward a uniform average, lower values concentrate weight on the lowest-loss samples, and 0 applies the lowest-loss candidate alone. Defaults to 1.0.
batched bool

Whether the n_samples candidate rollouts are computed with jax.vmap (default; requires dynamics.state_evolution to be vmap-compatible) or jax.lax.map (slower, but works for a dynamics.state_evolution that isn't vmap-compatible).
rng_key PRNGKeyArray | None

MPPI's own PRNG key, put in the policy state by initial_state (as s.key) and split internally on every call. If None (default), initial_state takes a key from the surrounding NumPyro seed handler (numpyro.prng_key()), and raises if there is none.

horizon_length: int property

Number of planning steps H in the grid horizon.

initial_state() -> MPPIState

Zero nominal control sequence plus MPPI's own PRNG key. Must be initialized explicitly. Override it to return an MPPIState subclass carrying more state.

The key is rng_key if given, otherwise one drawn from the surrounding NumPyro seed handler.

Raises:

Type Description
ValueError

If rng_key is None and there is no seed handler.

rollout_initial_condition(x_hat: Distribution, t_now: Real[Array, ''], s: MPPIState) -> Distribution

Distribution each candidate rollout's \(x_0\) is drawn from.

Default: a Delta at x_hat.mean, i.e. plan from the point estimate. Return x_hat itself to plan over the whole belief instead.

sample_controls(key: PRNGKeyArray, x_hat: Distribution, t_now: Real[Array, ''], s: MPPIState) -> Real[Array, 'n_samples horizon control_dim']

Candidate control sequences to roll out and score (the proposal).

Default: s.nominal_sequence plus n_samples noise_std-scaled draws from noise.

combine_sequences(losses: Real[Array, ' n_samples'], candidates: Real[Array, 'n_samples horizon control_dim'], x_hat: Distribution, t_now: Real[Array, ''], s: MPPIState) -> Real[Array, 'horizon control_dim']

Weight the candidates and combine them into this step's plan, whose first entry is applied. losses are always finite.

Default: the softmax weights softmax(-losses / temperature) and the weighted mean of the candidates. With temperature=0 all the weight goes to the lowest-loss candidate.

update_state(plan: Real[Array, 'horizon control_dim'], info: MPPIStepInfo, s: MPPIState) -> MPPIState

Next policy state, from this step's plan and info (an MPPIStepInfo: the belief, time, candidates, their finite losses and rollouts). Must return the same fields as s.

Default: s with nominal_sequence set to the plan shifted left by one with the last entry repeated (the receding-horizon warm start); other fields unchanged.

plan_step(x_hat: Distribution, t_now: Real[Array, ''], s: MPPIState) -> tuple[Real[Array, ' control_dim'], MPPIState, MPPIStepInfo]

Do MPPI's full planning step and also return what it computed, as the MPPIStepInfo, useful for debugging.

__call__ (used by DiscreteControlLoopSimulator) is a thin wrapper around this that drops the info.

Every field of info.results is shaped (n_samples, n_simulations, horizon, ...).

MPPIState

Bases: Module

MPPI's internal policy state.

To add fields, subclass it and override MPPI.initial_state to return the subclass.

To use them, override the MPPI hooks that need them. A key field is always required: MPPI uses it for its own randomness.

Update fields with s.replace(nominal_sequence=...).

Attributes:

Name Type Description
nominal_sequence Real[Array, 'horizon control_dim']

The sequence the candidates are sampled around, (horizon, control_dim).
key PRNGKeyArray

MPPI's own PRNG key, advanced by MPPI.plan_step on every call.

replace(**changes) -> Self

A copy with the given fields replaced, e.g. s.replace(nominal_sequence=new_nominal). Works for subclasses too.

MPPIStepInfo

Bases: NamedTuple

What one MPPI planning step computed, handed to MPPI.update_state.

Attributes:

Name Type Description
x_hat Distribution

The belief handed to the policy.
t_now Real[Array, '']

The current time.
candidates Real[Array, 'n_samples horizon control_dim']

The candidate control sequences, (n_samples, horizon, control_dim).
losses Real[Array, ' n_samples']

Their (always finite) losses, (n_samples,).
results SimulatedResult

Their rollouts, a SimulatedResult batched over candidates.

Noise distributions

NumPyro distributions for the perturbations MPPI adds to its nominal control sequence. Pass one as MPPI(noise=...), or any other NumPyro distribution whose samples have shape (horizon_length, control_dim).

WhiteNoise

Bases: Distribution

i.i.d. standard Gaussian perturbations, uncorrelated across the horizon. Only the number of planning steps matters, not their times.

Parameters:

Name Type Description Default
times Real[Array, ' horizon_plus_one']

The planning grid [0, t_1, ..., t_H] (MPPI.horizon).

 required
control_dim int

The control dimension.

 required

AR1Noise

Bases: Distribution

Ornstein-Uhlenbeck perturbations observed at the planning times, correlated as Cov(eps_h, eps_h') = rho ** |t_h - t_h'|, where t_h is the time perturbation h starts. Smoother than WhiteNoise; rho=0 is equivalent to WhiteNoise.

Parameters:

Name Type Description Default
times Real[Array, ' horizon_plus_one']

The planning grid [0, t_1, ..., t_H] (MPPI.horizon).

 required
control_dim int

The control dimension.

 required
rho float

Correlation between perturbations one time unit apart, in [0, 1]. Defaults to 0.5.

 0.5

ColoredNoise

Bases: Distribution

Power-law (1/f**beta) perturbations generated in the frequency domain. Smoother, low-frequency-dominated perturbations for larger beta. beta=0 = WhiteNoise.

The FFT assumes equally spaced planning times. On an uneven grid the spectrum is over the step index rather than time. A warning is raised for concrete uneven grids; the diagnostic is skipped when times are traced.

Parameters:

Name Type Description Default
times Real[Array, ' horizon_plus_one']

The planning grid [0, t_1, ..., t_H] (MPPI.horizon).

 required
control_dim int

The control dimension.

 required
beta float

Power-law exponent. 0 is white, 1 is "pink", 2 is Brownian-like. Defaults to 2.0.

 2.0