1import json, math, os
  2import numpy as np
  3from scipy.linalg import solve_discrete_are
  4
  5SEED = 7
  6rng_global = np.random.default_rng(SEED)
  7
  8# Small unstable LTI plant. The controller knows this model; process noise is external.
  9A = np.array([[1.10, 0.10], [0.00, 0.95]])
 10B = np.array([[0.0], [0.10]])
 11Q = np.diag([1.0, 0.2])
 12R = np.array([[0.08]])
 13P = solve_discrete_are(A, B, Q, R)
 14K = np.linalg.solve(R + B.T @ P @ B, B.T @ P @ A)
 15Acl = A - B @ K
 16D = P - Acl.T @ P @ Acl
 17alpha_P = np.min(np.linalg.eigvalsh(D))
 18
 19# A small learned residual surrogate: deliberately useful near the origin but imperfect.
 20def residual(x):
 21    return np.array([0.18 * np.tanh(1.5*x[0]) - 0.03*x[1]])
 22
 23def lqr(x):
 24    return (-K @ x).reshape(1)
 25
 26def rollout_model(x, u):
 27    return A @ x + (B @ np.asarray(u).reshape(1))[..., 0]
 28
 29def mppi_action(x, M, lam, H=8, sigma=0.55, rng=None):
 30    # Nominal residual policy around LQR, with receding-horizon sampled corrections.
 31    if rng is None: rng = np.random.default_rng()
 32    unom = float(lqr(x)[0] + residual(x)[0])
 33    # Constant nominal sequence is sufficient for the tiny sanity test.
 34    eps = rng.normal(0.0, sigma, size=(M, H))
 35    us = unom + eps
 36    xs = np.repeat(x[None, :], M, axis=0)
 37    costs = np.zeros(M)
 38    for t in range(H):
 39        costs += np.einsum('bi,ij,bj->b', xs, Q, xs) + R[0,0] * us[:,t]**2
 40        xs = xs @ A.T + us[:,t,None] * B.T
 41    costs += np.einsum('bi,ij,bj->b', xs, P, xs)
 42    # numerically stable importance weights
 43    z = -(costs - costs.min()) / lam
 44    w = np.exp(np.clip(z, -80, 0))
 45    sw = w.sum() + 1e-12
 46    vals = us[:,0]
 47    uhat = float(np.dot(w, vals) / sw)
 48    var = float(np.dot(w, (vals-uhat)**2) / sw)
 49    ess = float(sw*sw / (np.dot(w,w) + 1e-12))
 50    se = math.sqrt(max(var, 0.0) / (ess + 1e-12))
 51    return uhat, se, ess, unom
 52
 53def budgeted_action(x, lam, M0=16, Mmax=128, c=0.45, q=0.35, rng=None):
 54    if rng is None: rng = np.random.default_rng()
 55    # Conservative sample scaling from the supplied corollary (constant hidden factor omitted).
 56    scale = (np.linalg.norm(P)**2 * np.linalg.norm(B)**2 * np.linalg.norm(Acl)**2 /
 57             (alpha_P**2)) * math.log(2.0 / 0.05)
 58    # Start small, then spend computation if the estimated perturbation exceeds budget.
 59    M = M0
 60    attempts = 0
 61    while True:
 62        uhat, se, ess, unom = mppi_action(x, M, lam, rng=rng)
 63        ehat = se + abs(uhat-unom)*q
 64        budget = c * alpha_P * np.linalg.norm(x)
 65        accepted = np.linalg.norm(B[:,0]) * abs(ehat) <= budget
 66        attempts += 1
 67        if accepted or M >= Mmax:
 68            # At Mmax, safety takes precedence over refinement.
 69            if not accepted:
 70                return float(lqr(x)[0]), True, M, attempts, se, ess, scale
 71            return uhat, False, M, attempts, se, ess, scale
 72        M *= 2
 73
 74def one_episode(kind, rng, T=100, lam=0.35, fixed_M=64):
 75    x = np.array([1.8, 0.0]) + rng.normal(0, .05, 2)
 76    total = 0.; unstable = False; fallbacks = 0; samples = 0; sevals=[]
 77    for _ in range(T):
 78        if kind == 'lqr':
 79            u = float(lqr(x)[0]); Mused = 0
 80        elif kind == 'residual':
 81            u = float(lqr(x)[0] + residual(x)[0]); Mused = 0
 82        elif kind == 'fixed':
 83            u, se, _, _ = mppi_action(x, fixed_M, lam, rng=rng); Mused=fixed_M; sevals.append(se)
 84        else:
 85            u, fb, Mused, _, se, _, _ = budgeted_action(x, lam, rng=rng)
 86            fallbacks += int(fb); sevals.append(se)
 87        u = float(np.clip(u, -5, 5))
 88        total += float(x @ Q @ x + R[0,0]*u*u)
 89        x = A @ x + B[:,0]*u + rng.normal(0, .035, 2)
 90        samples += Mused
 91        if np.linalg.norm(x) > 8 or not np.all(np.isfinite(x)):
 92            unstable=True; total += 5000.; break
 93    return total, unstable, fallbacks, samples, np.mean(sevals) if sevals else 0.
 94
 95def scaling_check():
 96    # Repeated independent MPPI estimates at one state; report RMS error to a high-M reference.
 97    x=np.array([.2, -.05]); lam=.35
 98    ref=np.mean([mppi_action(x, 4096, lam, rng=np.random.default_rng(10000+i))[0] for i in range(20)])
 99    rows=[]
