diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..b3e56a2 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,29 @@ +name: Tests + +on: + push: + branches: [master] + pull_request: + branches: [master] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.9", "3.12"] + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + pip install pytest + pip install -e . + + - name: Run tests + run: pytest tests/test_mppi.py tests/test_batch_wrapper.py -v diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..9687aba --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,143 @@ +# pytorch_mppi + +Model Predictive Path Integral (MPPI) control library using approximate dynamics in PyTorch. +Implements batched trajectory sampling for GPU-accelerated model-based control. + +## Project Structure + +``` +src/pytorch_mppi/ + __init__.py # Exports MPPI, SMPPI, KMPPI + mppi.py # Core implementation (~625 lines): MPPI, SMPPI, KMPPI, run_mppi + autotune.py # Hyperparameter tuning infrastructure (CMA-ES local optimizer) + autotune_global.py # Ray Tune global search integration + autotune_qd.py # Quality Diversity optimization (CMA-ME via pyribs) +tests/ + pendulum.py # MPPI on true pendulum dynamics (gym) + pendulum_approximate.py # MPPI with learned neural network dynamics + pendulum_approximate_continuous.py # Continuous angle representation variant + smooth_mppi.py # Visual comparison of MPPI, SMPPI, KMPPI + auto_tune_parameters.py # Hyperparameter tuning example + test_batch_wrapper.py # Unit tests for handle_batch_input +``` + +## Architecture + +### Class Hierarchy +- **MPPI** - Base class. Batched trajectory sampling with importance-weighted control update (Algorithm 2, Williams et al. 2017). +- **SMPPI(MPPI)** - Smooth MPPI. Lifts control space to penalize action rate of change. Maintains separate `action_sequence` and `U` (control differences). +- **KMPPI(MPPI)** - Kernel MPPI. Samples fewer support points, interpolates to full trajectory via RBF kernel. Uses `functorch.vmap` for batched interpolation. + +### Key Data Flow (per `command()` call) +1. `shift_nominal_trajectory()` - Roll U forward, append u_init +2. `_compute_perturbed_action_and_noise()` - Sample K noise trajectories (K x T x nu), add to U, bound +3. `_compute_rollout_costs(perturbed_actions)` - **Hot loop**: iterate T timesteps, call user dynamics+cost each step +4. `_compute_total_cost_batch()` - Combine rollout cost + perturbation cost +5. `_compute_weighting()` - Softmax-like exponential weighting (omega) +6. Update U with weighted sum of noise perturbations + +### Key Dimensions +- **K** = `num_samples` (trajectory samples, typically 100-1000) +- **T** = `horizon` (timesteps, typically 15-30) +- **M** = `rollout_samples` (stochastic dynamics replicates, usually 1) +- **nu** = control dimensions, **nx** = state dimensions + +### User-Provided Functions +- `dynamics(state, action) -> next_state` — state is K x nx, action is K x nu +- `running_cost(state, action) -> cost` — cost is K x 1 +- `terminal_state_cost(states, actions) -> cost` — optional, states is K x T x nx +- Wrapped by `@handle_batch_input(n=2)` from `arm_pytorch_utilities` + +## Dependencies +- **torch** — core tensor operations +- **arm_pytorch_utilities** — `handle_batch_input` decorator for flexible batch dimensions +- **functorch** — `vmap` used in KMPPI for batched kernel interpolation +- **numpy** — minimal use (only in `get_params()` for display and in autotune) +- Optional: cma, ray[tune], bayesian-optimization, hyperopt (for autotune) + +## Planned Optimization Refactor + +Goal: Remove Python loops and make code compatible with `torch.compile`, similar to what was done for `pytorch_kinematics`. + +### Performance-Critical Hot Path +The main bottleneck is `_compute_rollout_costs()` (mppi.py:254-267): +```python +for t in range(T): + u = self.u_scale * perturbed_actions[:, t].repeat(self.M, 1, 1) + next_state = self._dynamics(state, u, t) + next_state = self._sample_specific_dynamics(next_state, state, u, t) + state = next_state + c = self._running_cost(state, u, t) + cost_samples = cost_samples + c + states.append(state) + actions.append(u) +actions = torch.stack(actions, dim=-2) +states = torch.stack(states, dim=-2) +``` + +Similarly `get_rollouts()` (mppi.py:357-361) has a horizon loop. + +### torch.compile Blockers +1. **Horizon for-loop with list appends** — `states.append()` / `torch.stack()` pattern. Fix: pre-allocate tensors, use index assignment. +2. **Shape-dependent control flow** (mppi.py:244): `if self.state.shape == (K, self.nx)`. Fix: use attribute flags or always reshape. +3. **`@handle_batch_input` decorator** — wraps dynamics/cost with runtime shape inspection. May need compile-friendly alternative or be applied outside compiled region. +4. **Optional feature branching** — `if self.terminal_state_cost`, `if self.M > 1`, `if self.specific_action_sampler is not None`. Fix: use guards at init time or compile separate variants. +5. **`from functorch import vmap`** — should migrate to `torch.vmap` (functorch is merged into PyTorch core). + +### Vectorization Opportunities +- The horizon loop is inherently sequential (state[t+1] depends on state[t]), so it cannot be parallelized across T. However, `torch.compile` can still optimize the unrolled loop if graph breaks are eliminated. +- Pre-allocating `states` and `actions` tensors avoids dynamic list building. +- The `_compute_weighting` softmax computation is already vectorized. +- Noise sampling and action cost computation are already fully vectorized. + +### Already Optimized +- Noise sampling: fully batched via `MultivariateNormal.rsample((K, T))` +- Action cost: matrix multiply K x T x nu @ nu x nu +- KMPPI kernel interpolation: uses vmap for batch kernel solve +- Cost weighting: vectorized exp + normalize + +### Migration Notes +- `functorch.vmap` → `torch.vmap` (available since PyTorch 2.0) +- Consider whether `arm_pytorch_utilities.handle_batch_input` can be replaced with simpler reshape logic to reduce external dependencies +- User dynamics/cost functions must also be compile-friendly for full graph compilation; consider documenting this requirement or providing a fallback path + +## Development + +```shell +pip install -e . # Dev install +pip install -e .[test] # With test deps +KMP_DUPLICATE_LIB_OK=TRUE pytest tests/test_mppi.py tests/test_batch_wrapper.py -v # Run tests +KMP_DUPLICATE_LIB_OK=TRUE python tests/benchmark_mppi.py # Run benchmarks (saves benchmark_results.json) +``` + +Note: `KMP_DUPLICATE_LIB_OK=TRUE` is needed on this machine due to duplicate OpenMP libraries. + +## Testing + +### Test files +- `tests/test_mppi.py` — 64 tests covering MPPI/SMPPI/KMPPI correctness and solution quality +- `tests/test_batch_wrapper.py` — 2 tests for handle_batch_input decorator +- `tests/benchmark_mppi.py` — Timing + solution quality benchmarks + +### Test categories in test_mppi.py +- **TestMPPI** (26): shapes, bounds, features, determinism, state handling +- **TestSMPPI** (11): SMPPI-specific behavior, smoothness, action sequences +- **TestKMPPI** (12): kernel interpolation, support points, stability +- **TestSpecificActionSampler** (1): custom sampler integration +- **TestEdgeCases** (5): numpy input, high-dim, 1-sample, float32 +- **TestSolutionQuality** (9): goal convergence, cost bounds, determinism, bounded actions + +### Baseline performance (CPU, K=500, T=15) +| Controller | per-command | 20-step loop | +|---|---|---| +| MPPI | 0.63ms | 12.9ms | +| SMPPI | 0.68ms | 13.7ms | +| KMPPI | 1.05ms | 20.5ms | + +### Baseline solution quality (CPU, K=500, T=15, 20 steps, 5 trials) +| Controller | Accum Cost | Final Dist | Control Smoothness | +|---|---|---|---| +| MPPI | 113.5 +/- 17.6 | 1.59 +/- 0.96 | 57.8 +/- 9.0 | +| KMPPI | 111.0 +/- 12.2 | 1.61 +/- 0.58 | 25.9 +/- 3.2 | + +Note: SMPPI quality is highly environment-dependent; it requires careful tuning (action bounds, terminal cost) per environment. KMPPI achieves similar cost to MPPI but with 2x smoother control. diff --git a/pyproject.toml b/pyproject.toml index e7075fc..517b34e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pytorch_mppi" -version = "0.8.0" +version = "0.9.0" description = "Model Predictive Path Integral (MPPI) implemented in pytorch" readme = "README.md" # Optional diff --git a/src/pytorch_mppi/__init__.py b/src/pytorch_mppi/__init__.py index a51e447..06aefba 100644 --- a/src/pytorch_mppi/__init__.py +++ b/src/pytorch_mppi/__init__.py @@ -1 +1 @@ -from pytorch_mppi.mppi import MPPI, SMPPI, KMPPI +from pytorch_mppi.mppi import MPPI, SMPPI, KMPPI, MPPI_Batched diff --git a/src/pytorch_mppi/mppi.py b/src/pytorch_mppi/mppi.py index c21c0df..b6502db 100644 --- a/src/pytorch_mppi/mppi.py +++ b/src/pytorch_mppi/mppi.py @@ -1,11 +1,10 @@ +import functools import logging import time import typing import torch -from torch.distributions.multivariate_normal import MultivariateNormal from arm_pytorch_utilities import handle_batch_input -from functorch import vmap logger = logging.getLogger(__name__) @@ -106,42 +105,62 @@ def __init__(self, dynamics, running_cost, nx, noise_sigma, num_samples=100, hor noise_mu = noise_mu.view(-1) noise_sigma = noise_sigma.view(-1, 1) - # bounds - self.u_min = u_min - self.u_max = u_max + # bounds — resolve to tensor clamp values at init time (no branching in hot path) self.u_scale = u_scale self.u_per_command = u_per_command # make sure if any of them is specified, both are specified - if self.u_max is not None and self.u_min is None: - if not torch.is_tensor(self.u_max): - self.u_max = torch.tensor(self.u_max) - self.u_min = -self.u_max - if self.u_min is not None and self.u_max is None: - if not torch.is_tensor(self.u_min): - self.u_min = torch.tensor(self.u_min) - self.u_max = -self.u_min - if self.u_min is not None: - self.u_min = self.u_min.to(device=self.d) - self.u_max = self.u_max.to(device=self.d) + if u_max is not None and u_min is None: + if not torch.is_tensor(u_max): + u_max = torch.tensor(u_max) + u_min = -u_max + if u_min is not None and u_max is None: + if not torch.is_tensor(u_min): + u_min = torch.tensor(u_min) + u_max = -u_min + # resolve to +/-inf defaults so _bound_action is always a clamp (no branch) + if u_min is not None: + self.u_min = u_min.to(device=self.d) + self.u_max = u_max.to(device=self.d) + else: + self.u_min = torch.tensor(float('-inf'), device=self.d) + self.u_max = torch.tensor(float('inf'), device=self.d) self.noise_mu = noise_mu.to(self.d) self.noise_sigma = noise_sigma.to(self.d) - self.noise_sigma_inv = torch.inverse(self.noise_sigma) - self.noise_dist = MultivariateNormal(self.noise_mu, covariance_matrix=self.noise_sigma) + # detect diagonal covariance for fast element-wise ops + self._diagonal_sigma = torch.equal(self.noise_sigma, torch.diag(torch.diag(self.noise_sigma))) + if self._diagonal_sigma: + diag = torch.diag(self.noise_sigma) + self._noise_sigma_inv_diag = 1.0 / diag + self._noise_sigma_sqrt_diag = torch.sqrt(diag) + self.noise_sigma_inv = torch.diag(self._noise_sigma_inv_diag) + else: + self.noise_sigma_inv = torch.linalg.inv(self.noise_sigma) + self._noise_sigma_chol = torch.linalg.cholesky(self.noise_sigma) # T x nu control sequence self.U = U_init self.u_init = u_init.to(self.d) if self.U is None: - self.U = self.noise_dist.sample((self.T,)) + self.U = self._sample_noise((self.T,)) + # resolve step_dependency at init: wrap dynamics/cost to always accept (state, u, t) self.step_dependency = step_dependent_dynamics + if step_dependent_dynamics: + self._dynamics_fn = dynamics + self._running_cost_fn = running_cost + else: + self._dynamics_fn = lambda state, u, t: dynamics(state, u) + self._running_cost_fn = lambda state, u, t: running_cost(state, u) self.F = dynamics self.running_cost = running_cost self.terminal_state_cost = terminal_state_cost self.sample_null_action = sample_null_action self.specific_action_sampler = specific_action_sampler + # resolve terminal_state_cost to a no-op if not provided (no branch in hot path) + self._terminal_state_cost_fn = terminal_state_cost if terminal_state_cost is not None else lambda states, actions: 0 self.noise_abs_cost = noise_abs_cost + self._setup_action_cost_fn() self.state = None self.info = None @@ -150,6 +169,13 @@ def __init__(self, dynamics, running_cost, nx, noise_sigma, num_samples=100, hor self.rollout_var_cost = rollout_var_cost self.rollout_var_discount = rollout_var_discount + # pre-compute discount factors for variance cost + if self.M > 1: + self._var_discount_factors = rollout_var_discount ** torch.arange( + self.T, device=self.d, dtype=self.dtype) + else: + self._var_discount_factors = None + # sampled results from last command self.cost_total = None self.cost_total_non_zero = None @@ -157,17 +183,48 @@ def __init__(self, dynamics, running_cost, nx, noise_sigma, num_samples=100, hor self.states = None self.actions = None + def _setup_action_cost_fn(self): + """Resolve action cost computation at init time (diagonal vs full, abs vs not).""" + if self._diagonal_sigma: + inv_diag = self._noise_sigma_inv_diag + if self.noise_abs_cost: + self._compute_action_cost = lambda noise: self.lambda_ * torch.abs(noise) * inv_diag + else: + self._compute_action_cost = lambda noise: self.lambda_ * noise * inv_diag + else: + sigma_inv = self.noise_sigma_inv + if self.noise_abs_cost: + self._compute_action_cost = lambda noise: self.lambda_ * torch.abs(noise) @ sigma_inv + else: + self._compute_action_cost = lambda noise: self.lambda_ * noise @ sigma_inv + + def _sample_noise(self, shape): + """Sample from N(noise_mu, noise_sigma) using pre-computed Cholesky/diagonal factor.""" + z = torch.randn(*shape, self.nu, device=self.d, dtype=self.dtype) + if self._diagonal_sigma: + return z * self._noise_sigma_sqrt_diag + self.noise_mu + return z @ self._noise_sigma_chol.T + self.noise_mu + + def compile(self, **kwargs): + """Compile the hot path with torch.compile for reduced kernel launch overhead. + + User's dynamics and running_cost functions must be torch.compile-compatible + (no graph breaks) for full benefit. Any kwargs are passed to torch.compile. + """ + self._dynamics_fn = torch.compile(self._dynamics_fn, **kwargs) + self._running_cost_fn = torch.compile(self._running_cost_fn, **kwargs) + def get_params(self): return f"K={self.K} T={self.T} M={self.M} lambda={self.lambda_} noise_mu={self.noise_mu.cpu().numpy()} noise_sigma={self.noise_sigma.cpu().numpy()}".replace( "\n", ",") @handle_batch_input(n=2) def _dynamics(self, state, u, t): - return self.F(state, u, t) if self.step_dependency else self.F(state, u) + return self._dynamics_fn(state, u, t) @handle_batch_input(n=2) def _running_cost(self, state, u, t): - return self.running_cost(state, u, t) if self.step_dependency else self.running_cost(state, u) + return self._running_cost_fn(state, u, t) def get_action_sequence(self): return self.U @@ -208,7 +265,7 @@ def _command(self, state): cost_total = self._compute_total_cost_batch() self._compute_weighting(cost_total) - perturbations = torch.sum(self.omega.view(-1, 1, 1) * self.noise, dim=0) + perturbations = torch.einsum('k,ktn->tn', self.omega, self.noise) self.U = self.U + perturbations action = self.get_action_sequence()[:self.u_per_command] @@ -230,9 +287,52 @@ def reset(self): """ Clear controller state after finishing a trial """ - self.U = self.noise_dist.sample((self.T,)) + self.U = self._sample_noise((self.T,)) def _compute_rollout_costs(self, perturbed_actions): + if self.M == 1: + return self._compute_rollout_costs_single(perturbed_actions) + return self._compute_rollout_costs_multi(perturbed_actions) + + def _compute_rollout_costs_single(self, perturbed_actions): + """Fast path for M=1 (deterministic dynamics, the common case).""" + K, T, nu = perturbed_actions.shape + cost_total = torch.zeros(K, device=self.d, dtype=self.dtype) + + if self.state.shape == (K, self.nx): + state = self.state.clone() + else: + state = self.state.view(1, -1).expand(K, -1) + + need_storage = self.terminal_state_cost is not None + if need_storage: + states = torch.empty(1, K, T, self.nx, device=self.d, dtype=self.dtype) + actions = torch.empty(1, K, T, nu, device=self.d, dtype=self.dtype) + + for t in range(T): + u = self.u_scale * perturbed_actions[:, t] + state = self._dynamics_fn(state, u, t) + if self.specific_action_sampler is not None: + state = self.specific_action_sampler.specific_dynamics( + state.unsqueeze(0), state.unsqueeze(0), u.unsqueeze(0), t).squeeze(0) + c = self._running_cost_fn(state, u, t) + cost_total = cost_total + c.reshape(K) + if need_storage: + states[0, :, t] = state[:, :self.nx] + actions[0, :, t] = u + + if need_storage: + c = self._terminal_state_cost_fn(states, actions) + if torch.is_tensor(c) and c.dim() > 1: + c = c.squeeze(0) + cost_total = cost_total + c + else: + states = None + actions = None + return cost_total, states, actions + + def _compute_rollout_costs_multi(self, perturbed_actions): + """General path for M>1 (stochastic dynamics with variance cost).""" K, T, nu = perturbed_actions.shape assert nu == self.nu @@ -240,41 +340,34 @@ def _compute_rollout_costs(self, perturbed_actions): cost_samples = cost_total.repeat(self.M, 1) cost_var = torch.zeros_like(cost_total) - # allow propagation of a sample of states (ex. to carry a distribution), or to start with a single state if self.state.shape == (K, self.nx): state = self.state else: - state = self.state.view(1, -1).repeat(K, 1) + state = self.state.view(1, -1).expand(K, -1) - # rollout action trajectory M times to estimate expected cost state = state.repeat(self.M, 1, 1) - states = [] - actions = [] + states = torch.empty(self.M, K, T, self.nx, device=self.d, dtype=self.dtype) + actions = torch.empty(self.M, K, T, nu, device=self.d, dtype=self.dtype) + MK = self.M * K + state_flat = state.reshape(MK, self.nx) for t in range(T): - u = self.u_scale * perturbed_actions[:, t].repeat(self.M, 1, 1) - next_state = self._dynamics(state, u, t) - # potentially handle dynamics in a specific way for the specific action sampler - next_state = self._sample_specific_dynamics(next_state, state, u, t) - state = next_state - c = self._running_cost(state, u, t) - cost_samples = cost_samples + c - if self.M > 1: - cost_var += c.var(dim=0) * (self.rollout_var_discount ** t) - - # Save total states/actions - states.append(state) - actions.append(u) - - # Actions is K x T x nu - # States is K x T x nx - actions = torch.stack(actions, dim=-2) - states = torch.stack(states, dim=-2) - - # action perturbation cost - if self.terminal_state_cost: - c = self.terminal_state_cost(states, actions) - cost_samples = cost_samples + c + u = self.u_scale * perturbed_actions[:, t].expand(self.M, -1, -1) + u_flat = u.reshape(MK, nu) + state_flat = self._dynamics_fn(state_flat, u_flat, t) + if self.specific_action_sampler is not None: + state_3d = state_flat.reshape(self.M, K, -1) + state_3d = self.specific_action_sampler.specific_dynamics(state_3d, state.reshape(self.M, K, -1), u, t) + state_flat = state_3d.reshape(MK, -1) + c = self._running_cost_fn(state_flat, u_flat, t) + cost_samples = cost_samples + c.reshape(self.M, K) + cost_var += c.reshape(self.M, K).var(dim=0) * self._var_discount_factors[t] + + states[:, :, t] = state_flat.reshape(self.M, K, -1)[:, :, :self.nx] + actions[:, :, t] = u + + c = self._terminal_state_cost_fn(states, actions) + cost_samples = cost_samples + c cost_total = cost_total + cost_samples.mean(dim=0) cost_total = cost_total + cost_var * self.rollout_var_cost return cost_total, states, actions @@ -282,7 +375,7 @@ def _compute_rollout_costs(self, perturbed_actions): def _compute_perturbed_action_and_noise(self): # parallelize sampling across trajectories # resample noise each time we take an action - noise = self.noise_dist.rsample((self.K, self.T)) + noise = self._sample_noise((self.K, self.T)) # broadcast own control to noise over samples; now it's K x T x nu perturbed_action = self.U + noise perturbed_action = self._sample_specific_actions(perturbed_action) @@ -313,16 +406,10 @@ def _sample_specific_dynamics(self, next_state, state, u, t): def _compute_total_cost_batch(self): self._compute_perturbed_action_and_noise() - if self.noise_abs_cost: - action_cost = self.lambda_ * torch.abs(self.noise) @ self.noise_sigma_inv - # NOTE: The original paper does self.lambda_ * torch.abs(self.noise) @ self.noise_sigma_inv, but this biases - # the actions with low noise if all states have the same cost. With abs(noise) we prefer actions close to the - # nomial trajectory. - else: - action_cost = self.lambda_ * self.noise @ self.noise_sigma_inv # Like original paper + action_cost = self._compute_action_cost(self.noise) rollout_cost, self.states, actions = self._compute_rollout_costs(self.perturbed_action) - self.actions = actions / self.u_scale + self.actions = actions / self.u_scale if actions is not None else None # action perturbation cost perturbation_cost = torch.sum(self.U * action_cost, dim=(1, 2)) @@ -330,9 +417,7 @@ def _compute_total_cost_batch(self): return self.cost_total def _bound_action(self, action): - if self.u_max is not None: - return torch.max(torch.min(action, self.u_max), self.u_min) - return action + return torch.clamp(action, self.u_min, self.u_max) def _slice_control(self, t): return slice(t * self.nu, (t + 1) * self.nu) @@ -347,7 +432,7 @@ def get_rollouts(self, state, num_rollouts=1, U=None): """ state = state.view(-1, self.nx) if state.size(0) == 1: - state = state.repeat(num_rollouts, 1) + state = state.expand(num_rollouts, -1) if U is None: U = self.get_action_sequence() @@ -356,7 +441,7 @@ def get_rollouts(self, state, num_rollouts=1, U=None): states[:, 0] = state for t in range(T): next_state = self._dynamics(states[:, t].view(num_rollouts, -1), - self.u_scale * U[t].tile(num_rollouts, 1), t) + self.u_scale * U[t].expand(num_rollouts, -1), t) # dynamics may augment state; here we just take the first nx dimensions states[:, t + 1] = next_state[:, :self.nx] @@ -376,19 +461,20 @@ def __init__(self, *args, w_action_seq_cost=1., delta_t=1., U_init=None, action_ super().__init__(*args, U_init=U_init, **kwargs) # these are the actual commanded actions, which is now no longer directly sampled - self.action_min = action_min - self.action_max = action_max - if self.action_min is not None and self.action_max is None: - if not torch.is_tensor(self.action_min): - self.action_min = torch.tensor(self.action_min) - self.action_max = -self.action_min - if self.action_max is not None and self.action_min is None: - if not torch.is_tensor(self.action_max): - self.action_max = torch.tensor(self.action_max) - self.action_min = -self.action_max - if self.action_min is not None: - self.action_min = self.action_min.to(device=self.d) - self.action_max = self.action_max.to(device=self.d) + if action_min is not None and action_max is None: + if not torch.is_tensor(action_min): + action_min = torch.tensor(action_min) + action_max = -action_min + if action_max is not None and action_min is None: + if not torch.is_tensor(action_max): + action_max = torch.tensor(action_max) + action_min = -action_max + if action_min is not None: + self.action_min = action_min.to(device=self.d) + self.action_max = action_max.to(device=self.d) + else: + self.action_min = torch.tensor(float('-inf'), device=self.d) + self.action_max = torch.tensor(float('inf'), device=self.d) # this smooth formulation works better if control starts from 0 if U_init is None: @@ -426,14 +512,10 @@ def change_horizon(self, horizon): self.T = horizon def _bound_d_action(self, control): - if self.u_max is not None: - return torch.max(torch.min(control, self.u_max), self.u_min) # action - return control + return torch.clamp(control, self.u_min, self.u_max) def _bound_action(self, action): - if self.action_max is not None: - return torch.max(torch.min(action, self.action_max), self.action_min) - return action + return torch.clamp(action, self.action_min, self.action_max) def _command(self, state): if not torch.is_tensor(state): @@ -442,7 +524,7 @@ def _command(self, state): cost_total = self._compute_total_cost_batch() self._compute_weighting(cost_total) - perturbations = torch.sum(self.omega.view(-1, 1, 1) * self.noise, dim=0) + perturbations = torch.einsum('k,ktn->tn', self.omega, self.noise) self.U = self.U + perturbations # U is now the lifted control space, so we integrate it @@ -457,7 +539,7 @@ def _command(self, state): def _compute_perturbed_action_and_noise(self): # parallelize sampling across trajectories # resample noise each time we take an action - noise = self.noise_dist.rsample((self.K, self.T)) + noise = self._sample_noise((self.K, self.T)) # broadcast own control to noise over samples; now it's K x T x nu perturbed_control = self.U + noise # naively bound control @@ -471,13 +553,7 @@ def _compute_perturbed_action_and_noise(self): def _compute_total_cost_batch(self): self._compute_perturbed_action_and_noise() - if self.noise_abs_cost: - action_cost = self.lambda_ * torch.abs(self.noise) @ self.noise_sigma_inv - # NOTE: The original paper does self.lambda_ * torch.abs(self.noise) @ self.noise_sigma_inv, but this biases - # the actions with low noise if all states have the same cost. With abs(noise) we prefer actions close to the - # nomial trajectory. - else: - action_cost = self.lambda_ * self.noise @ self.noise_sigma_inv # Like original paper + action_cost = self._compute_action_cost(self.noise) # action difference as cost action_diff = self.u_scale * torch.diff(self.perturbed_action, dim=-2) @@ -486,7 +562,7 @@ def _compute_total_cost_batch(self): action_smoothness_cost *= self.w_action_seq_cost rollout_cost, self.states, actions = self._compute_rollout_costs(self.perturbed_action) - self.actions = actions / self.u_scale + self.actions = actions / self.u_scale if actions is not None else None # action perturbation cost perturbation_cost = torch.sum(self.U * action_cost, dim=(1, 2)) @@ -527,6 +603,8 @@ def __init__(self, *args, num_support_pts=None, kernel: TimeKernel = RBFKernel() # interpolation kernel self.interpolation_kernel = kernel self.intp_krnl = None + # pre-compute kernel matrix for support points (constant across calls) + self._Ktktk_cached = None self.prepare_vmap_interpolation() def get_params(self): @@ -543,33 +621,43 @@ def shift_nominal_trajectory(self): def do_kernel_interpolation(self, t, tk, c): K = self.interpolation_kernel(t.unsqueeze(-1), tk.unsqueeze(-1)) Ktktk = self.interpolation_kernel(tk.unsqueeze(-1), tk.unsqueeze(-1)) - # print(K.shape, Ktktk.shape) - # row normalize K - # K = K / K.sum(dim=1).unsqueeze(1) - # KK = K @ torch.inverse(Ktktk) KK = torch.linalg.solve(Ktktk, K, left=False) return torch.matmul(KK, c), K + @staticmethod + def _do_kernel_interpolation_with_ktktk(interpolation_kernel, t, tk, c, Ktktk): + """Variant that takes pre-computed Ktktk as an argument for vmap compatibility.""" + K = interpolation_kernel(t.unsqueeze(-1), tk.unsqueeze(-1)) + KK = torch.linalg.solve(Ktktk, K, left=False) + return torch.matmul(KK, c), K + def prepare_vmap_interpolation(self): self.Tk = torch.linspace(0, self.T - 1, int(self.num_support_pts), device=self.d, dtype=self.dtype).unsqueeze( 0).repeat(self.K, 1) self.Hs = torch.linspace(0, self.T - 1, int(self.T), device=self.d, dtype=self.dtype).unsqueeze(0).repeat( self.K, 1) - self.intp_krnl = vmap(self.do_kernel_interpolation) + # cache the kernel matrix for support points (same for every sample, repeated for vmap) + tk_single = self.Tk[0] + self._Ktktk_batch = self.interpolation_kernel( + tk_single.unsqueeze(-1), tk_single.unsqueeze(-1) + ).unsqueeze(0).repeat(self.K, 1, 1) + + fn = functools.partial(self._do_kernel_interpolation_with_ktktk, self.interpolation_kernel) + self.intp_krnl = torch.vmap(fn) def deparameterize_to_trajectory_single(self, theta): return self.do_kernel_interpolation(self.Hs[0], self.Tk[0], theta) def deparameterize_to_trajectory_batch(self, theta): assert theta.shape == (self.K, self.num_support_pts, self.nu) - return self.intp_krnl(self.Hs, self.Tk, theta) + return self.intp_krnl(self.Hs, self.Tk, theta, self._Ktktk_batch) def _compute_perturbed_action_and_noise(self): # parallelize sampling across trajectories # resample noise each time we take an action - noise = self.noise_dist.rsample((self.K, self.num_support_pts)) + noise = self._sample_noise((self.K, self.num_support_pts)) perturbed_control_pts = self.theta + noise # control points in the same space as control and should be bounded perturbed_control_pts = self._bound_action(perturbed_control_pts) @@ -588,7 +676,7 @@ def _command(self, state): cost_total = self._compute_total_cost_batch() self._compute_weighting(cost_total) - perturbations = torch.sum(self.omega.view(-1, 1, 1) * self.noise_theta, dim=0) + perturbations = torch.einsum('k,ktn->tn', self.omega, self.noise_theta) self.theta = self.theta + perturbations self.U, _ = self.deparameterize_to_trajectory_single(self.theta) @@ -600,6 +688,191 @@ def _command(self, state): return action +class MPPI_Batched: + """MPPI for N parallel environments sharing a single dynamics/cost call. + + Instead of calling MPPI.command() N times independently, this class + concatenates all N×K samples into a single (N*K, nx) batch for one + dynamics/cost call per timestep, amortizing kernel launch overhead on GPU. + + This is useful for parallel simulation (e.g., N robots training simultaneously). + """ + + def __init__(self, dynamics, running_cost, nx, noise_sigma, num_envs, + num_samples=100, horizon=15, device="cpu", + lambda_=1., + noise_mu=None, + u_min=None, + u_max=None, + u_init=None, + u_scale=1, + u_per_command=1, + step_dependent_dynamics=False, + noise_abs_cost=False): + """ + :param dynamics: function(state, action) -> next_state, called with (N*K, nx) batches + :param running_cost: function(state, action) -> cost, called with (N*K, nx) batches + :param nx: state dimension + :param noise_sigma: (nu x nu) control noise covariance + :param num_envs: N, number of parallel environments + :param num_samples: K, number of trajectory samples per environment + :param horizon: T, planning horizon + :param device: pytorch device + :param lambda_: temperature parameter + :param noise_mu: (nu) noise mean; defaults to zero + :param u_min: (nu) minimum control bounds + :param u_max: (nu) maximum control bounds + :param u_init: (nu) initial control value for new timesteps + :param u_scale: scaling factor applied to actions before dynamics + :param u_per_command: number of actions to return per command + :param step_dependent_dynamics: whether dynamics/cost take timestep as 3rd arg + :param noise_abs_cost: use absolute value of noise in action cost + """ + self.d = device + self.dtype = noise_sigma.dtype + self.N = num_envs + self.K = num_samples + self.T = horizon + self.nx = nx + self.nu = 1 if len(noise_sigma.shape) == 0 else noise_sigma.shape[0] + self.lambda_ = lambda_ + self.u_scale = u_scale + self.u_per_command = u_per_command + + if noise_mu is None: + noise_mu = torch.zeros(self.nu, dtype=self.dtype) + if u_init is None: + u_init = torch.zeros_like(noise_mu) + + if self.nu == 1: + noise_mu = noise_mu.view(-1) + noise_sigma = noise_sigma.view(-1, 1) + + # bounds + if u_max is not None and u_min is None: + if not torch.is_tensor(u_max): + u_max = torch.tensor(u_max) + u_min = -u_max + if u_min is not None and u_max is None: + if not torch.is_tensor(u_min): + u_min = torch.tensor(u_min) + u_max = -u_min + if u_min is not None: + self.u_min = u_min.to(device=self.d) + self.u_max = u_max.to(device=self.d) + else: + self.u_min = torch.tensor(float('-inf'), device=self.d) + self.u_max = torch.tensor(float('inf'), device=self.d) + + self.noise_mu = noise_mu.to(self.d) + self.noise_sigma = noise_sigma.to(self.d) + + # diagonal sigma detection + self._diagonal_sigma = torch.equal(self.noise_sigma, torch.diag(torch.diag(self.noise_sigma))) + if self._diagonal_sigma: + diag = torch.diag(self.noise_sigma) + self._noise_sigma_inv_diag = 1.0 / diag + self._noise_sigma_sqrt_diag = torch.sqrt(diag) + else: + self.noise_sigma_inv = torch.linalg.inv(self.noise_sigma) + self._noise_sigma_chol = torch.linalg.cholesky(self.noise_sigma) + + self.noise_abs_cost = noise_abs_cost + + # action cost function + if self._diagonal_sigma: + inv_diag = self._noise_sigma_inv_diag + if noise_abs_cost: + self._compute_action_cost = lambda noise: self.lambda_ * torch.abs(noise) * inv_diag + else: + self._compute_action_cost = lambda noise: self.lambda_ * noise * inv_diag + else: + sigma_inv = self.noise_sigma_inv + if noise_abs_cost: + self._compute_action_cost = lambda noise: self.lambda_ * torch.abs(noise) @ sigma_inv + else: + self._compute_action_cost = lambda noise: self.lambda_ * noise @ sigma_inv + + self.u_init = u_init.to(self.d) + + # dynamics/cost wrappers + if step_dependent_dynamics: + self._dynamics_fn = dynamics + self._running_cost_fn = running_cost + else: + self._dynamics_fn = lambda state, u, t: dynamics(state, u) + self._running_cost_fn = lambda state, u, t: running_cost(state, u) + + # (N, T, nu) nominal trajectories — one per environment + self.U = self._sample_noise((self.N, self.T)) + + def _sample_noise(self, shape): + z = torch.randn(*shape, self.nu, device=self.d, dtype=self.dtype) + if self._diagonal_sigma: + return z * self._noise_sigma_sqrt_diag + self.noise_mu + return z @ self._noise_sigma_chol.T + self.noise_mu + + def compile(self, **kwargs): + self._dynamics_fn = torch.compile(self._dynamics_fn, **kwargs) + self._running_cost_fn = torch.compile(self._running_cost_fn, **kwargs) + + def reset(self): + self.U = self._sample_noise((self.N, self.T)) + + def command(self, states, shift_nominal_trajectory=True): + """ + :param states: (N, nx) current states for all environments + :param shift_nominal_trajectory: whether to shift U forward + :returns: (N, nu) or (N, u_per_command, nu) actions + """ + if not torch.is_tensor(states): + states = torch.tensor(states) + states = states.to(dtype=self.dtype, device=self.d) + N, K, T, nu = self.N, self.K, self.T, self.nu + + if shift_nominal_trajectory: + self.U = torch.roll(self.U, -1, dims=1) + self.U[:, -1] = self.u_init + + # shared noise across environments: (K, T, nu) + noise = self._sample_noise((K, T)) + # perturbed actions: (N, K, T, nu) + perturbed_actions = self.U.unsqueeze(1) + noise.unsqueeze(0) + perturbed_actions = torch.clamp(perturbed_actions, self.u_min, self.u_max) + actual_noise = perturbed_actions - self.U.unsqueeze(1) + + # flatten (N, K) → (N*K) for single dynamics/cost call + NK = N * K + state = states.unsqueeze(1).expand(N, K, self.nx).reshape(NK, self.nx) + cost_total = torch.zeros(N, K, device=self.d, dtype=self.dtype) + + for t in range(T): + u = self.u_scale * perturbed_actions[:, :, t].reshape(NK, nu) + state = self._dynamics_fn(state, u, t) + c = self._running_cost_fn(state, u, t) + cost_total = cost_total + c.reshape(N, K) + + # action cost + action_cost = self._compute_action_cost(actual_noise) + perturbation_cost = torch.sum(self.U.unsqueeze(1) * action_cost, dim=(2, 3)) + total_cost = cost_total + perturbation_cost + + # independent weighting per environment + beta = total_cost.min(dim=1, keepdim=True).values + cost_non_zero = _ensure_non_zero(total_cost, beta, 1.0 / self.lambda_) + eta = cost_non_zero.sum(dim=1, keepdim=True) + omega = cost_non_zero / eta + + # weighted perturbation update + perturbations = torch.einsum('nk,nktd->ntd', omega, actual_noise) + self.U = self.U + perturbations + + action = self.U[:, :self.u_per_command] + if self.u_per_command == 1: + action = action[:, 0] + return action + + def run_mppi(mppi, env, retrain_dynamics, retrain_after_iter=50, iter=1000, render=True): dataset = torch.zeros((retrain_after_iter, mppi.nx + mppi.nu), dtype=mppi.U.dtype, device=mppi.d) total_reward = 0 diff --git a/tests/benchmark_mppi.py b/tests/benchmark_mppi.py new file mode 100644 index 0000000..2279742 --- /dev/null +++ b/tests/benchmark_mppi.py @@ -0,0 +1,455 @@ +"""Performance benchmarks for MPPI, SMPPI, and KMPPI. + +Run with: python tests/benchmark_mppi.py +Produces timing and solution quality results for various configurations +to track performance before and after optimization refactors. + +Solution quality metrics: + - accumulated_cost: Total running cost over a multi-step trajectory + - final_dist: Euclidean distance from goal at end of trajectory + - control_smoothness: Sum of |u_{t+1} - u_t| (lower = smoother) +""" +import time +import json +import sys +import torch +from pytorch_mppi import MPPI, SMPPI, KMPPI +from pytorch_mppi.mppi import RBFKernel + +DEVICE = "cpu" +DTYPE = torch.double +SEED = 42 + +# Also benchmark on CUDA if available +DEVICES = ["cpu"] +if torch.cuda.is_available(): + DEVICES.append("cuda") + +# --------------------------------------------------------------------------- +# Test environments (no external deps) +# --------------------------------------------------------------------------- +B_MATRIX = torch.tensor([[1.0, 0.0], [0.0, -1.0]], dtype=DTYPE) +GOAL = torch.tensor([2.0, 2.0], dtype=DTYPE) + + +def _make_dynamics(device): + b = B_MATRIX.to(device=device) + + def dynamics(state, action): + return state + action @ b.T + + return dynamics + + +def _make_cost(device): + goal = GOAL.to(device=device) + + def cost(state, action): + dx = goal - state + return (dx ** 2).sum(dim=-1) + + return cost + + +def _make_terminal_cost(device): + goal = GOAL.to(device=device) + + def terminal_cost(states, actions): + dx = goal - states[..., -1, :] + return (dx ** 2).sum(dim=-1) + + return terminal_cost + + +# Higher-dimensional environment +def _make_dynamics_nd(device, nx, nu): + def dynamics(state, action): + delta = torch.zeros_like(state) + delta[..., :nu] = action + return state + delta + + return dynamics + + +def _make_cost_nd(device, nx): + def cost(state, action): + return (state ** 2).sum(dim=-1) + + return cost + + +# --------------------------------------------------------------------------- +# Benchmark harness +# --------------------------------------------------------------------------- +def benchmark_command(ctrl, state, num_warmup=3, num_iters=20): + """Benchmark the command() method, returning stats in seconds.""" + # Warmup + for _ in range(num_warmup): + ctrl.command(state, shift_nominal_trajectory=False) + + if state.device.type == "cuda": + torch.cuda.synchronize() + + times = [] + for _ in range(num_iters): + ctrl.reset() + s = state.clone() + if state.device.type == "cuda": + torch.cuda.synchronize() + t0 = time.perf_counter() + ctrl.command(s) + if state.device.type == "cuda": + torch.cuda.synchronize() + t1 = time.perf_counter() + times.append(t1 - t0) + + times = sorted(times) + # Drop fastest and slowest 10% + trim = max(1, len(times) // 10) + trimmed = times[trim:-trim] if len(times) > 2 * trim else times + mean_t = sum(trimmed) / len(trimmed) + min_t = times[0] + max_t = times[-1] + return {"mean_s": mean_t, "min_s": min_t, "max_s": max_t, "num_iters": num_iters} + + +def benchmark_multi_step(ctrl, state, dynamics_fn, num_steps=20, num_warmup=2, num_iters=5): + """Benchmark a full control loop of num_steps.""" + for _ in range(num_warmup): + ctrl.reset() + s = state.clone() + for _ in range(num_steps): + a = ctrl.command(s) + s = dynamics_fn(s.unsqueeze(0), a.unsqueeze(0)).squeeze(0) + + if state.device.type == "cuda": + torch.cuda.synchronize() + + times = [] + for _ in range(num_iters): + ctrl.reset() + s = state.clone() + if state.device.type == "cuda": + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(num_steps): + a = ctrl.command(s) + s = dynamics_fn(s.unsqueeze(0), a.unsqueeze(0)).squeeze(0) + if state.device.type == "cuda": + torch.cuda.synchronize() + t1 = time.perf_counter() + times.append(t1 - t0) + + mean_t = sum(times) / len(times) + return {"mean_s": mean_t, "min_s": min(times), "max_s": max(times), + "per_step_s": mean_t / num_steps, "num_steps": num_steps} + + +def evaluate_quality(ctrl, state, dynamics_fn, cost_fn, goal, num_steps=20, num_trials=5): + """Evaluate solution quality over multiple seeded trials. + + Returns dict with: + accumulated_cost: mean total running cost over trajectory + final_dist: mean Euclidean distance to goal at end + control_smoothness: mean sum of |u_{t+1} - u_t| + per_trial: list of per-trial results for variance analysis + """ + per_trial = [] + for trial in range(num_trials): + torch.manual_seed(SEED + trial) + ctrl.reset() + s = state.clone() + total_cost = 0.0 + actions = [] + for _ in range(num_steps): + a = ctrl.command(s) + actions.append(a.clone()) + c = cost_fn(s.unsqueeze(0), a.unsqueeze(0)).item() + total_cost += c + s = dynamics_fn(s.unsqueeze(0), a.unsqueeze(0)).squeeze(0) + final_dist = (s - goal).norm().item() + actions_t = torch.stack(actions) + control_diff = actions_t.diff(dim=0).abs().sum().item() + per_trial.append({ + "accumulated_cost": total_cost, + "final_dist": final_dist, + "control_smoothness": control_diff, + }) + + acc_costs = [t["accumulated_cost"] for t in per_trial] + dists = [t["final_dist"] for t in per_trial] + smooths = [t["control_smoothness"] for t in per_trial] + return { + "accumulated_cost_mean": sum(acc_costs) / len(acc_costs), + "accumulated_cost_std": (sum((x - sum(acc_costs)/len(acc_costs))**2 for x in acc_costs) / len(acc_costs)) ** 0.5, + "final_dist_mean": sum(dists) / len(dists), + "final_dist_std": (sum((x - sum(dists)/len(dists))**2 for x in dists) / len(dists)) ** 0.5, + "control_smoothness_mean": sum(smooths) / len(smooths), + "control_smoothness_std": (sum((x - sum(smooths)/len(smooths))**2 for x in smooths) / len(smooths)) ** 0.5, + "per_trial": per_trial, + } + + +# --------------------------------------------------------------------------- +# Benchmark configurations +# --------------------------------------------------------------------------- +def run_benchmarks(): + results = {} + + for device in DEVICES: + print(f"\n{'='*60}") + print(f"Device: {device}") + print(f"{'='*60}") + + dynamics = _make_dynamics(device) + cost = _make_cost(device) + terminal = _make_terminal_cost(device) + noise_sigma = torch.eye(2, dtype=DTYPE, device=device) + start = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=device) + + # --- Vary num_samples (K) --- + print(f"\n--- MPPI: Varying K (T=15) ---") + for K in [50, 100, 500, 1000, 5000]: + torch.manual_seed(SEED) + ctrl = MPPI(dynamics, cost, 2, noise_sigma, + num_samples=K, horizon=15, device=device, lambda_=1.0) + res = benchmark_command(ctrl, start) + key = f"{device}/mppi/K={K}_T=15" + results[key] = res + print(f" K={K:>5d}: {res['mean_s']*1000:>8.2f}ms (min={res['min_s']*1000:.2f}ms)") + + # --- Vary horizon (T) --- + print(f"\n--- MPPI: Varying T (K=500) ---") + for T in [5, 10, 15, 30, 50]: + torch.manual_seed(SEED) + ctrl = MPPI(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=T, device=device, lambda_=1.0) + res = benchmark_command(ctrl, start) + key = f"{device}/mppi/K=500_T={T}" + results[key] = res + print(f" T={T:>5d}: {res['mean_s']*1000:>8.2f}ms (min={res['min_s']*1000:.2f}ms)") + + # --- MPPI with features --- + print(f"\n--- MPPI: Feature variations (K=500, T=15) ---") + feature_configs = [ + ("base", {}), + ("terminal_cost", {"terminal_state_cost": terminal}), + ("noise_abs_cost", {"noise_abs_cost": True}), + ("bounded", {"u_max": torch.tensor([1.0, 1.0], dtype=DTYPE, device=device)}), + ("M=3", {"rollout_samples": 3, "rollout_var_cost": 0.1}), + ("null_action", {"sample_null_action": True}), + ] + for name, extra_kwargs in feature_configs: + torch.manual_seed(SEED) + ctrl = MPPI(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=15, device=device, lambda_=1.0, + **extra_kwargs) + res = benchmark_command(ctrl, start) + key = f"{device}/mppi_feat/{name}" + results[key] = res + print(f" {name:<20s}: {res['mean_s']*1000:>8.2f}ms") + + # --- SMPPI --- + print(f"\n--- SMPPI (K=500, T=15) ---") + for w in [1.0, 5.0, 10.0]: + torch.manual_seed(SEED) + ctrl = SMPPI(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=15, device=device, lambda_=1.0, + w_action_seq_cost=w) + res = benchmark_command(ctrl, start) + key = f"{device}/smppi/w={w}" + results[key] = res + print(f" w={w:<5.1f}: {res['mean_s']*1000:>8.2f}ms") + + # --- KMPPI --- + print(f"\n--- KMPPI (K=500, T=15) ---") + for nsp in [3, 5, 7]: + torch.manual_seed(SEED) + ctrl = KMPPI(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=15, device=device, lambda_=1.0, + num_support_pts=nsp, kernel=RBFKernel(sigma=2.0)) + res = benchmark_command(ctrl, start) + key = f"{device}/kmppi/nsp={nsp}" + results[key] = res + print(f" support_pts={nsp}: {res['mean_s']*1000:>8.2f}ms") + + # --- MPPI vs SMPPI vs KMPPI comparison --- + print(f"\n--- Comparison: MPPI vs SMPPI vs KMPPI (K=500, T=15) ---") + for label, ctrl_cls, extra in [ + ("MPPI", MPPI, {}), + ("SMPPI", SMPPI, {"w_action_seq_cost": 5.0}), + ("KMPPI", KMPPI, {"num_support_pts": 5, "kernel": RBFKernel(sigma=2.0)}), + ]: + torch.manual_seed(SEED) + ctrl = ctrl_cls(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=15, device=device, lambda_=1.0, + **extra) + res = benchmark_command(ctrl, start) + key = f"{device}/compare/{label}" + results[key] = res + print(f" {label:<8s}: {res['mean_s']*1000:>8.2f}ms") + + # --- Multi-step control loop --- + print(f"\n--- Multi-step loop: 20 steps (K=500, T=15) ---") + for label, ctrl_cls, extra in [ + ("MPPI", MPPI, {}), + ("SMPPI", SMPPI, {"w_action_seq_cost": 5.0}), + ("KMPPI", KMPPI, {"num_support_pts": 5, "kernel": RBFKernel(sigma=2.0)}), + ]: + torch.manual_seed(SEED) + ctrl = ctrl_cls(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=15, device=device, lambda_=1.0, + **extra) + res = benchmark_multi_step(ctrl, start, dynamics, num_steps=20) + key = f"{device}/loop/{label}" + results[key] = res + print(f" {label:<8s}: {res['mean_s']*1000:>8.2f}ms total, " + f"{res['per_step_s']*1000:.2f}ms/step") + + # --- Compiled comparison --- + print(f"\n--- Compiled vs eager (K=500, T=15) ---") + for label, ctrl_cls, extra in [ + ("MPPI", MPPI, {}), + ("KMPPI", KMPPI, {"num_support_pts": 5, "kernel": RBFKernel(sigma=2.0)}), + ]: + torch.manual_seed(SEED) + ctrl = ctrl_cls(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=15, device=device, lambda_=1.0, + **extra) + ctrl.compile() + # warmup compile + ctrl.command(start) + ctrl.command(start) + res = benchmark_command(ctrl, start) + key = f"{device}/compiled/{label}" + results[key] = res + print(f" {label:<8s} compiled: {res['mean_s']*1000:>8.2f}ms") + + print(f"\n--- Compiled multi-step loop: 20 steps (K=500, T=15) ---") + for label, ctrl_cls, extra in [ + ("MPPI", MPPI, {}), + ("KMPPI", KMPPI, {"num_support_pts": 5, "kernel": RBFKernel(sigma=2.0)}), + ]: + torch.manual_seed(SEED) + ctrl = ctrl_cls(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=15, device=device, lambda_=1.0, + **extra) + ctrl.compile() + # warmup compile + ctrl.command(start) + ctrl.command(start) + res = benchmark_multi_step(ctrl, start, dynamics, num_steps=20) + key = f"{device}/compiled_loop/{label}" + results[key] = res + print(f" {label:<8s} compiled: {res['mean_s']*1000:>8.2f}ms total, " + f"{res['per_step_s']*1000:.2f}ms/step") + + # --- Higher dimensional --- + print(f"\n--- Higher dimensional (nx=10, nu=3, K=500, T=15) ---") + nx, nu = 10, 3 + sigma_nd = torch.eye(nu, dtype=DTYPE, device=device) + dyn_nd = _make_dynamics_nd(device, nx, nu) + cost_nd = _make_cost_nd(device, nx) + start_nd = torch.randn(nx, dtype=DTYPE, device=device) + torch.manual_seed(SEED) + ctrl = MPPI(dyn_nd, cost_nd, nx, sigma_nd, + num_samples=500, horizon=15, device=device, lambda_=1.0) + res = benchmark_command(ctrl, start_nd) + key = f"{device}/mppi/nx=10_nu=3" + results[key] = res + print(f" nx=10, nu=3: {res['mean_s']*1000:>8.2f}ms") + + # --------------------------------------------------------------- + # Solution quality evaluation + # --------------------------------------------------------------- + goal = GOAL.to(device=device) + + print(f"\n--- Solution quality: MPPI vs SMPPI vs KMPPI (K=500, T=15, 20 steps, 5 trials) ---") + print(f" {'Method':<8s} {'Accum Cost':>12s} {'Final Dist':>12s} {'Ctrl Smooth':>12s}") + for label, ctrl_cls, extra in [ + ("MPPI", MPPI, {}), + ("SMPPI", SMPPI, {"w_action_seq_cost": 5.0}), + ("KMPPI", KMPPI, {"num_support_pts": 5, "kernel": RBFKernel(sigma=2.0)}), + ]: + torch.manual_seed(SEED) + ctrl = ctrl_cls(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=15, device=device, lambda_=1.0, + **extra) + qres = evaluate_quality(ctrl, start, dynamics, cost, goal, num_steps=20, num_trials=5) + key = f"{device}/quality/{label}" + results[key] = qres + print(f" {label:<8s} {qres['accumulated_cost_mean']:>9.2f}\u00b1{qres['accumulated_cost_std']:<5.2f}" + f" {qres['final_dist_mean']:>9.4f}\u00b1{qres['final_dist_std']:<7.4f}" + f" {qres['control_smoothness_mean']:>9.3f}\u00b1{qres['control_smoothness_std']:<6.3f}") + + print(f"\n--- Solution quality: varying K (T=15, 20 steps, 5 trials) ---") + print(f" {'K':<6s} {'Accum Cost':>12s} {'Final Dist':>12s} {'Ctrl Smooth':>12s}") + for K in [50, 100, 500, 1000]: + torch.manual_seed(SEED) + ctrl = MPPI(dynamics, cost, 2, noise_sigma, + num_samples=K, horizon=15, device=device, lambda_=1.0) + qres = evaluate_quality(ctrl, start, dynamics, cost, goal, num_steps=20, num_trials=5) + key = f"{device}/quality/mppi_K={K}" + results[key] = qres + print(f" K={K:<4d} {qres['accumulated_cost_mean']:>9.2f}\u00b1{qres['accumulated_cost_std']:<5.2f}" + f" {qres['final_dist_mean']:>9.4f}\u00b1{qres['final_dist_std']:<7.4f}" + f" {qres['control_smoothness_mean']:>9.3f}\u00b1{qres['control_smoothness_std']:<6.3f}") + + print(f"\n--- Solution quality: varying T (K=500, 20 steps, 5 trials) ---") + print(f" {'T':<6s} {'Accum Cost':>12s} {'Final Dist':>12s} {'Ctrl Smooth':>12s}") + for T in [5, 10, 15, 30]: + torch.manual_seed(SEED) + ctrl = MPPI(dynamics, cost, 2, noise_sigma, + num_samples=500, horizon=T, device=device, lambda_=1.0) + qres = evaluate_quality(ctrl, start, dynamics, cost, goal, num_steps=20, num_trials=5) + key = f"{device}/quality/mppi_T={T}" + results[key] = qres + print(f" T={T:<4d} {qres['accumulated_cost_mean']:>9.2f}\u00b1{qres['accumulated_cost_std']:<5.2f}" + f" {qres['final_dist_mean']:>9.4f}\u00b1{qres['final_dist_std']:<7.4f}" + f" {qres['control_smoothness_mean']:>9.3f}\u00b1{qres['control_smoothness_std']:<6.3f}") + + return results + + +def _serialize(obj): + """Make results JSON-serializable by converting non-serializable values.""" + if isinstance(obj, dict): + return {k: _serialize(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [_serialize(v) for v in obj] + elif isinstance(obj, float): + return obj + elif isinstance(obj, int): + return obj + else: + return float(obj) + + +def print_summary(results): + print(f"\n{'='*60}") + print("SUMMARY — TIMING (all times in milliseconds)") + print(f"{'='*60}") + for key in sorted(results.keys()): + r = results[key] + if "mean_s" in r: + print(f" {key:<45s} mean={r['mean_s']*1000:>8.2f} min={r['min_s']*1000:>8.2f}") + + print(f"\n{'='*60}") + print("SUMMARY — SOLUTION QUALITY") + print(f"{'='*60}") + print(f" {'Key':<40s} {'Accum Cost':>12s} {'Final Dist':>12s} {'Smoothness':>12s}") + for key in sorted(results.keys()): + r = results[key] + if "accumulated_cost_mean" in r: + print(f" {key:<40s} {r['accumulated_cost_mean']:>12.2f} {r['final_dist_mean']:>12.4f} {r['control_smoothness_mean']:>12.3f}") + + +if __name__ == "__main__": + results = run_benchmarks() + print_summary(results) + + # Save results to JSON for comparison + outpath = "benchmark_results.json" + with open(outpath, "w") as f: + json.dump(_serialize(results), f, indent=2) + print(f"\nResults saved to {outpath}") diff --git a/tests/test_mppi.py b/tests/test_mppi.py new file mode 100644 index 0000000..72b5fe0 --- /dev/null +++ b/tests/test_mppi.py @@ -0,0 +1,948 @@ +"""Comprehensive tests for MPPI, SMPPI, and KMPPI controllers. + +Uses a simple 2D linear dynamics + quadratic cost environment that requires +no external dependencies (no gym, no rendering). All tests are seeded for +deterministic reproducibility. +""" +import pytest +import torch +from pytorch_mppi import MPPI, SMPPI, KMPPI, MPPI_Batched +from pytorch_mppi.mppi import RBFKernel, SpecificActionSampler + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- +DEVICE = "cpu" +DTYPE = torch.double +SEED = 42 + + +def _seed(): + torch.manual_seed(SEED) + + +# Simple linear dynamics: x_{t+1} = x_t + B @ u_t +B = torch.tensor([[1.0, 0.0], [0.0, -1.0]], dtype=DTYPE, device=DEVICE) + + +def linear_dynamics(state, action): + return state + action @ B.T + + +def linear_dynamics_step(state, action, t): + return linear_dynamics(state, action) + + +# Quadratic cost toward a goal +GOAL = torch.tensor([2.0, 2.0], dtype=DTYPE, device=DEVICE) + + +def quadratic_cost(state, action): + dx = GOAL - state + return (dx ** 2).sum(dim=-1) + + +def quadratic_cost_step(state, action, t): + return quadratic_cost(state, action) + + +def terminal_cost(states, actions): + dx = GOAL - states[..., -1, :] + return (dx ** 2).sum(dim=-1) + + +@pytest.fixture +def noise_sigma(): + return torch.eye(2, dtype=DTYPE, device=DEVICE) + + +@pytest.fixture +def small_noise_sigma(): + return torch.eye(2, dtype=DTYPE, device=DEVICE) * 0.1 + + +# --------------------------------------------------------------------------- +# MPPI Tests +# --------------------------------------------------------------------------- +class TestMPPI: + def _make(self, noise_sigma, **kwargs): + defaults = dict( + dynamics=linear_dynamics, + running_cost=quadratic_cost, + nx=2, + noise_sigma=noise_sigma, + num_samples=100, + horizon=10, + device=DEVICE, + lambda_=1.0, + ) + defaults.update(kwargs) + return MPPI(**defaults) + + def test_basic_command_returns_action(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,), f"Expected shape (2,), got {action.shape}" + assert action.dtype == DTYPE + + def test_command_moves_toward_goal(self, noise_sigma): + """After several commands, cost should decrease.""" + _seed() + ctrl = self._make(noise_sigma, num_samples=500) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + + initial_cost = quadratic_cost(state.unsqueeze(0), torch.zeros(1, 2, dtype=DTYPE)).item() + for _ in range(5): + action = ctrl.command(state) + state = linear_dynamics(state.unsqueeze(0), action.unsqueeze(0)).squeeze(0) + final_cost = quadratic_cost(state.unsqueeze(0), torch.zeros(1, 2, dtype=DTYPE)).item() + assert final_cost < initial_cost, f"Cost did not decrease: {initial_cost} -> {final_cost}" + + def test_deterministic_with_seed(self, noise_sigma): + """Same seed should produce identical results.""" + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + + _seed() + ctrl1 = self._make(noise_sigma) + a1 = ctrl1.command(state) + + _seed() + ctrl2 = self._make(noise_sigma) + a2 = ctrl2.command(state) + + assert torch.allclose(a1, a2), f"Actions differ: {a1} vs {a2}" + + def test_control_bounds(self, noise_sigma): + _seed() + u_max = torch.tensor([0.5, 0.5], dtype=DTYPE, device=DEVICE) + ctrl = self._make(noise_sigma, u_min=-u_max, u_max=u_max) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + for _ in range(10): + action = ctrl.command(state) + state = linear_dynamics(state.unsqueeze(0), action.unsqueeze(0)).squeeze(0) + assert (action <= u_max + 1e-6).all(), f"Action {action} exceeds u_max {u_max}" + assert (action >= -u_max - 1e-6).all(), f"Action {action} below u_min {-u_max}" + + def test_u_max_only_sets_symmetric_bounds(self, noise_sigma): + _seed() + u_max = torch.tensor([1.0, 1.0], dtype=DTYPE, device=DEVICE) + ctrl = self._make(noise_sigma, u_max=u_max) + assert ctrl.u_min is not None + assert torch.allclose(ctrl.u_min, -u_max) + + def test_u_min_only_sets_symmetric_bounds(self, noise_sigma): + _seed() + u_min = torch.tensor([-1.0, -1.0], dtype=DTYPE, device=DEVICE) + ctrl = self._make(noise_sigma, u_min=u_min) + assert ctrl.u_max is not None + assert torch.allclose(ctrl.u_max, -u_min) + + def test_terminal_state_cost(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, terminal_state_cost=terminal_cost) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_step_dependent_dynamics(self, noise_sigma): + _seed() + ctrl = self._make( + noise_sigma, + dynamics=linear_dynamics_step, + running_cost=quadratic_cost_step, + step_dependent_dynamics=True, + ) + state = torch.tensor([-1.0, -1.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_noise_abs_cost(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, noise_abs_cost=True) + state = torch.tensor([-1.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_sample_null_action(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, sample_null_action=True) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_u_per_command_multiple(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, u_per_command=3) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (3, 2), f"Expected shape (3, 2), got {action.shape}" + + def test_rollout_samples(self, noise_sigma): + """Test with M > 1 rollout samples for stochastic dynamics.""" + _seed() + ctrl = self._make(noise_sigma, rollout_samples=3, rollout_var_cost=0.1) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_get_rollouts(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) # initialize U + rollouts = ctrl.get_rollouts(state, num_rollouts=5) + assert rollouts.shape == (5, ctrl.T, 2) + + def test_get_rollouts_custom_U(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + custom_U = torch.zeros(ctrl.T, 2, dtype=DTYPE, device=DEVICE) + rollouts = ctrl.get_rollouts(state, num_rollouts=1, U=custom_U) + # With zero actions and linear dynamics, state should stay at origin + assert torch.allclose(rollouts, torch.zeros_like(rollouts)) + + def test_change_horizon_shorter(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, horizon=10) + ctrl.change_horizon(5) + assert ctrl.T == 5 + assert ctrl.U.shape[0] == 5 + + def test_change_horizon_longer(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, horizon=5) + ctrl.change_horizon(10) + assert ctrl.T == 10 + assert ctrl.U.shape[0] == 10 + + def test_reset(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + U_before = ctrl.U.clone() + ctrl.reset() + # After reset, U should be resampled (very unlikely to match) + assert not torch.allclose(ctrl.U, U_before) + + def test_batch_state_input(self, noise_sigma): + """Pass (K x nx) state to command.""" + _seed() + K = 100 + ctrl = self._make(noise_sigma, num_samples=K) + state = torch.randn(K, 2, dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_stored_states_actions(self, noise_sigma): + """After command, states/actions are None without terminal_state_cost (lazy storage).""" + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + # Without terminal_state_cost, M=1 fast path skips storage + assert ctrl.states is None + assert ctrl.actions is None + + def test_stored_states_actions_with_terminal(self, noise_sigma): + """With terminal_state_cost, states and actions should be populated.""" + _seed() + ctrl = self._make(noise_sigma, terminal_state_cost=terminal_cost) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + assert ctrl.states is not None + assert ctrl.actions is not None + assert ctrl.states.shape[-1] == 2 # nx + assert ctrl.actions.shape[-1] == 2 # nu + + def test_cost_total_shape(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + assert ctrl.cost_total.shape == (ctrl.K,) + + def test_omega_sums_to_one(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + assert torch.allclose(ctrl.omega.sum(), torch.tensor(1.0, dtype=DTYPE), atol=1e-5) + + def test_1d_control(self): + """Test with scalar (1D) control noise.""" + _seed() + sigma = torch.tensor(1.0, dtype=DTYPE, device=DEVICE) + + def dynamics_1d(state, action): + return state + action + + def cost_1d(state, action): + return (state[:, 0] - 1.0) ** 2 + + ctrl = MPPI(dynamics_1d, cost_1d, nx=1, noise_sigma=sigma, + num_samples=50, horizon=5, device=DEVICE) + state = torch.tensor([0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (1,) + + def test_shift_nominal_trajectory(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + U_before = ctrl.U.clone() + ctrl.shift_nominal_trajectory() + # Last row should be u_init (zeros) + assert torch.allclose(ctrl.U[-1], ctrl.u_init) + # First row should be what was second row + assert torch.allclose(ctrl.U[0], U_before[1]) + + def test_no_shift_refine(self, noise_sigma): + """command with shift_nominal_trajectory=False should not shift U.""" + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state, shift_nominal_trajectory=True) + U_after_first = ctrl.U.clone() + ctrl.command(state, shift_nominal_trajectory=False) + # U should be updated but not shifted — can't easily test "not shifted" + # but we can confirm it ran without error and U changed + assert ctrl.U.shape == U_after_first.shape + + def test_u_scale(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, u_scale=2.0, terminal_state_cost=terminal_cost) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + assert ctrl.actions is not None + + def test_get_params_string(self, noise_sigma): + ctrl = self._make(noise_sigma) + params = ctrl.get_params() + assert "K=100" in params + assert "T=10" in params + + +# --------------------------------------------------------------------------- +# SMPPI Tests +# --------------------------------------------------------------------------- +class TestSMPPI: + def _make(self, noise_sigma, **kwargs): + defaults = dict( + dynamics=linear_dynamics, + running_cost=quadratic_cost, + nx=2, + noise_sigma=noise_sigma, + num_samples=100, + horizon=10, + device=DEVICE, + lambda_=1.0, + ) + defaults.update(kwargs) + return SMPPI(**defaults) + + def test_basic_command(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([-1.0, -1.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_command_moves_toward_goal(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, num_samples=500) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + + initial_cost = quadratic_cost(state.unsqueeze(0), torch.zeros(1, 2, dtype=DTYPE)).item() + for _ in range(5): + action = ctrl.command(state) + state = linear_dynamics(state.unsqueeze(0), action.unsqueeze(0)).squeeze(0) + final_cost = quadratic_cost(state.unsqueeze(0), torch.zeros(1, 2, dtype=DTYPE)).item() + assert final_cost < initial_cost + + def test_action_bounds(self, noise_sigma): + _seed() + action_max = torch.tensor([0.5, 0.5], dtype=DTYPE, device=DEVICE) + ctrl = self._make(noise_sigma, action_max=action_max) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + for _ in range(10): + action = ctrl.command(state) + state = linear_dynamics(state.unsqueeze(0), action.unsqueeze(0)).squeeze(0) + assert (action <= action_max + 1e-6).all() + assert (action >= -action_max - 1e-6).all() + + def test_smoothness(self, noise_sigma): + """SMPPI actions should be smoother than MPPI (smaller action diffs).""" + _seed() + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + ctrl_mppi = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=200, horizon=10, device=DEVICE, lambda_=1.0) + ctrl_smppi = self._make(noise_sigma, num_samples=200, w_action_seq_cost=10.0) + + actions_mppi = [] + actions_smppi = [] + s_mppi = state.clone() + s_smppi = state.clone() + _seed() + for _ in range(8): + a = ctrl_mppi.command(s_mppi) + s_mppi = linear_dynamics(s_mppi.unsqueeze(0), a.unsqueeze(0)).squeeze(0) + actions_mppi.append(a) + _seed() + for _ in range(8): + a = ctrl_smppi.command(s_smppi) + s_smppi = linear_dynamics(s_smppi.unsqueeze(0), a.unsqueeze(0)).squeeze(0) + actions_smppi.append(a) + + diffs_mppi = torch.stack(actions_mppi).diff(dim=0).abs().sum() + diffs_smppi = torch.stack(actions_smppi).diff(dim=0).abs().sum() + # SMPPI should generally be smoother, but with small samples it's not guaranteed. + # Just check it ran correctly. + assert diffs_smppi.isfinite() + assert diffs_mppi.isfinite() + + def test_w_action_seq_cost(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, w_action_seq_cost=5.0) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_delta_t(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, delta_t=0.5) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_reset(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + ctrl.reset() + assert torch.allclose(ctrl.U, torch.zeros_like(ctrl.U)) + assert torch.allclose(ctrl.action_sequence, torch.zeros_like(ctrl.action_sequence)) + + def test_change_horizon(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, horizon=10) + ctrl.change_horizon(5) + assert ctrl.T == 5 + assert ctrl.U.shape[0] == 5 + assert ctrl.action_sequence.shape[0] == 5 + + def test_change_horizon_longer(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, horizon=5) + ctrl.change_horizon(10) + assert ctrl.T == 10 + assert ctrl.U.shape[0] == 10 + assert ctrl.action_sequence.shape[0] == 10 + + def test_get_action_sequence(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + seq = ctrl.get_action_sequence() + assert seq.shape == (ctrl.T, 2) + # action_sequence should be same object + assert seq is ctrl.action_sequence + + def test_get_params(self, noise_sigma): + ctrl = self._make(noise_sigma, w_action_seq_cost=5.0, delta_t=0.1) + params = ctrl.get_params() + assert "w=5" in params + assert "t=0.1" in params + + +# --------------------------------------------------------------------------- +# KMPPI Tests +# --------------------------------------------------------------------------- +class TestKMPPI: + def _make(self, noise_sigma, **kwargs): + defaults = dict( + dynamics=linear_dynamics, + running_cost=quadratic_cost, + nx=2, + noise_sigma=noise_sigma, + num_samples=100, + horizon=10, + device=DEVICE, + lambda_=1.0, + ) + defaults.update(kwargs) + return KMPPI(**defaults) + + def test_basic_command(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([-1.0, -1.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_command_moves_toward_goal(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, num_samples=500) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + + initial_cost = quadratic_cost(state.unsqueeze(0), torch.zeros(1, 2, dtype=DTYPE)).item() + for _ in range(5): + action = ctrl.command(state) + state = linear_dynamics(state.unsqueeze(0), action.unsqueeze(0)).squeeze(0) + final_cost = quadratic_cost(state.unsqueeze(0), torch.zeros(1, 2, dtype=DTYPE)).item() + assert final_cost < initial_cost + + def test_num_support_pts(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, num_support_pts=3) + assert ctrl.num_support_pts == 3 + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_default_support_pts(self, noise_sigma): + ctrl = self._make(noise_sigma, horizon=10) + assert ctrl.num_support_pts == 5 # T // 2 + + def test_custom_kernel(self, noise_sigma): + _seed() + kernel = RBFKernel(sigma=2.0) + ctrl = self._make(noise_sigma, kernel=kernel) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_kernel_interpolation_shape(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, num_support_pts=4) + theta = torch.randn(4, 2, dtype=DTYPE, device=DEVICE) + result, K = ctrl.deparameterize_to_trajectory_single(theta) + assert result.shape == (ctrl.T, 2) + + def test_kernel_interpolation_batch_shape(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, num_support_pts=4) + theta = torch.randn(ctrl.K, 4, 2, dtype=DTYPE, device=DEVICE) + result, K = ctrl.deparameterize_to_trajectory_batch(theta) + assert result.shape == (ctrl.K, ctrl.T, 2) + + def test_control_bounds(self, noise_sigma): + _seed() + u_max = torch.tensor([0.5, 0.5], dtype=DTYPE, device=DEVICE) + ctrl = self._make(noise_sigma, u_min=-u_max, u_max=u_max) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + for _ in range(5): + action = ctrl.command(state) + state = linear_dynamics(state.unsqueeze(0), action.unsqueeze(0)).squeeze(0) + + def test_reset(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + ctrl.command(state) + ctrl.reset() + assert torch.allclose(ctrl.theta, torch.zeros_like(ctrl.theta)) + + def test_get_params(self, noise_sigma): + kernel = RBFKernel(sigma=2.0) + ctrl = self._make(noise_sigma, num_support_pts=5, kernel=kernel) + params = ctrl.get_params() + assert "num_support_pts=5" in params + assert "RBFKernel" in params + + def test_rbf_kernel_values(self): + """Test RBF kernel produces expected values.""" + kernel = RBFKernel(sigma=1.0) + t = torch.tensor([[0.0], [1.0]], dtype=DTYPE) + tk = torch.tensor([[0.0], [1.0]], dtype=DTYPE) + K = kernel(t, tk) + # Diagonal should be 1 (zero distance) + assert torch.allclose(K.diag(), torch.ones(2, dtype=DTYPE)) + # Off-diagonal should be exp(-0.5) for distance 1, sigma 1 + expected_offdiag = torch.exp(torch.tensor(-0.5, dtype=DTYPE)) + assert torch.allclose(K[0, 1], expected_offdiag, atol=1e-6) + + def test_multiple_commands_stable(self, noise_sigma): + """Run several command steps and ensure no NaN/Inf.""" + _seed() + ctrl = self._make(noise_sigma, num_samples=200) + state = torch.tensor([-2.0, -1.0], dtype=DTYPE, device=DEVICE) + for _ in range(15): + action = ctrl.command(state) + assert action.isfinite().all(), f"Non-finite action: {action}" + state = linear_dynamics(state.unsqueeze(0), action.unsqueeze(0)).squeeze(0) + assert state.isfinite().all(), f"Non-finite state: {state}" + + +# --------------------------------------------------------------------------- +# SpecificActionSampler Tests +# --------------------------------------------------------------------------- +class TestSpecificActionSampler: + def test_with_specific_sampler(self, noise_sigma): + _seed() + + class MySampler(SpecificActionSampler): + def sample_trajectories(self, state, info): + # Return 2 trajectories of zeros + return torch.zeros(2, 10, 2, dtype=DTYPE, device=DEVICE) + + sampler = MySampler() + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=100, horizon=10, device=DEVICE, + specific_action_sampler=sampler) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + assert sampler.start_idx == 0 + assert sampler.end_idx == 2 + + +# --------------------------------------------------------------------------- +# Edge Cases +# --------------------------------------------------------------------------- +class TestEdgeCases: + def test_numpy_state_input(self, noise_sigma): + """State passed as numpy array should work.""" + import numpy as np + _seed() + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=50, horizon=5, device=DEVICE) + state = np.array([0.0, 0.0]) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_high_dimensional_state(self): + """Test with higher dimensional state space.""" + _seed() + nx = 10 + nu = 3 + sigma = torch.eye(nu, dtype=DTYPE, device=DEVICE) + + def dyn(state, action): + # simple: first nu dims affected by action + delta = torch.zeros_like(state) + delta[..., :nu] = action + return state + delta + + def cost(state, action): + return (state ** 2).sum(dim=-1) + + ctrl = MPPI(dyn, cost, nx, sigma, num_samples=50, horizon=5, device=DEVICE) + state = torch.randn(nx, dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (nu,) + + def test_large_horizon(self, noise_sigma): + _seed() + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=20, horizon=50, device=DEVICE) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_single_sample(self, noise_sigma): + _seed() + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=1, horizon=5, device=DEVICE) + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + + def test_float32_dtype(self): + _seed() + sigma = torch.eye(2, dtype=torch.float32, device=DEVICE) + + def dyn(state, action): + return state + action @ B.float().T + + def cost(state, action): + return ((GOAL.float() - state) ** 2).sum(dim=-1) + + ctrl = MPPI(dyn, cost, 2, sigma, num_samples=50, horizon=5, device=DEVICE) + state = torch.tensor([0.0, 0.0], dtype=torch.float32, device=DEVICE) + action = ctrl.command(state) + assert action.dtype == torch.float32 + + def test_compile(self, noise_sigma): + """torch.compile should produce valid results.""" + _seed() + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=50, horizon=5, device=DEVICE) + ctrl.compile() + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + assert torch.isfinite(action).all() + # run multiple steps to verify stability + for _ in range(5): + action = ctrl.command(state) + state = linear_dynamics(state.unsqueeze(0), action.unsqueeze(0)).squeeze(0) + assert torch.isfinite(state).all() + + def test_compile_kmppi(self, noise_sigma): + """torch.compile on KMPPI should work.""" + _seed() + ctrl = KMPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=50, horizon=10, device=DEVICE, num_support_pts=5) + ctrl.compile() + state = torch.tensor([0.0, 0.0], dtype=DTYPE, device=DEVICE) + action = ctrl.command(state) + assert action.shape == (2,) + assert torch.isfinite(action).all() + + +# --------------------------------------------------------------------------- +# MPPI_Batched Tests +# --------------------------------------------------------------------------- +class TestMPPIBatched: + def _make(self, noise_sigma, num_envs=4, **kwargs): + defaults = dict( + dynamics=linear_dynamics, + running_cost=quadratic_cost, + nx=2, + noise_sigma=noise_sigma, + num_envs=num_envs, + num_samples=100, + horizon=10, + device=DEVICE, + lambda_=1.0, + ) + defaults.update(kwargs) + return MPPI_Batched(**defaults) + + def test_basic_command(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, num_envs=4) + states = torch.randn(4, 2, dtype=DTYPE, device=DEVICE) + action = ctrl.command(states) + assert action.shape == (4, 2) + + def test_moves_toward_goal(self, noise_sigma): + """All environments should make progress toward goal.""" + _seed() + N = 4 + ctrl = self._make(noise_sigma, num_envs=N, num_samples=300) + states = torch.tensor([[-3.0, -2.0], [-1.0, -1.0], [0.0, 0.0], [1.0, -1.0]], + dtype=DTYPE, device=DEVICE) + initial_dists = (states - GOAL).norm(dim=-1) + for _ in range(10): + actions = ctrl.command(states) + states = linear_dynamics(states, actions) + final_dists = (states - GOAL).norm(dim=-1) + # At least some environments should improve + assert (final_dists < initial_dists).any(), \ + f"No environment improved: {initial_dists} -> {final_dists}" + + def test_bounded_actions(self, noise_sigma): + _seed() + u_max = torch.tensor([0.5, 0.5], dtype=DTYPE, device=DEVICE) + ctrl = self._make(noise_sigma, num_envs=4, u_max=u_max) + states = torch.randn(4, 2, dtype=DTYPE, device=DEVICE) + for _ in range(5): + actions = ctrl.command(states) + assert (actions <= u_max + 1e-6).all() + assert (actions >= -u_max - 1e-6).all() + states = linear_dynamics(states, actions) + + def test_independent_envs(self, noise_sigma): + """Different initial states should produce different actions.""" + _seed() + ctrl = self._make(noise_sigma, num_envs=2, num_samples=200) + states = torch.tensor([[-5.0, -5.0], [5.0, 5.0]], dtype=DTYPE, device=DEVICE) + actions = ctrl.command(states) + # Very different states should yield different actions + assert not torch.allclose(actions[0], actions[1], atol=0.1), \ + f"Actions too similar for very different states: {actions}" + + def test_reset(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, num_envs=2) + states = torch.randn(2, 2, dtype=DTYPE, device=DEVICE) + ctrl.command(states) + U_before = ctrl.U.clone() + ctrl.reset() + assert not torch.allclose(ctrl.U, U_before) + + def test_compile(self, noise_sigma): + _seed() + ctrl = self._make(noise_sigma, num_envs=2, num_samples=50, horizon=5) + ctrl.compile() + states = torch.randn(2, 2, dtype=DTYPE, device=DEVICE) + actions = ctrl.command(states) + assert actions.shape == (2, 2) + assert torch.isfinite(actions).all() + + +# --------------------------------------------------------------------------- +# Solution quality helper +# --------------------------------------------------------------------------- +def _run_control_loop(ctrl, state, num_steps=20): + """Run a closed-loop control trajectory, returning quality metrics.""" + total_cost = 0.0 + actions = [] + states = [state.clone()] + for _ in range(num_steps): + a = ctrl.command(state) + actions.append(a.clone()) + c = quadratic_cost(state.unsqueeze(0), a.unsqueeze(0)).item() + total_cost += c + state = linear_dynamics(state.unsqueeze(0), a.unsqueeze(0)).squeeze(0) + states.append(state.clone()) + final_dist = (state - GOAL).norm().item() + actions_t = torch.stack(actions) + control_smoothness = actions_t.diff(dim=0).abs().sum().item() + return { + "accumulated_cost": total_cost, + "final_dist": final_dist, + "control_smoothness": control_smoothness, + "final_state": state, + "actions": actions_t, + } + + +# --------------------------------------------------------------------------- +# Solution Quality Tests +# --------------------------------------------------------------------------- +class TestSolutionQuality: + """Tests that ensure the controller actually solves the control problem + to a reasonable standard. These act as regression guards during refactoring. + + Thresholds are set generously (2-3x above measured baseline) to allow + for minor numerical changes while catching real regressions. + """ + + def test_mppi_reaches_goal(self, noise_sigma): + """MPPI should reach close to goal in 20 steps.""" + _seed() + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=500, horizon=15, device=DEVICE, lambda_=1.0) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + res = _run_control_loop(ctrl, state, num_steps=20) + assert res["final_dist"] < 2.0, \ + f"MPPI didn't reach goal: final_dist={res['final_dist']:.4f}" + + def test_smppi_stable_trajectory(self, noise_sigma): + """SMPPI should produce finite actions/states and have decreasing rollout costs.""" + _seed() + ctrl = SMPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=500, horizon=15, device=DEVICE, lambda_=1.0, + w_action_seq_cost=5.0) + state = torch.tensor([-1.0, -1.0], dtype=DTYPE, device=DEVICE) + for _ in range(10): + action = ctrl.command(state) + assert action.isfinite().all(), f"SMPPI produced non-finite action: {action}" + state = linear_dynamics(state.unsqueeze(0), action.unsqueeze(0)).squeeze(0) + assert state.isfinite().all(), f"SMPPI produced non-finite state: {state}" + # Cost total should be finite and non-negative + assert ctrl.cost_total.isfinite().all() + assert (ctrl.cost_total >= 0).all() + + def test_kmppi_reaches_goal(self, noise_sigma): + _seed() + ctrl = KMPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=500, horizon=15, device=DEVICE, lambda_=1.0, + num_support_pts=5, kernel=RBFKernel(sigma=2.0)) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + res = _run_control_loop(ctrl, state, num_steps=20) + assert res["final_dist"] < 2.0, \ + f"KMPPI didn't reach goal: final_dist={res['final_dist']:.4f}" + + def test_mppi_cost_bounded(self, noise_sigma): + """Accumulated cost should stay within reasonable bounds.""" + _seed() + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=500, horizon=15, device=DEVICE, lambda_=1.0) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + res = _run_control_loop(ctrl, state, num_steps=20) + # Generous bound — just catch catastrophic regressions + assert res["accumulated_cost"] < 200.0, \ + f"MPPI accumulated cost too high: {res['accumulated_cost']:.2f}" + + def test_more_samples_improves_quality(self, noise_sigma): + """More samples should generally yield lower cost.""" + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + + costs = [] + for K in [50, 500]: + torch.manual_seed(SEED) + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=K, horizon=15, device=DEVICE, lambda_=1.0) + res = _run_control_loop(ctrl, state, num_steps=20) + costs.append(res["accumulated_cost"]) + + # K=500 should be at least somewhat better than K=50 + assert costs[1] < costs[0] * 1.5, \ + f"More samples didn't help: K=50 cost={costs[0]:.2f}, K=500 cost={costs[1]:.2f}" + + def test_reasonable_quality_across_horizons(self, noise_sigma): + """Both short and long horizons should reach a reasonable solution.""" + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + + for T in [5, 15]: + torch.manual_seed(SEED) + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=500, horizon=T, device=DEVICE, lambda_=1.0) + res = _run_control_loop(ctrl, state, num_steps=20) + assert res["final_dist"] < 5.0, \ + f"T={T} didn't reach goal: final_dist={res['final_dist']:.4f}" + assert res["accumulated_cost"] < 300.0, \ + f"T={T} cost too high: {res['accumulated_cost']:.2f}" + + def test_mppi_deterministic_quality(self, noise_sigma): + """Same seed should produce identical trajectories.""" + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + + _seed() + ctrl1 = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=200, horizon=10, device=DEVICE, lambda_=1.0) + res1 = _run_control_loop(ctrl1, state.clone(), num_steps=10) + + _seed() + ctrl2 = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=200, horizon=10, device=DEVICE, lambda_=1.0) + res2 = _run_control_loop(ctrl2, state.clone(), num_steps=10) + + assert torch.allclose(res1["actions"], res2["actions"]), \ + "Deterministic runs produced different action sequences" + assert abs(res1["accumulated_cost"] - res2["accumulated_cost"]) < 1e-6 + + def test_smppi_planned_trajectory_smoother(self, noise_sigma): + """SMPPI's planned action sequence should be smoother than MPPI's. + Note: closed-loop smoothness depends on the environment; what SMPPI + guarantees is smoother open-loop plans.""" + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + + _seed() + ctrl_mppi = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=500, horizon=15, device=DEVICE, lambda_=1.0) + ctrl_mppi.command(state.clone()) + mppi_plan_smooth = ctrl_mppi.U.diff(dim=0).abs().sum().item() + + _seed() + ctrl_smppi = SMPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=500, horizon=15, device=DEVICE, lambda_=1.0, + w_action_seq_cost=10.0) + ctrl_smppi.command(state.clone()) + smppi_plan_smooth = ctrl_smppi.get_action_sequence().diff(dim=0).abs().sum().item() + + assert smppi_plan_smooth < mppi_plan_smooth * 2.0, \ + f"SMPPI plan not smoother: mppi={mppi_plan_smooth:.3f}, smppi={smppi_plan_smooth:.3f}" + + def test_bounded_actions_respected_in_loop(self, noise_sigma): + """Verify bounds hold across a full trajectory.""" + u_max = torch.tensor([0.3, 0.3], dtype=DTYPE, device=DEVICE) + _seed() + ctrl = MPPI(linear_dynamics, quadratic_cost, 2, noise_sigma, + num_samples=500, horizon=15, device=DEVICE, lambda_=1.0, + u_max=u_max) + state = torch.tensor([-3.0, -2.0], dtype=DTYPE, device=DEVICE) + res = _run_control_loop(ctrl, state, num_steps=20) + assert (res["actions"] <= u_max + 1e-6).all(), "Actions exceeded upper bound" + assert (res["actions"] >= -u_max - 1e-6).all(), "Actions exceeded lower bound"