1import sys, json, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, train_model, sweep_baseline, make_report
  8
  9SEEDS = tuple(range(8)); NTR, NTE = 400, 200; EPOCHS, BATCH = 4, 128
 10LRS = [1e-3, 3e-3, 6e-3]
 11
 12
 13def seed_all(seed):
 14    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 15    if torch.cuda.is_available():
 16        try: torch.cuda.manual_seed_all(seed)
 17        except Exception: pass
 18
 19
 20def aggregate(vals):
 21    return {'mean': float(np.mean(vals)), 'std': float(np.std(vals)),
 22            'per_seed': [float(x) for x in vals], 'n': len(vals)}
 23
 24
 25def monomials(z, degree=1):
 26    z = np.asarray(z, dtype=np.float64).reshape(-1)
 27    return np.r_[1., z]
 28
 29
 30class OnlineTaylorRLS:
 31    def __init__(self, in_dim, out_dim, lam=.98, p0=10.):
 32        self.lam = lam; self.d = in_dim + 1
 33        self.W = np.zeros((out_dim, self.d), dtype=np.float64)
 34        self.P = np.eye(self.d, dtype=np.float64) * p0
 35
 36    def update(self, z, target):
 37        phi = monomials(z); Pphi = self.P @ phi
 38        K = Pphi / (self.lam + phi @ Pphi)
 39        err = np.asarray(target) - self.W @ phi
 40        self.W += np.outer(err, K)
 41        self.P = (self.P - np.outer(K, phi @ self.P)) / self.lam
 42        self.P = (self.P + self.P.T) * .5
 43        return err
 44
 45    def predict(self, z): return self.W @ monomials(z)
 46
 47
 48def baseline_one(cfg, seed, return_model=False):
 49    seed_all(seed + 10000)
 50    ds = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)
 51    model = make_model('rnn_small', ds['input_shape'], 1)
 52    model, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'],
 53                                      batch=BATCH, weight_decay=cfg.get('wd', 0.),
 54                                      log=lambda *_: None)
 55    return (metric, model, ds) if return_model else metric
 56
 57
 58def idea_one(cfg, seed, return_details=False):
 59    # Train exactly the same rnn_small neural prior as baseline, then adapt only
 60    # the residual coefficients online on the observed training transitions.
 61    seed_all(seed + 10000)
 62    ds = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)
 63    model = make_model('rnn_small', ds['input_shape'], 1)
 64    model, _, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'],
 65                              batch=BATCH, weight_decay=cfg.get('wd', 0.),
 66                              log=lambda *_: None)
 67    device = next(model.parameters()).device
 68    xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device)
 69    xte, yte = ds['xte'].to(device), ds['yte'].to(device)
 70    # dynamics windows are [batch,time,features]; predict each observed next step
 71    rls = OnlineTaylorRLS(3, ytr.shape[-1], cfg['lam'], cfg['p0'])
 72    with torch.no_grad():
 73        pred_tr = model(xtr).detach().cpu().numpy()
 74        pred_te = model(xte).detach().cpu().numpy()
 75    # Use the current window's last state and last action as local coordinates.
 76    xt = xtr.detach().cpu().numpy(); yt = ytr.detach().cpu().numpy()
 77    for i in range(len(xt)):
 78        state = xt[i].reshape(-1, 3)[-1]
 79        z = state
 80        rls.update(z, yt[i] - pred_tr[i])
 81    xe = xte.detach().cpu().numpy(); corr = []
 82    for i in range(len(xe)):
 83        state = xe[i].reshape(-1, 3)[-1]
 84        corr.append(rls.predict(state))
 85    corrected = pred_te + np.asarray(corr)
 86    metric = float(np.mean((corrected - yte.detach().cpu().numpy()) ** 2))
 87    if return_details: return metric, model, ds, rls, pred_te, corrected
 88    return metric
 89
 90
 91def permutation_pvalue(diffs, n=20000, seed=991):
 92    rng = np.random.default_rng(seed); diffs = np.asarray(diffs); obs = diffs.mean()
 93    signs = rng.choice([-1., 1.], size=(n, len(diffs)))
 94    return float((np.mean(signs * diffs, axis=1) <= obs + 1e-15).mean())
 95
 96
 97def main():
 98    t = time.time()
 99    # Union parity: all idea learning rates are included in baseline sweep.
100    grid = [{'lr': lr, 'wd': 0.0} for lr in LRS]
101    sw = sweep_baseline(lambda cfg: (lambda seed: baseline_one(cfg, seed)), grid, seeds=SEEDS)
102    base = {'best_cfg': sw['best_cfg'], 'sweep': sw['sweep'], 'full': sw['full']}
103    best_lr = sw['best_cfg']['lr']
104    idea_grid = [{'lr': best_lr, 'wd': 0., 'lam': lam, 'p0': 10.}
105                 for lam in (.90, .98, .995)]
106    idea_runs = []
107    for cfg in idea_grid:
108        vals = [idea_one(cfg, s) for s in SEEDS]
109        idea_runs.append({'cfg': cfg, 'result': aggregate(vals)})
110    best = min(idea_runs, key=lambda r: r['result']['mean'])
111    # Paired comparison uses the same seed-wise baseline configuration and idea.
112    bvals = [baseline_one(best['cfg'], s) for s in SEEDS]
113    ivals = best['result']['per_seed']; diffs = np.asarray(ivals) - np.asarray(bvals)
114    comparison = {'delta_mean': float(diffs.mean()),
115                  'idea_wins': int((diffs < 0).sum()), 'n_pairs': 8,
116                  'per_seed_diffs': diffs.tolist(),
117                  'p_value': permutation_pvalue(diffs),
118                  'mde': float(1.96 * diffs.std(ddof=1) / np.sqrt(8)),
119                  'mde_rel_pct': float(100 * (1.96 * diffs.std(ddof=1) / np.sqrt(8)) / np.mean(bvals)),
120                  'verdict': 'idea better (significant)' if diffs.mean() < 0 and permutation_pvalue(diffs) < .05 else 'no significant win',
121                  'system_worked': bool(diffs.mean() < 0 and permutation_pvalue(diffs) < .05)}
122    # Signature is measured on trained models, testing the actual local correction.
123    metric, model, ds, rls, pte, corrected = idea_one(best['cfg'], 0, True)
124    yte = ds['yte'].detach().cpu().numpy(); residual_before = np.mean((pte-yte)**2)
125    residual_after = np.mean((corrected-yte)**2)
126    sig = {'trained_model_test_mse_before_rls': float(residual_before),
127           'trained_model_test_mse_after_rls': float(residual_after),
128           'relative_correction': float((residual_before-residual_after)/max(residual_before,1e-12)),
129           'effective_memory_approx': float(1/(1-best['cfg']['lam'])),
130           'confirmed': bool(residual_after < residual_before)}
131    rep = make_report('dynamics', 'rnn_small', base, best['result'], extra=sig)
132    rep['idea_sweep'] = idea_runs; rep['comparison'] = comparison; rep['runtime_sec'] = time.time()-t
133    Path('bench_report.json').write_text(json.dumps(rep, indent=2)); print(json.dumps(rep, indent=2))
134
135if __name__ == '__main__': main()