import sys, json, random, math import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report SEED0 = 1215 EPOCHS = 12 BATCH = 128 THREADS = 4 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 js_batch(p, q): m = (p + q) / 2 return 0.5 * ((p * (p.clamp_min(1e-8).log() - m.clamp_min(1e-8).log())).sum(-1) + (q * (q.clamp_min(1e-8).log() - m.clamp_min(1e-8).log())).sum(-1)) def math_check(): p = torch.tensor([[.9, .1], [.9, .1], [.2, .8]], dtype=torch.float64) z = float(js_batch(p[:1], p[1:2])[0]) unequal = float(js_batch(p[:1], p[2:3])[0]) sets = [{0, 1}, {0}, {0, 2}] empty = [{0}, {1}] return { "js_equal": z, "js_unequal": unequal, "js_symmetric_error": abs(float(js_batch(p[:1], p[2:3])[0] - js_batch(p[2:3], p[:1])[0])), "intersection_feasible": bool(set.intersection(*map(set, sets))), "intersection_empty": not bool(set.intersection(*map(set, empty))), "passed": z < 1e-12 and unequal > 0 and not bool(set.intersection(*map(set, empty))), } def grouped_consistency(pred, x): # Dynamics observations are aliased by a coarse theta projection. The # scalar forecast induces a distribution over three forcing-mask choices. theta_obs = x[:, -3] keys = torch.round(theta_obs / 0.35).to(torch.int64) centers = pred.new_tensor([-0.35, 0.0, 0.35]) q = F.softmax(-((pred - centers[None, :]) / 0.30) ** 2, dim=1) vals = [] for k in torch.unique(keys): ix = keys == k if int(ix.sum()) > 1: bar = q[ix].mean(0, keepdim=True) vals.append(js_batch(q[ix], bar.expand_as(q[ix])).mean()) return torch.stack(vals).mean() if vals else pred.new_zeros(()) def robust_safety(pred, x): # Candidate forcing corrections; require one shared correction to keep all # histories in each observation class inside the safe angular region. theta_obs = x[:, -3] keys = torch.round(theta_obs / 0.35).to(torch.int64) deltas = pred.new_tensor([-0.20, 0.0, 0.20]) margins = torch.sigmoid((1.20 - (pred[:, None] + deltas[None, :]).abs()) / 0.12) losses = [] for k in torch.unique(keys): ix = keys == k if int(ix.sum()) > 1: worst = margins[ix].min(0).values best = worst.max() losses.append(F.relu(pred.new_tensor(0.75) - best)) return torch.stack(losses).mean() if losses else pred.new_zeros(()) def train_one(seed, lr, weight_decay=0.0, fc_weight=0.0, safety_weight=0.0, return_sig=False): seed_all(seed) ds = get_dataset("dynamics", seed, n_train=400, n_test=200) model = make_model("rnn_small", ds["input_shape"], ds["out_dim"]) device = "cuda" if torch.cuda.is_available() else "cpu" try: model.to(device) opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) x, y = ds["xtr"].to(device), ds["ytr"].to(device) model.train() for _ in range(EPOCHS): order = torch.randperm(len(x), device=device) for start in range(0, len(x), BATCH): ix = order[start:start+BATCH] pred = model(x[ix]) task = F.mse_loss(pred, y[ix]) loss = task if fc_weight: loss = loss + fc_weight * grouped_consistency(pred, x[ix]) if safety_weight: loss = loss + safety_weight * robust_safety(pred, x[ix]) opt.zero_grad(set_to_none=True); loss.backward(); opt.step() model.eval() with torch.no_grad(): pred = model(ds["xte"].to(device)) metric = float(F.mse_loss(pred, ds["yte"].to(device)).cpu()) # Signature measured on the trained model, not an analytic toy. kappa = float(grouped_consistency(pred, ds["xte"].to(device)).cpu()) safety = float(robust_safety(pred, ds["xte"].to(device)).cpu()) if return_sig: return metric, {"observed_kappa": kappa, "observed_safety_hinge": safety} return metric except RuntimeError: if device != "cpu": torch.cuda.empty_cache() return train_one(seed, lr, weight_decay, fc_weight, safety_weight, return_sig) raise def main(): check = math_check() assert check["passed"] # Baseline decisive knob (weight decay) and learning rate are both swept. grid = [{"lr": lr, "weight_decay": wd} for lr in (1e-2, 3e-2, 1e-1) for wd in (0.0, 1e-4)] base = sweep_baseline( lambda c: lambda s: train_one(s, c["lr"], c["weight_decay"]), grid) best_lr = base["best_cfg"]["lr"] idea_grid = [{"lr": lr, "weight_decay": 0.0, "fc_weight": 0.08, "safety_weight": 0.08} for lr in (1e-2, best_lr, 1e-1)] idea_trials = [] for cfg in idea_grid: r = evaluate(lambda s, c=cfg: train_one(s, c["lr"], c["weight_decay"], c["fc_weight"], c["safety_weight"]), (0,1,2,3)) idea_trials.append({"cfg": cfg, "mean": r["mean"]}) best_idea_cfg = min(idea_trials, key=lambda z: z["mean"])["cfg"] idea = evaluate(lambda s: train_one(s, best_idea_cfg["lr"], best_idea_cfg["weight_decay"], best_idea_cfg["fc_weight"], best_idea_cfg["safety_weight"])) _, sigs = zip(*(train_one(s, best_idea_cfg["lr"], best_idea_cfg["weight_decay"], best_idea_cfg["fc_weight"], best_idea_cfg["safety_weight"], True) for s in range(8))) base_sig = [train_one(s, base["best_cfg"]["lr"], base["best_cfg"]["weight_decay"], return_sig=True)[1] for s in range(8)] signature = { "prediction": "FC should lower observational kappa on trained recurrent models", "baseline_kappa_mean": float(np.mean([z["observed_kappa"] for z in base_sig])), "idea_kappa_mean": float(np.mean([z["observed_kappa"] for z in sigs])), "baseline_safety_hinge_mean": float(np.mean([z["observed_safety_hinge"] for z in base_sig])), "idea_safety_hinge_mean": float(np.mean([z["observed_safety_hinge"] for z in sigs])), "confirmed": float(np.mean([z["observed_kappa"] for z in sigs])) < float(np.mean([z["observed_kappa"] for z in base_sig])) } report = make_report("dynamics", "rnn_small", base, idea, { "math_check": check, "idea_sweep": idea_trials, "best_idea_cfg": best_idea_cfg, "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()