1import json, os, sys, 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, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11EPOCHS = 12
 12NTRAIN, NTEST = 400, 200
 13BATCH = 128
 14# Shared union: all lrs and all fixed/dual pressure values are evaluated on baseline.
 15LRS = (1e-3, 3e-3, 1e-2)
 16PRESSURES = (0.0, 0.05, 0.2)
 17HGOAL = 0.55
 18ETA_LAMBDA = 0.08
 19LAMBDA_MAX = 1.5
 20
 21
 22def seed_all(seed):
 23    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 24    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 25
 26
 27def device_try():
 28    return 'cuda' if torch.cuda.is_available() else 'cpu'
 29
 30
 31def features(net, x):
 32    # Uses the exact bench rnn_small GRU and head; h is the learned latent state.
 33    seq = x.view(x.shape[0], -1, 3)
 34    _, h = net.rnn(seq)
 35    z = h[-1]
 36    return z, net.head(z)
 37
 38
 39def train_one(seed, lr, pressure, dual, eta_lambda=ETA_LAMBDA):
 40    seed_all(seed)
 41    ds = get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)
 42    net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 43    # Auxiliary conditional and action-only Gaussian mean decoders. They estimate
 44    # I(Y;Z|U) by log q(y|z,u)-log q(y|u); y is the observed future target.
 45    cond = nn.Linear(64 + 8, 1)
 46    marginal = nn.Linear(1, 1)
 47    dev = device_try()
 48    try:
 49        net, cond, marginal = net.to(dev), cond.to(dev), marginal.to(dev)
 50        opt = torch.optim.Adam(list(net.parameters()) + list(cond.parameters()) + list(marginal.parameters()), lr=lr)
 51        x, y = ds['xtr'].to(dev), ds['ytr'].to(dev)
 52        lam = 0.0; ema = 0.0; last_I = 0.0
 53        for ep in range(EPOCHS):
 54            net.train(); cond.train(); marginal.train()
 55            perm = torch.randperm(len(x), device=dev)
 56            for i in range(0, len(x), BATCH):
 57                ix = perm[i:i+BATCH]; xb, yb = x[ix], y[ix]
 58                z, pred = features(net, xb)
 59                # actions are the eight u coordinates (indices 2,5,...)
 60                u = xb.view(-1, 8, 3)[:, :, 2]
 61                us = u.mean(1, keepdim=True)
 62                mu_c = cond(torch.cat([z, u], 1))
 63                mu_m = marginal(us)
 64                # Equal-variance Gaussian log-density difference.
 65                logc = -0.5 * (yb - mu_c).pow(2)
 66                logm = -0.5 * (yb - mu_m).pow(2)
 67                I = (logc - logm).mean()
 68                task = (pred - yb).pow(2).mean()
 69                # Baseline is fixed information pressure; dual is the intervention.
 70                if dual:
 71                    loss = task + lam * (HGOAL - I)
 72                else:
 73                    loss = task + pressure * (HGOAL - I)
 74                opt.zero_grad(); loss.backward()
 75                torch.nn.utils.clip_grad_norm_(list(net.parameters()) + list(cond.parameters()) + list(marginal.parameters()), 5.0)
 76                opt.step()
 77                with torch.no_grad():
 78                    last_I = float(I.detach()); ema = 0.9 * ema + 0.1 * last_I
 79                    if dual:
 80                        lam = float(np.clip(lam + eta_lambda * (HGOAL - ema), 0., LAMBDA_MAX))
 81        net.eval(); cond.eval(); marginal.eval()
 82        with torch.no_grad():
 83            z, pred = features(net, ds['xte'].to(dev)); yt = ds['yte'].to(dev)
 84            u = ds['xte'].to(dev).view(-1, 8, 3)[:, :, 2]
 85            I_test = float((-0.5*(yt-cond(torch.cat([z,u],1))).pow(2) + 0.5*(yt-marginal(u.mean(1,keepdim=True))).pow(2)).mean())
 86            mse = float((pred-yt).pow(2).mean())
 87        return mse, {'I_train': last_I, 'I_test': I_test, 'lambda': lam, 'H_goal': HGOAL, 'pressure': pressure}
 88    except RuntimeError:
 89        if dev == 'cuda':
 90            torch.cuda.empty_cache()
 91            # deterministic CPU fallback
 92            return train_one_cpu(seed, lr, pressure, dual, eta_lambda)
 93        raise
 94
 95
 96def train_one_cpu(seed, lr, pressure, dual, eta_lambda=ETA_LAMBDA):
 97    old = torch.cuda.is_available
 98    # Same implementation, forced by a temporary explicit helper path.
 99    torch.cuda.is_available = lambda: False
100    try: return train_one(seed, lr, pressure, dual, eta_lambda)
101    finally: torch.cuda.is_available = old
102
103
104def eval_cfg(lr, pressure, dual, seeds=SEEDS, collect=False, eta_lambda=ETA_LAMBDA):
105    vals=[]; details=[]
106    for s in seeds:
107        v, d = train_one(s, lr, pressure, dual, eta_lambda); vals.append(v); details.append(d)
108    out={'mean':float(np.mean(vals)), 'std':float(np.std(vals)), 'per_seed':vals, 'n':len(vals)}
109    if collect: out['details']=details
110    return out
111
112
113def main():
114    # Baseline sweep includes every lr and pressure used by the idea-side grid.
115    grid=[{'lr':lr,'pressure':p} for lr in LRS for p in PRESSURES]
116    base=sweep_baseline(lambda c: lambda s: eval_cfg(c['lr'], c['pressure'], False, (s,))['mean'], grid)
117    best=base['best_cfg']
118    idea_grid=[{'lr':best['lr'],'pressure':best['pressure'],'eta_lambda':e} for e in (0.04, 0.08, 0.16)]
119    idea_candidates=[]
120    for c in idea_grid:
121        r=eval_cfg(c['lr'], c['pressure'], True, SEEDS, collect=True, eta_lambda=c['eta_lambda'])
122        idea_candidates.append({'cfg':c,'result':r})
123    chosen=min(idea_candidates, key=lambda q:q['result']['mean'])
124    idea=chosen['result']; idea['best_cfg']=chosen['cfg']; idea['candidates']=[{'cfg':q['cfg'],'mean':q['result']['mean']} for q in idea_candidates]
125    # Signature is measured from trained systems, not the toy identity.
126    base_sig=eval_cfg(best['lr'], best['pressure'], False, SEEDS, collect=True)
127    bd=base_sig['details']; idd=idea['details']
128    sig={'H_goal':HGOAL, 'baseline_fixed_pressure':best['pressure'],
129         'baseline_mean_I_test':float(np.mean([d['I_test'] for d in bd])),
130         'idea_mean_I_test':float(np.mean([d['I_test'] for d in idd])),
131         'idea_mean_lambda':float(np.mean([d['lambda'] for d in idd])),
132         'predicted': 'dual pressure increases when I_test/EMA is below H_goal and should reduce deficit',
133         'observed_deficit_baseline':float(HGOAL-np.mean([d['I_test'] for d in bd])),
134         'observed_deficit_idea':float(HGOAL-np.mean([d['I_test'] for d in idd])),
135         'confirmed': bool(np.mean([d['I_test'] for d in idd]) > np.mean([d['I_test'] for d in bd]))}
136    rep=make_report('dynamics','rnn_small',base,idea,{'bench_report':{'track_match':'controlled dynamics / latent predictive information','custom_track':None},'signature':sig})
137    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
138    print(json.dumps(rep,indent=2))
139
140if __name__=='__main__': main()