1import json, sys, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11# Shared union: every idea learning rate is also evaluated for baseline.
 12LR_GRID = [0.0015, 0.003, 0.006]
 13EPOCHS = 15
 14NTRAIN, NTEST = 1200, 400
 15LAMBDA_D = 1e-3
 16TAU = 0.8
 17
 18
 19def seed_all(seed):
 20    random.seed(seed)
 21    np.random.seed(seed)
 22    torch.manual_seed(seed)
 23    if torch.cuda.is_available():
 24        torch.cuda.manual_seed_all(seed)
 25
 26
 27def baseline_fn(cfg):
 28    def run(seed):
 29        seed_all(seed)
 30        ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
 31        net = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
 32        _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=cfg["lr"], batch=128)
 33        return float(metric)
 34    return run
 35
 36
 37class ActuatorRNN(nn.Module):
 38    """Same rnn_small GRU/head, with d(u)=alpha*tanh(u/tau) before it."""
 39    def __init__(self, base):
 40        super().__init__()
 41        self.rnn = base.rnn
 42        self.head = base.head
 43        self.alpha = nn.Parameter(torch.zeros(1))
 44        self.tau = TAU
 45
 46    def disturbance(self, x):
 47        z = x.view(x.shape[0], -1, 3).clone()
 48        u = z[:, :, 2]
 49        d = self.alpha * torch.tanh(u / self.tau)
 50        z[:, :, 2] = u + d
 51        return z.reshape(x.shape)
 52
 53    def forward(self, x):
 54        seq = self.disturbance(x).view(x.shape[0], -1, 3)
 55        try:
 56            _, h = self.rnn(seq)
 57        except RuntimeError:
 58            old = torch.backends.cudnn.enabled
 59            torch.backends.cudnn.enabled = False
 60            try:
 61                _, h = self.rnn(seq)
 62            finally:
 63                torch.backends.cudnn.enabled = old
 64        return self.head(h[-1])
 65
 66
 67def train_idea(seed, cfg, return_model=False):
 68    seed_all(seed)
 69    ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
 70    base = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
 71    net = ActuatorRNN(base)
 72    # This is a custom loop because the intervention explicitly changes the loss.
 73    errors = []
 74    for device in (["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"]):
 75        try:
 76            net = net.to(device)
 77            xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device)
 78            opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"])
 79            for _ in range(EPOCHS):
 80                net.train()
 81                perm = torch.randperm(len(xtr), device=device)
 82                for i in range(0, len(xtr), 128):
 83                    ix = perm[i:i+128]
 84                    pred = net(xtr[ix])
 85                    d = net.disturbance(xtr[ix]).view(len(ix), -1, 3)[:, :, 2] - xtr[ix].view(len(ix), -1, 3)[:, :, 2]
 86                    loss = ((pred - ytr[ix]) ** 2).mean() + LAMBDA_D * (d ** 2).mean()
 87                    opt.zero_grad(); loss.backward(); opt.step()
 88            net.eval()
 89            with torch.no_grad():
 90                metric = float(((net(ds["xte"].to(device)) - ds["yte"].to(device)) ** 2).mean())
 91            return (metric, net, ds) if return_model else metric
 92        except RuntimeError as e:
 93            errors.append(str(e))
 94            if device == "cpu":
 95                raise
 96            # Recreate clean CPU tensors/model after any CUDA failure.
 97            base = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
 98            net = ActuatorRNN(base)
 99    raise RuntimeError("training failed: " + repr(errors))
100
101
102def idea_fn(cfg):
103    return lambda seed: train_idea(seed, cfg)
104
105
106def mechanism_signature():
107    # Measurements are from trained benchmark models, not an analytic toy graph.
108    vals = []
109    corrs = []
110    for seed in SEEDS:
111        metric, net, ds = train_idea(seed, {"lr": 0.003}, True)
112        with torch.no_grad():
113            dev = next(net.parameters()).device
114            x = ds["xte"].to(dev)
115            u = x.view(len(x), -1, 3)[:, :, 2]
116            d = net.alpha.detach().cpu().item() * torch.tanh(u / TAU)
117            vals.append(abs(net.alpha.detach().cpu().item()) / TAU)
118            # In a null-distortion track, residual/action correlation should be small.
119            pred = net(x).detach().cpu().numpy().reshape(-1)
120            residual = ds["yte"].numpy() - pred
121            signal = d.detach().cpu().numpy().mean(axis=1)
122            corrs.append(float(np.corrcoef(residual, signal)[0, 1]) if np.std(signal) > 1e-10 else 0.0)
123    observed = float(np.mean(vals))
124    return {"predicted_max_jacobian": 0.0, "observed_mean_abs_alpha_over_tau": observed,
125            "observed_residual_disturbance_corr_abs": float(np.mean(np.abs(corrs))),
126            "n_trained_models": 8, "confirmed": bool(observed < 0.05),
127            "interpretation": "The matched bench has no latent actuator error; the simplicity prior should identify zero disturbance."}
128
129
130def main():
131    base = sweep_baseline(baseline_fn, [{"lr": x} for x in LR_GRID])
132    # Same 3-point grid for the idea; best selected on the same four sweep seeds.
133    idea_trials = []
134    for cfg in [{"lr": x} for x in LR_GRID]:
135        r = evaluate(idea_fn(cfg), seeds=(0, 1, 2, 3))
136        idea_trials.append({"cfg": cfg, "mean": r["mean"]})
137    best_cfg = min(idea_trials, key=lambda z: z["mean"])["cfg"]
138    idea = evaluate(idea_fn(best_cfg), seeds=SEEDS)
139    report = make_report("dynamics", "rnn_small", base, idea,
140                         {"idea_lr_sweep": idea_trials, "chosen_idea_cfg": best_cfg,
141                          **mechanism_signature()})
142    Path("bench_report.json").write_text(json.dumps(report, indent=2))
143    print(json.dumps(report, indent=2))
144
145if __name__ == "__main__":
146    main()