1import json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6import sys
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, train_model, sweep_baseline, evaluate, make_report
  9
 10TRACK='dynamics'; MODEL='rnn_small'
 11# Equal-budget settings; the union is used on both sides.
 12LR_GRID=[1e-3, 3e-3, 6e-3]
 13EPOCHS=12
 14N_TRAIN=400
 15N_POOL=1600
 16SEEDS=tuple(range(8))
 17
 18
 19def seed_all(s):
 20    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 21    if torch.cuda.is_available():
 22        try: torch.cuda.manual_seed_all(s)
 23        except Exception: pass
 24
 25
 26def state_view(x):
 27    # Last observed pendulum state in each 8-step input window.
 28    return np.asarray(x, dtype=np.float64).reshape(len(x), 8, 3)[:, -1, :]
 29
 30
 31def residual(x):
 32    # Conservative, task-independent safety margin: theta and angular velocity.
 33    z=state_view(x)
 34    return np.minimum(1.25-np.abs(z[:,0]), 2.0-0.35*np.abs(z[:,1]))
 35
 36
 37def cover_radius(query, samples):
 38    # Small pool: exact vectorized nearest-neighbor radius.
 39    q=np.asarray(query); s=np.asarray(samples)
 40    best=np.full(len(q), np.inf)
 41    for j in range(0,len(s),128):
 42        d=((q[:,None,:]-s[None,j:j+128,:])**2).sum(2)
 43        best=np.minimum(best, np.sqrt(d.min(1)))
 44    return float(best.max()), best
 45
 46
 47def adaptive_indices(pool_x, n, seed, gamma=0.35):
 48    z=state_view(pool_x)
 49    rng=np.random.default_rng(seed)
 50    # Four seed states, then q = normalized distance + gamma*low-margin.
 51    chosen=list(rng.choice(len(pool_x), size=8, replace=False))
 52    chosen_set=set(chosen)
 53    scale=np.std(z,axis=0)+1e-6
 54    zn=z/scale
 55    r=residual(pool_x)
 56    eps=0.15
 57    while len(chosen)<n:
 58        sel=np.asarray(chosen)
 59        best=np.full(len(pool_x), np.inf)
 60        for j in range(0,len(sel),128):
 61            d=((zn[:,None,:]-zn[sel[j:j+128]][None,:,:])**2).sum(2)
 62            best=np.minimum(best, np.sqrt(d.min(1)))
 63        q=best + gamma*np.maximum(0.,eps-r)
 64        q[list(chosen_set)]=-np.inf
 65        k=int(np.argmax(q)); chosen.append(k); chosen_set.add(k)
 66    return np.asarray(chosen,dtype=int)
 67
 68
 69def prepare(seed, mode):
 70    # Pool and test set are generated from the same benchmark dynamics task.
 71    base=get_dataset(TRACK, seed, n_train=N_TRAIN, n_test=1000)
 72    pool=get_dataset(TRACK, seed+10000, n_train=N_POOL, n_test=1000)
 73    rng=np.random.default_rng(seed+991)
 74    if mode=='baseline': idx=rng.choice(N_POOL,size=N_TRAIN,replace=False)
 75    else: idx=adaptive_indices(pool['xtr'].numpy(),N_TRAIN,seed)
 76    ds=dict(base)
 77    ds['xtr']=pool['xtr'][idx].clone(); ds['ytr']=pool['ytr'][idx].clone()
 78    return ds, pool, idx
 79
 80
 81def run_one(seed, mode, cfg, return_net=False):
 82    seed_all(seed + (0 if mode=='baseline' else 100000))
 83    ds,pool,idx=prepare(seed,mode)
 84    net=make_model(MODEL, ds['input_shape'], ds['out_dim'])
 85    net,metric,hist=train_model(net,ds,epochs=EPOCHS,lr=float(cfg['lr']),batch=128,log=lambda *_:None)
 86    if return_net:
 87        return float(metric),net,ds,pool,idx
 88    return float(metric)
 89
 90
 91def math_check():
 92    rng=np.random.default_rng(4); cloud=rng.uniform(-1,1,(500,3)); samples=cloud[rng.choice(500,40,False)]
 93    # r(x)=0.2+L||x-c|| is L-Lipschitz; verify the stated lower bound.
 94    L=0.7; c=np.array([0.1,-0.2,0.3]); r=lambda a:0.2+L*np.linalg.norm(a-c,axis=1)
 95    d=cover_radius(cloud,samples)[0]; assert r(cloud).min() >= r(samples).min()-L*d-1e-10
 96    return {'checked':True,'L':L,'delta':d,'sample_min':float(r(samples).min()),'cloud_min':float(r(cloud).min()),'bound':float(r(samples).min()-L*d)}
 97
 98
 99def main():
