import json, sys, time from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import train_model, evaluate, sweep_baseline, make_report, get_dataset as bench_get_dataset from compat_track import META TRACK = 'multi_agent_critic_compatibility' MODEL = 'smooth_mlp_tiny' SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) LR_GRID = [0.001, 0.003, 0.006] LAMBDA_GRID = [0.0, 0.01, 0.05] EPOCHS = 18 BATCH = 64 records = {} class SmoothMLP(nn.Module): """Shared smooth base architecture; tanh is necessary for nonzero Hessians.""" def __init__(self, input_dim=6, out_dim=2): super().__init__() self.net = nn.Sequential(nn.Linear(input_dim, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, out_dim)) def forward(self, x): return self.net(x) def as_torch_ds(d): out = dict(d) for k in ('xtr', 'ytr', 'xte', 'yte'): out[k] = torch.as_tensor(out[k], dtype=torch.float32) # bench's generic regression adapter flattens multi-output custom targets # to [N*2,1]; restore the registered track's two-agent [N,2] target. for k in ('ytr', 'yte'): if out[k].shape[0] != out['x' + k[1:]].shape[0] and out[k].numel() % 2 == 0: out[k] = out[k].reshape(-1, 2) out['input_shape'] = tuple(out['xtr'].shape[1:]) return out def seed_all(seed): np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def new_model(): return SmoothMLP(6, 2) def cross_penalty(net, z, create_graph=True): q = net(z) g1 = torch.autograd.grad(q[:, 0].sum(), z, create_graph=True)[0][:, 2] g2 = torch.autograd.grad(q[:, 1].sum(), z, create_graph=True, retain_graph=True)[0][:, 4] h12 = torch.autograd.grad(g1.sum(), z, create_graph=create_graph, retain_graph=True)[0][:, 4] h21 = torch.autograd.grad(g2.sum(), z, create_graph=create_graph, retain_graph=True)[0][:, 2] return ((h12 - h21) ** 2).mean() def residual(net, x): dev = next(net.parameters()).device z = x[:128].to(dev).detach().clone().requires_grad_(True) return float(cross_penalty(net, z, create_graph=False).detach().cpu()) def baseline_fn(cfg): lr = float(cfg['lr']) def run(seed): seed_all(seed); ds = as_torch_ds(bench_get_dataset(TRACK, seed=seed, n_train=400, n_test=160)) net, metric, _ = train_model(new_model(), ds, epochs=EPOCHS, lr=lr, batch=BATCH, weight_decay=0.0, log=lambda *_: None) if net is None: return float('nan') with torch.enable_grad(): records[('baseline', lr, seed)] = residual(net, ds['xte']) return metric return run def idea_fn(cfg): lr, lam = float(cfg['lr']), float(cfg['lambda']) def run(seed): seed_all(seed); ds = as_torch_ds(bench_get_dataset(TRACK, seed=seed, n_train=400, n_test=160)) device = 'cuda' if torch.cuda.is_available() else 'cpu' try: net = new_model().to(device); xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device) opt = torch.optim.Adam(net.parameters(), lr=lr) for _ in range(EPOCHS): perm = torch.randperm(len(xtr), device=device); net.train() for i in range(0, len(xtr), BATCH): idx = perm[i:i+BATCH]; z = xtr[idx].detach().clone().requires_grad_(True) q = net(z); fit = ((q-ytr[idx])**2).mean() loss = fit + lam * cross_penalty(net, z, create_graph=True) opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): metric = float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean().cpu()) with torch.enable_grad(): records[('idea', lr, lam, seed)] = residual(net, ds['xte']) return metric except RuntimeError: # CPU fallback keeps the same architecture, optimizer, data and updates. net = new_model(); xtr, ytr = ds['xtr'], ds['ytr']; opt = torch.optim.Adam(net.parameters(), lr=lr) for _ in range(EPOCHS): perm = torch.randperm(len(xtr)) for i in range(0, len(xtr), BATCH): z=xtr[perm[i:i+BATCH]].detach().clone().requires_grad_(True); q=net(z) fit=((q-ytr[perm[i:i+BATCH]])**2).mean() loss=fit+lam*cross_penalty(net,z,create_graph=True) opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): metric=float(((net(ds['xte'])-ds['yte'])**2).mean()) with torch.enable_grad(): records[('idea',lr,lam,seed)] = residual(net,ds['xte']) return metric return run def main(): t=time.time() base_grid=[{'lr':lr, 'weight_decay':0.0} for lr in LR_GRID] base=sweep_baseline(baseline_fn, base_grid, seeds=SWEEP_SEEDS) idea_trials=[] for lr in LR_GRID: for lam in LAMBDA_GRID: cfg={'lr':lr,'lambda':lam}; r=evaluate(idea_fn(cfg), SWEEP_SEEDS) idea_trials.append({'cfg':cfg,'mean':r['mean']}) best_cfg=min(idea_trials,key=lambda x:x['mean'])['cfg'] idea_full=evaluate(idea_fn(best_cfg), SEEDS) base_final=evaluate(baseline_fn(base['best_cfg']), SEEDS) bres=np.mean([records[('baseline',float(base['best_cfg']['lr']),s)] for s in SEEDS]) ires=np.mean([records[('idea',float(best_cfg['lr']),float(best_cfg['lambda']),s)] for s in SEEDS]) report=make_report(TRACK, MODEL, {**base,'full':base_final}, idea_full, { 'prediction':'compatibility penalty lowers trained-NN cross-partial residual; task MSE may trade off against fit', 'baseline_mean_residual':float(bres), 'idea_mean_residual':float(ires), 'predicted_vs_observed':{'predicted_direction':'lower residual','observed_direction':'lower residual' if ires < bres else 'not lower'}, 'confirmed':bool(ires < bres) }) report['baseline']['union_grid']=base_grid; report['idea_sweep']=idea_trials report['custom_track']={'name':META['name'],'file':'compat_track.py','domain':META['domain']} report['runtime_sec']=time.time()-t Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2)) if __name__=='__main__': main()