import sys, json, random import numpy as np import torch import torch.nn as nn sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) LR_GRID = (1e-3, 3e-3, 6e-3) EPOCHS = 12 NTRAIN, NTEST = 800, 300 BATCH = 128 LAMBDA = 0.01 GAMMA = 0.50 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) def baseline_fn(cfg): def run(seed): seed_all(seed) ds = get_dataset("dynamics", seed=seed, n_train=NTRAIN, n_test=NTEST) net = make_model("rnn_small", ds["input_shape"], ds["out_dim"]) _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, log=lambda *_: None) return float(metric) return run def observer_train(seed, lr, return_signature=False): """Train the identical GRU with an empirical bounded-real sensitivity penalty. Input perturbations are the disturbance channel. For q equal to the scalar prediction, ||dq/dx|| is the local induced gain. The penalty is the positive part of gain^2-gamma^2, a differentiable finite-dimensional proxy for the bounded-real inequality; the task MSE remains the primary objective. """ seed_all(seed) ds = get_dataset("dynamics", seed=seed, n_train=NTRAIN, n_test=NTEST) net = make_model("rnn_small", ds["input_shape"], ds["out_dim"]) device = "cuda" if torch.cuda.is_available() else "cpu" try: net = net.to(device) x = ds["xtr"].to(device) y = ds["ytr"].to(device) opt = torch.optim.Adam(net.parameters(), lr=lr) lossf = nn.MSELoss() for _ in range(EPOCHS): net.train() perm = torch.randperm(len(x), device=device) for start in range(0, len(x), BATCH): idx = perm[start:start+BATCH] xb = x[idx].detach().requires_grad_(True) pred = net(xb) mse = lossf(pred, y[idx]) grad = torch.autograd.grad(pred.sum(), xb, create_graph=True)[0] local_gain_sq = grad.reshape(len(idx), -1).pow(2).sum(1).mean() penalty = torch.relu(local_gain_sq - GAMMA ** 2) loss = mse + LAMBDA * penalty opt.zero_grad(set_to_none=True) loss.backward() opt.step() net.eval() with torch.no_grad(): metric = float(((net(ds["xte"].to(device)) - ds["yte"].to(device)) ** 2).mean()) if return_signature: # Measured on the trained model, not an analytical identity. xt = ds["xte"][:64].to(device).detach().requires_grad_(True) out = net(xt) g = torch.autograd.grad(out.sum(), xt)[0] gains = g.reshape(len(xt), -1).norm(dim=1).detach().cpu().numpy() return metric, float(np.mean(gains)), float(np.quantile(gains, .95)) return metric except RuntimeError: if device == "cuda": torch.cuda.empty_cache() # Explicit CPU fallback with same seed and settings. return observer_train_cpu(seed, lr, return_signature) raise def observer_train_cpu(seed, lr, return_signature=False): seed_all(seed) ds = get_dataset("dynamics", seed=seed, n_train=NTRAIN, n_test=NTEST) net = make_model("rnn_small", ds["input_shape"], ds["out_dim"]) opt = torch.optim.Adam(net.parameters(), lr=lr) for _ in range(EPOCHS): perm = torch.randperm(len(ds["xtr"])) for start in range(0, len(perm), BATCH): idx = perm[start:start+BATCH] xb = ds["xtr"][idx].detach().requires_grad_(True) pred = net(xb) mse = ((pred - ds["ytr"][idx]) ** 2).mean() grad = torch.autograd.grad(pred.sum(), xb, create_graph=True)[0] gain2 = grad.reshape(len(idx), -1).pow(2).sum(1).mean() loss = mse + LAMBDA * torch.relu(gain2 - GAMMA ** 2) opt.zero_grad(set_to_none=True); loss.backward(); opt.step() net.eval() xt = ds["xte"][:64].detach().requires_grad_(True) out = net(xt) metric = float(((out - ds["yte"][:64]) ** 2).mean()) if return_signature: g = torch.autograd.grad(out.sum(), xt)[0] gains = g.reshape(len(xt), -1).norm(dim=1).detach().numpy() return metric, float(gains.mean()), float(np.quantile(gains, .95)) return metric def main(): grid = [{"lr": x} for x in LR_GRID] base = sweep_baseline(baseline_fn, grid, seeds=(0, 1, 2, 3)) idea_cfg = base["best_cfg"] # Same union of learning rates as baseline; report best idea configuration. idea_results = [] for lr in LR_GRID: r = evaluate(lambda s, lr=lr: observer_train(s, lr), seeds=SEEDS) idea_results.append((r, lr)) idea, idea_lr = min(idea_results, key=lambda z: z[0]["mean"]) sig = [observer_train(s, idea_lr, True) for s in SEEDS] def baseline_gain(seed): seed_all(seed) ds = get_dataset("dynamics", seed=seed, n_train=NTRAIN, n_test=NTEST) net = make_model("rnn_small", ds["input_shape"], ds["out_dim"]) net, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=idea_lr, batch=BATCH, log=lambda *_: None) dev = next(net.parameters()).device xt = ds["xte"][:64].to(dev).detach().requires_grad_(True) out = net(xt) g = torch.autograd.grad(out.sum(), xt)[0] gains = g.reshape(len(xt), -1).norm(dim=1).detach().cpu().numpy() return float(metric), float(gains.mean()), float(np.quantile(gains, .95)) base_beh = [baseline_gain(s) for s in SEEDS] signature = { "claim": "bounded-real attenuation reduces local input-disturbance gain", "gamma": GAMMA, "idea_lr": idea_lr, "baseline_observed_mean_local_gain": float(np.mean([x[1] for x in base_beh])), "idea_observed_mean_local_gain": float(np.mean([x[1] for x in sig])), "baseline_observed_p95_local_gain": float(np.mean([x[2] for x in base_beh])), "idea_observed_p95_local_gain": float(np.mean([x[2] for x in sig])), "baseline_task_mse_at_same_lr": float(np.mean([x[0] for x in base_beh])), "idea_task_mse_at_same_lr": float(np.mean([x[0] for x in sig])), "confirmed": bool(np.mean([x[1] for x in sig]) < np.mean([x[1] for x in base_beh]) * 0.95) } report = make_report("dynamics", "rnn_small", base, idea, {"mechanism_signature": signature, "idea_sweep": [{"lr": lr, "mean": r["mean"]} for r, lr in idea_results], "structural_match": "controlled damped pendulum; stability/control track"}) report["mechanism_signature"] = signature with open("bench_report.json", "w") as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2)) if __name__ == "__main__": main()