1import json, random, sys
  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, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = tuple(range(4))
 12EPOCHS = 15
 13BATCH = 128
 14# Shared union: every idea lr is also evaluated by baseline sweep.
 15LRS = [1e-3, 3e-3, 1e-2]
 16UBAR = 1.5
 17DT = 0.05
 18
 19
 20def seed_all(seed):
 21    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 22    if torch.cuda.is_available():
 23        torch.cuda.manual_seed_all(seed)
 24
 25
 26def data(seed):
 27    return get_dataset('dynamics', seed, n_train=400, n_test=400)
 28
 29
 30def baseline_metric(cfg, seed, keep=False):
 31    seed_all(seed)
 32    ds = data(seed)
 33    net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 34    net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_: None)
 35    return float(metric)
 36
 37
 38def infer_rho(ds):
 39    # State-only demonstration geometry: observed one-step target-minus-current-state errors.
 40    x = ds['xtr'][:, -3]
 41    y = ds['ytr'].reshape(-1)
 42    raw = torch.abs(y - x)
 43    return float(torch.quantile(raw, 0.90).item() + 0.02)
 44
 45
 46def idea_train(cfg, seed, return_details=False):
 47    seed_all(seed)
 48    ds = data(seed)
 49    rho = infer_rho(ds)
 50    # Same rnn_small weights as baseline. The only change is bounded, actuator-aware
 51    # state correction readout and its funnel/authority training penalties.
 52    net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 53    xtr, ytr = ds['xtr'], ds['ytr'].reshape(-1)
 54    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 55    try:
 56        net = net.to(device); xtr=xtr.to(device); ytr=ytr.to(device)
 57        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'])
 58        mse = nn.MSELoss()
 59        for _ in range(EPOCHS):
 60            net.train(); perm=torch.randperm(len(xtr), device=device)
 61            for i in range(0, len(xtr), BATCH):
 62                ix=perm[i:i+BATCH]; raw=net(xtr[ix]).reshape(-1)
 63                current=xtr[ix, -3]
 64                # bounded actuator-like state increment over one sample interval
 65                pred=current + UBAR*DT*torch.tanh(raw)
 66                err=pred-ytr[ix]
 67                funnel=torch.relu(torch.abs(err)/rho-1.0).pow(2).mean()
 68                # requested feedback gain times funnel radius <= actuator authority
 69                gain=torch.abs(pred-current)/(torch.abs(current-ytr[ix])+1e-3)
 70                authority=torch.relu(gain*rho-UBAR).pow(2).mean()
 71                loss=mse(pred,ytr[ix]) + 0.20*funnel + 0.05*authority
 72                opt.zero_grad(); loss.backward(); opt.step()
 73        net.eval()
 74        with torch.no_grad():
 75            xt=ds['xte'].to(device); yt=ds['yte'].reshape(-1).to(device)
 76            pred=xt[:, -3] + UBAR*DT*torch.tanh(net(xt).reshape(-1))
 77            metric=float(((pred-yt)**2).mean().cpu())
 78            err=torch.abs(pred-yt)
 79            violation=float((err>rho).float().mean().cpu())
 80            gain=torch.abs(pred-xt[:, -3])/(torch.abs(xt[:, -3]-yt)+1e-3)
 81            requested=float((gain*rho).mean().cpu())
 82            observed=float(torch.clamp(gain*rho, max=UBAR).mean().cpu())
 83        if return_details:
 84            return metric, {'rho':rho, 'violation_rate':violation,
 85                            'requested_authority':requested,
 86                            'observed_authority':observed,
 87                            'authority_ratio':observed/(requested+1e-9)}
 88        return metric
 89    except RuntimeError:
 90        # Explicit shared-GPU fallback, rebuilding on CPU after any CUDA failure.
 91        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 92        torch.set_default_device('cpu')
 93        return idea_train_cpu(cfg, seed, return_details)
 94
 95
 96def idea_train_cpu(cfg, seed, return_details=False):
 97    seed_all(seed); ds=data(seed); rho=infer_rho(ds)
 98    net=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 99    x,y=ds['xtr'],ds['ytr'].reshape(-1); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
100    for _ in range(EPOCHS):
101        for i in range(0,len(x),BATCH):
102            raw=net(x[i:i+BATCH]).reshape(-1); cur=x[i:i+BATCH,-3]; yy=y[i:i+BATCH]
103            pred=cur+UBAR*DT*torch.tanh(raw); gain=torch.abs(pred-cur)/(torch.abs(cur-yy)+1e-3)
104            loss=((pred-yy)**2).mean()+.20*torch.relu(torch.abs(pred-yy)/rho-1).pow(2).mean()+.05*torch.relu(gain*rho-UBAR).pow(2).mean()
105            opt.zero_grad();loss.backward();opt.step()
106    with torch.no_grad():
107        xt,yt=ds['xte'],ds['yte'].reshape(-1); pred=xt[:,-3]+UBAR*DT*torch.tanh(net(xt).reshape(-1)); err=torch.abs(pred-yt); gain=torch.abs(pred-xt[:,-3])/(torch.abs(xt[:,-3]-yt)+1e-3)
108        metric=float(((pred-yt)**2).mean()); det={'rho':rho,'violation_rate':float((err>rho).float().mean()),'requested_authority':float((gain*rho).mean()),'observed_authority':float(torch.clamp(gain*rho,max=UBAR).mean())}
109        det['authority_ratio']=det['observed_authority']/(det['requested_authority']+1e-9)
110    return (metric,det) if return_details else metric
111
112
113def main():
114    grid=[{'lr':v} for v in LRS]
115    base=sweep_baseline(lambda cfg: lambda seed: baseline_metric(cfg,seed), grid, seeds=SWEEP_SEEDS)
116    # Evaluate the idea at all shared settings; choose by the same four-seed selection.
117    idea_sweep=[]
118    for cfg in grid:
119        vals=[idea_train(cfg,s) for s in SWEEP_SEEDS]
120        idea_sweep.append({'cfg':cfg,'mean':float(np.mean(vals))})
121    best=min(idea_sweep,key=lambda z:z['mean'])['cfg']
122    ivals=[idea_train(best,s) for s in SEEDS]
123    idea={'mean':float(np.mean(ivals)),'std':float(np.std(ivals)),'per_seed':ivals,'n':len(ivals), 'best_cfg':best,'sweep':idea_sweep}
124    sig=[]
125    for s in SEEDS:
126        _,d=idea_train(best,s,True); sig.append(d)
127    signature={'predicted_kmax':UBAR,'observed_authority_mean':float(np.mean([x['observed_authority'] for x in sig])), 'requested_authority_mean':float(np.mean([x['requested_authority'] for x in sig])), 'authority_ratio_mean':float(np.mean([x['authority_ratio'] for x in sig])), 'funnel_violation_rate_mean':float(np.mean([x['violation_rate'] for x in sig])), 'confirmed': bool(np.mean([x['authority_ratio'] for x in sig]) <= 1.01)}
128    report=make_report('dynamics','rnn_small',base,idea,{'custom_track':None, **signature})
129    Path('bench_report.json').write_text(json.dumps(report,indent=2))
130    print(json.dumps(report,indent=2))
131
132if __name__=='__main__': main()