import os, sys, json, random 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, train_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS = 18 BATCH = 128 # Union is shared: baseline is evaluated at every lr considered by the idea. LRS = [1e-3, 2e-3, 3e-3] def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass class VectorCell(nn.Module): """Unconstrained recurrent vector field, used only as a matched reference.""" def __init__(self, hidden=64): super().__init__() self.inp = nn.Linear(3, hidden) self.field = nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, hidden)) self.head = nn.Linear(hidden, 1) def forward(self, x): seq = x.view(x.shape[0], -1, 3) z = torch.tanh(self.inp(seq[:, 0])) for k in range(seq.shape[1]): u = torch.tanh(self.inp(seq[:, k])) z = z + 0.12 * self.field(z + u) return self.head(z) class EnergyCell(nn.Module): """Same recurrent scaffold, but the state update is -grad_z E(z,input).""" def __init__(self, hidden=64, mu=0.03): super().__init__() self.inp = nn.Linear(3, hidden) self.energy_net = nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1)) self.head = nn.Linear(hidden, 1) self.mu = mu def energy(self, z, u): # Conditioning is fixed during each state update; mu gives coercivity. return self.energy_net(torch.tanh(z + u)).squeeze(-1) + self.mu * (z*z).sum(-1) / 2 def forward(self, x, collect=False): seq = x.view(x.shape[0], -1, 3) z = torch.tanh(self.inp(seq[:, 0])) energies, grads, states = [], [], [] for k in range(seq.shape[1]): u = torch.tanh(self.inp(seq[:, k])) z = z.requires_grad_(True) e = self.energy(z, u) g = torch.autograd.grad(e.sum(), z, create_graph=True)[0] z = z - 0.12 * g if collect: energies.append(e.detach()); grads.append(g.detach()); states.append(z.detach()) out = self.head(z) if collect: return out, energies, grads, states return out def train_energy(ds, epochs, lr, seed, collect=False): seed_all(seed) model = EnergyCell().to('cuda' if torch.cuda.is_available() else 'cpu') # Explicit fallback mirrors train_model's robust CUDA->CPU behavior. devices = ['cuda', 'cpu'] if torch.cuda.is_available() else ['cpu'] last = None for device in devices: try: model = model.to(device) x, y = ds['xtr'].to(device), ds['ytr'].to(device) opt = torch.optim.Adam(model.parameters(), lr=lr) for _ in range(epochs): model.train(); perm = torch.randperm(len(x), device=device) for i in range(0, len(x), BATCH): idx = perm[i:i+BATCH] loss = ((model(x[idx]) - y[idx]) ** 2).mean() opt.zero_grad(); loss.backward(); opt.step() model.eval() with torch.enable_grad(): pred = model(ds['xte'].to(device)) metric = float(((pred - ds['yte'].to(device)) ** 2).mean().detach().cpu()) return metric, model, device except RuntimeError as exc: last = exc if device == 'cuda': torch.cuda.empty_cache() continue raise raise last def idea_metric(cfg, seed, keep=False): ds = get_dataset('dynamics', seed) metric, model, device = train_energy(ds, EPOCHS, cfg['lr'], seed, keep) if keep: return metric, model, device, ds return metric def signature(model, ds, device): model.eval(); x = ds['xte'][:64].to(device) with torch.enable_grad(): _, es, gs, zs = model(x, collect=True) # Re-test the learned model's actual discrete dissipation, not an analytic toy. e = torch.stack(es, 1).mean(0).detach().cpu().numpy() gn = torch.stack([g.norm(dim=1) for g in gs], 1).mean(0).detach().cpu().numpy() diffs = np.diff(e) return {'predicted': 'energy should not increase under sufficiently small Euler steps', 'observed_energy_nonincreasing_fraction': float(np.mean(diffs <= 1e-7)), 'observed_energy_first_last': [float(e[0]), float(e[-1])], 'observed_grad_norm_first_last': [float(gn[0]), float(gn[-1])], 'predicted_boundary_eta_L': 2.0, 'observed_local_energy_boundary': 'not estimated (learned Hessian unavailable in budget)', 'confirmed': bool(np.all(diffs <= 1e-7) and gn[-1] < gn[0])} def main(): # Baseline sweep uses the standard bench train_model and the full lr union. def base_fn(cfg): def run(seed): seed_all(seed); ds = get_dataset('dynamics', seed) net = VectorCell(hidden=64).to('cpu') _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_: None) return metric return run base = sweep_baseline(base_fn, [{'lr': v} for v in LRS], seeds=SWEEP_SEEDS) # Evaluate each idea setting on the same tuning seeds, then full 8 for best. idea_sweep = [] for lr in LRS: r = evaluate(lambda s, lr=lr: idea_metric({'lr': lr}, s), seeds=SWEEP_SEEDS) idea_sweep.append({'cfg': {'lr': lr}, 'mean': r['mean']}) best_lr = min(idea_sweep, key=lambda q: q['mean'])['cfg']['lr'] idea = evaluate(lambda s: idea_metric({'lr': best_lr}, s), seeds=SEEDS) # Signature from one of the trained benchmark models, seed 0. _, trained, dev, ds0 = idea_metric({'lr': best_lr}, 0, keep=True) sig = signature(trained, ds0, dev) report = make_report('dynamics', 'rnn_small', {'best_cfg': base['best_cfg'], 'sweep': base['sweep'], 'full': base['full']}, idea, {'idea_sweep': idea_sweep, **sig}) report['protocol_notes'] = {'track_choice': 'dynamics matches stability/control/Lyapunov structure', 'shared_lr_union': LRS, 'epochs': EPOCHS, 'batch': BATCH, 'baseline_architecture': 'matched VectorCell recurrent scaffold (vector field)', 'idea_architecture': 'matched recurrent hidden width with scalar energy gradient'} with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()