1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5# Reproducible MVP for PAC transition-cover monitoring.
  6SEED = 2077
  7
  8class PACCellMonitor:
  9    def __init__(self, n_cells, epsilon=0.1, delta=0.1):
 10        self.n_cells = n_cells
 11        self.epsilon = epsilon
 12        self.delta = delta
 13        self.required = int(math.ceil(math.log(1.0 / delta) / epsilon))
 14        self.counts = np.zeros(n_cells, dtype=int)
 15        self.successes = np.zeros(n_cells, dtype=int)
 16
 17    def observe(self, cell, hit):
 18        self.counts[cell] += 1
 19        self.successes[cell] += int(hit)
 20
 21    def deficient_cells(self):
 22        # A cell is not PAC-ready until it has the theorem's sample count.
 23        return np.flatnonzero(self.counts < self.required).tolist()
 24
 25    def report(self):
 26        return {"required_n": self.required,
 27                "counts": self.counts.tolist(),
 28                "deficient": self.deficient_cells()}
 29
 30
 31def math_sanity(rng):
 32    # Prediction 1: missed-region probability follows (1-eps)^n and its
 33    # exponential upper bound. Prediction 2: n >= log(1/delta)/eps crosses
 34    # below delta. Each row uses independent Bernoulli region hits.
 35    rows = []
 36    reps = 40000
 37    for eps in (0.02, 0.05, 0.10, 0.20):
 38        for n in (10, 25, 50, 100, 200):
 39            hits = rng.random((reps, n)) < eps
 40            observed = float(np.mean(~np.any(hits, axis=1)))
 41            exact = (1.0 - eps) ** n
 42            bound = math.exp(-n * eps)
 43            rows.append({"epsilon": eps, "n": n, "observed_miss": observed,
 44                         "exact": exact, "exp_bound": bound,
 45                         "abs_error_exact": abs(observed-exact)})
 46    threshold_rows = []
 47    for eps in (0.02, 0.05, 0.10, 0.20):
 48        delta = 0.1
 49        n_req = math.ceil(math.log(1/delta) / eps)
 50        hits = rng.random((reps, n_req)) < eps
 51        observed = float(np.mean(~np.any(hits, axis=1)))
 52        threshold_rows.append({"epsilon": eps, "delta": delta, "n_required": n_req,
 53                               "observed_miss": observed, "target_delta": delta,
 54                               "exact_miss": (1-eps)**n_req})
 55    max_err = max(x["abs_error_exact"] for x in rows)
 56    threshold_pass = all(x["observed_miss"] <= x["target_delta"] + 0.01
 57                         for x in threshold_rows)
 58    return {"scaling_rows": rows, "threshold_rows": threshold_rows,
 59            "max_abs_error_vs_exact": max_err, "threshold_pass": threshold_pass}
 60
 61
 62def dynamics(x, rng):
 63    # A nonlinear noisy transition with a rare, dynamically important branch.
 64    # The branch is concentrated in cell 3 and is observable from x.
 65    cell = np.clip(((x + 1.0) / 2.0 * 8).astype(int), 0, 7)
 66    rare = (cell == 3) & (rng.random(len(x)) < 0.12)
 67    y = 0.72*x + 0.22*np.sin(3*x) + rng.normal(0, 0.025, len(x))
 68    y = y + rare * (0.85 + 0.08*np.sin(5*x))
 69    return np.clip(y, -1.5, 1.5), cell, rare
 70
 71
 72def collect_uniform(rng, budget):
 73    x = rng.uniform(-1, 1, budget)
 74    y, cells, rare = dynamics(x, rng)
 75    return x, y, cells, rare
 76
 77
 78def collect_pac(rng, budget, n_cells=8, epsilon=0.1, delta=0.1):
 79    mon = PACCellMonitor(n_cells, epsilon, delta)
 80    xs, ys, cs, rs = [], [], [], []
 81    # Deficit-first allocation: this is the proposed extra-rollout policy.
 82    for t in range(budget):
 83        deficient = mon.deficient_cells()
 84        if deficient:
 85            c = deficient[t % len(deficient)]
 86        else:
 87            c = int(rng.integers(n_cells))
 88        lo, hi = -1 + 2*c/n_cells, -1 + 2*(c+1)/n_cells
 89        x = np.array([rng.uniform(lo, hi)])
 90        y, actual_cell, rare = dynamics(x, rng)
 91        # For this toy, a hit means observing the rare successor component;
 92        # ordinary transitions are also retained in the training buffer.
 93        mon.observe(c, bool(rare[0]))
 94        xs.append(x[0]); ys.append(y[0]); cs.append(c); rs.append(bool(rare[0]))
 95    return np.array(xs), np.array(ys), np.array(cs), np.array(rs), mon.report()
 96
 97
 98def fit_and_score(x, y, rng, steps=350):
 99    # Small torch MLP; CPU fallback is automatic and CUDA is optional.
