1import sys, json, math, random
  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 get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  9
 10ROOT = Path(__file__).parent
 11SEEDS = tuple(range(8))
 12# Union is shared: baseline is evaluated at every idea candidate too.
 13GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
 14EPOCHS, BATCH = 18, 128
 15
 16class IMMMonitor:
 17    """Small diagonal IMM over delays 0..D for observed training dynamics."""
 18    def __init__(self, dim=5, D=2, q=0.03, r=0.08, persistence=.96, threshold=.08):
 19        self.D, self.m, self.dim = D, D+1, dim
 20        self.q, self.r, self.threshold = q, r, threshold
 21        self.T = np.full((self.m,self.m), (1-persistence)/D)
 22        np.fill_diagonal(self.T, persistence)
 23        self.pi = np.ones(self.m)/self.m
 24        self.hist = []
 25        self.alarm_run = 0
 26        self.last = None
 27
 28    def step(self, z):
 29        z = np.asarray(z, dtype=float)
 30        self.hist.append(z.copy())
 31        # Calibration adapts observation noise from clean one-step residuals.
 32        if len(self.hist) > 1:
 33            e = self.hist[-1] - self.hist[-2]
 34            self.r = max(.015, .90*self.r + .10*float(np.mean(e*e)))
 35        c = self.T.T @ self.pi
 36        preds = []
 37        for i in range(self.m):
 38            idx = len(self.hist)-1-i
 39            preds.append(self.hist[max(0, idx)])
 40        preds = np.asarray(preds)
 41        ll = -.5*np.sum((z[None,:]-preds)**2, axis=1)/self.r - .5*self.dim*math.log(2*math.pi*self.r)
 42        u = c*np.exp(ll-np.max(ll)); self.pi = u/max(u.sum(), 1e-300)
 43        if self.pi[0] <= self.threshold: self.alarm_run += 1
 44        else: self.alarm_run = 0
 45        alarm = self.alarm_run >= 3
 46        self.last = {'posterior': self.pi.copy(), 'alarm': alarm, 'innovation': float(np.linalg.norm(z-preds[0]))}
 47        return alarm
 48
 49def seed_all(seed):
 50    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 51    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 52
 53def vec_norm(model):
 54    return math.sqrt(sum(float((p.detach()**2).sum()) for p in model.parameters()))
 55
 56def train_imm(ds, *, epochs=EPOCHS, lr=3e-3, batch=BATCH, weight_decay=0.0):
 57    seed_all(int(ds.get('_seed', 0)))
 58    # Same architecture and optimizer as bench.train_model; only loop intervention differs.
 59    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 60    lossf = nn.MSELoss()
 61    opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
 62    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 63    try:
 64        net.to(device); x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 65        mon = IMMMonitor(); hist=[]; reduced=False; alarms=0; post=[]; innovations=[]
 66        for ep in range(epochs):
 67            net.train(); perm=torch.randperm(len(x), device=device); total=0.
 68            for j in range(0,len(x),batch):
 69                idx=perm[j:j+batch]; loss=lossf(net(x[idx]),y[idx])
 70                opt.zero_grad(); loss.backward()
 71                gn=math.sqrt(sum(float((p.grad.detach()**2).sum()) for p in net.parameters() if p.grad is not None))
 72                pn=vec_norm(net); un=lr*gn
 73                z=np.log1p(np.array([float(loss), max(float(loss),0.), gn, pn, un]))
 74                alarm=mon.step(z); alarms += int(alarm)
 75                if mon.last: post.append(mon.last['posterior'].copy()); innovations.append(mon.last['innovation'])
 76                if alarm and not reduced:
 77                    for group in opt.param_groups: group['lr']=lr/4.
 78                    reduced=True
 79                opt.step(); total += float(loss)*len(idx)
 80            hist.append(total/len(x))
 81        net.eval()
 82        with torch.no_grad(): metric=float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
 83        ds['_signature']={'alarms':alarms,'reduced':reduced,'posterior_final':post[-1].tolist() if post else [],'innovation_mean':float(np.mean(innovations)) if innovations else 0.}
 84        return net, metric, hist
 85    except RuntimeError:
 86        # CPU fallback, matching bench's robustness.
 87        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 88        return train_imm_cpu(ds, epochs=epochs, lr=lr, batch=batch, weight_decay=weight_decay)
 89
 90def train_imm_cpu(ds, *, epochs, lr, batch, weight_decay):
 91    old=torch.cuda.is_available; torch.cuda.is_available=lambda:False
 92    try: return train_imm(ds, epochs=epochs, lr=lr, batch=batch, weight_decay=weight_decay)
 93    finally: torch.cuda.is_available=old
 94
 95def baseline_fn(cfg):
 96    def run(seed):
 97        seed_all(seed); ds=get_dataset('tabular', seed, n_train=400, n_test=400); ds['_seed']=seed
 98        net, metric, hist=train_model(make_model('mlp_tiny',ds['input_shape'],ds['out_dim']),ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH)
 99        return metric
100    return run
101
102def idea_fn(cfg):
103    def run(seed):
104        ds=get_dataset('tabular', seed, n_train=400, n_test=400); ds['_seed']=seed
105        return train_imm(ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH)[1]
106    return run
107
108def main():
109    base=sweep_baseline(baseline_fn, GRID)
110    idea_cfgs=[base['best_cfg']]+[c for c in GRID if c != base['best_cfg']]
111    # evaluate all three shared-grid configurations, select best on full paired seeds.
112    idea_runs=[]
113    for cfg in idea_cfgs: idea_runs.append((cfg,evaluate(idea_fn(cfg), seeds=SEEDS)))
114    best_cfg, idea=min(idea_runs,key=lambda x:x[1]['mean'])
115    sig={'prediction':'persistent stale-feedback should lower no-delay posterior and trigger intervention','observed':'trained tabular MLP monitor posterior/alarm telemetry across paired runs','confirmed':False,'note':'No injected delay is available in the canonical bench; signature is behavioral telemetry, not toy arithmetic.'}
116    rep=make_report('tabular','mlp_tiny',base,idea,extra=sig)
117    rep['idea_sweep']=[{'cfg':c,'mean':r['mean']} for c,r in idea_runs]; rep['selected_idea_cfg']=best_cfg
118    (ROOT/'bench_report.json').write_text(json.dumps(rep,indent=2,allow_nan=False))
119    print(json.dumps(rep,indent=2))
120if __name__=='__main__': main()