mppx

mppx implements Model Predictive Path Integral (MPPI) control using JAX and Equinox. The main features of mppx include:

• MPPI control via solve, for deterministic dynamics models
• batch_solve, for stochastic or ensemble dynamics models
• JIT-compiled solvers built on JAX and Equinox
• Automatic temperature tuning via effective sample size
• Control-sequence smoothing, including a built-in Savitzky-Golay filter

Installation

mppx can be installed using pip or uv,

# pip installation
pip install git+https://github.com/OSU-LRAM/mppx.git

# uv installation
uv pip install git+https://github.com/OSU-LRAM/mppx.git

# or add to your uv project using
uv add git+ssh://git@github.com/OSU-LRAM/mppx.git

Usage

See the following example for usage with learned dynamics models (e.g., neural ODEs).

import diffrax
import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
from mppx import solve
from lrax.common.samplers import ZeroOrderSampler

class VectorField(eqx.Module):
    mlp: eqx.nn.MLP

    def __init__(self, state_dim, control_dim, key):
        in_size = state_dim + control_dim
        self.mlp = eqx.nn.MLP(in_size, state_dim, width_size=32, depth=2, key=key)

    def __call__(self, t, y, sampler):
        return self.mlp(jnp.concatenate([y, sampler.evaluate(t)]))

class NeuralODE(eqx.Module):
    """A basic neural ODE, demonstrating learning-based control."""

    vf: VectorField

    def __init__(self, state_dim, control_dim, key):
        self.vf = VectorField(state_dim, control_dim, key)

    def __call__(self, ts, y0, us, dt):
        sol = diffrax.diffeqsolve(
            diffrax.ODETerm(self.vf),
            diffrax.Euler(),
            t0=0.0,
            t1=ts[-1],
            dt0=dt,
            y0=y0,
            saveat=diffrax.SaveAt(ts=ts),
            args=ZeroOrderSampler(ts, us),
        )
        return sol.ys

def cost_fn(t, y, u, args):
    return jnp.sum(y**2) + jnp.sum(u**2)

min_us, max_us = -jnp.ones(2), jnp.ones(2)
cov = jnp.eye(2) * 0.1
model = NeuralODE(state_dim=2, control_dim=2, key=jr.key(0))

solution = solve(
    model, 0.0, jnp.zeros(2), min_us, max_us, cov, cost_fn,
    horizon=5, num_samples=20, key=jr.key(1),
)

solution is a Solution instance containing the next control to apply (u0), the updated control sequence (us), the predicted next state (y0), and diagnostics from the solve (ys, rollouts, costs, temperature).

References

G. Williams, P. Drews, B. Goldfain, J. M. Rehg and E. A. Theodorou, "Information-Theoretic Model Predictive Control: Theory and Applications to Autonomous Driving," in IEEE Transactions on Robotics, vol. 34, no. 6, pp. 1603-1622, 2018.

License

mppx is released under the MIT license.