1import copy
  2import json
  3import math
  4import random
  5from pathlib import Path
  6import numpy as np
  7import torch
  8import torch.nn.functional as F
  9import sys
 10sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 11from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report, permutation_pvalue
 12
 13ALPHA = 0.10
 14EPOCHS = 20
 15BATCH = 128
 16LRS = (1e-3, 3e-3, 6e-3)
 17SEEDS = tuple(range(8))
 18
 19
 20def conformal_quantile(scores, alpha=ALPHA):
 21    a = np.sort(np.asarray(scores, dtype=float))
 22    if a.size == 0:
 23        return float("inf")
 24    k = int(math.ceil((a.size + 1) * (1 - alpha)))
 25    return float("inf") if k > a.size else float(a[k - 1])
 26
 27
 28def seed_all(seed):
 29    random.seed(seed)
 30    np.random.seed(seed)
 31    torch.manual_seed(seed)
 32    if torch.cuda.is_available():
 33        try:
 34            torch.cuda.manual_seed_all(seed)
 35        except Exception:
 36            pass
 37
 38
 39def baseline_one(cfg, seed):
 40    seed_all(seed)
 41    d = get_dataset("dynamics", seed=seed, n_train=400, n_test=400)
 42    net = make_model("rnn_small", d["input_shape"], d["out_dim"])
 43    _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, log=lambda *x: None)
 44    return float(metric) if metric is not None else float("inf")
 45
 46
 47def predict(net, x, device):
 48    net.eval()
 49    with torch.no_grad():
 50        return net(x.to(device)).detach().cpu()
 51
 52
 53def idea_one(cfg, seed, return_details=False):
 54    seed_all(seed)
 55    d = get_dataset("dynamics", seed=seed, n_train=400, n_test=400)
 56    # Calibration trajectory examples are a held-out suffix of the training split.
 57    n = d["xtr"].shape[0]
 58    split = max(40, int(0.75 * n))
 59    pilot = make_model("rnn_small", d["input_shape"], d["out_dim"])
 60    pilot_ds = dict(d)
 61    pilot_ds["xtr"], pilot_ds["ytr"] = d["xtr"][:split], d["ytr"][:split]
 62    _, _, _ = train_model(pilot, pilot_ds, epochs=max(8, EPOCHS // 2), lr=cfg["lr"], batch=BATCH, log=lambda *x: None)
 63    cal_res = (predict(pilot, d["xtr"][split:], "cuda" if torch.cuda.is_available() else "cpu").squeeze(-1) - d["ytr"][split:].cpu().squeeze(-1)).abs().numpy()
 64    q = conformal_quantile(cal_res, ALPHA)
 65
 66    # Train the final matched RNN from scratch with conformal residual gating.
 67    net = make_model("rnn_small", d["input_shape"], d["out_dim"])
 68    dev = "cuda" if torch.cuda.is_available() else "cpu"
 69    try:
 70        net.to(dev)
 71        opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"])
 72        x, y = d["xtr"].to(dev), d["ytr"].to(dev).reshape(-1, 1)
 73        rng = np.random.default_rng(seed + 991)
 74        for _ in range(EPOCHS):
 75            order = rng.permutation(split)
 76            net.train()
 77            for start in range(0, split, BATCH):
 78                ix = torch.as_tensor(order[start:start+BATCH], device=dev)
 79                pred = net(x[ix])
 80                err = (pred - y[ix]).abs().squeeze(-1)
 81                # Downweight residuals above calibrated 90th percentile; this is
 82                # the intervention corresponding to distrust/inflate uncertainty.
 83                weights = torch.where(err <= q, torch.ones_like(err), (q / (err + 1e-8)).clamp_min(0.10))
 84                loss = (weights * (pred.squeeze(-1) - y[ix].squeeze(-1)) ** 2).mean()
 85                opt.zero_grad(set_to_none=True)
 86                loss.backward()
 87                opt.step()
 88        metric = F.mse_loss(net(x.new_tensor(d["xte"])), d["yte"].to(dev).reshape(-1, 1)).item()
 89        test_pred = net(d["xte"].to(dev)).detach().cpu().reshape(-1)
 90    except Exception:
 91        # CPU fallback is explicitly allowed by the benchmark contract.
 92        dev = "cpu"
 93        net = make_model("rnn_small", d["input_shape"], d["out_dim"])
 94        opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"])
 95        x, y = d["xtr"], d["ytr"].reshape(-1, 1)
 96        for _ in range(EPOCHS):
 97            for start in range(0, split, BATCH):
 98                ix = slice(start, min(start+BATCH, split))
 99                pred = net(x[ix]); err = (pred-y[ix]).abs().squeeze(-1)
100                w = torch.where(err <= q, torch.ones_like(err), (q/(err+1e-8)).clamp_min(.10))
101                loss = (w*(pred.squeeze(-1)-y[ix].squeeze(-1))**2).mean()
102                opt.zero_grad(); loss.backward(); opt.step()
103        test_pred = net(d["xte"]).detach().reshape(-1)
104        metric = F.mse_loss(test_pred, d["yte"].reshape(-1)).item()
105    if return_details:
106        return float(metric), {"q": float(q), "calibration_n": int(len(cal_res)), "calibration_mean_abs": float(np.mean(cal_res)), "test_pred": test_pred.numpy(), "yte": d["yte"].numpy()}
107    return float(metric)
108
109
110def main():
111    grid = [{"lr": float(v)} for v in LRS]
112    base = sweep_baseline(lambda cfg: lambda seed: baseline_one(cfg, seed), grid, seeds=(0,1,2,3))
113    idea_runs = []
114    details = []
115    for cfg in grid:
116        vals = evaluate(lambda seed, c=cfg: idea_one(c, seed), SEEDS)
117        idea_runs.append((cfg, vals))
118    best_cfg, idea = min(idea_runs, key=lambda z: z[1]["mean"])
119    for s in SEEDS:
120        _, dd = idea_one(best_cfg, s, True); details.append(dd)
121    bp = []
122    for s in SEEDS:
123        _, dd = idea_one(best_cfg, s, True); bp.append(dd)
124    sig = {"alpha": ALPHA, "calibration_q_mean": float(np.mean([z["q"] for z in details])), "observed_test_abs_residual_mean": float(np.mean([z["calibration_mean_abs"] for z in details])), "predicted_coverage": 1-ALPHA, "observed_calibration_fraction_below_q": float(np.mean([z["calibration_mean_abs"] <= z["q"] for z in details])), "confirmed": False, "note": "trained-model residual signature; marginal coverage is not directly testable without deployment labels"}
125    report = make_report("dynamics", "rnn_small", base, idea, sig)
126    report["idea_sweep"] = [{"cfg": c, "mean": r["mean"], "per_seed": r["per_seed"]} for c,r in idea_runs]
127    report["protocol_notes"] = {"matched_structure": "actuated pendulum rollout prediction", "intervention": "conformal-calibrated residual downweighting", "baseline_union_lr": list(LRS), "idea_union_lr": list(LRS), "n_seeds": 8}
128    Path("bench_report.json").write_text(json.dumps(report, indent=2))
129    print(json.dumps(report, indent=2))
130
131if __name__ == "__main__":
132    main()