import json, sys, random import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, '/home/maxwelhelp/all/math2nn') import bench SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS = 12 BATCH = 128 THETA_LIMIT = 1.5 HORIZON = 0.20 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 def baseline_train(cfg, seed): seed_all(seed) ds = bench.get_dataset('dynamics', seed, 400, 100) model = bench.make_model('rnn_small', ds['input_shape'], ds['out_dim']) _, metric, _ = bench.train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *a, **k: None) return float(metric) def certificate_terms(x, pred): last = x[:, -3:] theta, omega, action = last[:, 0], last[:, 1], last[:, 2] phi0 = THETA_LIMIT - torch.abs(theta) gamma = torch.abs(omega) + 0.20 * (torch.abs(action) + 1.0) + 0.50 phi_pred = THETA_LIMIT - torch.abs(pred[:, 0]) contract = phi0 - gamma * HORIZON return phi0, gamma, phi_pred, contract def idea_train(cfg, seed, return_signature=False): seed_all(seed) ds = bench.get_dataset('dynamics', seed, 400, 100) model = bench.make_model('rnn_small', ds['input_shape'], ds['out_dim']) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') try: model.to(device) xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device) xte, yte = ds['xte'].to(device), ds['yte'].to(device) opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) model.train() for _ in range(EPOCHS): perm = torch.randperm(len(xtr), device=device) for j in range(0, len(xtr), BATCH): xb, yb = xtr[perm[j:j+BATCH]], ytr[perm[j:j+BATCH]] pred = model(xb) phi0, gamma, phi_pred, contract = certificate_terms(xb, pred) safety_penalty = F.relu(-contract).pow(2).mean() loss = F.mse_loss(pred, yb) + cfg['cert_weight'] * safety_penalty opt.zero_grad(); loss.backward(); opt.step() model.eval() with torch.no_grad(): pred = model(xte) mse = F.mse_loss(pred, yte).item() phi0, gamma, phi_pred, contract = certificate_terms(xte, pred) predicted_drop = (gamma * HORIZON).cpu().numpy() observed_drop = (phi0 - phi_pred).cpu().numpy() residual = observed_drop - predicted_drop permitted = (contract >= 0).cpu().numpy() violations = ((phi_pred < 0) & permitted).sum() result = float(mse) if return_signature: return result, { 'n_test': int(len(xte)), 'predicted_certificate_drop_mean': float(np.mean(predicted_drop)), 'observed_certificate_drop_mean': float(np.mean(observed_drop)), 'predicted_vs_observed_drop_ratio': float(np.mean(observed_drop) / (np.mean(predicted_drop) + 1e-12)), 'contract_permitted_fraction': float(np.mean(permitted)), 'permitted_certificate_violations': int(violations), 'confirmed': bool(np.mean(observed_drop) <= np.mean(predicted_drop) * 1.20 + 1e-8) } return result except Exception: if device.type == 'cuda': torch.cuda.empty_cache() # Re-run entirely on CPU after any CUDA/runtime failure. torch.set_default_device('cpu') return idea_train_cpu(cfg, seed, return_signature) def idea_train_cpu(cfg, seed, return_signature=False): seed_all(seed) ds = bench.get_dataset('dynamics', seed, 400, 100) model = bench.make_model('rnn_small', ds['input_shape'], ds['out_dim']) xtr, ytr, xte, yte = ds['xtr'], ds['ytr'], ds['xte'], ds['yte'] opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) for _ in range(EPOCHS): for j in range(0, len(xtr), BATCH): pred = model(xtr[j:j+BATCH]); _, _, _, c = certificate_terms(xtr[j:j+BATCH], pred) loss = F.mse_loss(pred, ytr[j:j+BATCH]) + cfg['cert_weight'] * F.relu(-c).pow(2).mean() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): pred = model(xte); mse = F.mse_loss(pred, yte).item() phi0, gamma, phi_pred, contract = certificate_terms(xte, pred) pd = (gamma*HORIZON).numpy(); od = (phi0-phi_pred).numpy(); permitted = (contract >= 0).numpy() if return_signature: return mse, {'n_test': len(xte), 'predicted_certificate_drop_mean': float(pd.mean()), 'observed_certificate_drop_mean': float(od.mean()), 'predicted_vs_observed_drop_ratio': float(od.mean()/(pd.mean()+1e-12)), 'contract_permitted_fraction': float(permitted.mean()), 'permitted_certificate_violations': int(((phi_pred < 0).numpy() & permitted).sum()), 'confirmed': bool(od.mean() <= pd.mean()*1.20+1e-8)} return mse def main(): # Union of all idea learning rates is included in the baseline sweep. grid = [{'lr': lr, 'weight_decay': wd} for lr in (0.0015, 0.003, 0.006) for wd in (0.0, 1e-4)] base = bench.sweep_baseline(lambda c: lambda s: baseline_train(c, s), grid, seeds=SWEEP_SEEDS) idea_cfgs = [dict(base['best_cfg'], cert_weight=w) for w in (0.01, 0.05, 0.20)] idea_trials = [] for c in idea_cfgs: vals = [idea_train(c, s) for s in SWEEP_SEEDS] idea_trials.append({'cfg': c, 'mean': float(np.mean(vals))}) best = min(idea_trials, key=lambda z: z['mean']) best_cfg = best['cfg'] vals = [idea_train(best_cfg, s) for s in SEEDS] idea = {'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': vals, 'n': len(vals), 'best_cfg': best_cfg, 'sweep': idea_trials} signatures = [idea_train(best_cfg, s, True)[1] for s in SEEDS] sig = {k: (float(np.mean([x[k] for x in signatures])) if isinstance(signatures[0][k], (float, int)) and k not in ('n_test','permitted_certificate_violations') else (int(np.sum([x[k] for x in signatures])) if k in ('n_test','permitted_certificate_violations') else bool(all(x[k] for x in signatures)))) for k in signatures[0]} report = bench.make_report('dynamics', 'rnn_small', base, idea, {'predicted_vs_observed': sig, 'confirmed': sig['confirmed']}) with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()