1import sys, json, random, itertools
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6import torch.nn.functional as F
 7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 8from bench import make_model, train_model, evaluate, sweep_baseline, make_report, get_dataset
 9import route_track
10
11SEEDS=tuple(range(8)); EPOCHS=15; BATCH=128
12LRS=[1e-3,3e-3,1e-2]; WDS=[0.0,1e-4]
13
14def seed_all(s):
15    random.seed(s); np.random.seed(s); torch.manual_seed(s)
16    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
17
18def ds(seed):
19    # Load the accepted track through bench; restore its documented 4-edge target
20    # shape because the read-only adapter flattens all custom regression targets.
21    d=get_dataset('budgeted_route_costs', seed, n_train=400, n_test=160)
22    d['ytr']=d['ytr'].reshape(400,4); d['yte']=d['yte'].reshape(160,4)
23    d['out_dim']=4
24    return d
25
26class Net(nn.Module):
27    def __init__(self,shape):
28        super().__init__(); self.m=make_model('mlp_tiny',shape,4)
29    def forward(self,x): return F.softplus(self.m(x))
30
31def base(seed,cfg,ret=False):
32    seed_all(1000+seed); d=ds(seed); m=Net(d['input_shape'])
33    m,metric,_=train_model(m,d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['wd'],log=lambda x:None)
34    return (float(metric),m,d) if ret else float(metric)
35
36def subsets(B=1): return [()] + [(i,) for i in range(4)]
37
38def adv_loss(pred,c,tau=.10):
39    # Parallel routes: path A=(0,1), B=(2,3); enumerate all budget-1 delays.
40    delays=torch.tensor([[0.,0.,0.,0.],[0.,0.,1.2,0.],[0.,1.2,0.,0.],[1.2,0.,0.,0.],[0.,0.,0.,1.2]],device=pred.device)
41    vals=[]
42    for d in delays:
43        ct=c+d; cp=pred+d
44        true=torch.stack((ct[:,0]+ct[:,1],ct[:,2]+ct[:,3]),1)
45        pp=torch.softmax(-torch.stack((cp[:,0]+cp[:,1],cp[:,2]+cp[:,3]),1)/tau,1)
46        vals.append((pp*true).sum(1)-true.min(1).values)
47    return torch.stack(vals,1).max(1).values.mean()
48
49def idea(seed,cfg,ret=False):
50    seed_all(1000+seed); d=ds(seed); m=Net(d['input_shape']); opt=torch.optim.Adam(m.parameters(),lr=cfg['lr'],weight_decay=cfg['wd'])
51    x,y=d['xtr'],d['ytr']
52    for ep in range(EPOCHS):
53        p=torch.randperm(len(x))
54        for j in range(0,len(x),BATCH):
55            q=m(x[p[j:j+BATCH]]); c=y[p[j:j+BATCH]]
56            loss=adv_loss(q,c)+.01*((q-c)**2).mean()
57            opt.zero_grad(); loss.backward(); opt.step()
58    m.eval()
59    with torch.no_grad(): metric=float(((m(d['xte'])-d['yte'])**2).mean())
60    return (metric,m,d) if ret else metric
61
62def graph_metrics(m,d):
63    m = m.cpu()
64    with torch.no_grad(): p=m(d['xte']).numpy(); c=d['yte'].numpy()
65    worst=[]; flips=[]; nominal=[]
66    for a,b in zip(p,c):
67        vals=[]; fs=[]
68        for sub in subsets():
69            dd=np.zeros(4); dd[list(sub)]=1.2
70            ta=float(b[0]+b[1]+dd[0]+dd[1]); tb=float(b[2]+b[3]+dd[2]+dd[3])
71            pa=float(a[0]+a[1]+dd[0]+dd[1]); pb=float(a[2]+a[3]+dd[2]+dd[3])
72            true_i=0 if ta <= tb else 1; pred_i=0 if pa <= pb else 1
73            true_cost=(ta,tb); vals.append(max(0., true_cost[pred_i]-true_cost[true_i])); fs.append(int(true_i!=pred_i))
74        worst.append(max(vals)); flips.append(max(fs)); nominal.append(vals[0])
75    return {'worst_regret':float(np.mean(worst)),'flip_rate':float(np.mean(flips)),'nominal_regret':float(np.mean(nominal))}
76
77def main():
78    grid=[{'lr':lr,'wd':wd} for lr in LRS for wd in WDS]
79    baseblock=sweep_baseline(lambda cfg: lambda s:base(s,cfg),grid)
80    best=baseblock['best_cfg']
81    idea_grid=[best,{'lr':1e-3,'wd':best['wd']},{'lr':1e-2,'wd':best['wd']}]
82    # de-duplicate while preserving the required shared union.
83    ir=[]; chosen=None
84    for cfg in idea_grid:
85        r=evaluate(lambda s,cfg=cfg:idea(s,cfg),SEEDS); ir.append({'cfg':cfg,'result':r})
86        if chosen is None or r['mean']<chosen['result']['mean']: chosen=ir[-1]
87    ideares=chosen['result']
88    bm,im=[],[]
89    for s in SEEDS:
90        _,mb,db=base(s,best,True); _,mi,di=idea(s,chosen['cfg'],True); bm.append(graph_metrics(mb,db)); im.append(graph_metrics(mi,di))
91    sig={'baseline_worst_regret_mean':float(np.mean([x['worst_regret'] for x in bm])),'idea_worst_regret_mean':float(np.mean([x['worst_regret'] for x in im])),'predicted_effect':'adversarial training reduces budgeted path flips/regret','confirmed':float(np.mean([x['worst_regret'] for x in im]))<float(np.mean([x['worst_regret'] for x in bm]))}
92    rep=make_report('budgeted_route_costs','mlp_tiny',baseblock,ideares,{'custom_track':{'name':'budgeted_route_costs','file':'route_track.py','domain':'graph-decision'},'idea_sweep':ir,'mechanism_signature':sig})
93    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
94if __name__=='__main__': main()