1import json
  2import random
  3from pathlib import Path
  4import numpy as np
  5import torch
  6from torch import nn
  7import sys
  8sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  9from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
 10
 11EPOCHS = 10
 12BATCH = 128
 13WEIGHT_DECAY = 0.0
 14
 15
 16def seed_all(seed):
 17    random.seed(seed)
 18    np.random.seed(seed)
 19    torch.manual_seed(seed)
 20    if torch.cuda.is_available():
 21        torch.cuda.manual_seed_all(seed)
 22
 23
 24def sign_transform(x):
 25    # Pendulum equations are equivariant under (theta, omega, u) -> -(theta, omega, u).
 26    return -x
 27
 28
 29def output_transform(y):
 30    return -y
 31
 32
 33def equivariance_penalty(net, x):
 34    pred = net(x)
 35    pred_transformed = net(sign_transform(x))
 36    return ((pred_transformed - output_transform(pred)) ** 2).mean()
 37
 38
 39def train_idea(net, ds, epochs, lr, lam):
 40    # Same Adam/batch/epochs as bench.train_model; only the proposed loss is added.
 41    devices = [("cuda", False), ("cuda", True), ("cpu", False)] if torch.cuda.is_available() else [("cpu", False)]
 42    errors = []
 43    for devname, no_cudnn in devices:
 44        try:
 45            device = torch.device(devname)
 46            old_cudnn = torch.backends.cudnn.enabled
 47            if no_cudnn:
 48                torch.backends.cudnn.enabled = False
 49            net = net.to(device)
 50            opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=WEIGHT_DECAY)
 51            xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device)
 52            for _ in range(epochs):
 53                net.train()
 54                perm = torch.randperm(len(xtr), device=device)
 55                for i in range(0, len(xtr), BATCH):
 56                    idx = perm[i:i+BATCH]
 57                    pred = net(xtr[idx])
 58                    loss = ((pred - ytr[idx]) ** 2).mean() + lam * equivariance_penalty(net, xtr[idx])
 59                    opt.zero_grad(set_to_none=True)
 60                    loss.backward()
 61                    opt.step()
 62            net.eval()
 63            with torch.no_grad():
 64                metric = float(((net(ds["xte"].to(device)) - ds["yte"].to(device)) ** 2).mean().cpu())
 65            torch.backends.cudnn.enabled = old_cudnn
 66            return net, metric
 67        except RuntimeError as exc:
 68            errors.append(str(exc)[:160])
 69            if devname == "cuda":
 70                continue
 71            raise
 72    raise RuntimeError("training failed: " + " | ".join(errors))
 73
 74
 75def fit_baseline(cfg, seed):
 76    seed_all(seed)
 77    ds = get_dataset("dynamics", seed, n_train=400, n_test=400)
 78    net = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
 79    _, metric, _ = train_model(net, ds, epochs=cfg["epochs"], lr=cfg["lr"], batch=BATCH, weight_decay=WEIGHT_DECAY, log=lambda *_: None)
 80    return metric
 81
 82
 83def fit_idea(cfg, seed, capture=False):
 84    seed_all(seed)
 85    ds = get_dataset("dynamics", seed, n_train=400, n_test=400)
 86    net = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
 87    trained, metric = train_idea(net, ds, cfg["epochs"], cfg["lr"], cfg["lam"])
 88    if not capture:
 89        return metric
 90    device = next(trained.parameters()).device
 91    x = ds["xte"].to(device)
 92    with torch.no_grad():
 93        p = trained(x)
 94        pt = trained(-x)
 95        eq = float(((pt + p) ** 2).mean().cpu())
 96        # Behavioural signature: compare observed transformed prediction to the
 97        # transformed prediction expected from the trained model's original output.
 98        pred_scale = float(p.abs().mean().cpu())
 99    return metric, {"mean_prediction_abs": pred_scale, "trained_model_c2_penalty": eq}
100
101
102def math_check():
103    x = torch.randn(64, 24)
104    y = torch.randn(64, 1)
105    return {"c2_input_composition_max_abs": float((sign_transform(sign_transform(x)) - x).abs().max()),
106            "c2_output_composition_max_abs": float((output_transform(output_transform(y)) - y).abs().max()),
107            "exact_transform_identity": bool(torch.equal(sign_transform(sign_transform(x)), x))}
108
109
110def main():
111    # Union parity: every lr and lambda considered for the idea is also run by
112    # baseline; baseline lambda is the standard no-penalty value 0.
113    lrs = [1e-3, 3e-3, 1e-2]
114    base_grid = [{"lr": lr, "epochs": EPOCHS, "lam": 0.0} for lr in lrs]
115    idea_grid = [{"lr": lr, "epochs": EPOCHS, "lam": lam} for lr in lrs for lam in [0.05, 0.2, 1.0]]
116    def base_fn(cfg):
117        return lambda seed: fit_baseline(cfg, seed)
118    base = sweep_baseline(base_fn, base_grid)
119    # Evaluate all idea settings on sweep seeds for selection, then full 8 seeds.
120    idea_trials = []
121    for cfg in idea_grid:
122        r = evaluate(lambda seed, c=cfg: fit_idea(c, seed), seeds=(0, 1, 2, 3))
123        idea_trials.append({"cfg": cfg, "mean": r["mean"]})
124    best_cfg = min(idea_trials, key=lambda z: z["mean"])["cfg"]
125    idea = evaluate(lambda seed: fit_idea(best_cfg, seed), seeds=tuple(range(8)))
126    sig_metric, sig = fit_idea(best_cfg, 0, capture=True)
127    sig.update({"prediction_metric_seed0": sig_metric,
128                "expected_c2_penalty_direction": "lower is more equivariant",
129                "confirmed": sig["trained_model_c2_penalty"] < 1e-3})
130    report = make_report("dynamics", "rnn_small", {"best_cfg": base["best_cfg"], "sweep": base["sweep"], "full": base["full"]}, idea, {"mechanism_signature": sig, "math_check": math_check(), "idea_sweep": idea_trials, "structural_match": "controlled pendulum rollout has sign-equivariant local dynamics"})
131    Path("bench_report.json").write_text(json.dumps(report, indent=2))
132    print(json.dumps(report, indent=2))
133
134
135if __name__ == "__main__":
136    main()