1import sys, json, random, math
 2import numpy as np
 3import torch
 4from torch import nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6from bench import get_dataset as bench_get_dataset, train_model, evaluate, sweep_baseline, make_report
 7from custom_incidence_actions import MAX_ATOMS, MAX_ACTIONS, D
 8
 9SEEDS = tuple(range(8)); SWEEP = tuple(range(4))
10
11def seed_all(s):
12    random.seed(s); np.random.seed(s); torch.manual_seed(s)
13    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
14
15class IncidenceHead(nn.Module):
16    def __init__(self, width=32, alpha=0.0):
17        super().__init__(); self.alpha=alpha
18        self.atom = nn.Sequential(nn.Linear(D,width), nn.ReLU(), nn.Linear(width,1))
19    def forward(self, x):
20        atoms=x[:, :MAX_ATOMS*D].view(-1,MAX_ATOMS,D)
21        A=x[:, MAX_ATOMS*D:].view(-1,MAX_ACTIONS,MAX_ATOMS)
22        s=self.atom(atoms).squeeze(-1)
23        counts=A.sum(-1)
24        z=torch.bmm(A,s.unsqueeze(-1)).squeeze(-1)
25        z=z/(counts.clamp_min(1e-6)**self.alpha)
26        return z.masked_fill(counts.eq(0), -1e9)
27
28class PaddedActionMLP(nn.Module):
29    def __init__(self, width=32):
30        super().__init__()
31        # Standard padded action head: each legal row receives its padded aggregate
32        # atom representation and a count, then a shared MLP emits its action logit.
33        self.net=nn.Sequential(nn.Linear(D+1,width),nn.ReLU(),nn.Linear(width,1))
34    def forward(self,x):
35        atoms=x[:, :MAX_ATOMS*D].view(-1,MAX_ATOMS,D)
36        A=x[:, MAX_ATOMS*D:].view(-1,MAX_ACTIONS,MAX_ATOMS)
37        counts=A.sum(-1)
38        feat=torch.bmm(A,atoms)
39        feat=feat/(counts.unsqueeze(-1).clamp_min(1.0))
40        z=self.net(torch.cat([feat, counts.unsqueeze(-1)],-1)).squeeze(-1)
41        return z.masked_fill(counts.eq(0), -1e9)
42
43def run(kind,cfg,seed,return_model=False):
44    seed_all(seed); ds=bench_get_dataset('incidence_actions', seed, 400, 160)
45    model=IncidenceHead(alpha=cfg.get('alpha',0.0)) if kind=='idea' else PaddedActionMLP()
46    model, metric, hist=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,weight_decay=0.0)
47    if return_model: return metric,model,ds
48    return metric
49
50def trial(kind): return lambda cfg: (lambda seed: run(kind,cfg,seed))
51
52def main():
53    # Union parity: every idea lr is included in this baseline sweep.
54    grid=[{'lr':lr,'epochs':12} for lr in (0.001,0.003,0.006)]
55    base=sweep_baseline(trial('base'),grid,seeds=SWEEP)
56    best=base['best_cfg']
57    idea_grid=[{'lr':best['lr'],'epochs':best['epochs'],'alpha':a} for a in (0.0,0.5,1.0)]
58    tried=[]
59    for c in idea_grid:
60        r=evaluate(trial('idea')(c),seeds=SWEEP)
61        tried.append({'cfg':c,'mean':r['mean']})
62    bestidea=min(idea_grid,key=lambda c: next(q['mean'] for q in tried if q['cfg']==c))
63    idea=evaluate(trial('idea')(bestidea),seeds=SEEDS)
64    idea.update({'best_cfg':bestidea,'sweep':tried})
65    # Re-test the quantitative identity on outputs of one actually trained model.
66    metric,m,ds=run('idea',bestidea,0,True); m.eval()
67    with torch.no_grad():
68        dev=next(m.parameters()).device
69        x=ds['xte'][:64].to(dev); atoms=x[:,:MAX_ATOMS*D].view(-1,MAX_ATOMS,D)
70        A=x[:,MAX_ATOMS*D:].view(-1,MAX_ACTIONS,MAX_ATOMS)
71        s=m.atom(atoms).squeeze(-1); counts=A.sum(-1)
72        observed=m(x); predicted=torch.bmm(A,s.unsqueeze(-1)).squeeze(-1)/(counts.clamp_min(1e-6)**bestidea['alpha'])
73        valid=counts.gt(0)
74        err=(observed[valid]-predicted[valid]).abs()
75    sig={'prediction':'trained incidence logits equal A times trained atomic scores (alpha-normalized)','n_observed_logits':int(valid.sum()),'predicted_max_abs_error':float(err.max()),'observed_mean_abs_error':float(err.mean()),'confirmed':bool(float(err.max())<1e-5)}
76    report=make_report('incidence_actions','custom_incidence_head',base,idea,{'custom_track':{'name':'incidence_actions','file':'incidence_actions.py','domain':'structured_action_classification'},'baseline_sweep_union_lrs':[c['lr'] for c in grid],'idea_sweep':tried,'mechanism_signature':sig})
77    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
78    print(json.dumps(report,indent=2))
79if __name__=='__main__': main()