100    for M in [16,32,64,128,256]:
101        vals=np.array([mppi_action(x,M,lam,rng=np.random.default_rng(20000+100*M+j))[0]
102                       for j in range(30)])
103        rms=float(np.sqrt(np.mean((vals-ref)**2)))
104        rows.append((M,rms))
105    slope=float(np.polyfit(np.log([r[0] for r in rows]), np.log([r[1] for r in rows]), 1)[0])
106    return {'reference_action':float(ref), 'rows':rows, 'loglog_slope':slope}
107
108def main():
109    print('geometry', json.dumps({'P_norm':float(np.linalg.norm(P)), 'K':K.tolist(),
110      'Acl':Acl.tolist(), 'alpha_P':float(alpha_P),
111      'rho_Acl':float(max(abs(np.linalg.eigvals(Acl))))}, sort_keys=True))
112    scaling=scaling_check(); print('scaling',json.dumps(scaling))
113    n=40; results={}
114    for kind in ['lqr','residual','fixed','budgeted']:
115        vals=[]
116        for i in range(n): vals.append(one_episode(kind,np.random.default_rng(5000+i),lam=.35))
117        arr=np.array(vals)
118        results[kind]={'mean_cost':float(arr[:,0].mean()), 'std_cost':float(arr[:,0].std()),
119          'escape_rate':float(arr[:,1].mean()), 'fallback_rate_per_step':float(arr[:,2].sum()/(n*100)),
120          'mean_rollout_samples_per_step':float(arr[:,3].mean()/100), 'mean_action_se':float(arr[:,4].mean())}
121    # Temperature sweep for fixed and budgeted, same states/noise seeds per temperature.
122    temps={}
123    for lam in [.15,.35,.7]:
124        for kind in ['fixed','budgeted']:
125            vals=[one_episode(kind,np.random.default_rng(9000+i),lam=lam) for i in range(25)]
126            a=np.array(vals); temps[kind+str(lam)]={'cost':float(a[:,0].mean()),'escape':float(a[:,1].mean()),'samples':float(a[:,3].mean()/100)}
127    out={'geometry':{'P':P.tolist(),'K':K.tolist(),'Acl':Acl.tolist(),'alpha_P':float(alpha_P)},
128         'corollary_scale_without_hidden_constant':float((np.linalg.norm(P)**2*np.linalg.norm(B)**2*np.linalg.norm(Acl)**2/alpha_P**2)*math.log(2/.05)),
129         'scaling_check':scaling,'episodes':results,'temperature_sweep':temps}
130    with open('results.json','w') as f: json.dump(out,f,indent=2)
131    print(json.dumps(out,indent=2))
132if __name__=='__main__': main()