1import json, sys, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import train_model, evaluate, sweep_baseline, make_report, get_dataset as bench_get_dataset
  9from compat_track import META
 10
 11TRACK = 'multi_agent_critic_compatibility'
 12MODEL = 'smooth_mlp_tiny'
 13SEEDS = tuple(range(8))
 14SWEEP_SEEDS = (0, 1, 2, 3)
 15LR_GRID = [0.001, 0.003, 0.006]
 16LAMBDA_GRID = [0.0, 0.01, 0.05]
 17EPOCHS = 18
 18BATCH = 64
 19records = {}
 20
 21class SmoothMLP(nn.Module):
 22    """Shared smooth base architecture; tanh is necessary for nonzero Hessians."""
 23    def __init__(self, input_dim=6, out_dim=2):
 24        super().__init__()
 25        self.net = nn.Sequential(nn.Linear(input_dim, 64), nn.Tanh(),
 26                                 nn.Linear(64, 64), nn.Tanh(),
 27                                 nn.Linear(64, out_dim))
 28    def forward(self, x):
 29        return self.net(x)
 30
 31def as_torch_ds(d):
 32    out = dict(d)
 33    for k in ('xtr', 'ytr', 'xte', 'yte'):
 34        out[k] = torch.as_tensor(out[k], dtype=torch.float32)
 35    # bench's generic regression adapter flattens multi-output custom targets
 36    # to [N*2,1]; restore the registered track's two-agent [N,2] target.
 37    for k in ('ytr', 'yte'):
 38        if out[k].shape[0] != out['x' + k[1:]].shape[0] and out[k].numel() % 2 == 0:
 39            out[k] = out[k].reshape(-1, 2)
 40    out['input_shape'] = tuple(out['xtr'].shape[1:])
 41    return out
 42
 43def seed_all(seed):
 44    np.random.seed(seed); torch.manual_seed(seed)
 45    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 46
 47def new_model():
 48    return SmoothMLP(6, 2)
 49
 50def cross_penalty(net, z, create_graph=True):
 51    q = net(z)
 52    g1 = torch.autograd.grad(q[:, 0].sum(), z, create_graph=True)[0][:, 2]
 53    g2 = torch.autograd.grad(q[:, 1].sum(), z, create_graph=True, retain_graph=True)[0][:, 4]
 54    h12 = torch.autograd.grad(g1.sum(), z, create_graph=create_graph, retain_graph=True)[0][:, 4]
 55    h21 = torch.autograd.grad(g2.sum(), z, create_graph=create_graph, retain_graph=True)[0][:, 2]
 56    return ((h12 - h21) ** 2).mean()
 57
 58def residual(net, x):
 59    dev = next(net.parameters()).device
 60    z = x[:128].to(dev).detach().clone().requires_grad_(True)
 61    return float(cross_penalty(net, z, create_graph=False).detach().cpu())
 62
 63def baseline_fn(cfg):
 64    lr = float(cfg['lr'])
 65    def run(seed):
 66        seed_all(seed); ds = as_torch_ds(bench_get_dataset(TRACK, seed=seed, n_train=400, n_test=160))
 67        net, metric, _ = train_model(new_model(), ds, epochs=EPOCHS, lr=lr,
 68                                     batch=BATCH, weight_decay=0.0, log=lambda *_: None)
 69        if net is None: return float('nan')
 70        with torch.enable_grad(): records[('baseline', lr, seed)] = residual(net, ds['xte'])
 71        return metric
 72    return run
 73
 74def idea_fn(cfg):
 75    lr, lam = float(cfg['lr']), float(cfg['lambda'])
 76    def run(seed):
 77        seed_all(seed); ds = as_torch_ds(bench_get_dataset(TRACK, seed=seed, n_train=400, n_test=160))
 78        device = 'cuda' if torch.cuda.is_available() else 'cpu'
 79        try:
 80            net = new_model().to(device); xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device)
 81            opt = torch.optim.Adam(net.parameters(), lr=lr)
 82            for _ in range(EPOCHS):
 83                perm = torch.randperm(len(xtr), device=device); net.train()
 84                for i in range(0, len(xtr), BATCH):
 85                    idx = perm[i:i+BATCH]; z = xtr[idx].detach().clone().requires_grad_(True)
 86                    q = net(z); fit = ((q-ytr[idx])**2).mean()
 87                    loss = fit + lam * cross_penalty(net, z, create_graph=True)
 88                    opt.zero_grad(); loss.backward(); opt.step()
 89            net.eval()
 90            with torch.no_grad(): metric = float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean().cpu())
 91            with torch.enable_grad(): records[('idea', lr, lam, seed)] = residual(net, ds['xte'])
 92            return metric
 93        except RuntimeError:
 94            # CPU fallback keeps the same architecture, optimizer, data and updates.
 95            net = new_model(); xtr, ytr = ds['xtr'], ds['ytr']; opt = torch.optim.Adam(net.parameters(), lr=lr)
 96            for _ in range(EPOCHS):
 97                perm = torch.randperm(len(xtr))
 98                for i in range(0, len(xtr), BATCH):
 99                    z=xtr[perm[i:i+BATCH]].detach().clone().requires_grad_(True); q=net(z)
100                    fit=((q-ytr[perm[i:i+BATCH]])**2).mean()
101                    loss=fit+lam*cross_penalty(net,z,create_graph=True)
102                    opt.zero_grad(); loss.backward(); opt.step()
103            with torch.no_grad(): metric=float(((net(ds['xte'])-ds['yte'])**2).mean())
104            with torch.enable_grad(): records[('idea',lr,lam,seed)] = residual(net,ds['xte'])
105            return metric
106    return run
107
108def main():
109    t=time.time()
110    base_grid=[{'lr':lr, 'weight_decay':0.0} for lr in LR_GRID]
111    base=sweep_baseline(baseline_fn, base_grid, seeds=SWEEP_SEEDS)
112    idea_trials=[]
113    for lr in LR_GRID:
114        for lam in LAMBDA_GRID:
115            cfg={'lr':lr,'lambda':lam}; r=evaluate(idea_fn(cfg), SWEEP_SEEDS)
116            idea_trials.append({'cfg':cfg,'mean':r['mean']})
117    best_cfg=min(idea_trials,key=lambda x:x['mean'])['cfg']
118    idea_full=evaluate(idea_fn(best_cfg), SEEDS)
119    base_final=evaluate(baseline_fn(base['best_cfg']), SEEDS)
120    bres=np.mean([records[('baseline',float(base['best_cfg']['lr']),s)] for s in SEEDS])
121    ires=np.mean([records[('idea',float(best_cfg['lr']),float(best_cfg['lambda']),s)] for s in SEEDS])
122    report=make_report(TRACK, MODEL, {**base,'full':base_final}, idea_full, {
123        'prediction':'compatibility penalty lowers trained-NN cross-partial residual; task MSE may trade off against fit',
124        'baseline_mean_residual':float(bres), 'idea_mean_residual':float(ires),
125        'predicted_vs_observed':{'predicted_direction':'lower residual','observed_direction':'lower residual' if ires < bres else 'not lower'},
126        'confirmed':bool(ires < bres)
127    })
128    report['baseline']['union_grid']=base_grid; report['idea_sweep']=idea_trials
129    report['custom_track']={'name':META['name'],'file':'compat_track.py','domain':META['domain']}
130    report['runtime_sec']=time.time()-t
131    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
132
133if __name__=='__main__': main()