1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 1653
  6
  7def simulate_bandit(mu, sigma, T, c, reps=220, seed=SEED):
  8    rng = np.random.default_rng(seed)
  9    K = len(mu)
 10    zs = np.zeros((reps, K)); means = np.zeros((reps, K)); counts = np.zeros((reps, K), dtype=int)
 11    for rep in range(reps):
 12        sums = np.zeros(K); sums2 = np.zeros(K); n = np.zeros(K, dtype=int)
 13        # forced initialization makes every empirical statistic defined
 14        for a in range(K):
 15            r = rng.normal(mu[a], sigma[a]); sums[a] += r; sums2[a] += r*r; n[a] += 1
 16        for t in range(K, T):
 17            f = c * math.sqrt(math.log(t + 1.0))
 18            bonus = f / np.sqrt(n)
 19            arm = int(np.argmax(sums / n + bonus))
 20            r = rng.normal(mu[arm], sigma[arm]); sums[arm] += r; sums2[arm] += r*r; n[arm] += 1
 21        m = sums / n
 22        # population variance estimate is used here to avoid tiny-n instability
 23        sd = np.sqrt(np.maximum(sums2 / n - m*m, 1e-12))
 24        means[rep] = m; counts[rep] = n
 25        zs[rep] = np.sqrt(n) * (m - mu) / sigma
 26    fT = c * math.sqrt(math.log(T + 1.0))
 27    return {"z_mean": zs.mean(0), "z_se": zs.std(0, ddof=1)/math.sqrt(reps),
 28            "mean_bias": (means-mu).mean(0), "counts": counts.mean(0), "fT": fT}
 29
 30def mechanism_sweep():
 31    # Arm 1 is suboptimal and therefore covered by coefficient 1 in the theorem.
 32    mu = np.array([0.50, 0.00, 0.50])
 33    sigma = np.ones(3)
 34    rows=[]
 35    for c in [0.5, 1.0, 2.0, 4.0]:
 36        out = simulate_bandit(mu, sigma, T=1400, c=c)
 37        z = float(out["z_mean"][1])
 38        rows.append({"c":c, "fT":out["fT"], "z_suboptimal":z,
 39                     "predicted_sign":"negative", "scaled_fT_times_z":out["fT"]*z,
 40                     "predicted_scaled_limit":-1.0,
 41                     "mean_count_suboptimal":float(out["counts"][1])})
 42    # T sweep checks the predicted 1/sqrt(log T) UCB1 decay.
 43    trows=[]
 44    for T in [400, 800, 1400, 2400]:
 45        out=simulate_bandit(mu, sigma, T=T, c=1.0, reps=260, seed=SEED+T)
 46        z=float(out["z_mean"][1]); f=out["fT"]
 47        trows.append({"T":T,"fT":f,"z_suboptimal":z,"fT_times_z":f*z,
 48                      "predicted_z":-1.0/f})
 49    # Direct target correction, using g=1 for the known non-unique-optimal test arm.
 50    out=simulate_bandit(mu,sigma,T=1400,c=1.0,reps=350,seed=SEED+99)
 51    raw=float(out["z_mean"][1]); corrected=raw+1.0/out["fT"]
 52    return {"exploration_sweep":rows,"horizon_sweep":trows,
 53            "correction":{"raw_z":raw,"corrected_z":corrected,
 54                           "predicted_raw_z":-1/out["fT"],"fT":out["fT"]}}
 55
 56def neural_comparison():
 57    # Small contextual bandit: nonlinear arm means, Gaussian reward. Torch is optional.
 58    try:
 59        import torch
 60        import torch.nn as nn
 61        torch.manual_seed(SEED); np.random.seed(SEED); random.seed(SEED)
 62        device = "cuda" if torch.cuda.is_available() else "cpu"
 63        try:
 64            torch.tensor([1.], device=device)
 65        except Exception:
 66            device="cpu"
 67        rng=np.random.default_rng(SEED); K=6; d=4; T=1100
 68        x=rng.normal(size=(T,d));
 69        def true_mean(xx):
 70            vals=[]
 71            for a in range(K):
 72                vals.append(.55*xx[:,0]*np.sin(.7*(a+1))+ .35*xx[:,1]**2/(a+1)
 73                            +.25*xx[:,2]*(-1)**a + .08*a)
 74            return np.stack(vals,1)
 75        mt=true_mean(x); noise=.6
 76        counts=np.zeros(K,int); sums=np.zeros(K); sumsq=np.zeros(K); data=[]
 77        for a in range(K):
 78            r=mt[a,a] + rng.normal(0,noise); counts[a]+=1; sums[a]+=r; sumsq[a]+=r*r; data.append((x[a],a,r))
 79        c=1.0
 80        for t in range(K,T):
 81            # policy uses the known functional oracle plus empirical arm residual; this isolates sampling bias
 82            est=sums/counts
 83            arm=int(np.argmax(mt[t]+c*math.sqrt(math.log(t+1))/np.sqrt(counts)))
 84            r=mt[t,arm]+rng.normal(0,noise)
 85            counts[arm]+=1; sums[arm]+=r; sumsq[arm]+=r*r; data.append((x[t],arm,r))
 86        X=np.stack([z[0] for z in data]); A=np.array([z[1] for z in data]); Y=np.array([z[2] for z in data])
 87        f=c*math.sqrt(math.log(T+1)); m=sums/counts
 88        sd=np.sqrt(np.maximum(sumsq/counts-m*m, .05)); u=m+sd/np.sqrt(counts)
 89        gates=np.zeros(K)
 90        for a in range(K):
 91            competitor=np.max(np.delete(u,a)); gates[a]=1/(1+np.exp(-(competitor-u[a])/.15))
 92        corr=gates[A]*sd[A]/(np.sqrt(counts[A])*f)
 93        class Net(nn.Module):
 94            def __init__(self):
 95                super().__init__(); self.net=nn.Sequential(nn.Linear(d+K,32),nn.Tanh(),nn.Linear(32,1))
 96            def forward(self,z): return self.net(z).squeeze(-1)
 97        inp=np.concatenate([X,np.eye(K)[A]],1).astype("float32")
 98        tx=torch.tensor(inp,device=device); ty=torch.tensor(Y.astype("float32"),device=device)
 99        def fit(target):
