import sys, json, random from pathlib import Path import numpy as np import torch import torch.nn as nn import torch.nn.functional as F sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, train_model, sweep_baseline, make_report SEEDS = tuple(range(8)) TRACK = 'dynamics' MODEL = 'rnn_small' EPOCHS = 18 BATCH = 64 NTRAIN, NTEST = 1000, 300 # All learning rates are shared by baseline and idea, satisfying search parity. GRID = [ {'lr': 1e-3, 'weight_decay': 0.0}, {'lr': 3e-3, 'weight_decay': 0.0}, {'lr': 1e-2, 'weight_decay': 0.0}, ] # Safety constants fixed before running: large predicted angle is unsafe. SAFETY_SCALE = 0.25 SAFETY_LIMIT = 0.80 BETA = 1.0 TEMP = 0.12 TOP_Q = 8 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass torch.set_num_threads(4) def baseline_train(seed, cfg, keep=False): seed_all(seed) ds = get_dataset(TRACK, seed, n_train=NTRAIN, n_test=NTEST) net = make_model(MODEL, ds['input_shape'], ds['out_dim']) net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None) if keep: return metric, net, ds return metric def pick_train(seed, cfg, keep=False): seed_all(seed) ds = get_dataset(TRACK, seed, n_train=NTRAIN, n_test=NTEST) net = make_model(MODEL, ds['input_shape'], ds['out_dim']) # The custom loop is the intervention: append worst violating trajectories # and optimize task MSE plus a differentiable surrogate safety penalty. try: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') net = net.to(device) x, y = ds['xtr'].to(device), ds['ytr'].to(device) opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) buffer = [] for ep in range(EPOCHS): net.train(); perm = torch.randperm(len(x), device=device) for start in range(0, len(x), BATCH): idx = perm[start:start+BATCH] pred = net(x[idx]).reshape(-1) task = F.mse_loss(pred, y[idx].reshape(-1)) # normalized violation v=[(|prediction|-limit)/scale]_+ v = F.relu((pred.abs() - SAFETY_LIMIT) / SAFETY_SCALE) # Pick worst current trajectories, as prescribed by the idea. k = min(TOP_Q, len(idx)) worst = torch.topk(v.detach(), k=k).indices buffer.extend(idx[worst].detach().cpu().tolist()) if len(buffer) > 64: buffer = buffer[-64:] bx = x[torch.as_tensor(buffer, device=device)] bp = net(bx).reshape(-1) bv = F.relu((bp.abs() - SAFETY_LIMIT) / SAFETY_SCALE) safe = (F.softplus(bv / TEMP) * TEMP).pow(2).mean() loss = task + BETA * safe opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): pred = net(ds['xte'].to(device)).reshape(-1) metric = float(F.mse_loss(pred, ds['yte'].to(device).reshape(-1)).cpu()) if keep: return metric, net, ds return metric except RuntimeError: # Robust CPU fallback for a shared/fragile CUDA slot. seed_all(seed); torch.cuda.empty_cache() if torch.cuda.is_available() else None old = torch.cuda.is_available # Re-run identical intervention on CPU by temporarily selecting device locally. net = make_model(MODEL, ds['input_shape'], ds['out_dim']) x, y = ds['xtr'], ds['ytr']; opt = torch.optim.Adam(net.parameters(), lr=cfg['lr']) buffer=[] for ep in range(EPOCHS): for start in range(0,len(x),BATCH): idx=torch.randperm(len(x))[start:start+BATCH]; pred=net(x[idx]).reshape(-1) v=F.relu((pred.abs()-SAFETY_LIMIT)/SAFETY_SCALE); k=min(TOP_Q,len(idx)) buffer.extend(idx[torch.topk(v.detach(),k).indices].tolist()); buffer=buffer[-64:] bp=net(x[torch.tensor(buffer)]).reshape(-1); bv=F.relu((bp.abs()-SAFETY_LIMIT)/SAFETY_SCALE) loss=F.mse_loss(pred,y[idx].reshape(-1))+BETA*(F.softplus(bv/TEMP)*TEMP).pow(2).mean() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): metric=float(F.mse_loss(net(ds['xte']).reshape(-1),ds['yte'].reshape(-1))) return (metric,net,ds) if keep else metric def behavior_signature(base_cfg, idea_cfg): rows=[] for s in SEEDS: bm,bn,bd=baseline_train(s,base_cfg,True); im,inn,idd=pick_train(s,idea_cfg,True) with torch.no_grad(): bp=bn(bd['xte'].to(next(bn.parameters()).device)).reshape(-1).cpu().numpy() ip=inn(idd['xte'].to(next(inn.parameters()).device)).reshape(-1).cpu().numpy() true=bd['yte'].numpy().reshape(-1) rows.append({'seed':s,'baseline_pred_rate':float(np.mean(np.abs(bp)>SAFETY_LIMIT)), 'idea_pred_rate':float(np.mean(np.abs(ip)>SAFETY_LIMIT)), 'observed_rate':float(np.mean(np.abs(true)>SAFETY_LIMIT)), 'baseline_pred_max_margin':float(np.max(np.maximum(np.abs(bp)-SAFETY_LIMIT,0))), 'idea_pred_max_margin':float(np.max(np.maximum(np.abs(ip)-SAFETY_LIMIT,0)))}) obs=float(np.mean([r['observed_rate'] for r in rows])) pred_b=float(np.mean([r['baseline_pred_rate'] for r in rows])); pred_i=float(np.mean([r['idea_pred_rate'] for r in rows])) # Prediction tested at NN scale: adaptive training should reduce predicted and # observed rare-event rates; confirmed only if both decrease by >=20%. return {'safety_limit':SAFETY_LIMIT,'scale':SAFETY_SCALE,'rows':rows, 'observed_rate_mean':obs,'predicted_rate_baseline_mean':pred_b, 'predicted_rate_idea_mean':pred_i, 'predicted_reduction_fraction':float((pred_b-pred_i)/max(pred_b,1e-9)), 'observed_reduction_fraction': 0.0, 'confirmed': False} def main(): base = sweep_baseline(lambda cfg: (lambda seed: baseline_train(seed,cfg)), GRID, seeds=(0,1,2,3)) # Three idea settings: baseline-best plus two nearby settings, all in GRID. idea_cfgs = GRID idea_runs=[] for cfg in idea_cfgs: r={'cfg':cfg,'result':{'mean':0,'std':0,'per_seed':[],'n':0}} vals=[pick_train(s,cfg) for s in SEEDS] r['result']={'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':[float(v) for v in vals],'n':len(vals)} idea_runs.append(r) best=min(idea_runs,key=lambda z:z['result']['mean']) sig=behavior_signature(base['best_cfg'],best['cfg']) report=make_report(TRACK,MODEL,base,best['result'],sig) report['idea_sweep']=idea_runs report['protocol']={'seeds':list(SEEDS),'n_train':NTRAIN,'n_test':NTEST,'epochs':EPOCHS,'batch':BATCH, 'structural_match':'dynamics: controlled pendulum rollout and stability/safety violations', 'baseline_method':'uniform per-example MSE','idea_method':'top-8 normalized predicted-angle violation replay'} Path('bench_report.json').write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__=='__main__': main()