Optimize MPPI hot path for GPU performance and add torch.compile support

[LemonPi]

Summary

• Optimize the MPPI hot path for ~1.5-2x speedup on CUDA (eager mode), with no breaking API changes
• Add optional torch.compile support via mppi.compile() (beneficial at large K)
• Add MPPI_Batched class for running N parallel environments in a single dynamics call
• Add comprehensive test suite (75 tests) and benchmarks
• Bump version to 0.9.0

Optimizations

• Pre-allocate tensors instead of list + torch.stack in rollout loop
• Direct Cholesky noise sampling replacing MultivariateNormal.rsample
• Diagonal covariance detection at init — element-wise ops instead of matmul when sigma is diagonal (common case)
• M=1 fast path — skip .repeat(), variance computation, and M dimension when rollout_samples=1 (default)
• Lazy state/action storage — skip allocating (M, K, T, nx) tensors when no terminal_state_cost
• torch.einsum for weighted perturbation sum (avoids materializing K×T×nu broadcast)
• Bypass handle_batch_input in hot path — pre-flatten to (M*K, nx) and call dynamics/cost directly
• Resolve all branches at init — bounds default to ±inf, dynamics/cost wrapped to uniform signature, action cost function pre-selected
• Cache KMPPI kernel matrix — Ktktk computed once at init, not every call
• functorch.vmap → torch.vmap, torch.inverse → torch.linalg.inv

CUDA benchmark results (RTX 4070)

Config	Before	After	Speedup
MPPI K=500 T=15	1.38ms	0.92ms	1.50x
MPPI K=5000 T=15	3.78ms	2.02ms	1.87x
MPPI 20-step loop	27.5ms	18.5ms	1.49x

Notes on torch.compile

The mppi.compile() API is provided for users with large batch sizes and compile-friendly dynamics/cost functions, but the gains are situational. At K=5000 compiled mode matches the eager speedup (~1.95x vs original), but at typical K=500 it actually adds overhead (0.88x) and hurts in multi-step loops (0.81x). This is because the sequential horizon loop (state[t+1] depends on state[t]) prevents meaningful graph fusion — compile just adds tracing and guard-check overhead. The bulk of the speedup comes from the eager-mode optimizations listed above.

New features

• mppi.compile(**kwargs) — compiles dynamics/cost with torch.compile
• MPPI_Batched — runs N environments with shared noise and a single (N*K, nx) dynamics call per timestep

Test plan

• 75 unit tests passing on CPU and CUDA
• Solution quality regression tests verify controllers still reach goals within bounds
• Benchmarks track timing + solution quality (accumulated cost, final distance, control smoothness)
• Verify downstream users are unaffected (no API changes to existing classes)

🤖 Generated with Claude Code