import sys, json, random, time from pathlib import Path import numpy as np import torch 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)); NTR, NTE = 400, 200; EPOCHS, BATCH = 4, 128 LRS = [1e-3, 3e-3, 6e-3] 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 aggregate(vals): return {'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': [float(x) for x in vals], 'n': len(vals)} def monomials(z, degree=1): z = np.asarray(z, dtype=np.float64).reshape(-1) return np.r_[1., z] class OnlineTaylorRLS: def __init__(self, in_dim, out_dim, lam=.98, p0=10.): self.lam = lam; self.d = in_dim + 1 self.W = np.zeros((out_dim, self.d), dtype=np.float64) self.P = np.eye(self.d, dtype=np.float64) * p0 def update(self, z, target): phi = monomials(z); Pphi = self.P @ phi K = Pphi / (self.lam + phi @ Pphi) err = np.asarray(target) - self.W @ phi self.W += np.outer(err, K) self.P = (self.P - np.outer(K, phi @ self.P)) / self.lam self.P = (self.P + self.P.T) * .5 return err def predict(self, z): return self.W @ monomials(z) def baseline_one(cfg, seed, return_model=False): seed_all(seed + 10000) ds = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE) model = make_model('rnn_small', ds['input_shape'], 1) model, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg.get('wd', 0.), log=lambda *_: None) return (metric, model, ds) if return_model else metric def idea_one(cfg, seed, return_details=False): # Train exactly the same rnn_small neural prior as baseline, then adapt only # the residual coefficients online on the observed training transitions. seed_all(seed + 10000) ds = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE) model = make_model('rnn_small', ds['input_shape'], 1) model, _, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg.get('wd', 0.), log=lambda *_: None) device = next(model.parameters()).device xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device) xte, yte = ds['xte'].to(device), ds['yte'].to(device) # dynamics windows are [batch,time,features]; predict each observed next step rls = OnlineTaylorRLS(3, ytr.shape[-1], cfg['lam'], cfg['p0']) with torch.no_grad(): pred_tr = model(xtr).detach().cpu().numpy() pred_te = model(xte).detach().cpu().numpy() # Use the current window's last state and last action as local coordinates. xt = xtr.detach().cpu().numpy(); yt = ytr.detach().cpu().numpy() for i in range(len(xt)): state = xt[i].reshape(-1, 3)[-1] z = state rls.update(z, yt[i] - pred_tr[i]) xe = xte.detach().cpu().numpy(); corr = [] for i in range(len(xe)): state = xe[i].reshape(-1, 3)[-1] corr.append(rls.predict(state)) corrected = pred_te + np.asarray(corr) metric = float(np.mean((corrected - yte.detach().cpu().numpy()) ** 2)) if return_details: return metric, model, ds, rls, pred_te, corrected return metric def permutation_pvalue(diffs, n=20000, seed=991): rng = np.random.default_rng(seed); diffs = np.asarray(diffs); obs = diffs.mean() signs = rng.choice([-1., 1.], size=(n, len(diffs))) return float((np.mean(signs * diffs, axis=1) <= obs + 1e-15).mean()) def main(): t = time.time() # Union parity: all idea learning rates are included in baseline sweep. grid = [{'lr': lr, 'wd': 0.0} for lr in LRS] sw = sweep_baseline(lambda cfg: (lambda seed: baseline_one(cfg, seed)), grid, seeds=SEEDS) base = {'best_cfg': sw['best_cfg'], 'sweep': sw['sweep'], 'full': sw['full']} best_lr = sw['best_cfg']['lr'] idea_grid = [{'lr': best_lr, 'wd': 0., 'lam': lam, 'p0': 10.} for lam in (.90, .98, .995)] idea_runs = [] for cfg in idea_grid: vals = [idea_one(cfg, s) for s in SEEDS] idea_runs.append({'cfg': cfg, 'result': aggregate(vals)}) best = min(idea_runs, key=lambda r: r['result']['mean']) # Paired comparison uses the same seed-wise baseline configuration and idea. bvals = [baseline_one(best['cfg'], s) for s in SEEDS] ivals = best['result']['per_seed']; diffs = np.asarray(ivals) - np.asarray(bvals) comparison = {'delta_mean': float(diffs.mean()), 'idea_wins': int((diffs < 0).sum()), 'n_pairs': 8, 'per_seed_diffs': diffs.tolist(), 'p_value': permutation_pvalue(diffs), 'mde': float(1.96 * diffs.std(ddof=1) / np.sqrt(8)), 'mde_rel_pct': float(100 * (1.96 * diffs.std(ddof=1) / np.sqrt(8)) / np.mean(bvals)), 'verdict': 'idea better (significant)' if diffs.mean() < 0 and permutation_pvalue(diffs) < .05 else 'no significant win', 'system_worked': bool(diffs.mean() < 0 and permutation_pvalue(diffs) < .05)} # Signature is measured on trained models, testing the actual local correction. metric, model, ds, rls, pte, corrected = idea_one(best['cfg'], 0, True) yte = ds['yte'].detach().cpu().numpy(); residual_before = np.mean((pte-yte)**2) residual_after = np.mean((corrected-yte)**2) sig = {'trained_model_test_mse_before_rls': float(residual_before), 'trained_model_test_mse_after_rls': float(residual_after), 'relative_correction': float((residual_before-residual_after)/max(residual_before,1e-12)), 'effective_memory_approx': float(1/(1-best['cfg']['lam'])), 'confirmed': bool(residual_after < residual_before)} rep = make_report('dynamics', 'rnn_small', base, best['result'], extra=sig) rep['idea_sweep'] = idea_runs; rep['comparison'] = comparison; rep['runtime_sec'] = time.time()-t Path('bench_report.json').write_text(json.dumps(rep, indent=2)); print(json.dumps(rep, indent=2)) if __name__ == '__main__': main()