100            torch.manual_seed(SEED); net=Net().to(device); opt=torch.optim.Adam(net.parameters(),lr=.01)
101            for _ in range(18):
102                perm=torch.randperm(len(ty),device=device)
103                for j in range(0,len(ty),128):
104                    q=perm[j:j+128]; loss=((net(tx[q])-target[q])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
105            # uniformly randomized held-out contexts/arms
106            xe=rng.normal(size=(1200,d)); ae=rng.integers(K,size=1200); truth=true_mean(xe)[np.arange(1200),ae]
107            ii=np.concatenate([xe,np.eye(K)[ae]],1).astype("float32")
108            with torch.no_grad(): pred=net(torch.tensor(ii,device=device)).cpu().numpy()
109            mse=float(np.mean((pred-truth)**2)); rare=np.argsort(counts)[:3]
110            rare_mask=np.isin(ae,rare); rare_bias=float(np.mean(pred[rare_mask]-truth[rare_mask]))
111            return mse,rare_bias
112        b=fit(ty); idea=fit(ty+torch.tensor(corr.astype("float32"),device=device))
113        return {"device":device,"counts":counts.tolist(),"gates":gates.tolist(),
114                "baseline":{"heldout_mse":b[0],"rare_signed_bias":b[1]},
115                "idea":{"heldout_mse":idea[0],"rare_signed_bias":idea[1]}}
116    except Exception as e:
117        return {"error":repr(e)}
118
119if __name__ == "__main__":
120    result={"seed":SEED,"mechanism":mechanism_sweep(),"neural":neural_comparison()}
121    Path("results.json").write_text(json.dumps(result,indent=2))
122    print(json.dumps(result,indent=2))