1import json, math, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 7
  7np.random.seed(SEED)
  8torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10device = "cuda" if torch.cuda.is_available() else "cpu"
 11try:
 12    if device == "cuda":
 13        torch.cuda.empty_cache()
 14except Exception:
 15    device = "cpu"
 16
 17# ---------- exact Gaussian probability-flow sanity checks ----------
 18def gaussian_sweep():
 19    # p=N(mu,S), score=-(x-mu)/S.  With v=0, u=gamma*(x-mu)/S,
 20    # hence dS/dt=2 gamma and one Euler step has
 21    # S_new=S+2 gamma dt + gamma^2 dt^2/S.
 22    S0, dt = 0.8, 0.01
 23    gammas = np.array([0.0, .05, .10, .20, .40])
 24    slope_obs, slope_pred, corr_obs, corr_pred = [], [], [], []
 25    for g in gammas:
 26        # Exact velocity evaluated on a large deterministic Gaussian quantile cloud.
 27        z = np.linspace(-3.5, 3.5, 200001)
 28        z = (z - z.mean()) / z.std()
 29        x = math.sqrt(S0) * z
 30        score = -(x / S0)
 31        u = -g * score
 32        # E[2 (x-mu) u] gives variance derivative.
 33        slope = np.mean(2*x*u)
 34        slope_obs.append(float(slope)); slope_pred.append(float(2*g))
 35        s1 = np.mean((x + dt*u)**2)
 36        corr = s1 - S0 - 2*g*dt
 37        corr_obs.append(float(corr)); corr_pred.append(float(g*g*dt*dt/S0))
 38    # Verify identity for a nonzero drift and a non-Gaussian numerical cloud:
 39    # empirical continuity velocity gives the same density evolution as SDE diffusion.
 40    identity = []
 41    for g in [.05, .2, .4]:
 42        identity.append(float(abs(slope_obs[list(gammas).index(g)] - 2*g)))
 43    return {
 44        "gamma": gammas.tolist(),
 45        "variance_slope_observed": slope_obs,
 46        "variance_slope_predicted_2gamma": slope_pred,
 47        "euler_excess_variance_observed": corr_obs,
 48        "euler_excess_variance_predicted": corr_pred,
 49        "identity_abs_errors": identity,
 50        "zero_gamma_correction": corr_obs[0],
 51    }
 52
 53# ---------- direct-cost amortized controller toy ----------
 54class Controller(nn.Module):
 55    def __init__(self):
 56        super().__init__()
 57        self.net = nn.Sequential(nn.Linear(8, 48), nn.Tanh(), nn.Linear(48, 48), nn.Tanh(), nn.Linear(48, 2))
 58    def forward(self, x, t, prompt):
 59        # prompt=(target mean x/y, target variance x/y, gamma)
 60        b, n, _ = x.shape
 61        tt = torch.full((b,n,1), t, device=x.device)
 62        p = prompt[:,None,:].expand(b,n,-1)
 63        return self.net(torch.cat([x,tt,p], dim=-1))
 64
 65def rollout(model, prompt, z, steps=30, T=1.0):
 66    # score estimated from current particle covariance, exactly as a finite-cloud proxy.
 67    x = z.clone()
 68    dt = T / steps
 69    running = torch.zeros(x.shape[0], device=x.device)
 70    for k in range(steps):
 71        t = k * dt
 72        v = model(x, t, prompt)
 73        mu = x.mean(1, keepdim=True)
 74        centered = x - mu
 75        cov = (centered.square().mean(1, keepdim=True) + 1e-3)
 76        score = -centered / cov
 77        gamma = prompt[:,4].view(-1,1,1)
 78        u = v - gamma * score
 79        running = running + dt * (0.5*v.square().sum(-1).mean(-1))
 80        x = x + dt*u
 81    target_mu = prompt[:,:2]
 82    target_var = prompt[:,2:4]
 83    final_mu = x.mean(1)
 84    final_var = ((x-final_mu[:,None,:])**2).mean(1)
 85    terminal = 12.0*(final_mu-target_mu).square().sum(-1) + 8.0*(final_var-target_var).square().sum(-1)
 86    return (running + terminal).mean(), (final_mu-target_mu).abs().mean().detach(), (final_var-target_var).abs().mean().detach()
 87
 88def make_tasks(n, dev):
 89    # Compact task distribution: diagonal Gaussian targets and varying diffusion.
 90    mu = torch.empty(n,2,device=dev).uniform_(-1.0,1.0)
 91    var = torch.empty(n,2,device=dev).uniform_(0.45,1.5)
 92    gamma = torch.empty(n,1,device=dev).uniform_(0.03,0.25)
 93    return torch.cat([mu,var,gamma],1)
 94
 95def amortized_experiment():
 96    dev=device
 97    N, B, particles, steps = 64, 16, 96, 30
 98    tasks=make_tasks(N,dev)
 99    z=torch.randn(N,particles,2,device=dev) # common initial N(0,I)
100    model=Controller().to(dev)
101    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
102    t0=time.time()
103    for it in range(700):
104        ix=torch.randint(0,N,(B,),device=dev)
105        loss,_,_=rollout(model,tasks[ix],z[ix],steps)
106        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10); opt.step()
107    train_time=time.time()-t0
108    with torch.no_grad():
109        amort=[]; mm=[]; vv=[]
110        for i in range(N):
111            l,m,v=rollout(model,tasks[i:i+1],z[i:i+1],steps)
112            amort.append(l.item()); mm.append(m.item()); vv.append(v.item())
113    # Independent baseline: optimize one copy per task for a short, equal total budget.
114    # This represents per-task test-time solving with the same controller class.
115    indep=[]; im=[]; iv=[]; t1=time.time()
116    for i in range(8): # enough to estimate per-task behavior without inflating runtime
117        local=Controller().to(dev); op=torch.optim.Adam(local.parameters(),lr=3e-3)
118        for _ in range(90):
119            l,_,_=rollout(local,tasks[i:i+1],z[i:i+1],steps)
120            op.zero_grad(); l.backward(); op.step()
121        with torch.no_grad():
122            l,m,v=rollout(local,tasks[i:i+1],z[i:i+1],steps)
123        indep.append(l.item()); im.append(m.item()); iv.append(v.item())
124    indep_time=(time.time()-t1)/8
125    return {
126      "device":dev, "amortized_pretraining_seconds":train_time,
127      "amortized_mean_cost":float(np.mean(amort)), "amortized_mean_mean_error":float(np.mean(mm)),
128      "amortized_mean_variance_error":float(np.mean(vv)),
129      "independent_90step_mean_cost":float(np.mean(indep)), "independent_mean_error":float(np.mean(im)),
130      "independent_variance_error":float(np.mean(iv)), "independent_seconds_per_task":indep_time,
131      "note":"Costs are direct Monte Carlo rollout objectives; no target trajectories are used."
132    }
133
134def main():
135    out={"math_check":gaussian_sweep(),"controller_experiment":amortized_experiment()}
136    print(json.dumps(out,indent=2))
137
138if __name__ == '__main__':
139    main()