import os, sys, json, random 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, sweep_baseline, make_report SEEDS = tuple(range(8)); TUNE = (0,1,2,3) LRS = [1e-3, 3e-3, 1e-2] MUS = [0.05, 0.2, 0.5] EPOCHS = 12; BATCH = 128 def seed_all(s): random.seed(s); np.random.seed(s); torch.manual_seed(s) try: if torch.cuda.is_available(): torch.cuda.manual_seed_all(s) except Exception: pass def get_residuals(ds, seed): seed_all(seed+10000) m=make_model('rnn_small',ds['input_shape'],ds['out_dim']) m,_,_=train_model(m,ds,epochs=3,lr=3e-3,batch=BATCH,log=lambda *a:None) if m is None: return None dev=next(m.parameters()).device with torch.no_grad(): p=m(ds['xtr'].to(dev)).cpu() r=(ds['ytr']-p).reshape(-1) scale=max(float(r.std()),1e-5) return r.float(), scale def scenario_train(ds, seed, lr, mu): seed_all(seed); m=make_model('rnn_small',ds['input_shape'],ds['out_dim']) rb=get_residuals(ds,seed) if rb is None: return float('nan'), {} residuals,scale=rb try: dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu') m=m.to(dev); x=ds['xtr'].to(dev); y=ds['ytr'].to(dev) r=residuals.to(dev); lo=float(y.min())-0.05; hi=float(y.max())+0.05 opt=torch.optim.Adam(m.parameters(),lr=lr) gen=torch.Generator(device=dev); gen.manual_seed(seed+20000) for _ in range(EPOCHS): m.train(); perm=torch.randperm(len(x),device=dev) for j in range(0,len(x),BATCH): ix=perm[j:j+BATCH]; pred=m(x[ix]); nom=((pred-y[ix])**2).mean() ind=torch.randint(len(r),(8,len(ix)),generator=gen,device=dev) ys=pred[:,None,:]+r[ind].T[:,:,None]*scale h=torch.relu(lo-ys)+torch.relu(ys-hi) loss=nom+mu*h.mean(); opt.zero_grad(); loss.backward(); opt.step() m.eval() with torch.no_grad(): metric=float(((m(ds['xte'].to(dev))-ds['yte'].to(dev))**2).mean()) pred=m(x).cpu(); obs=(pred-ds['ytr']).reshape(-1) sg=torch.Generator().manual_seed(seed+30000) draws=obs[torch.randint(len(obs),(8,len(obs)),generator=sg)] yy=pred[:,None,:]+draws.T[:,:,None] slack=torch.relu(lo-yy)+torch.relu(yy-hi) return metric, {'residual_std':float(obs.std()),'scenario_violation_rate':float((slack>0).float().mean()),'mean_slack':float(slack.mean())} except RuntimeError: return float('nan'), {} def baseline_train(ds,seed,lr): seed_all(seed); m=make_model('rnn_small',ds['input_shape'],ds['out_dim']) try: m,metric,_=train_model(m,ds,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *a:None) if m is None:return float('nan'),{} dev=next(m.parameters()).device with torch.no_grad(): pred=m(ds['xtr'].to(dev)).cpu(); obs=(pred-ds['ytr']).reshape(-1) lo=float(ds['ytr'].min())-0.05; hi=float(ds['ytr'].max())+0.05 sg=torch.Generator().manual_seed(seed+30000) draws=obs[torch.randint(len(obs),(8,len(obs)),generator=sg)] slack=torch.relu(lo-(pred[:,None,:]+draws.T[:,:,None]))+torch.relu((pred[:,None,:]+draws.T[:,:,None])-hi) return float(metric),{'residual_std':float(obs.std()),'scenario_violation_rate':float((slack>0).float().mean()),'mean_slack':float(slack.mean())} except RuntimeError:return float('nan'),{} def main(): # Baseline grid includes every lr and every method-side mu value (mu is inert for MSE). grid=[{'lr':lr,'mu':mu} for lr in LRS for mu in [0.0]+MUS] def base_factory(cfg): return lambda s: baseline_train(get_dataset('dynamics',s,400,200),s,cfg['lr'])[0] base=sweep_baseline(base_factory,grid,seeds=TUNE) # Tune the idea on exactly the same four seeds, then only the selected config is scored full. idea_sweep=[] for lr in LRS: for mu in MUS: vals=[scenario_train(get_dataset('dynamics',s,400,200),s,lr,mu)[0] for s in TUNE] idea_sweep.append({'cfg':{'lr':lr,'mu':mu},'mean':float(np.nanmean(vals))}) best=min(idea_sweep,key=lambda z:z['mean'])['cfg'] rows=[]; vals=[]; sig=[] for s in SEEDS: ds=get_dataset('dynamics',s,400,200) b,bs=baseline_train(ds,s,base['best_cfg']['lr']) v,isig=scenario_train(ds,s,best['lr'],best['mu']) vals.append(v); sig.append((bs,isig)) idea={'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':[float(v) for v in vals],'n':8,'best_cfg':best,'sweep':idea_sweep} br={'mean':float(np.mean([baseline_train(get_dataset('dynamics',s,400,200),s,base['best_cfg']['lr'])[0] for s in SEEDS])),'std':base['full']['std'],'per_seed':[float(baseline_train(get_dataset('dynamics',s,400,200),s,base['best_cfg']['lr'])[0]) for s in SEEDS],'n':8} # Use the already paired baseline values from rows, avoiding any synthetic signature. bvals=[] for s in SEEDS: bvals.append(baseline_train(get_dataset('dynamics',s,400,200),s,base['best_cfg']['lr'])[0]) br={'mean':float(np.mean(bvals)),'std':float(np.std(bvals)),'per_seed':[float(v) for v in bvals],'n':8} signature={'baseline_residual_std':float(np.mean([x[0]['residual_std'] for x in sig])),'idea_residual_std':float(np.mean([x[1]['residual_std'] for x in sig])),'baseline_mean_slack':float(np.mean([x[0]['mean_slack'] for x in sig])),'idea_mean_slack':float(np.mean([x[1]['mean_slack'] for x in sig])),'prediction':'scenario penalty reduces empirical residual-boundary slack','confirmed':bool(np.mean([x[1]['mean_slack'] for x in sig]) < np.mean([x[0]['mean_slack'] for x in sig]))} rep=make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':br},idea,{'mechanism_signature':signature,'track_justification':'Dynamics is structurally matched: the task is controlled pendulum rollout and the intervention constrains predicted outputs under empirical error scenarios.','idea_sweep':idea_sweep,'custom_track':None}) os.makedirs('artifacts',exist_ok=True) json.dump(rep,open('artifacts/bench_report.json','w'),indent=2); print(json.dumps(rep,indent=2)) if __name__=='__main__': main()