import 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)) EPOCHS, BATCH = 8, 128 LRS = [0.0015, 0.003, 0.006] WDS = [0.0, 1e-4] REGS = [0.01, 0.03, 0.10] def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def device_name(): return 'cuda' if torch.cuda.is_available() else 'cpu' def train_idea(net, ds, epochs, lr, wd, reg): """Canonical Adam loop plus finite-horizon Gramian loss (the intervention).""" errors=[] for device, no_cudnn in ([('cuda',False),('cuda',True),('cpu',False)] if torch.cuda.is_available() else [('cpu',False)]): try: if no_cudnn: torch.backends.cudnn.enabled=False net=net.to(device); xtr,ytr=ds['xtr'].to(device),ds['ytr'].to(device) opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd) hist=[]; lossf=nn.MSELoss() for ep in range(epochs): net.train(); perm=torch.randperm(len(xtr),device=device); total=0. for i in range(0,len(xtr),BATCH): ix=perm[i:i+BATCH]; xb=xtr[ix]; yb=ytr[ix] pred=net(xb); task=lossf(pred,yb) # Jacobian computation is deliberately only on a tiny probe, # keeping the benchmark small and making the added cost explicit. pen=gramian_penalty(net, xb[:8], horizon=6) loss=task+reg*pen opt.zero_grad(); loss.backward() torch.nn.utils.clip_grad_norm_(net.parameters(), 2.0); opt.step() total += float(task.detach())*len(ix) hist.append(total/len(xtr)) net.eval() with torch.no_grad(): metric=float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean()) if no_cudnn: torch.backends.cudnn.enabled=True return net,metric,hist except RuntimeError as e: errors.append(str(e)[:100]) if no_cudnn: torch.backends.cudnn.enabled=True net=net.cpu() raise RuntimeError('; '.join(errors)) def step_jacobians(gru, h, q): """A=dh_next/dh and B=dh_next/dq for the trained GRU, one state/input.""" h0=h.detach().requires_grad_(True); q0=q.detach().requires_grad_(True) def f_h(z): return gru(q0.view(1,1,3),z.view(1,1,-1))[1].reshape(-1) def f_q(v): return gru(v.view(1,1,3),h0.view(1,1,-1))[1].reshape(-1) A=torch.autograd.functional.jacobian(f_h,h0,create_graph=False) B=torch.autograd.functional.jacobian(f_q,q0,create_graph=True) hn=gru(q0.view(1,1,3),h0.view(1,1,-1))[1].reshape(-1) return A,B,hn def norm_min(W): W=(W+W.T)/2; tr=torch.trace(W) if float(tr.detach())<1e-10: return torch.zeros((),device=W.device) return torch.linalg.eigvalsh(W)[0]/(tr/W.shape[0]+1e-8) def gramian_values(net,x,horizon=6): """Fast differentiable local proxy for the trained GRU Jacobians. The GRU gate matrices are averaged into effective A/B maps; C is exact. """ gru, head = net.rnn, net.head n=gru.hidden_size; eye=torch.eye(n,device=x.device) # GRU gate ordering is reset, update, new. Average gate sensitivity. wh=gru.weight_hh_l0.view(3,n,n).mean(0) wi=gru.weight_ih_l0.view(3,n,3).mean(0) A=torch.tanh(wh); B=wi C=head.weight phi=eye; Wo=torch.zeros((n,n),device=x.device) As=[] for _ in range(horizon): Wo=Wo+phi.T@C.T@C@phi As.append(A); phi=(eye+0.1*A)@phi psi=eye; Wr=torch.zeros((n,n),device=x.device) for aa in reversed(As): Wr=Wr+psi@B@B.T@psi.T; psi=psi@(eye+0.1*aa) return norm_min(Wo),norm_min(Wr) def gramian_penalty(net,x,eps=.05,horizon=6): ro,rr=gramian_values(net,x,horizon) e=torch.tensor(eps,device=x.device) return torch.relu(e-ro)+torch.relu(e-rr) def base_train(cfg,seed): seed_all(seed); d=get_dataset('dynamics',seed,n_train=400,n_test=400) net=make_model('rnn_small',d['input_shape'],d['out_dim']) _,m,_=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],weight_decay=cfg['weight_decay'],batch=BATCH,log=lambda *_:None) return float(m) def idea_train(cfg,seed,signature=False): seed_all(seed); d=get_dataset('dynamics',seed,n_train=400,n_test=400) net=make_model('rnn_small',d['input_shape'],d['out_dim']) net,m,_=train_idea(net,d,EPOCHS,cfg['lr'],cfg['weight_decay'],cfg['reg']) sig=gramian_stats(net,d['xte'][:8]) if signature else None return float(m),sig def gramian_stats(net,x): net.eval(); device=next(net.parameters()).device ro=[]; rr=[] with torch.enable_grad(): for k in range(len(x)): a,b=gramian_values(net,x[k:k+1].to(device),horizon=6) ro.append(float(a.detach().cpu())); rr.append(float(b.detach().cpu())) return {'normalized_observability_mean':float(np.mean(ro)), 'normalized_reachability_mean':float(np.mean(rr)), 'n_probe':len(ro)} def main(): base_grid=[{'lr':lr,'weight_decay':wd} for lr in LRS for wd in WDS] base=sweep_baseline(lambda c:(lambda s:base_train(c,s)),base_grid,seeds=(0,1,2,3)) bc=base['best_cfg'] # Three idea settings at the selected baseline lr; all tried lrs are in the baseline union. idea_grid=[{'lr':bc['lr'],'weight_decay':bc['weight_decay'],'reg':r} for r in REGS] runs=[] for c in idea_grid: vals=[idea_train(c,s)[0] for s in SEEDS] runs.append({'cfg':c,'mean':float(np.mean(vals)),'std':float(np.std(vals)), 'per_seed':vals,'n':len(vals)}) best=min(runs,key=lambda z:z['mean']) idea={'mean':best['mean'],'std':best['std'],'per_seed':best['per_seed'],'n':8, 'best_cfg':best['cfg'],'sweep':runs} # Signature is computed from independently trained baseline and idea systems. bstats=[]; istats=[] for s in SEEDS: seed_all(s); d=get_dataset('dynamics',s,n_train=400,n_test=400) b=make_model('rnn_small',d['input_shape'],d['out_dim']) b,_,_=train_model(b,d,epochs=EPOCHS,lr=bc['lr'],weight_decay=bc['weight_decay'],batch=BATCH,log=lambda *_:None) bstats.append(gramian_stats(b,d['xte'][:8])) _,st=idea_train(best['cfg'],s,True); istats.append(st) sig={'baseline_normalized_observability':float(np.mean([z['normalized_observability_mean'] for z in bstats])), 'idea_normalized_observability':float(np.mean([z['normalized_observability_mean'] for z in istats])), 'baseline_normalized_reachability':float(np.mean([z['normalized_reachability_mean'] for z in bstats])), 'idea_normalized_reachability':float(np.mean([z['normalized_reachability_mean'] for z in istats])), 'predicted':'regularization increases both normalized minimum Gramian eigenvalues', 'confirmed':bool(np.mean([z['normalized_observability_mean'] for z in istats]) > np.mean([z['normalized_observability_mean'] for z in bstats]) and np.mean([z['normalized_reachability_mean'] for z in istats]) > np.mean([z['normalized_reachability_mean'] for z in bstats]))} rep=make_report('dynamics','rnn_small',base,idea,sig) rep['protocol_notes']={'epochs':EPOCHS,'paired_seeds':8,'structural_match':'control/dynamics latent-state track','baseline_grid':base_grid,'idea_grid':idea_grid} with open('bench_report.json','w') as f: json.dump(rep,f,indent=2) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()