1import sys, json, math, random
  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
  8from bench.protocol import evaluate
  9
 10ALPHA = 0.10
 11GAMMA = 0.25
 12K = 8
 13EPOCHS = 12
 14CAL_FRAC = 0.75
 15# This is the complete shared union: every idea lr is also a baseline lr.
 16GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 6e-3}]
 17SEEDS = tuple(range(8))
 18
 19
 20def cvar_lower(z, gamma):
 21    z = np.sort(np.asarray(z, dtype=float))
 22    h = gamma * len(z)
 23    m = int(math.floor(h))
 24    total = z[:m].sum()
 25    if h - m > 1e-12 and m < len(z):
 26        total += (h - m) * z[m]
 27    return float(total / h)
 28
 29
 30def cvar_eta(z, gamma):
 31    z = np.asarray(z, dtype=float)
 32    cand = np.r_[z, z.min() - 1., z.max() + 1.]
 33    vals = cand - np.maximum(cand[:, None] - z[None, :], 0).sum(1) / (gamma * len(z))
 34    return float(vals.max())
 35
 36
 37def conformal_q(r, alpha):
 38    r = np.sort(np.asarray(r, dtype=float))
 39    k = int(math.ceil((len(r) + 1) * (1 - alpha)))
 40    return float("inf") if k > len(r) else float(r[k - 1])
 41
 42
 43def seed_all(seed):
 44    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 45    if torch.cuda.is_available():
 46        try: torch.cuda.manual_seed_all(seed)
 47        except Exception: pass
 48
 49
 50def fit(seed, lr, idea):
 51    seed_all(seed)
 52    # The built-in dynamics task is the structural match: controlled pendulum
 53    # rollout and stability/safety margins. Use a held-out calibration slice.
 54    ds = get_dataset("dynamics", seed, n_train=400, n_test=400)
 55    ncal = int(len(ds["xtr"]) * (1 - CAL_FRAC))
 56    train_ds = dict(ds)
 57    train_ds["xtr"] = ds["xtr"][:len(ds["xtr"]) - ncal]
 58    train_ds["ytr"] = ds["ytr"][:len(ds["ytr"]) - ncal]
 59    net = make_model("rnn_small", ds["input_shape"], K)
 60    net, _, _ = train_model(net, train_ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
 61    if net is None: return float("nan"), {}
 62    device = next(net.parameters()).device
 63    net.eval()
 64    with torch.no_grad():
 65        pc = net(ds["xtr"][-ncal:].to(device)).detach().cpu().numpy()
 66        pt = net(ds["xte"].to(device)).detach().cpu().numpy()
 67    yc = ds["ytr"][-ncal:].numpy().reshape(-1)
 68    yt = ds["yte"].numpy().reshape(-1)
 69    # Clearance is signed distance to the pendulum angle limit |theta|=1.5.
 70    # It is computed on the selected rollout output, not analytically.
 71    mc = 1.5 - np.abs(pc)
 72    mt = 1.5 - np.abs(pt)
 73    clearance_cvar = np.array([cvar_lower(row, GAMMA) for row in mc])
 74    clearance_test_cvar = np.array([cvar_lower(row, GAMMA) for row in mt])
 75    q = conformal_q(clearance_cvar - (1.5 - np.abs(yc)), ALPHA)
 76    cert = clearance_test_cvar - q
 77    # Preserve predicted direction while enforcing the lower clearance bound.
 78    sign = np.sign(pt.mean(1)); sign[sign == 0] = 1
 79    pred = sign * (1.5 - cert)
 80    pred = np.clip(pred, -3., 3.)
 81    if not idea:
 82        pred = pt.mean(1)
 83    mse = float(np.mean((pred - yt) ** 2))
 84    # Signature is measured from this trained model's predictions and observations.
 85    realized_clear = 1.5 - np.abs(yt)
 86    coverage = float(np.mean(realized_clear >= cert)) if idea else float(np.mean(realized_clear >= clearance_test_cvar))
 87    return mse, {"q": q, "coverage": coverage, "certificate_nonnegative": float(np.mean(cert >= 0)),
 88                 "predicted_clearance_mean": float(np.mean(clearance_test_cvar)),
 89                 "observed_clearance_mean": float(np.mean(realized_clear))}
 90
 91
 92def baseline_factory(cfg):
 93    return lambda seed: fit(seed, cfg["lr"], False)[0]
 94
 95
 96def idea_factory(cfg):
 97    return lambda seed: fit(seed, cfg["lr"], True)[0]
 98
 99
100def main():
101    # Core math sanity check comes before any neural training.
102    rng = np.random.default_rng(2747)
103    errs = [abs(cvar_eta((z := rng.normal(size=K)), GAMMA) - cvar_lower(z, GAMMA)) for _ in range(20)]
104    rr = np.sort(rng.normal(size=31)); q = conformal_q(rr, ALPHA)
105    math_check = {"cvar_max_abs_error": float(max(errs)), "q_is_order_statistic": bool(q in rr),
106                  "q_rank": int(np.where(rr == q)[0][0] + 1)}
107
108    base = sweep_baseline(baseline_factory, GRID, seeds=(0, 1, 2, 3))
109    idea_cfg = base["best_cfg"]
110    # Full idea sweep at exactly the same three learning rates; best is selected
111    # on the four sweep seeds, then rerun on all eight paired seeds.
112    idea_sweep = []
113    for cfg in GRID:
114        r = evaluate(idea_factory(cfg), seeds=(0, 1, 2, 3))
115        idea_sweep.append({"cfg": cfg, "mean": r["mean"]})
116    idea_cfg = min(idea_sweep, key=lambda x: x["mean"])["cfg"]
117    idea_full = evaluate(idea_factory(idea_cfg), seeds=SEEDS)
118    # Mechanism signature is recomputed using trained models on all paired seeds.
119    sig = [fit(s, idea_cfg["lr"], True)[1] for s in SEEDS]
120    signature = {"definition": "trained RNN predicted-vs-observed clearance coverage",
121                 "alpha": ALPHA, "gamma": GAMMA, "K": K,
122                 "coverage_mean": float(np.mean([x["coverage"] for x in sig])),
123                 "coverage_per_seed": [x["coverage"] for x in sig],
124                 "q_mean": float(np.mean([x["q"] for x in sig])),
125                 "certificate_nonnegative_mean": float(np.mean([x["certificate_nonnegative"] for x in sig])),
126                 "target_coverage": 1 - ALPHA,
127                 "confirmed": bool(abs(float(np.mean([x["coverage"] for x in sig])) - (1-ALPHA)) <= 0.06),
128                 "math_check": math_check,
129                 "idea_sweep": idea_sweep}
130    report = make_report("dynamics", "rnn_small", {**base, "sweep_union": GRID}, idea_full, signature)
131    report["idea"]["selected_cfg"] = idea_cfg
132    report["math_check"] = math_check
133    Path("bench_report.json").write_text(json.dumps(report, indent=2))
134    print(json.dumps(report, indent=2))
135
136if __name__ == "__main__": main()