100    check=math_check()
101    grid=[{'lr':x} for x in LR_GRID]
102    base=sweep_baseline(lambda cfg: lambda s: run_one(s,'baseline',cfg),grid,seeds=(0,1,2,3))
103    # Explicitly evaluate idea at all three settings; best is selected on the same tuning seeds.
104    idea_cfg_results=[]
105    for cfg in grid:
106        rr=evaluate(lambda s,cfg=cfg: run_one(s,'idea',cfg),seeds=(0,1,2,3))
107        idea_cfg_results.append((rr['mean'],cfg))
108    idea_cfg=min(idea_cfg_results,key=lambda x:x[0])[1]
109    idea=evaluate(lambda s: run_one(s,'idea',idea_cfg),seeds=SEEDS)
110    # Signature uses trained models, not a toy identity: compare pool coverage and
111    # boundary/general-region prediction errors for one paired seed.
112    bm,bnet,bds,bpool,bidx=run_one(0,'baseline',base['best_cfg'],True)
113    im,inet,ids,ipool,iidx=run_one(0,'idea',idea_cfg,True)
114    with torch.no_grad():
115        bx=bpool['xtr']; by=bpool['ytr']; ix=ipool['xtr']; iy=ipool['ytr']
116        bdev=next(bnet.parameters()).device; idev=next(inet.parameters()).device
117        bp=bnet(bx.to(bdev)).detach().cpu().numpy().reshape(-1); ip=inet(ix.to(idev)).detach().cpu().numpy().reshape(-1)
118    bz=state_view(bpool['xtr'].numpy()); iz=state_view(ipool['xtr'].numpy())
119    bdelta=cover_radius(bz,state_view(bpool['xtr'][bidx].numpy()))[0]
120    idelta=cover_radius(iz,state_view(ipool['xtr'][iidx].numpy()))[0]
121    br=residual(bpool['xtr'].numpy()); ir=residual(ipool['xtr'].numpy())
122    sig={'prediction':'adaptive reachable-state selection should reduce empirical coverage radius at equal simulator budget',
123         'observed':{'baseline_delta':bdelta,'idea_delta':idelta,'baseline_pool_mse':float(np.mean((bp-by.numpy())**2)),'idea_pool_mse':float(np.mean((ip-iy.numpy())**2)), 'baseline_min_sample_residual':float(br[bidx].min()),'idea_min_sample_residual':float(ir[iidx].min())},
124         'relative_delta_reduction_pct':float(100*(bdelta-idelta)/bdelta),
125         'confirmed':bool(idelta < bdelta)}
126    report=make_report(TRACK,MODEL,base,idea,extra={'math_check':check,'mechanism_signature':sig,'idea_grid_results':idea_cfg_results,'selected_idea_cfg':idea_cfg,'budget':{'epochs':EPOCHS,'train_samples':N_TRAIN,'pool_samples':N_POOL}})
127    Path('bench_report.json').write_text(json.dumps(report,indent=2))
128    print(json.dumps(report,indent=2))
129
130if __name__=='__main__': main()