Contraction-budgeted MPPI policy head

(1) Integration point: use a neural policy \(\pi_\theta(x)\) inside a model-based RL or world-model controller. At every state, generate action perturbations around the network's proposed action, roll them through a learned dynamics model, score each trajectory, and form an MPPI-weighted action. The nominal reference is \(\pi_\theta\); the output is the corrected action \(u^{\mathrm{MPPI}}\). This is primarily an inference-time module, with optional training losses for improving the contraction margin.

(2) Pseudocode:

input x, policy pi_theta, dynamics f_phi, temperature lambda
a0 = pi_theta(x)
for j = 1,...,M:
    sample delta_j[0:H-1]
    set a_j[t] = pi_theta(x_j[t]) + delta_j[t]
    roll out x_j[t+1] = f_phi(x_j[t], a_j[t])
    compute trajectory cost J_j
weights w_j = softmax(-J_j / lambda)
u = pi_theta(x) + sum_j w_j * delta_j[0]
estimate beta_inf, L_u, and beta from recent transitions
e_M = beta0_hat(lambda) + c_hat * sqrt(log(1/eta) / M)
if sqrt(mu_bar/mu)*L_u*beta_inf_hat + e_M > (1-beta_hat)/2:
    increase M, or decrease lambda if beta0_hat decreases
execute u

(3) Estimate the mathematical quantities empirically. Estimate \(L_u\) with finite differences through the learned dynamics: \(\widehat L_u=\max_i\|f_\phi(x_i,u_i+\delta u_i)-f_\phi(x_i,u_i)\|/\|\delta u_i\|\). Estimate \(\widehat\beta_\infty\) as a robust upper-quantile slope of \(\|u^{\mathrm{MPPI}}-\pi_\theta(x)\|\) versus \(\|x-x^*\|\). Estimate \(\widehat\beta\) from the 95th-percentile ratio of successive state distances under the nominal policy. Estimate \(\widehat c\) by running independent rollout batches and fitting action-estimate standard deviation against \(1/\sqrt M\). Estimate \(\widehat\beta_0(\lambda)\) using a very large rollout batch at each candidate temperature. Initially use Euclidean state distance, so \(\bar\mu/\mu=1\); later replace it with bounds from a learned positive-definite metric.

(4) First cheap experiment: use a 2-layer MLP policy and a small learned ensemble dynamics model on Pendulum-v1 or CartPole swing-up. Compare fixed-sample MPPI, fixed-temperature MPPI, and this adaptive controller at equal average rollout FLOPs. Measure success rate, action variance, catastrophic departures from a safe state set, reward, and executed rollouts. The expected signal is fewer unstable episodes and lower action variance at the same average sample count, with adaptive rollout increases concentrated near difficult states. At equal failure rate, it should require fewer rollouts than a conservative fixed-sample baseline.