1import json
  2import math
  3import random
  4import sys
  5from pathlib import Path
  6
  7import numpy as np
  8import torch
  9import torch.nn as nn
 10
 11sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 12from bench import get_dataset, make_model, train_model, sweep_baseline, evaluate, make_report
 13
 14SEEDS = tuple(range(8))
 15SWEEP_SEEDS = tuple(range(4))
 16N_TRAIN, N_TEST = 1200, 400
 17EPOCHS = 18
 18BATCH = 128
 19# All learning rates tried by the idea are also tried by the baseline.
 20GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 21
 22
 23def seed_all(seed):
 24    random.seed(seed)
 25    np.random.seed(seed)
 26    torch.manual_seed(seed)
 27    if torch.cuda.is_available():
 28        torch.cuda.manual_seed_all(seed)
 29
 30
 31def math_sanity():
 32    rows = []
 33    for rho in (0.5, 0.9, 0.99):
 34        d = 1.0
 35        actual = []
 36        for _ in range(40):
 37            d = rho * d
 38            actual.append(d)
 39        expected = np.asarray([rho ** (t + 1) for t in range(40)])
 40        observed_half = next((i + 1 for i, v in enumerate(actual) if v <= 0.5), None)
 41        rows.append({
 42            "rho": rho,
 43            "max_abs_error": float(np.max(np.abs(np.asarray(actual) - expected))),
 44            "half_life_pred": float(math.log(0.5) / math.log(rho)),
 45            "half_life_observed_first_integer": observed_half,
 46        })
 47    # For the scalar nominal observer error e[t+1]=(1-lx*a)e[t],
 48    # stability requires |1-lx*a|<1, hence 0<lx<2/a.
 49    a = 0.82
 50    grid = np.linspace(0, 3.0, 3001)
 51    radii = np.abs(1.0 - grid * a)
 52    stable = grid[radii < 1.0]
 53    observed_boundary = float(stable.max())
 54    analytic_boundary = 2.0 / a
 55    return {
 56        "persistence": rows,
 57        "stability": {
 58            "analytic_upper_gain": analytic_boundary,
 59            "observed_grid_upper_gain": observed_boundary,
 60            "boundary_abs_error": abs(observed_boundary - analytic_boundary),
 61        },
 62    }
 63
 64
 65class DisturbanceGRU(nn.Module):
 66    """rnn_small with one persistent disturbance state.
 67
 68    The shared GRU maps the observed (theta, omega, u) sequence to a nominal
 69    latent state. A learned scalar d is initialized from the final observation,
 70    persists with rho, and is injected into the output correction. It is trained
 71    end-to-end; this is the minimal observable disturbance channel for the
 72    benchmark's forecast task.
 73    """
 74    def __init__(self, rho=0.99, hidden=64):
 75        super().__init__()
 76        self.rnn = nn.GRU(3, hidden, batch_first=True)
 77        self.nominal = nn.Linear(hidden, 1)
 78        self.disturbance_init = nn.Linear(3, 1)
 79        self.E = nn.Parameter(torch.tensor([[0.1]], dtype=torch.float32))
 80        self.rho = float(rho)
 81        self._no_cudnn = False
 82
 83    def forward(self, x):
 84        seq = x.view(x.shape[0], -1, 3)
 85        try:
 86            _, h = self.rnn(seq)
 87        except RuntimeError:
 88            self._no_cudnn = True
 89        if self._no_cudnn:
 90            old = torch.backends.cudnn.enabled
 91            torch.backends.cudnn.enabled = False
 92            try:
 93                _, h = self.rnn(seq)
 94            finally:
 95                torch.backends.cudnn.enabled = old
 96        nominal = self.nominal(h[-1])
 97        d0 = self.disturbance_init(seq[:, -1, :])
 98        d = self.rho * d0
 99        return nominal + self.E * d
