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, evaluate, sweep_baseline, make_report SEEDS=tuple(range(8)) # Same union of learning rates is swept for both systems. GRID=[{'lr':1e-3,'epochs':8},{'lr':3e-3,'epochs':8},{'lr':1e-2,'epochs':8}] NTR,NTE=400,200 BATCH=128 class MeanVarianceRNN(nn.Module): """The bench rnn_small backbone with the required mean and variance heads.""" def __init__(self, input_shape): super().__init__() self.rnn=nn.GRU(3,64,batch_first=True) self.mean=nn.Linear(64,1) self.rawvar=nn.Linear(64,1) def encode(self,x): _,h=self.rnn(x.view(x.shape[0],-1,3)); return h[-1] def forward(self,x): h=self.encode(x) return self.mean(h) def both(self,x): h=self.encode(x) return self.mean(h), torch.nn.functional.softplus(self.rawvar(h))+1e-4 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) def device_run(fn): # Small batches/models, with explicit GPU fallback as required by the harness. dev='cuda' if torch.cuda.is_available() else 'cpu' try: return fn(dev) except RuntimeError: return fn('cpu') def train_base(seed,cfg, capture=False): seed_all(seed); d=get_dataset('dynamics',seed,NTR,NTE) def run(dev): net=make_model('rnn_small',d['input_shape'],1).to(dev) opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']) x,y=d['xtr'].to(dev),d['ytr'].to(dev) for _ in range(cfg['epochs']): net.train(); p=torch.randperm(len(x),device=dev) for i in range(0,len(x),BATCH): z=net(x[p[i:i+BATCH]]); loss=((z-y[p[i:i+BATCH]])**2).mean() opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): pred=net(d['xte'].to(dev)); err=(pred-d['yte'].to(dev))**2 metric=float(err.mean()) if capture: return metric, {'pred':pred.cpu().numpy().ravel(),'err':err.cpu().numpy().ravel()} return metric return device_run(run) def recursive_q(net,x): """One-step Bellman variance estimate from four stochastic child branches. Branch noise represents rollout/transition uncertainty in the final observed (theta, omega, action); q_rec=E[q_child]+Var(m_child), with stop-gradient.""" b=x.shape[0]; k=4 scales=x.new_tensor([.01,.03,.15]).view(1,1,3) eps=torch.randn((b,k,3),device=x.device)*scales branches=x[:,None,:].repeat(1,k,1).view(b*k,-1) branches[:,-3:]=branches[:,-3:]+eps.view(b*k,3) mu,q=net.both(branches); mu=mu.view(b,k); q=q.view(b,k) return (q.mean(1)+(mu-mu.mean(1,keepdim=True)).pow(2).mean(1)).detach() def train_idea(seed,cfg,capture=False): seed_all(seed); d=get_dataset('dynamics',seed,NTR,NTE) def run(dev): net=MeanVarianceRNN(d['input_shape']).to(dev) opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']) x,y=d['xtr'].to(dev),d['ytr'].to(dev) for _ in range(cfg['epochs']): net.train(); p=torch.randperm(len(x),device=dev) for i in range(0,len(x),BATCH): ix=p[i:i+BATCH]; mu,q=net.both(x[ix]); qr=recursive_q(net,x[ix]) # recursive target is detached when weighting mean loss; q head is # separately regressed to the recursive target. meanloss=.5*((y[ix]-mu).pow(2)/(qr+1e-3)+torch.log(qr+1e-3)).mean() varloss=.05*(torch.log(q+1e-4)-torch.log(qr+1e-4)).pow(2).mean() loss=meanloss+varloss opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): mu,q=net.both(d['xte'].to(dev)); err=(mu-d['yte'].to(dev))**2 # Signature is measured on this trained model, not an analytical toy. qr=recursive_q(net,d['xte'].to(dev)) metric=float(err.mean()) if capture: return metric, {'pred':mu.cpu().numpy().ravel(),'err':err.cpu().numpy().ravel(),'q':q.cpu().numpy().ravel(),'qr':qr.cpu().numpy().ravel()} return metric return device_run(run) def base_factory(cfg): return lambda seed: train_base(seed,cfg) def idea_factory(cfg): return lambda seed: train_idea(seed,cfg) def main(): # Baseline tuning uses the harness sweep (4 seeds), then final paired 8-seed result. base=sweep_baseline(base_factory,GRID) idea_runs=[] for cfg in GRID: idea_runs.append({'cfg':cfg,'result':evaluate(idea_factory(cfg),SEEDS)}) best=min(idea_runs,key=lambda z:z['result']['mean']) idea=best['result']; idea['selected_cfg']=best['cfg'] # Capture the same selected systems once for a trained-model signature. bvals=[]; ivals=[]; ratios=[]; correlations=[] for s in SEEDS: bm,bd=train_base(s,base['best_cfg'],True); im,idat=train_idea(s,best['cfg'],True) bvals.append(bm); ivals.append(im) ratios.append(float(np.mean(idat['err'])/(np.mean(idat['qr'])+1e-12))) correlations.append(float(np.corrcoef(idat['err'],idat['qr'])[0,1])) base['full']={'mean':float(np.mean(bvals)),'std':float(np.std(bvals)),'per_seed':bvals,'n':8} idea={'mean':float(np.mean(ivals)),'std':float(np.std(ivals)),'per_seed':ivals,'n':8,'selected_cfg':best['cfg']} sig={'prediction':'recursive predicted variance should calibrate observed squared error near 1 and track heteroscedasticity', 'predicted_vs_observed':{'mean_q_recursive':float(np.mean([np.mean(idat['qr'])])), 'mean_observed_squared_error':float(np.mean([np.mean(idat['err'])])), 'calibration_ratio_mean':float(np.mean(ratios)), 'calibration_ratio_per_seed':ratios,'q_error_correlation_mean':float(np.mean(correlations))}, 'confirmed':bool(.8 <= np.mean(ratios) <= 1.2)} rep=make_report('dynamics','rnn_small',base,idea,sig) rep['protocol_notes']={'structural_match':'multi-step stochastic actuated pendulum rollout; recursive uncertainty applies to transition branches', 'baseline':'standard MSE training', 'idea':'mean/variance RNN with detached recursive child q plus between-child variance', 'n_train':NTR,'n_test':NTE,'branch_count':4,'grid_union':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()