import json, os, sys, random 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 get_dataset, make_model, sweep_baseline, make_report SEEDS = tuple(range(8)) EPOCHS = 12 NTRAIN, NTEST = 400, 200 BATCH = 128 # Shared union: all lrs and all fixed/dual pressure values are evaluated on baseline. LRS = (1e-3, 3e-3, 1e-2) PRESSURES = (0.0, 0.05, 0.2) HGOAL = 0.55 ETA_LAMBDA = 0.08 LAMBDA_MAX = 1.5 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def device_try(): return 'cuda' if torch.cuda.is_available() else 'cpu' def features(net, x): # Uses the exact bench rnn_small GRU and head; h is the learned latent state. seq = x.view(x.shape[0], -1, 3) _, h = net.rnn(seq) z = h[-1] return z, net.head(z) def train_one(seed, lr, pressure, dual, eta_lambda=ETA_LAMBDA): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST) net = make_model('rnn_small', ds['input_shape'], ds['out_dim']) # Auxiliary conditional and action-only Gaussian mean decoders. They estimate # I(Y;Z|U) by log q(y|z,u)-log q(y|u); y is the observed future target. cond = nn.Linear(64 + 8, 1) marginal = nn.Linear(1, 1) dev = device_try() try: net, cond, marginal = net.to(dev), cond.to(dev), marginal.to(dev) opt = torch.optim.Adam(list(net.parameters()) + list(cond.parameters()) + list(marginal.parameters()), lr=lr) x, y = ds['xtr'].to(dev), ds['ytr'].to(dev) lam = 0.0; ema = 0.0; last_I = 0.0 for ep in range(EPOCHS): net.train(); cond.train(); marginal.train() perm = torch.randperm(len(x), device=dev) for i in range(0, len(x), BATCH): ix = perm[i:i+BATCH]; xb, yb = x[ix], y[ix] z, pred = features(net, xb) # actions are the eight u coordinates (indices 2,5,...) u = xb.view(-1, 8, 3)[:, :, 2] us = u.mean(1, keepdim=True) mu_c = cond(torch.cat([z, u], 1)) mu_m = marginal(us) # Equal-variance Gaussian log-density difference. logc = -0.5 * (yb - mu_c).pow(2) logm = -0.5 * (yb - mu_m).pow(2) I = (logc - logm).mean() task = (pred - yb).pow(2).mean() # Baseline is fixed information pressure; dual is the intervention. if dual: loss = task + lam * (HGOAL - I) else: loss = task + pressure * (HGOAL - I) opt.zero_grad(); loss.backward() torch.nn.utils.clip_grad_norm_(list(net.parameters()) + list(cond.parameters()) + list(marginal.parameters()), 5.0) opt.step() with torch.no_grad(): last_I = float(I.detach()); ema = 0.9 * ema + 0.1 * last_I if dual: lam = float(np.clip(lam + eta_lambda * (HGOAL - ema), 0., LAMBDA_MAX)) net.eval(); cond.eval(); marginal.eval() with torch.no_grad(): z, pred = features(net, ds['xte'].to(dev)); yt = ds['yte'].to(dev) u = ds['xte'].to(dev).view(-1, 8, 3)[:, :, 2] 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()) mse = float((pred-yt).pow(2).mean()) return mse, {'I_train': last_I, 'I_test': I_test, 'lambda': lam, 'H_goal': HGOAL, 'pressure': pressure} except RuntimeError: if dev == 'cuda': torch.cuda.empty_cache() # deterministic CPU fallback return train_one_cpu(seed, lr, pressure, dual, eta_lambda) raise def train_one_cpu(seed, lr, pressure, dual, eta_lambda=ETA_LAMBDA): old = torch.cuda.is_available # Same implementation, forced by a temporary explicit helper path. torch.cuda.is_available = lambda: False try: return train_one(seed, lr, pressure, dual, eta_lambda) finally: torch.cuda.is_available = old def eval_cfg(lr, pressure, dual, seeds=SEEDS, collect=False, eta_lambda=ETA_LAMBDA): vals=[]; details=[] for s in seeds: v, d = train_one(s, lr, pressure, dual, eta_lambda); vals.append(v); details.append(d) out={'mean':float(np.mean(vals)), 'std':float(np.std(vals)), 'per_seed':vals, 'n':len(vals)} if collect: out['details']=details return out def main(): # Baseline sweep includes every lr and pressure used by the idea-side grid. grid=[{'lr':lr,'pressure':p} for lr in LRS for p in PRESSURES] base=sweep_baseline(lambda c: lambda s: eval_cfg(c['lr'], c['pressure'], False, (s,))['mean'], grid) best=base['best_cfg'] idea_grid=[{'lr':best['lr'],'pressure':best['pressure'],'eta_lambda':e} for e in (0.04, 0.08, 0.16)] idea_candidates=[] for c in idea_grid: r=eval_cfg(c['lr'], c['pressure'], True, SEEDS, collect=True, eta_lambda=c['eta_lambda']) idea_candidates.append({'cfg':c,'result':r}) chosen=min(idea_candidates, key=lambda q:q['result']['mean']) idea=chosen['result']; idea['best_cfg']=chosen['cfg']; idea['candidates']=[{'cfg':q['cfg'],'mean':q['result']['mean']} for q in idea_candidates] # Signature is measured from trained systems, not the toy identity. base_sig=eval_cfg(best['lr'], best['pressure'], False, SEEDS, collect=True) bd=base_sig['details']; idd=idea['details'] sig={'H_goal':HGOAL, 'baseline_fixed_pressure':best['pressure'], 'baseline_mean_I_test':float(np.mean([d['I_test'] for d in bd])), 'idea_mean_I_test':float(np.mean([d['I_test'] for d in idd])), 'idea_mean_lambda':float(np.mean([d['lambda'] for d in idd])), 'predicted': 'dual pressure increases when I_test/EMA is below H_goal and should reduce deficit', 'observed_deficit_baseline':float(HGOAL-np.mean([d['I_test'] for d in bd])), 'observed_deficit_idea':float(HGOAL-np.mean([d['I_test'] for d in idd])), 'confirmed': bool(np.mean([d['I_test'] for d in idd]) > np.mean([d['I_test'] for d in bd]))} rep=make_report('dynamics','rnn_small',base,idea,{'bench_report':{'track_match':'controlled dynamics / latent predictive information','custom_track':None},'signature':sig}) Path('bench_report.json').write_text(json.dumps(rep,indent=2)) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()