import sys, json, random from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) EPOCHS = 6 BATCH = 128 EPS = 0.02 FAIL_THR = 1.5 def seed_all(s): random.seed(s); np.random.seed(s); torch.manual_seed(s) if torch.cuda.is_available(): torch.cuda.manual_seed_all(s) def proposal(c, alpha): a = (EPS + np.clip(c, .01, .99)) ** alpha z = float(a.mean()) q = a / a.sum() w = z / a ess = float((w.sum() ** 2) / (w @ w) / len(w)) return q, w, ess, z def math_check(): rng = np.random.default_rng(13) c = rng.uniform(.01,.99,1000); y = rng.binomial(1,c) g = rng.normal(size=(1000,7)); rows=[] for alpha in (.5,1.,2.): q,w,ess,z = proposal(c,alpha) exact = (q[:,None]*w[:,None]*g).sum(0) uniform = g.mean(0) pred = (q*y).sum()/y.mean() obs = float(y[rng.choice(len(y),20000,p=q)].mean()/y.mean()) rows.append({'alpha':alpha,'predicted_enrichment':float(pred), 'sampled_enrichment':obs,'identity_l2':float(np.linalg.norm(exact-uniform)), 'ess_fraction':ess}) return {'rows':rows,'max_identity_l2':max(x['identity_l2'] for x in rows), 'confirmed_identity':max(x['identity_l2'] for x in rows)<1e-12} def critic_scores(x,y,seed): seed_all(seed+10000) # auxiliary predictor sees state/window only; labels are eventual high-energy proxy net=nn.Sequential(nn.Linear(x.shape[1],32),nn.Tanh(),nn.Linear(32,1)) opt=torch.optim.Adam(net.parameters(),lr=.01) for _ in range(80): loss=nn.functional.binary_cross_entropy_with_logits(net(x).squeeze(1),y) opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): c=torch.sigmoid(net(x)).squeeze(1).numpy() return np.clip(c,.01,.99), float(loss) def idea_train(seed,cfg,capture=False): seed_all(seed) d=get_dataset('dynamics',seed,n_train=400,n_test=200) dev='cuda' if torch.cuda.is_available() else 'cpu' # failure labels are held-out rollout outcome proxies from the training targets labels=(d['ytr'].abs().max(dim=1).values.numpy()>FAIL_THR).astype(np.float32) c,bce=critic_scores(d['xtr'],torch.tensor(labels),seed) q,w,ess,z=proposal(c,cfg['alpha']) try: net=make_model('rnn_small',d['input_shape'],d['out_dim']).to(dev) opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']) x,y=d['xtr'].to(dev),d['ytr'].to(dev) for _ in range(EPOCHS): net.train() for _step in range((len(x)+BATCH-1)//BATCH): ids=np.random.choice(len(x),BATCH,replace=True,p=q) ix=torch.as_tensor(ids,device=dev) per=(net(x[ix])-y[ix]).pow(2).mean(dim=1) # finite-population p/q correction, self-normalized for stability ww=torch.as_tensor(w[ids],dtype=per.dtype,device=dev) loss=(per*ww).sum()/(ww.sum()+1e-8) if cfg['weighted'] else per.mean() opt.zero_grad();loss.backward();opt.step() net.eval() with torch.no_grad(): metric=float((net(d['xte'].to(dev))-d['yte'].to(dev)).pow(2).mean()) except RuntimeError: # explicit CPU fallback for tight/shared CUDA allocations net=make_model('rnn_small',d['input_shape'],d['out_dim']) opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']); x,y=d['xtr'],d['ytr'] for _ in range(EPOCHS): for _step in range((len(x)+BATCH-1)//BATCH): ids=np.random.choice(len(x),BATCH,replace=True,p=q); ix=torch.as_tensor(ids) per=(net(x[ix])-y[ix]).pow(2).mean(1); ww=torch.tensor(w[ids],dtype=per.dtype) loss=(per*ww).sum()/ww.sum() if cfg['weighted'] else per.mean() opt.zero_grad();loss.backward();opt.step() with torch.no_grad(): metric=float((net(d['xte'])-d['yte']).pow(2).mean()) if not capture:return metric ids=np.random.choice(len(x),20000,replace=True,p=q) observed=float(labels[ids].mean()/labels.mean()) # model-behaviour check: weighted and unweighted train-pool prediction losses with torch.no_grad(): pred=(net(x).detach().cpu()-y.detach().cpu()).pow(2).mean(1).numpy() weighted_mean=float((q*w*pred).sum()); uniform_mean=float(pred.mean()) return {'metric':metric,'critic_bce':bce,'failure_rate':float(labels.mean()), 'predicted_enrichment':float((q*labels).sum()/labels.mean()), 'observed_enrichment':observed,'ess_fraction':ess, 'weighted_pool_loss':weighted_mean,'uniform_pool_loss':uniform_mean} def baseline_train(seed,cfg): seed_all(seed); d=get_dataset('dynamics',seed,n_train=400,n_test=200) _,metric,_=train_model(make_model('rnn_small',d['input_shape'],d['out_dim']),d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,log=lambda *_:None) return metric def main(): check=math_check() grid=[{'lr':v,'epochs':EPOCHS,'alpha':0.,'weighted':False} for v in (.001,.003,.006)] base=sweep_baseline(lambda cfg: (lambda s: baseline_train(s,cfg)),grid,seeds=SEEDS) best_lr=base['best_cfg']['lr'] idea_cfgs=[{'lr':lr,'epochs':EPOCHS,'alpha':a,'weighted':True} for lr in (best_lr,.001,.006) for a in (.5,1.,2.)] # keep a comparable 3-setting intervention sweep: best baseline lr, three replay strengths idea_cfgs=[{'lr':best_lr,'epochs':EPOCHS,'alpha':a,'weighted':True} for a in (.5,1.,2.)] idea_runs=[(c,evaluate(lambda s,c=c:idea_train(s,c),seeds=SEEDS)) for c in idea_cfgs] cfg,idea=min(idea_runs,key=lambda z:z[1]['mean']) sig=idea_train(0,cfg,True) report=make_report('dynamics','rnn_small',base,idea,extra={ 'mechanism_signature':sig,'math_check':check, 'prediction':'criticality proposal enrichment equals E_q[y]/E_p[y] while weighted expectation preserves uniform loss', 'confirmed': bool(check['confirmed_identity'] and abs(sig['observed_enrichment']-sig['predicted_enrichment'])/max(sig['predicted_enrichment'],1e-9)<.20), 'idea_sweep':[{'cfg':c,'result':r} for c,r in idea_runs], 'track_justification':'dynamics is the built-in structural match for rollout failure/control states'}) Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2)) if __name__=='__main__': main()