import json, math, random from pathlib import Path import numpy as np SEED = 1653 def simulate_bandit(mu, sigma, T, c, reps=220, seed=SEED): rng = np.random.default_rng(seed) K = len(mu) zs = np.zeros((reps, K)); means = np.zeros((reps, K)); counts = np.zeros((reps, K), dtype=int) for rep in range(reps): sums = np.zeros(K); sums2 = np.zeros(K); n = np.zeros(K, dtype=int) # forced initialization makes every empirical statistic defined for a in range(K): r = rng.normal(mu[a], sigma[a]); sums[a] += r; sums2[a] += r*r; n[a] += 1 for t in range(K, T): f = c * math.sqrt(math.log(t + 1.0)) bonus = f / np.sqrt(n) arm = int(np.argmax(sums / n + bonus)) r = rng.normal(mu[arm], sigma[arm]); sums[arm] += r; sums2[arm] += r*r; n[arm] += 1 m = sums / n # population variance estimate is used here to avoid tiny-n instability sd = np.sqrt(np.maximum(sums2 / n - m*m, 1e-12)) means[rep] = m; counts[rep] = n zs[rep] = np.sqrt(n) * (m - mu) / sigma fT = c * math.sqrt(math.log(T + 1.0)) return {"z_mean": zs.mean(0), "z_se": zs.std(0, ddof=1)/math.sqrt(reps), "mean_bias": (means-mu).mean(0), "counts": counts.mean(0), "fT": fT} def mechanism_sweep(): # Arm 1 is suboptimal and therefore covered by coefficient 1 in the theorem. mu = np.array([0.50, 0.00, 0.50]) sigma = np.ones(3) rows=[] for c in [0.5, 1.0, 2.0, 4.0]: out = simulate_bandit(mu, sigma, T=1400, c=c) z = float(out["z_mean"][1]) rows.append({"c":c, "fT":out["fT"], "z_suboptimal":z, "predicted_sign":"negative", "scaled_fT_times_z":out["fT"]*z, "predicted_scaled_limit":-1.0, "mean_count_suboptimal":float(out["counts"][1])}) # T sweep checks the predicted 1/sqrt(log T) UCB1 decay. trows=[] for T in [400, 800, 1400, 2400]: out=simulate_bandit(mu, sigma, T=T, c=1.0, reps=260, seed=SEED+T) z=float(out["z_mean"][1]); f=out["fT"] trows.append({"T":T,"fT":f,"z_suboptimal":z,"fT_times_z":f*z, "predicted_z":-1.0/f}) # Direct target correction, using g=1 for the known non-unique-optimal test arm. out=simulate_bandit(mu,sigma,T=1400,c=1.0,reps=350,seed=SEED+99) raw=float(out["z_mean"][1]); corrected=raw+1.0/out["fT"] return {"exploration_sweep":rows,"horizon_sweep":trows, "correction":{"raw_z":raw,"corrected_z":corrected, "predicted_raw_z":-1/out["fT"],"fT":out["fT"]}} def neural_comparison(): # Small contextual bandit: nonlinear arm means, Gaussian reward. Torch is optional. try: import torch import torch.nn as nn torch.manual_seed(SEED); np.random.seed(SEED); random.seed(SEED) device = "cuda" if torch.cuda.is_available() else "cpu" try: torch.tensor([1.], device=device) except Exception: device="cpu" rng=np.random.default_rng(SEED); K=6; d=4; T=1100 x=rng.normal(size=(T,d)); def true_mean(xx): vals=[] for a in range(K): vals.append(.55*xx[:,0]*np.sin(.7*(a+1))+ .35*xx[:,1]**2/(a+1) +.25*xx[:,2]*(-1)**a + .08*a) return np.stack(vals,1) mt=true_mean(x); noise=.6 counts=np.zeros(K,int); sums=np.zeros(K); sumsq=np.zeros(K); data=[] for a in range(K): r=mt[a,a] + rng.normal(0,noise); counts[a]+=1; sums[a]+=r; sumsq[a]+=r*r; data.append((x[a],a,r)) c=1.0 for t in range(K,T): # policy uses the known functional oracle plus empirical arm residual; this isolates sampling bias est=sums/counts arm=int(np.argmax(mt[t]+c*math.sqrt(math.log(t+1))/np.sqrt(counts))) r=mt[t,arm]+rng.normal(0,noise) counts[arm]+=1; sums[arm]+=r; sumsq[arm]+=r*r; data.append((x[t],arm,r)) 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]) f=c*math.sqrt(math.log(T+1)); m=sums/counts sd=np.sqrt(np.maximum(sumsq/counts-m*m, .05)); u=m+sd/np.sqrt(counts) gates=np.zeros(K) for a in range(K): competitor=np.max(np.delete(u,a)); gates[a]=1/(1+np.exp(-(competitor-u[a])/.15)) corr=gates[A]*sd[A]/(np.sqrt(counts[A])*f) class Net(nn.Module): def __init__(self): super().__init__(); self.net=nn.Sequential(nn.Linear(d+K,32),nn.Tanh(),nn.Linear(32,1)) def forward(self,z): return self.net(z).squeeze(-1) inp=np.concatenate([X,np.eye(K)[A]],1).astype("float32") tx=torch.tensor(inp,device=device); ty=torch.tensor(Y.astype("float32"),device=device) def fit(target): torch.manual_seed(SEED); net=Net().to(device); opt=torch.optim.Adam(net.parameters(),lr=.01) for _ in range(18): perm=torch.randperm(len(ty),device=device) for j in range(0,len(ty),128): q=perm[j:j+128]; loss=((net(tx[q])-target[q])**2).mean(); opt.zero_grad(); loss.backward(); opt.step() # uniformly randomized held-out contexts/arms xe=rng.normal(size=(1200,d)); ae=rng.integers(K,size=1200); truth=true_mean(xe)[np.arange(1200),ae] ii=np.concatenate([xe,np.eye(K)[ae]],1).astype("float32") with torch.no_grad(): pred=net(torch.tensor(ii,device=device)).cpu().numpy() mse=float(np.mean((pred-truth)**2)); rare=np.argsort(counts)[:3] rare_mask=np.isin(ae,rare); rare_bias=float(np.mean(pred[rare_mask]-truth[rare_mask])) return mse,rare_bias b=fit(ty); idea=fit(ty+torch.tensor(corr.astype("float32"),device=device)) return {"device":device,"counts":counts.tolist(),"gates":gates.tolist(), "baseline":{"heldout_mse":b[0],"rare_signed_bias":b[1]}, "idea":{"heldout_mse":idea[0],"rare_signed_bias":idea[1]}} except Exception as e: return {"error":repr(e)} if __name__ == "__main__": result={"seed":SEED,"mechanism":mechanism_sweep(),"neural":neural_comparison()} Path("results.json").write_text(json.dumps(result,indent=2)) print(json.dumps(result,indent=2))