1import json
  2import random
  3import sys
  4from pathlib import Path
  5
  6import numpy as np
  7import torch
  8from torch import nn
  9
 10sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 11from bench import get_dataset
 12from bench.protocol import evaluate, make_report, sweep_baseline
 13
 14SEEDS = tuple(range(8))
 15SWEEP_SEEDS = tuple(range(4))
 16EPOCHS = 12
 17BATCH = 128
 18
 19
 20def seed_all(seed):
 21    random.seed(seed)
 22    np.random.seed(seed)
 23    torch.manual_seed(seed)
 24    if torch.cuda.is_available():
 25        torch.cuda.manual_seed_all(seed)
 26
 27
 28class MatchedRNN(nn.Module):
 29    """The bench rnn_small architecture, with optional trajectory exposure."""
 30    def __init__(self, hidden=64):
 31        super().__init__()
 32        self.rnn = nn.GRU(3, hidden, batch_first=True)
 33        self.head = nn.Linear(hidden, 1)
 34
 35    def forward(self, x, return_state=False):
 36        seq = x.view(x.shape[0], -1, 3)
 37        hidden, _ = self.rnn(seq)
 38        yseq = self.head(hidden).squeeze(-1)
 39        pred = self.head(hidden[:, -1])
 40        if return_state:
 41            return pred, yseq, hidden
 42        return pred
 43
 44
 45def normalized_dataset(ds):
 46    d = dict(ds)
 47    mu = ds["xtr"].mean(0, keepdim=True)
 48    sd = ds["xtr"].std(0, keepdim=True).clamp_min(1e-4)
 49    for k in ("xtr", "xte"):
 50        d[k] = (ds[k] - mu) / sd
 51    ym = ds["ytr"].mean()
 52    ys = ds["ytr"].std().clamp_min(1e-4)
 53    d["ytr"] = (ds["ytr"] - ym) / ys
 54    d["yte"] = (ds["yte"] - ym) / ys
 55    return d
 56
 57
 58def passivity_loss(x, out, h, gamma=0.1, paired=None, eta=0.1):
 59    # h0=0, and output energy is assigned to each recurrent step.
 60    h0 = torch.zeros(x.shape[0], 1, h.shape[-1], device=x.device, dtype=h.dtype)
 61    hall = torch.cat((h0, h), dim=1)
 62    u = x.view(x.shape[0], -1, 3)
 63    y = out
 64    r = (hall[:, 1:].square().sum(-1) + y.square()
 65         - hall[:, :-1].square().sum(-1) - u.square().sum(-1))
 66    loss = torch.relu(r).mean() + gamma * torch.relu(r.sum(1)).mean()
 67    stats = {
 68        "max_positive_residual": float(torch.relu(r).max().detach().cpu()),
 69        "mean_cumulative_residual": float(r.sum(1).mean().detach().cpu()),
 70        "mean_hidden_norm": float(h.norm(dim=-1).mean().detach().cpu()),
 71    }
 72    if paired is not None:
 73        x2, out2, h2 = paired
 74        dh0 = torch.zeros(x.shape[0], 1, h.shape[-1], device=x.device, dtype=h.dtype)
 75        dh = torch.cat((dh0, h2 - h), dim=1)
 76        du = (x2 - x).view(x.shape[0], -1, 3)
 77        dy = out2 - out
 78        ri = (dh[:, 1:].square().sum(-1) + dy.square()
 79              - dh[:, :-1].square().sum(-1) - du.square().sum(-1))
 80        inc = torch.relu(ri).mean()
 81        loss = loss + eta * inc
 82        stats["incremental_penalty"] = float(inc.detach().cpu())
 83    return loss, stats
 84
 85
 86def run_one(seed, lr, lam, paired, return_stats=False):
 87    seed_all(seed)
 88    device = "cuda" if torch.cuda.is_available() else "cpu"
 89    try:
 90        ds = normalized_dataset(get_dataset("dynamics", seed, n_train=400, n_test=200))
 91        model = MatchedRNN().to(device)
 92        opt = torch.optim.Adam(model.parameters(), lr=lr)
 93        xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device)
 94        for _ in range(EPOCHS):
 95            model.train()
 96            perm = torch.randperm(len(xtr), device=device)
 97            for start in range(0, len(xtr), BATCH):
 98                ix = perm[start:start+BATCH]
 99                x = xtr[ix]; target = ytr[ix]