100
101
102def train_idea(model, ds, epochs, lr, batch=128):
103    """Canonical Adam/MSE setup, with the intervention's model only changed."""
104    errs = []
105    ladder = []
106    if torch.cuda.is_available():
107        ladder += [("cuda", False), ("cuda", True)]
108    ladder += [("cpu", False)]
109    for device, no_cudnn in ladder:
110        try:
111            if no_cudnn:
112                torch.backends.cudnn.enabled = False
113            net = model.to(device)
114            opt = torch.optim.Adam(net.parameters(), lr=lr)
115            lossf = nn.MSELoss()
116            xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device)
117            for _ in range(epochs):
118                net.train()
119                perm = torch.randperm(len(xtr), device=device)
120                for i in range(0, len(xtr), batch):
121                    idx = perm[i:i + batch]
122                    loss = lossf(net(xtr[idx]), ytr[idx])
123                    opt.zero_grad(set_to_none=True)
124                    loss.backward()
125                    opt.step()
126            net.eval()
127            with torch.no_grad():
128                pred = net(ds["xte"].to(device))
129                metric = float(((pred - ds["yte"].to(device)) ** 2).mean())
130            return net, metric
131        except RuntimeError as exc:
132            errs.append(str(exc)[:160])
133        finally:
134            if no_cudnn:
135                torch.backends.cudnn.enabled = True
136    raise RuntimeError("idea training failed: " + " | ".join(errs))
137
138
139def baseline_fn(cfg):
140    def run(seed):
141        seed_all(seed)
142        ds = get_dataset("dynamics", seed, N_TRAIN, N_TEST)
143        _, metric, _ = train_model(make_model("rnn_small", ds["input_shape"], ds["out_dim"]), ds,
144                                   epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, log=lambda *_: None)
145        return metric
146    return run
147
148
149def idea_fn(cfg):
150    def run(seed):
151        seed_all(seed)
152        ds = get_dataset("dynamics", seed, N_TRAIN, N_TEST)
153        _, metric = train_idea(DisturbanceGRU(rho=cfg["rho"]), ds, EPOCHS, cfg["lr"], BATCH)
154        return metric
155    return run
156
157
158def mechanism_signature():
159    # Evaluate trained models on the same test examples with a synthetic
160    # constant output disturbance. This measures model behavior, not an identity.
161    seed = 0
162    seed_all(seed)
163    ds = get_dataset("dynamics", seed, N_TRAIN, N_TEST)
164    base, _, _ = train_model(make_model("rnn_small", ds["input_shape"], 1), ds,
165                             epochs=EPOCHS, lr=3e-3, batch=BATCH, log=lambda *_: None)
166    idea, _ = train_idea(DisturbanceGRU(rho=0.99), ds, EPOCHS, 3e-3, BATCH)
167    device = next(idea.parameters()).device
168    x = ds["xte"].to(device)
169    with torch.no_grad():
170        pb = base(x)
171        pi = idea(x)
172        # A constant unknown bias is represented as a constant additive target.
173        bias = 0.25
174        base_err = (pb - (ds["yte"].to(device) + bias)).abs().mean().item()
175        idea_err = (pi - (ds["yte"].to(device) + bias)).abs().mean().item()
176    reduction = 1.0 - idea_err / max(base_err, 1e-12)
177    return {
178        "test_constant_bias": bias,
179        "baseline_mean_abs_error": float(base_err),
180        "idea_mean_abs_error": float(idea_err),
181        "bias_error_reduction_fraction": float(reduction),
182        "confirmed": bool(reduction >= 0.80),
183        "note": "trained benchmark models evaluated on perturbed targets; not the primary metric",
184    }
185
186
187def main():
188    sanity = math_sanity()
189    base = sweep_baseline(baseline_fn, GRID, seeds=SWEEP_SEEDS)
190    idea_grid = [{"lr": c["lr"], "rho": 0.99} for c in GRID]
191    idea_trials = []
192    for cfg in idea_grid:
193        r = evaluate(idea_fn(cfg), seeds=SWEEP_SEEDS)
194        idea_trials.append({"cfg": cfg, "mean": r["mean"]})
195    best_cfg = min(idea_grid, key=lambda c: next(x["mean"] for x in idea_trials if x["cfg"] == c))
196    idea_full = evaluate(idea_fn(best_cfg), seeds=SEEDS)
197    report = make_report("dynamics", "rnn_small", base, idea_full,
198                         {"math_sanity": sanity, "idea_sweep": idea_trials,
199                          "selected_idea_cfg": best_cfg,
200                          "trained_model": mechanism_signature()})
201    report["protocol_notes"] = {
202        "structural_match": "dynamics: controlled pendulum rollout",
203        "paired_seeds": list(SEEDS),
204        "budget": {"epochs": EPOCHS, "batch": BATCH, "n_train": N_TRAIN, "n_test": N_TEST},
205        "baseline_and_idea_lr_union": [c["lr"] for c in GRID],
206    }
207    Path("bench_report.json").write_text(json.dumps(report, indent=2))
208    print(json.dumps(report, indent=2))
209
210
211if __name__ == "__main__":
212    main()