import sys, json, random, math 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 SEED0 = 2686 EPOCHS = 18 BATCH = 128 # Union of baseline and idea grids: each side is evaluated at all learning rates. LRS = [1e-3, 3e-3, 1e-2] WEIGHT_DECAYS = [0.0, 1e-4] # baseline method knob, swept fairly def math_check(): # Exact scalar zonotope radius recurrence r+=|lambda|r+gamma. r0, gamma, n = .2, .03, 8 rows = [] for lam in [0., .2, .5, .8, 1., 1.2]: observed = abs(lam)**n*r0 + gamma*sum(abs(lam)**i for i in range(n)) propagated = r0 for _ in range(n): propagated = abs(lam)*propagated + gamma rows.append(abs(observed-propagated)) return {'max_radius_recurrence_error': float(max(rows)), 'confirmed': bool(max(rows) < 1e-12)} def device_net(model): return model def interval_reach_loss(net, x, horizon=8, disturbance=.025): """Differentiable interval-zonotope certificate around each training window. State is (theta, omega); the learned RNN predicts terminal theta. We propagate the known local pendulum bounds with interval Euler dynamics and penalize state and terminal-set support violations plus non-contraction.""" b = x.shape[0] # initial state uncertainty induced by a bounded set around first observed state lo = x[:, :2] - torch.tensor([.04, .04], device=x.device) hi = x[:, :2] + torch.tensor([.04, .04], device=x.device) total = torch.zeros((), device=x.device) radii = [] dt = .05 for k in range(horizon): # conservative interval for acceleration: -sin(theta)-.1 omega + u + w thl, thh = lo[:, 0], hi[:, 0] oml, omh = lo[:, 1], hi[:, 1] ul = x[:, 3*k+2] - .08 uh = x[:, 3*k+2] + .08 # sin range via endpoint plus extrema (bounds here are small enough, but sound) candidates = [torch.sin(thl), torch.sin(thh)] twopi = 2*math.pi for m in range(-3, 4): p = m*math.pi + math.pi/2 candidates.append(torch.where((thl <= p) & (p <= thh), torch.ones_like(thl)*math.sin(p), candidates[0])) sl = torch.stack(candidates).amin(0); sh = torch.stack(candidates).amax(0) # a=-sin(theta)-.1 omega+u plus bounded model disturbance al = -sh - .1*omh + ul - disturbance ah = -sl - .1*oml + uh + disturbance nlo = lo.clone(); nhi = hi.clone() nlo[:, 0] = lo[:, 0] + dt*oml; nhi[:, 0] = hi[:, 0] + dt*omh nlo[:, 1] = lo[:, 1] + dt*al; nhi[:, 1] = hi[:, 1] + dt*ah lo, hi = nlo, nhi rad = (hi-lo).mean(0) radii.append(rad) # support of box against |theta|<=1.5, |omega|<=4 total = total + torch.relu(hi[:,0]-1.5).mean() + torch.relu(-lo[:,0]-1.5).mean() total = total + .15*(torch.relu(hi[:,1]-4).mean()+torch.relu(-lo[:,1]-4).mean()) # terminal set and contraction deficit: radius should not grow beyond initial radius terminal_rad = (hi-lo).abs().mean() total = total + 2.0*(torch.relu(terminal_rad - .25))**2 if len(radii) >= 2: total = total + .2*torch.relu(radii[-1].mean() - radii[0].mean())**2 # Couple certificate to the actual trained network's prediction on the set center. center = (lo+hi)/2 pred = net(x) total = total + .05*torch.relu(pred.abs().mean() - 1.5)**2 return total def train_idea(model, ds, epochs, lr, weight_decay, lam=.15): """Only custom loop because the method changes the training loss.""" errs=[] for dev in (['cuda','cpu'] if torch.cuda.is_available() else ['cpu']): try: net=model.to(dev); opt=torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay) lossf=nn.MSELoss(); xtr,ytr=ds['xtr'].to(dev),ds['ytr'].to(dev) for ep in range(epochs): net.train(); perm=torch.randperm(len(xtr),device=dev) for i in range(0,len(xtr),BATCH): ix=perm[i:i+BATCH]; xb,yb=xtr[ix],ytr[ix] loss=lossf(net(xb),yb)+lam*interval_reach_loss(net,xb) opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step() net.eval() with torch.no_grad(): metric=float(((net(ds['xte'].to(dev))-ds['yte'].to(dev))**2).mean()) return metric except RuntimeError as e: errs.append(str(e)); continue raise RuntimeError(';'.join(errs)) def baseline_fn(cfg): def run(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) ds=get_dataset('dynamics',seed,n_train=400,n_test=200) net=make_model('rnn_small', ds['input_shape'], ds['out_dim']) _, metric, _=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['weight_decay'],log=lambda *_:None) return metric return run def idea_fn(cfg): def run(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) ds=get_dataset('dynamics',seed,n_train=400,n_test=200) net=make_model('rnn_small', ds['input_shape'], ds['out_dim']) return train_idea(net,ds,EPOCHS,cfg['lr'],cfg['weight_decay'],cfg['lambda']) return run if __name__=='__main__': print(json.dumps({'math_check':math_check()}, indent=2)) base_grid=[{'lr':lr,'weight_decay':wd} for lr in LRS for wd in WEIGHT_DECAYS] base=sweep_baseline(baseline_fn,base_grid) # Same 3 lr choices and the baseline's best weight-decay, plus the same penalty sweep. idea_grid=[{'lr':lr,'weight_decay':base['best_cfg']['weight_decay'],'lambda':lam} for lr in LRS for lam in [.05,.15,.3]] # Evaluate idea configs on sweep seeds, choose best; full result only for winner. tried=[] for cfg in idea_grid: r=evaluate(idea_fn(cfg),seeds=(0,1,2,3)); tried.append({'cfg':cfg,'mean':r['mean']}) best=min(tried,key=lambda z:z['mean'])['cfg']; idea=evaluate(idea_fn(best)) base['idea_grid']=tried rep=make_report('dynamics','rnn_small',base,idea,extra={ 'prediction':'adding the interval-zonotope loss should change the trained model task metric', 'observed_nn':{ 'baseline_mean_test_mse':base['full']['mean'], 'idea_mean_test_mse':idea['mean'], 'observed_delta':idea['mean']-base['full']['mean'], 'best_cfg':best}, 'toy_numeric_check':math_check(), 'confirmed':False, 'note':'The arithmetic radius recurrence is exact, but this does not establish the mechanism on the trained benchmark models.'}) Path('bench_report.json').write_text(json.dumps(rep,indent=2)) print(json.dumps(rep,indent=2))