100                pred, out, h = model(x, True)
101                pair = None
102                if paired:
103                    x2 = x + 0.02 * torch.randn_like(x)
104                    pred2, out2, h2 = model(x2, True)
105                    pair = (x2, out2, h2)
106                task = nn.functional.mse_loss(pred, target)
107                pl, _ = passivity_loss(x, out, h, paired=pair)
108                loss = task + lam * pl
109                opt.zero_grad(); loss.backward()
110                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
111                opt.step()
112        model.eval()
113        with torch.no_grad():
114            pred, out, h = model(ds["xte"].to(device), True)
115            metric = float(nn.functional.mse_loss(pred, ds["yte"].to(device)).cpu())
116            _, st = passivity_loss(ds["xte"].to(device), out, h)
117            x = ds["xte"].to(device)
118            x2 = x + 0.02 * torch.randn_like(x)
119            _, _, h2 = model(x2, True)
120            amp = ((h2[:, -1] - h[:, -1]).norm(dim=1) /
121                   (x2 - x).view(x.shape[0], -1).norm(dim=1).clamp_min(1e-8)).mean()
122            st["amplification"] = float(amp.cpu())
123        return (metric, st) if return_stats else metric
124    except RuntimeError:
125        if device == "cuda":
126            torch.cuda.empty_cache()
127            old = torch.cuda.is_available
128            # Explicit CPU retry without changing the benchmark architecture.
129            torch.set_default_device("cpu")
130            try:
131                return run_one_cpu(seed, lr, lam, paired, return_stats)
132            finally:
133                torch.set_default_device("cpu")
134        raise
135
136
137def run_one_cpu(seed, lr, lam, paired, return_stats=False):
138    seed_all(seed)
139    ds = normalized_dataset(get_dataset("dynamics", seed, n_train=400, n_test=200))
140    model = MatchedRNN()
141    opt = torch.optim.Adam(model.parameters(), lr=lr)
142    xtr, ytr = ds["xtr"], ds["ytr"]
143    for _ in range(EPOCHS):
144        perm = torch.randperm(len(xtr))
145        for start in range(0, len(xtr), BATCH):
146            x = xtr[perm[start:start+BATCH]]; target = ytr[perm[start:start+BATCH]]
147            pred, out, h = model(x, True); pair = None
148            if paired:
149                x2 = x + 0.02 * torch.randn_like(x)
150                _, out2, h2 = model(x2, True); pair = (x2, out2, h2)
151            task = nn.functional.mse_loss(pred, target)
152            pl, _ = passivity_loss(x, out, h, paired=pair)
153            opt.zero_grad(); (task + lam * pl).backward()
154            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
155    model.eval()
156    with torch.no_grad():
157        x = ds["xte"]; pred, out, h = model(x, True)
158        metric = float(nn.functional.mse_loss(pred, ds["yte"]))
159        _, st = passivity_loss(x, out, h)
160        x2 = x + 0.02 * torch.randn_like(x); _, _, h2 = model(x2, True)
161        st["amplification"] = float(((h2[:, -1]-h[:, -1]).norm(dim=1) /
162            (x2-x).view(x.shape[0],-1).norm(dim=1).clamp_min(1e-8)).mean())
163    return (metric, st) if return_stats else metric
164
165
166def main():
167    # Search-space parity: every idea learning rate is also swept by baseline.
168    lrs = (1e-3, 3e-3, 1e-2)
169    grid = [{"lr": lr} for lr in lrs]
170    base = sweep_baseline(
171        lambda cfg: lambda seed: run_one(seed, cfg["lr"], 0.0, False),
172        grid, seeds=SWEEP_SEEDS)
173    best_lr = base["best_cfg"]["lr"]
174    # Three regularizer settings, with and without the optional paired term.
175    idea_grid = [{"lr": lr, "lam": lam, "paired": paired}
176                 for lr in lrs for lam in (1e-3, 1e-2, 1e-1)
177                 for paired in (False, True)]
178    scored = []
179    for cfg in idea_grid:
180        vals = [run_one(s, cfg["lr"], cfg["lam"], cfg["paired"])
181                for s in SWEEP_SEEDS]
182        scored.append((float(np.mean(vals)), cfg))
183    idea_cfg = min(scored, key=lambda z: z[0])[1]
184    idea = evaluate(lambda s: run_one(s, idea_cfg["lr"], idea_cfg["lam"], idea_cfg["paired"]),
185                    seeds=SEEDS)
186
187    # Signature is measured on separately trained baseline and idea systems.
188    base_stats = [run_one(s, best_lr, 0.0, False, True)[1] for s in SEEDS]
189    idea_stats = [run_one(s, idea_cfg["lr"], idea_cfg["lam"], idea_cfg["paired"], True)[1]
190                  for s in SEEDS]
191    b_res = float(np.mean([z["mean_cumulative_residual"] for z in base_stats]))
192    i_res = float(np.mean([z["mean_cumulative_residual"] for z in idea_stats]))
193    b_amp = float(np.mean([z["amplification"] for z in base_stats]))
194    i_amp = float(np.mean([z["amplification"] for z in idea_stats]))
195    signature = {
196        "baseline_cumulative_residual": b_res,
197        "idea_cumulative_residual": i_res,
198        "baseline_amplification": b_amp,
199        "idea_amplification": i_amp,
200        "predicted_direction": "passivity reduces residual and amplification",
201        "confirmed": bool(i_res < b_res and i_amp < b_amp),
202    }
203    report = make_report(
204        "dynamics", "rnn_small", {**base, "selected_cfg": base["best_cfg"]}, idea,
205        {"selected_idea_cfg": idea_cfg, "mechanism_signature": signature})
206    Path("bench_report.json").write_text(json.dumps(report, indent=2))
207    print(json.dumps(report, indent=2))
208
209
210if __name__ == "__main__":
211    main()