import os, sys, json, math, 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, evaluate, sweep_baseline, make_report SEEDS=tuple(range(8)) # Union is shared: baseline and idea both evaluated at all lr/epochs settings. GRID=[{'lr':1e-3,'epochs':10},{'lr':3e-3,'epochs':10},{'lr':6e-3,'epochs':10}] BATCH=128 def sanity(): # A monotone map, independently of learned models, checks the claimed arithmetic. def F(z): return np.tanh(np.array([[.7,.2],[.1,.8]]) @ z) rng=np.random.default_rng(1907); rates=[]; widths=[] for h in [.2,.1,.05,.025]: ok=[]; ww=[] for _ in range(3000): lo=rng.uniform(-.5,.3,2); hi=lo+rng.uniform(.01,.15,2) yl=F(lo); yh=F(hi) lower=h*np.floor((yl-.01)/h) upper=h*np.floor((yh+.01+h)/h) z=lo+rng.random(2)*(hi-lo); q=h*np.floor(F(z)/h) ok.append(np.all((q>=lower-1e-9)&(q<=upper+1e-9))) ww.append(np.mean(upper-lower)) rates.append(float(np.mean(ok))); widths.append(float(np.mean(ww))) return {'containment':rates,'h': [.2,.1,.05,.025], 'mean_width':widths, 'all_contained':bool(min(rates)>=.999), 'width_decreases':bool(all(widths[i+1] <= widths[i]+1e-9 for i in range(3)))} 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 train_base(seed,cfg,return_model=False): seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=120) net=make_model('rnn_small',ds['input_shape'],ds['out_dim']) net,metric,hist=train_model(net,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=BATCH,weight_decay=0.0,log=lambda x:None) if return_model: return float(metric),net,ds return float(metric) def interval_loss(net,x,y,h=.05,delta=.015): # Local state-action cell: opposite corners are formed by coordinatewise +/- h/2. # This is a training surrogate for T- / T+ containment; target residual is inflated. lo=x-h/2; hi=x+h/2 yl=net(lo); yh=net(hi) lower=torch.minimum(yl,yh)-delta upper=torch.maximum(yl,yh)+delta+h contain=torch.relu(lower-y)+torch.relu(y-upper) width=torch.relu(upper-lower) # retain task accuracy while discouraging unnecessarily broad certificates return (contain**2).mean() + .02*width.mean() def train_idea(seed,cfg,return_model=False): seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=120) net=make_model('rnn_small',ds['input_shape'],ds['out_dim']) # Own loop is required because the idea changes the loss; all other settings match train_model. device='cuda' if torch.cuda.is_available() else 'cpu' try: net=net.to(device); xtr,ytr=ds['xtr'].to(device),ds['ytr'].to(device) opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=0.0) for ep in range(cfg['epochs']): net.train(); perm=torch.randperm(len(xtr),device=device) for i in range(0,len(xtr),BATCH): ix=perm[i:i+BATCH]; x=xtr[ix]; y=ytr[ix] pred=net(x); mse=((pred-y)**2).mean() loss=mse + .15*interval_loss(net,x,y) opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): metric=float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean()) if return_model:return metric,net,ds return metric except RuntimeError: # CPU fallback from a fresh identical seed/model. seed_all(seed); net=make_model('rnn_small',ds['input_shape'],ds['out_dim']).cpu() opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']) for ep in range(cfg['epochs']): perm=torch.randperm(len(ds['xtr'])) for i in range(0,len(perm),BATCH): ix=perm[i:i+BATCH]; pred=net(ds['xtr'][ix]); y=ds['ytr'][ix] loss=((pred-y)**2).mean()+.15*interval_loss(net,ds['xtr'][ix],y) opt.zero_grad();loss.backward();opt.step() with torch.no_grad(): metric=float(((net(ds['xte'])-ds['yte'])**2).mean()) if return_model:return metric,net,ds return metric def signature(cfg): vals=[]; widths=[]; violations=[] for s in SEEDS: metric,net,ds=train_idea(s,cfg,True) x,y=ds['xte'],ds['yte']; dev=next(net.parameters()).device with torch.no_grad(): z=x.to(dev); yl=net(z-.025); yh=net(z+.025) lower=torch.minimum(yl,yh)-.015; upper=torch.maximum(yl,yh)+.015+.05 yy=y.to(dev); vals.append(float(((yy>=lower)&(yy<=upper)).all(1).float().mean())) widths.append(float((upper-lower).mean())) violations.append(float(torch.relu(lower-yy).mean()+torch.relu(yy-upper).mean())) return {'predicted':'smaller h should reduce interval width while residual inflation preserves containment', 'observed_containment_mean':float(np.mean(vals)), 'observed_width_mean':float(np.mean(widths)), 'observed_violation_mean':float(np.mean(violations)), 'confirmed':bool(np.mean(vals)>=.95 and np.mean(violations)<1e-5)} def main(): san=sanity() # Baseline sweep uses all three configs on four seeds; each union lr is also idea-tested. base=sweep_baseline(lambda c: (lambda s: train_base(s,c)),GRID,seeds=(0,1,2,3)) idea_cfgs=[base['best_cfg']]+[c for c in GRID if c!=base['best_cfg']] idea_runs=[] for c in idea_cfgs: r=evaluate(lambda s,c=c: train_idea(s,c),SEEDS) idea_runs.append({'cfg':c,'result':r}) best=min(idea_runs,key=lambda q:q['result']['mean']) report=make_report('dynamics','rnn_small',base,best['result'],{'mechanism_signature':signature(best['cfg']), 'sanity_check':san,'idea_config_sweep':idea_runs, 'protocol_note':'Baseline and idea share rnn_small, dataset, lr/epoch union, batch and seeds; idea differs only by interval training loss.'}) report['custom_track']=None with open('bench_report.json','w') as f: json.dump(report,f,indent=2) print(json.dumps(report,indent=2)) if __name__=='__main__': main()