100    try:
101        import torch
102        import torch.nn as nn
103        device = "cuda" if torch.cuda.is_available() else "cpu"
104        try:
105            torch.manual_seed(SEED)
106            model = nn.Sequential(nn.Linear(1, 24), nn.Tanh(), nn.Linear(24, 24),
107                                  nn.Tanh(), nn.Linear(24, 1)).to(device)
108            opt = torch.optim.Adam(model.parameters(), lr=0.008)
109            X = torch.tensor(x[:,None], dtype=torch.float32, device=device)
110            Y = torch.tensor(y[:,None], dtype=torch.float32, device=device)
111            for _ in range(steps):
112                idx = torch.randint(0, len(x), (min(128, len(x)),), device=device)
113                loss = ((model(X[idx])-Y[idx])**2).mean()
114                opt.zero_grad(); loss.backward(); opt.step()
115            # Balanced test plus explicit rare-branch probe.
116            tx = np.linspace(-1, 1, 1600)
117            ty, tc, tr = dynamics(tx, np.random.default_rng(SEED+99))
118            with torch.no_grad(): pred = model(torch.tensor(tx[:,None], dtype=torch.float32, device=device)).cpu().numpy()[:,0]
119            mse = float(np.mean((pred-ty)**2))
120            rare_mask = (tc == 3) & tr
121            common_mask = (tc == 3) & ~tr
122            rare_mse = float(np.mean((pred[rare_mask]-ty[rare_mask])**2)) if rare_mask.any() else float('nan')
123            cell_mse = float(np.mean((pred[tc==3]-ty[tc==3])**2))
124            return {"test_mse": mse, "cell3_mse": cell_mse, "rare_branch_mse": rare_mse,
125                    "device": device}
126        except Exception:
127            pass
128    except Exception:
129        pass
130    # Deterministic least-squares fallback if torch/CUDA fails.
131    z = np.stack([x**k for k in range(6)], axis=1)
132    coef = np.linalg.lstsq(z, y, rcond=None)[0]
133    tx = np.linspace(-1,1,1600); ty,tc,tr = dynamics(tx, np.random.default_rng(SEED+99))
134    pred = np.stack([tx**k for k in range(6)], axis=1) @ coef
135    return {"test_mse": float(np.mean((pred-ty)**2)), "cell3_mse": float(np.mean((pred[tc==3]-ty[tc==3])**2)),
136            "rare_branch_mse": float(np.mean((pred[(tc==3)&tr]-ty[(tc==3)&tr])**2)), "device":"cpu-fallback"}
137
138
139def mini_experiment(rng):
140    budget = 640
141    ux, uy, uc, ur = collect_uniform(rng, budget)
142    px, py, pc, pr, report = collect_pac(rng, budget)
143    baseline = fit_and_score(ux, uy, rng)
144    idea = fit_and_score(px, py, rng)
145    return {"budget": budget, "baseline": baseline, "idea": idea,
146            "training_cell_counts": {"uniform": np.bincount(uc, minlength=8).tolist(),
147                                     "pac": np.bincount(pc, minlength=8).tolist()},
148            "pac_monitor": report,
149            "rare_samples": {"uniform": int(ur.sum()), "pac": int(pr.sum())}}
150
151
152def main():
153    random.seed(SEED); np.random.seed(SEED)
154    rng = np.random.default_rng(SEED)
155    out = {"seed": SEED, "math": math_sanity(rng), "mini_experiment": mini_experiment(rng)}
156    Path("results.json").write_text(json.dumps(out, indent=2))
157    print(json.dumps(out, indent=2))
158
159if __name__ == "__main__":
160    main()