1import json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7import sys
  8sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  9from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
 10
 11SEEDS = tuple(range(8))
 12# The same union is evaluated by both systems; baseline tuning uses the official 4 seeds.
 13GRID = [
 14    {"lr": 1e-3, "epochs": 10},
 15    {"lr": 3e-3, "epochs": 10},
 16    {"lr": 6e-3, "epochs": 10},
 17]
 18
 19class CompositionalRNN(nn.Module):
 20    """Shared recurrent primitive critic with a scalar aggregator."""
 21    def __init__(self, monotone=False, hidden=64):
 22        super().__init__()
 23        self.monotone = monotone
 24        self.rnn = nn.GRU(3, hidden, batch_first=True)
 25        self.primitive = nn.Linear(hidden, 2)
 26        self.raw_w = nn.Parameter(torch.zeros(2))
 27        self.bias = nn.Parameter(torch.zeros(1))
 28        # Both variants start at the same effective positive weights.
 29        nn.init.constant_(self.raw_w, float(np.log(np.expm1(0.5))))
 30
 31    def forward(self, x):
 32        seq = x.view(x.shape[0], -1, 3)
 33        _, h = self.rnn(seq)
 34        z = self.primitive(h[-1])
 35        if self.monotone:
 36            w = torch.nn.functional.softplus(self.raw_w)
 37        else:
 38            w = self.raw_w
 39        return (self.bias + z @ w).view(-1, 1)
 40
 41
 42def seed_all(seed):
 43    np.random.seed(seed); random.seed(seed); torch.manual_seed(seed)
 44    if torch.cuda.is_available():
 45        torch.cuda.manual_seed_all(seed)
 46    torch.set_num_threads(4)
 47
 48
 49def run_one(seed, cfg, monotone, return_model=False):
 50    seed_all(seed)
 51    # Reduced but fixed dataset size keeps the complete paired protocol quick.
 52    ds = get_dataset("dynamics", int(seed), n_train=400, n_test=200)
 53    model = CompositionalRNN(monotone=monotone)
 54    net, metric, _ = train_model(model, ds, epochs=int(cfg["epochs"]),
 55                                 lr=float(cfg["lr"]), batch=128, log=lambda *_: None)
 56    if metric is None:
 57        return (float("nan"), None, ds)
 58    return (float(metric), net, ds) if return_model else float(metric)
 59
 60
 61def main():
 62    # Official sweep helper is used as the default baseline training path.
 63    base = sweep_baseline(
 64        lambda cfg: (lambda seed: run_one(seed, cfg, False)), GRID)
 65
 66    # Required idea sweep: baseline-best and two nearby settings; all are in baseline union.
 67    idea_trials = []
 68    for cfg in GRID:
 69        r = evaluate(lambda seed, c=cfg: run_one(seed, c, True), seeds=SEEDS)
 70        idea_trials.append({"cfg": cfg, "mean": r["mean"], "full": r})
 71    best_trial = min(idea_trials, key=lambda q: q["mean"])
 72    idea = best_trial["full"]
 73
 74    # Refit one paired seed at each selected best configuration for behavior-based signature.
 75    bval, bnet, bds = run_one(0, base["best_cfg"], False, True)
 76    ival, inet, ids = run_one(0, best_trial["cfg"], True, True)
 77    def behavior(net, ds):
 78        # Signature evaluation is deliberately CPU-safe after GPU training.
 79        net = net.cpu()
 80        net.eval()
 81        device = torch.device("cpu")
 82        x = ds["xte"][:128].to(device).clone().requires_grad_(True)
 83        seq = x.view(x.shape[0], -1, 3)
 84        _, h = net.rnn(seq)
 85        z = net.primitive(h[-1])
 86        w = torch.nn.functional.softplus(net.raw_w) if net.monotone else net.raw_w
 87        out = (net.bias + z @ w).view(-1, 1)
 88        # Derivatives are measured on outputs of the trained model, not synthetic weights.
 89        dg = torch.autograd.grad(out.sum(), z, retain_graph=True)[0]
 90        # Finite primitive perturbation: change each trained primitive output and reapply
 91        # the trained aggregator, measuring the observed output response.
 92        zp = z.detach().clone(); zp[:, 0] += 0.1
 93        response = ((net.bias.detach() + zp @ w.detach()) - out.detach().view(-1)).mean().item()
 94        return {
 95            "samples": int(len(x)),
 96            "derivative_nonnegative_fraction": float((dg >= -1e-7).all(1).float().mean()),
 97            "mean_derivatives": dg.mean(0).tolist(),
 98            "positive_primitive_perturbation_mean_response": float(response),
 99        }
100    sig_b = behavior(bnet, bds); sig_i = behavior(inet, ids)
101    signature = {
102        "prediction": "monotone aggregator has nonnegative derivatives; increasing a primitive cannot lower composite value",
103        "baseline_observed": sig_b,
104        "idea_observed": sig_i,
105        "confirmed": bool(sig_i["derivative_nonnegative_fraction"] == 1.0 and
106                            sig_i["positive_primitive_perturbation_mean_response"] >= -1e-7 and
107                            sig_b["derivative_nonnegative_fraction"] < 1.0),
108        "measurement_note": "All quantities were computed from trained benchmark models on held-out dynamics inputs."
109    }
110    report = make_report("dynamics", "rnn_small", base, idea, signature)
111    report["idea"]["sweep"] = [{"cfg": q["cfg"], "mean": q["mean"]} for q in idea_trials]
112    report["custom_track"] = None
113    Path("bench_report.json").write_text(json.dumps(report, indent=2))
114    print(json.dumps(report, indent=2))
115
116if __name__ == "__main__":
117    main()