1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 2237
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_default_dtype(torch.float64)
  9DEVICE = 'cpu'
 10T, D, EPS = 18, 5, 1e-3
 11LAMBDA = 0.015
 12
 13# Bounded open-loop signal u_k=tanh(theta_k); this is a minimal policy
 14# parameterization that remains differentiable when environment gradients exist.
 15def rollout(theta, a):
 16    u = torch.tanh(theta)
 17    x = torch.zeros((), dtype=theta.dtype, device=theta.device)
 18    rows = []
 19    for k in range(T + 1):
 20        rows.append(torch.stack((u[k], x, u[k]**2, x**2, u[k] * x)))
 21        if k < T:
 22            # candidate nonlinear world models
 23            x = 0.78 * x + a * torch.tanh(x)**3 + u[k]
 24    return torch.stack(rows)
 25
 26def phi(theta, a):
 27    W = rollout(theta, a)
 28    G = EPS * torch.eye(D) + W.T @ W / (T + 1)
 29    sign, ld = torch.linalg.slogdet(G)
 30    return ld
 31
 32def objective(theta, models, lam=LAMBDA):
 33    vals = torch.stack([phi(theta, float(a)) for a in models])
 34    # mild action penalty; the robust mechanism is the expectation over models
 35    return vals.mean() - lam * torch.mean(torch.tanh(theta)**2)
 36
 37INIT = 0.35 * torch.sin(torch.arange(T + 1, dtype=torch.float64) * 1.37 + 0.21)
 38
 39def optimize(models, steps=700, lr=0.035):
 40    theta = INIT.clone().requires_grad_()
 41    opt = torch.optim.Adam([theta], lr=lr)
 42    for _ in range(steps):
 43        opt.zero_grad(); loss = -objective(theta, models); loss.backward(); opt.step()
 44    with torch.no_grad():
 45        vals = np.array([phi(theta, float(a)).item() for a in models])
 46        return theta.detach(), vals, objective(theta, models).item()
 47
 48def finite_difference_check():
 49    theta = torch.linspace(-.4, .4, T + 1, requires_grad=True)
 50    a = 0.08
 51    y = phi(theta, a); g = torch.autograd.grad(y, theta)[0].detach().numpy()
 52    h = 1e-5; fd=[]
 53    for i in range(T + 1):
 54        tp=theta.detach().clone(); tm=theta.detach().clone(); tp[i]+=h; tm[i]-=h
 55        fd.append((phi(tp,a).item()-phi(tm,a).item())/(2*h))
 56    fd=np.array(fd)
 57    return float(np.max(np.abs(g-fd))), float(np.linalg.norm(g-fd)/max(np.linalg.norm(fd),1e-12))
 58
 59def spread_sweep():
 60    # Nominal design optimizes only center model; robust design optimizes expectation.
 61    out=[]
 62    for s in [0.0, .03, .06, .10, .14, .18]:
 63        models=np.array([-.10-s, -.10, -.10+s])
 64        tn, _vn_nom, _ = optimize([-.10])
 65        vn = np.array([phi(tn, float(a)).item() for a in models])
 66        tr, vr, _ = optimize(models)
 67        out.append({'spread':s, 'single_min':float(vn.min()),
 68                    'robust_min':float(vr.min()),
 69                    'gain_pct':float(100*(vr.min()-vn.min())/max(abs(vn.min()),1e-9)),
 70                    'single_mean':float(vn.mean()), 'robust_mean':float(vr.mean())})
 71    return out
 72
 73def variance_sweep():
 74    # At a fixed theta, independently sample model indices and estimate grad Phi.
 75    theta,_v,_=optimize([-.10,-.16,-.04])
 76    models=np.array([-.16,-.10,-.04]); probs=np.ones(3)/3
 77    grads=[]
 78    for a in models:
 79        t=theta.clone().requires_grad_(True); g=torch.autograd.grad(phi(t,float(a)),t)[0]
 80        grads.append(g.detach().numpy())
 81    grads=np.array(grads); true=grads.mean(0)
 82    single_var=float(np.mean((grads-true)**2))
 83    rows=[]
 84    rng=np.random.default_rng(SEED+9)
 85    for B in [1,2,4,8,16,32,64]:
 86        errs=[]
 87        for _ in range(3000):
 88            ix=rng.integers(0,3,size=B); est=grads[ix].mean(0)
 89            errs.append(np.mean((est-true)**2))
 90        v=float(np.mean(errs)); rows.append({'B':B,'mse':v,'B_times_mse':B*v,'ratio_to_B1':v/max(single_var, 1e-30)})
 91    return rows
 92
 93def main():
 94    fd_abs, fd_rel=finite_difference_check()
 95    spread=spread_sweep(); variance=variance_sweep()
 96    result={'seed':SEED,'device':DEVICE,'T':T,'epsilon':EPS,
 97            'finite_difference_max_abs':fd_abs,'finite_difference_relative_l2':fd_rel,
 98            'spread_sweep':spread,'variance_sweep':variance}
 99    Path('results.json').write_text(json.dumps(result, indent=2))
100    print(json.dumps(result, indent=2))
101
102if __name__=='__main__': main()