1import sys, json, random, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  9
 10SEED0 = 2686
 11EPOCHS = 18
 12BATCH = 128
 13# Union of baseline and idea grids: each side is evaluated at all learning rates.
 14LRS = [1e-3, 3e-3, 1e-2]
 15WEIGHT_DECAYS = [0.0, 1e-4]  # baseline method knob, swept fairly
 16
 17
 18def math_check():
 19    # Exact scalar zonotope radius recurrence r+=|lambda|r+gamma.
 20    r0, gamma, n = .2, .03, 8
 21    rows = []
 22    for lam in [0., .2, .5, .8, 1., 1.2]:
 23        observed = abs(lam)**n*r0 + gamma*sum(abs(lam)**i for i in range(n))
 24        propagated = r0
 25        for _ in range(n): propagated = abs(lam)*propagated + gamma
 26        rows.append(abs(observed-propagated))
 27    return {'max_radius_recurrence_error': float(max(rows)),
 28            'confirmed': bool(max(rows) < 1e-12)}
 29
 30
 31def device_net(model):
 32    return model
 33
 34
 35def interval_reach_loss(net, x, horizon=8, disturbance=.025):
 36    """Differentiable interval-zonotope certificate around each training window.
 37    State is (theta, omega); the learned RNN predicts terminal theta. We propagate
 38    the known local pendulum bounds with interval Euler dynamics and penalize state
 39    and terminal-set support violations plus non-contraction."""
 40    b = x.shape[0]
 41    # initial state uncertainty induced by a bounded set around first observed state
 42    lo = x[:, :2] - torch.tensor([.04, .04], device=x.device)
 43    hi = x[:, :2] + torch.tensor([.04, .04], device=x.device)
 44    total = torch.zeros((), device=x.device)
 45    radii = []
 46    dt = .05
 47    for k in range(horizon):
 48        # conservative interval for acceleration: -sin(theta)-.1 omega + u + w
 49        thl, thh = lo[:, 0], hi[:, 0]
 50        oml, omh = lo[:, 1], hi[:, 1]
 51        ul = x[:, 3*k+2] - .08
 52        uh = x[:, 3*k+2] + .08
 53        # sin range via endpoint plus extrema (bounds here are small enough, but sound)
 54        candidates = [torch.sin(thl), torch.sin(thh)]
 55        twopi = 2*math.pi
 56        for m in range(-3, 4):
 57            p = m*math.pi + math.pi/2
 58            candidates.append(torch.where((thl <= p) & (p <= thh), torch.ones_like(thl)*math.sin(p), candidates[0]))
 59        sl = torch.stack(candidates).amin(0); sh = torch.stack(candidates).amax(0)
 60        # a=-sin(theta)-.1 omega+u plus bounded model disturbance
 61        al = -sh - .1*omh + ul - disturbance
 62        ah = -sl - .1*oml + uh + disturbance
 63        nlo = lo.clone(); nhi = hi.clone()
 64        nlo[:, 0] = lo[:, 0] + dt*oml; nhi[:, 0] = hi[:, 0] + dt*omh
 65        nlo[:, 1] = lo[:, 1] + dt*al; nhi[:, 1] = hi[:, 1] + dt*ah
 66        lo, hi = nlo, nhi
 67        rad = (hi-lo).mean(0)
 68        radii.append(rad)
 69        # support of box against |theta|<=1.5, |omega|<=4
 70        total = total + torch.relu(hi[:,0]-1.5).mean() + torch.relu(-lo[:,0]-1.5).mean()
 71        total = total + .15*(torch.relu(hi[:,1]-4).mean()+torch.relu(-lo[:,1]-4).mean())
 72    # terminal set and contraction deficit: radius should not grow beyond initial radius
 73    terminal_rad = (hi-lo).abs().mean()
 74    total = total + 2.0*(torch.relu(terminal_rad - .25))**2
 75    if len(radii) >= 2:
 76        total = total + .2*torch.relu(radii[-1].mean() - radii[0].mean())**2
 77    # Couple certificate to the actual trained network's prediction on the set center.
 78    center = (lo+hi)/2
 79    pred = net(x)
 80    total = total + .05*torch.relu(pred.abs().mean() - 1.5)**2
 81    return total
 82
 83
 84def train_idea(model, ds, epochs, lr, weight_decay, lam=.15):
 85    """Only custom loop because the method changes the training loss."""
 86    errs=[]
 87    for dev in (['cuda','cpu'] if torch.cuda.is_available() else ['cpu']):
 88        try:
 89            net=model.to(dev); opt=torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
 90            lossf=nn.MSELoss(); xtr,ytr=ds['xtr'].to(dev),ds['ytr'].to(dev)
 91            for ep in range(epochs):
 92                net.train(); perm=torch.randperm(len(xtr),device=dev)
 93                for i in range(0,len(xtr),BATCH):
 94                    ix=perm[i:i+BATCH]; xb,yb=xtr[ix],ytr[ix]
 95                    loss=lossf(net(xb),yb)+lam*interval_reach_loss(net,xb)
 96                    opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step()
 97            net.eval()
 98            with torch.no_grad(): metric=float(((net(ds['xte'].to(dev))-ds['yte'].to(dev))**2).mean())
 99            return metric
100        except RuntimeError as e: errs.append(str(e)); continue
101    raise RuntimeError(';'.join(errs))
102
103
104def baseline_fn(cfg):
105    def run(seed):
106        random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
107        ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
108        net=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
109        _, metric, _=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['weight_decay'],log=lambda *_:None)
110        return metric
111    return run
112
113def idea_fn(cfg):
114    def run(seed):
115        random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
116        ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
117        net=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
118        return train_idea(net,ds,EPOCHS,cfg['lr'],cfg['weight_decay'],cfg['lambda'])
119    return run
120
121if __name__=='__main__':
122    print(json.dumps({'math_check':math_check()}, indent=2))
123    base_grid=[{'lr':lr,'weight_decay':wd} for lr in LRS for wd in WEIGHT_DECAYS]
124    base=sweep_baseline(baseline_fn,base_grid)
125    # Same 3 lr choices and the baseline's best weight-decay, plus the same penalty sweep.
126    idea_grid=[{'lr':lr,'weight_decay':base['best_cfg']['weight_decay'],'lambda':lam} for lr in LRS for lam in [.05,.15,.3]]
127    # Evaluate idea configs on sweep seeds, choose best; full result only for winner.
128    tried=[]
129    for cfg in idea_grid:
130        r=evaluate(idea_fn(cfg),seeds=(0,1,2,3)); tried.append({'cfg':cfg,'mean':r['mean']})
131    best=min(tried,key=lambda z:z['mean'])['cfg']; idea=evaluate(idea_fn(best))
132    base['idea_grid']=tried
133    rep=make_report('dynamics','rnn_small',base,idea,extra={
134        'prediction':'adding the interval-zonotope loss should change the trained model task metric',
135        'observed_nn':{
136            'baseline_mean_test_mse':base['full']['mean'],
137            'idea_mean_test_mse':idea['mean'],
138            'observed_delta':idea['mean']-base['full']['mean'],
139            'best_cfg':best},
140        'toy_numeric_check':math_check(),
141        'confirmed':False,
142        'note':'The arithmetic radius recurrence is exact, but this does not establish the mechanism on the trained benchmark models.'})
143    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
144    print(json.dumps(rep,indent=2))