import json import random import sys from pathlib import Path import numpy as np import torch from torch import nn sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import get_dataset from bench.protocol import evaluate, make_report, sweep_baseline SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS = 12 BATCH = 128 def seed_all(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) class MatchedRNN(nn.Module): """The bench rnn_small architecture, with optional trajectory exposure.""" def __init__(self, hidden=64): super().__init__() self.rnn = nn.GRU(3, hidden, batch_first=True) self.head = nn.Linear(hidden, 1) def forward(self, x, return_state=False): seq = x.view(x.shape[0], -1, 3) hidden, _ = self.rnn(seq) yseq = self.head(hidden).squeeze(-1) pred = self.head(hidden[:, -1]) if return_state: return pred, yseq, hidden return pred def normalized_dataset(ds): d = dict(ds) mu = ds["xtr"].mean(0, keepdim=True) sd = ds["xtr"].std(0, keepdim=True).clamp_min(1e-4) for k in ("xtr", "xte"): d[k] = (ds[k] - mu) / sd ym = ds["ytr"].mean() ys = ds["ytr"].std().clamp_min(1e-4) d["ytr"] = (ds["ytr"] - ym) / ys d["yte"] = (ds["yte"] - ym) / ys return d def passivity_loss(x, out, h, gamma=0.1, paired=None, eta=0.1): # h0=0, and output energy is assigned to each recurrent step. h0 = torch.zeros(x.shape[0], 1, h.shape[-1], device=x.device, dtype=h.dtype) hall = torch.cat((h0, h), dim=1) u = x.view(x.shape[0], -1, 3) y = out r = (hall[:, 1:].square().sum(-1) + y.square() - hall[:, :-1].square().sum(-1) - u.square().sum(-1)) loss = torch.relu(r).mean() + gamma * torch.relu(r.sum(1)).mean() stats = { "max_positive_residual": float(torch.relu(r).max().detach().cpu()), "mean_cumulative_residual": float(r.sum(1).mean().detach().cpu()), "mean_hidden_norm": float(h.norm(dim=-1).mean().detach().cpu()), } if paired is not None: x2, out2, h2 = paired dh0 = torch.zeros(x.shape[0], 1, h.shape[-1], device=x.device, dtype=h.dtype) dh = torch.cat((dh0, h2 - h), dim=1) du = (x2 - x).view(x.shape[0], -1, 3) dy = out2 - out ri = (dh[:, 1:].square().sum(-1) + dy.square() - dh[:, :-1].square().sum(-1) - du.square().sum(-1)) inc = torch.relu(ri).mean() loss = loss + eta * inc stats["incremental_penalty"] = float(inc.detach().cpu()) return loss, stats def run_one(seed, lr, lam, paired, return_stats=False): seed_all(seed) device = "cuda" if torch.cuda.is_available() else "cpu" try: ds = normalized_dataset(get_dataset("dynamics", seed, n_train=400, n_test=200)) model = MatchedRNN().to(device) opt = torch.optim.Adam(model.parameters(), lr=lr) xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device) for _ in range(EPOCHS): model.train() perm = torch.randperm(len(xtr), device=device) for start in range(0, len(xtr), BATCH): ix = perm[start:start+BATCH] x = xtr[ix]; target = ytr[ix] pred, out, h = model(x, True) pair = None if paired: x2 = x + 0.02 * torch.randn_like(x) pred2, out2, h2 = model(x2, True) pair = (x2, out2, h2) task = nn.functional.mse_loss(pred, target) pl, _ = passivity_loss(x, out, h, paired=pair) loss = task + lam * pl opt.zero_grad(); loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) opt.step() model.eval() with torch.no_grad(): pred, out, h = model(ds["xte"].to(device), True) metric = float(nn.functional.mse_loss(pred, ds["yte"].to(device)).cpu()) _, st = passivity_loss(ds["xte"].to(device), out, h) x = ds["xte"].to(device) x2 = x + 0.02 * torch.randn_like(x) _, _, h2 = model(x2, True) amp = ((h2[:, -1] - h[:, -1]).norm(dim=1) / (x2 - x).view(x.shape[0], -1).norm(dim=1).clamp_min(1e-8)).mean() st["amplification"] = float(amp.cpu()) return (metric, st) if return_stats else metric except RuntimeError: if device == "cuda": torch.cuda.empty_cache() old = torch.cuda.is_available # Explicit CPU retry without changing the benchmark architecture. torch.set_default_device("cpu") try: return run_one_cpu(seed, lr, lam, paired, return_stats) finally: torch.set_default_device("cpu") raise def run_one_cpu(seed, lr, lam, paired, return_stats=False): seed_all(seed) ds = normalized_dataset(get_dataset("dynamics", seed, n_train=400, n_test=200)) model = MatchedRNN() opt = torch.optim.Adam(model.parameters(), lr=lr) xtr, ytr = ds["xtr"], ds["ytr"] for _ in range(EPOCHS): perm = torch.randperm(len(xtr)) for start in range(0, len(xtr), BATCH): x = xtr[perm[start:start+BATCH]]; target = ytr[perm[start:start+BATCH]] pred, out, h = model(x, True); pair = None if paired: x2 = x + 0.02 * torch.randn_like(x) _, out2, h2 = model(x2, True); pair = (x2, out2, h2) task = nn.functional.mse_loss(pred, target) pl, _ = passivity_loss(x, out, h, paired=pair) opt.zero_grad(); (task + lam * pl).backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step() model.eval() with torch.no_grad(): x = ds["xte"]; pred, out, h = model(x, True) metric = float(nn.functional.mse_loss(pred, ds["yte"])) _, st = passivity_loss(x, out, h) x2 = x + 0.02 * torch.randn_like(x); _, _, h2 = model(x2, True) st["amplification"] = float(((h2[:, -1]-h[:, -1]).norm(dim=1) / (x2-x).view(x.shape[0],-1).norm(dim=1).clamp_min(1e-8)).mean()) return (metric, st) if return_stats else metric def main(): # Search-space parity: every idea learning rate is also swept by baseline. lrs = (1e-3, 3e-3, 1e-2) grid = [{"lr": lr} for lr in lrs] base = sweep_baseline( lambda cfg: lambda seed: run_one(seed, cfg["lr"], 0.0, False), grid, seeds=SWEEP_SEEDS) best_lr = base["best_cfg"]["lr"] # Three regularizer settings, with and without the optional paired term. idea_grid = [{"lr": lr, "lam": lam, "paired": paired} for lr in lrs for lam in (1e-3, 1e-2, 1e-1) for paired in (False, True)] scored = [] for cfg in idea_grid: vals = [run_one(s, cfg["lr"], cfg["lam"], cfg["paired"]) for s in SWEEP_SEEDS] scored.append((float(np.mean(vals)), cfg)) idea_cfg = min(scored, key=lambda z: z[0])[1] idea = evaluate(lambda s: run_one(s, idea_cfg["lr"], idea_cfg["lam"], idea_cfg["paired"]), seeds=SEEDS) # Signature is measured on separately trained baseline and idea systems. base_stats = [run_one(s, best_lr, 0.0, False, True)[1] for s in SEEDS] idea_stats = [run_one(s, idea_cfg["lr"], idea_cfg["lam"], idea_cfg["paired"], True)[1] for s in SEEDS] b_res = float(np.mean([z["mean_cumulative_residual"] for z in base_stats])) i_res = float(np.mean([z["mean_cumulative_residual"] for z in idea_stats])) b_amp = float(np.mean([z["amplification"] for z in base_stats])) i_amp = float(np.mean([z["amplification"] for z in idea_stats])) signature = { "baseline_cumulative_residual": b_res, "idea_cumulative_residual": i_res, "baseline_amplification": b_amp, "idea_amplification": i_amp, "predicted_direction": "passivity reduces residual and amplification", "confirmed": bool(i_res < b_res and i_amp < b_amp), } report = make_report( "dynamics", "rnn_small", {**base, "selected_cfg": base["best_cfg"]}, idea, {"selected_idea_cfg": idea_cfg, "mechanism_signature": signature}) Path("bench_report.json").write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == "__main__": main()