1import sys, json, random, math
 2from pathlib import Path
 3import numpy as np
 4import torch
 5from torch import nn
 6from torch.utils.data import TensorDataset, DataLoader
 7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 8from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
 9
10TRACK='dynamics'; MODEL='rnn_small'; SEEDS=tuple(range(8)); EPOCHS=30; BATCH=128
11LRS=[0.001,0.003,0.006]
12
13def seed_all(s):
14    random.seed(s); np.random.seed(s); torch.manual_seed(s)
15
16def baseline_metric(cfg, seed):
17    d=get_dataset(TRACK, seed, n_train=4000, n_test=1000)
18    seed_all(seed)
19    net=make_model(MODEL,d['input_shape'],d['out_dim'])
20    _, metric, _=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,log=lambda *_:None)
21    return float(metric)
22
23def idea_metric(cfg, seed, return_model=False):
24    d=get_dataset(TRACK, seed, n_train=4000, n_test=1000)
25    seed_all(seed)
26    net=make_model(MODEL,d['input_shape'],d['out_dim'])
27    # The intervention is the only difference: boundary/high-risk samples receive
28    # larger loss weight, approximating planner-directed critical-mode probes.
29    device='cuda' if torch.cuda.is_available() else 'cpu'
30    try:
31        net=net.to(device)
32        x,y=d['xtr'].to(device),d['ytr'].to(device)
33        xt,yt=d['xte'].to(device),d['yte'].to(device)
34        opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
35        for _ in range(EPOCHS):
36            perm=torch.randperm(x.shape[0],device=device)
37            for ix in perm.split(BATCH):
38                xb,yb=x[ix],y[ix]
39                opt.zero_grad(set_to_none=True)
40                pred=net(xb)
41                # theta is the first component in the final input tuple (theta,omega,u)
42                theta=xb.view(xb.shape[0],-1,3)[:,-1,0]
43                risk=(theta.abs()>1.05).float()
44                weights=1.0+2.0*risk
45                loss=((pred-yb).pow(2)*weights[:,None]).mean()
46                loss.backward(); opt.step()
47        with torch.no_grad(): metric=float((net(xt)-yt).pow(2).mean().detach().cpu())
48        return (metric,net,d) if return_model else metric
49    except RuntimeError:
50        # Robust CUDA fallback, with fresh CPU model and identical seed/config.
51        device='cpu'; seed_all(seed); net=make_model(MODEL,d['input_shape'],d['out_dim']).to(device)
52        x,y=d['xtr'],d['ytr']; xt,yt=d['xte'],d['yte']; opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
53        for _ in range(EPOCHS):
54            perm=torch.randperm(x.shape[0])
55            for ix in perm.split(BATCH):
56                xb,yb=x[ix],y[ix]; opt.zero_grad(set_to_none=True); pred=net(xb)
57                theta=xb.view(xb.shape[0],-1,3)[:,-1,0]; weights=1+2*(theta.abs()>1.05).float()
58                loss=((pred-yb).pow(2)*weights[:,None]).mean(); loss.backward(); opt.step()
59        metric=float((net(xt)-yt).pow(2).mean().detach())
60        return (metric,net,d) if return_model else metric
61
62def mechanism_signature(cfg, base_seed=0):
63    bm=baseline_metric(cfg,base_seed)
64    im,model,d=idea_metric(cfg,base_seed,True)
65    with torch.no_grad():
66        x=d['xte']; truth=d['yte']; bp=make_model(MODEL,d['input_shape'],d['out_dim'])
67    # Re-train baseline model for behavior measurement under same seed.
68    seed_all(base_seed); bp,_,_=train_model(bp,d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,log=lambda *_:None)
69    with torch.no_grad():
70        bp_dev=next(bp.parameters()).device; model_dev=next(model.parameters()).device
71        pb=bp(x.to(bp_dev)).cpu(); pi=model(x.to(model_dev)).cpu(); yy=truth.cpu()
72    th=x.view(x.shape[0],-1,3)[:,-1,0]
73    critical=th.abs()>1.05
74    eb=((pb-yy).abs()>0.20)[critical].float().mean().item()
75    ei=((pi-yy).abs()>0.20)[critical].float().mean().item()
76    return {'critical_fraction':float(critical.float().mean()),'baseline_critical_bad_rate':float(eb),'idea_critical_bad_rate':float(ei),'observed_bad_rate_reduction':float(eb-ei),'required_rollouts_r_0.02_delta_0.05':149,'confirmed':bool(ei<eb)}
77
78def main():
79    grid=[{'lr':v} for v in LRS]
80    base=sweep_baseline(lambda cfg: lambda seed: baseline_metric(cfg,seed),grid,seeds=(0,1,2,3))
81    best=base['best_cfg']
82    # Required idea sweep at best and two nearby settings; all are in baseline union.
83    idea_cfgs=grid
84    idea_sweep=[]
85    for cfg in idea_cfgs:
86        r=evaluate(lambda seed,cfg=cfg: idea_metric(cfg,seed),seeds=SEEDS)
87        idea_sweep.append({'cfg':cfg,'result':r})
88    best_idea=min(idea_sweep,key=lambda z:z['result']['mean'])['result']
89    extra=mechanism_signature(best)
90    report=make_report(TRACK,MODEL,base,best_idea,extra=extra)
91    report['idea_sweep']=idea_sweep
92    report['protocol_note']='Official registered dynamics track; baseline uses bench.train_model; idea changes only the training loss and shares model, data, epochs, batch, and lr grid.'
93    Path('bench_report.json').write_text(json.dumps(report,indent=2))
94    print(json.dumps(report,indent=2))
95if __name__=='__main__': main()