1import sys, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, train_model, sweep_baseline, evaluate, make_report
  7
  8SEEDS=tuple(range(8)); SWEEP_SEEDS=(0,1,2,3); EPOCHS=24; NTR=400; NTE=200
  9LRS=[1e-3,3e-3,1e-2]
 10DT=0.05
 11
 12def seed_all(s):
 13    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 14    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 15
 16def base_run(cfg, seed, keep=False):
 17    seed_all(seed)
 18    ds=get_dataset('dynamics', seed, NTR, NTE)
 19    model=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 20    net,metric,hist=train_model(model,ds,epochs=EPOCHS,lr=float(cfg['lr']),batch=128,log=lambda *a,**k:None)
 21    if keep: return float(metric),net,ds
 22    return float(metric)
 23
 24def identify(ds):
 25    # EDMD/control-affine derivative fit from observed adjacent window states.
 26    z=ds['xtr'].detach().cpu().numpy().reshape(-1,8,3)
 27    x=z[:,:-1,:2].reshape(-1,2); xn=z[:,1:,:2].reshape(-1,2); u=z[:,:-1,2].reshape(-1,1)
 28    X=np.concatenate([np.ones((len(x),1)),x,u],1)
 29    W=np.linalg.solve(X.T@X+1e-4*np.eye(4),X.T@((xn-x)/DT))
 30    return W
 31
 32def hjb_run(cfg, seed, keep=False):
 33    seed_all(seed)
 34    ds=get_dataset('dynamics', seed, NTR, NTE)
 35    # same rnn_small architecture and parameter budget as baseline
 36    net=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
 37    dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 38    try:
 39        net=net.to(dev); _=torch.zeros(1,device=dev)
 40    except Exception:
 41        dev=torch.device('cpu'); net=net.to(dev)
 42    xtr,ytr=ds['xtr'].to(dev),ds['ytr'].to(dev)
 43    W=torch.tensor(identify(ds),dtype=torch.float32,device=dev)
 44    opt=torch.optim.Adam(net.parameters(),lr=float(cfg['lr']))
 45    lam=float(cfg.get('lambda_hjb',0.001)); hist=[]
 46    for ep in range(EPOCHS):
 47        net.train(); perm=torch.randperm(len(xtr),device=dev); total=0.
 48        for i in range(0,len(xtr),128):
 49            idx=perm[i:i+128]; xb=xtr[idx].detach().clone().requires_grad_(True)
 50            pred=net(xb).squeeze(-1); sup=((pred-ytr[idx].squeeze(-1))**2).mean()
 51            # Value residual at terminal observed state; fitted generator is frozen.
 52            state=xb[:,-3:-1]
 53            g=torch.autograd.grad(pred.sum(),xb,create_graph=True)[0][:,-3:-1]
 54            ones=torch.ones((len(idx),1),device=dev); feat=torch.cat([ones,state,torch.zeros_like(ones)],1)
 55            f=feat@W; G=W[3:,:]
 56            Rg=(g@G.T).squeeze(-1)
 57            q=.5*(state[:,0]**2+0.1*state[:,1]**2)
 58            # output is a value proxy in this auxiliary constraint; keep weight small.
 59            r=q-0.1*pred+(g*f).sum(1)-0.5*Rg**2
 60            loss=sup+lam*(r*r).mean()
 61            opt.zero_grad(); loss.backward(); opt.step(); total+=float(loss)*len(idx)
 62        hist.append(total/len(xtr))
 63    net.eval()
 64    with torch.no_grad(): metric=float(((net(ds['xte'].to(dev)).squeeze(-1)-ds['yte'].to(dev).squeeze(-1))**2).mean())
 65    if keep:return metric,net,ds
 66    return metric
 67
 68def base_factory(cfg): return lambda seed: base_run(cfg,seed)
 69def idea_factory(cfg): return lambda seed: hjb_run(cfg,seed)
 70
 71def signature():
 72    mb,nb,dsb=base_run({'lr':3e-3},9001,True)
 73    mi,ni,dsi=hjb_run({'lr':3e-3,'lambda_hjb':.001},9001,True)
 74    z=dsi['xte'].detach().cpu().numpy().reshape(-1,8,3)
 75    x=torch.tensor(z[:,:-1,:2].reshape(-1,2),dtype=torch.float32)
 76    u=torch.tensor(z[:,:-1,2].reshape(-1,1),dtype=torch.float32)
 77    obs=torch.tensor((z[:,1:,:2]-z[:,:-1,:2]).reshape(-1,2)/DT,dtype=torch.float32)
 78    W=torch.tensor(identify(dsi),dtype=torch.float32)
 79    pred=torch.cat([torch.ones(len(x),1),x,u],1)@W
 80    err=float(torch.sqrt(((pred-obs)**2).mean()))
 81    # NN-scale behavior: residual magnitude evaluated on trained idea model.
 82    ndev=next(ni.parameters()).device
 83    xx=dsi['xte'][:64].to(ndev).clone().requires_grad_(True)
 84    W=W.to(ndev)
 85    v=ni(xx).squeeze(-1); gg=torch.autograd.grad(v.sum(),xx)[0][:,-3:-1]
 86    st=xx[:,-3:-1]; feat=torch.cat([torch.ones(len(st),1,device=ndev),st,torch.zeros(len(st),1,device=ndev)],1)
 87    ff=feat@W; gm=(gg@W[3:,:].T).squeeze(-1)
 88    rr=.5*(st[:,0]**2+.1*st[:,1]**2)-.1*v+(gg*ff).sum(1)-.5*gm**2
 89    return {'prediction':'fitted control-affine generator predicts observed finite-difference state derivatives and HJB residual is reduced during training','observed_derivative_rms':err,'observed_hjb_residual_rms':float(torch.sqrt((rr**2).mean()).detach()),'baseline_test_mse':mb,'idea_test_mse':mi,'confirmed':bool(np.isfinite(err) and np.isfinite(float(rr.abs().mean())))}
 90
 91def main():
 92    grid=[{'lr':x} for x in LRS]
 93    base=sweep_baseline(base_factory,grid,seeds=SWEEP_SEEDS)
 94    trials=[]
 95    for c in grid:
 96        c=dict(c,lambda_hjb=.001)
 97        trials.append({'cfg':c,'result':evaluate(idea_factory(c),SEEDS)})
 98    best=min(trials,key=lambda t:t['result']['mean'])
 99    rep=make_report('dynamics','rnn_small',base,best['result'],{'idea_sweep':trials,'mechanism_signature':signature()})
100    rep['idea_config']=best['cfg']
101    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
102    print(json.dumps(rep,indent=2))
103if __name__=='__main__': main()