import os, sys, json, math, 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, sweep_baseline, evaluate, make_report SEEDS = tuple(range(8)) # Union is shared: baseline is evaluated at every idea lr. LRS = [1e-3, 3e-3, 6e-3] DELAYS = [0, 2, 4] EPOCHS = 12 BATCH = 128 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 flat_params(model): return torch.cat([p.detach().reshape(-1) for p in model.parameters()]) def metric_eval(model, ds, device): model.eval() with torch.no_grad(): pred = model(ds['xte'].to(device)) return float(((pred - ds['yte'].to(device)) ** 2).mean().cpu()) def adam_run(seed, lr): seed_all(seed); ds = get_dataset('dynamics', seed, n_train=400, n_test=400) model = make_model('rnn_small', ds['input_shape'], ds['out_dim']) # Explicitly use the canonical baseline optimization, with the same budget. opt = torch.optim.Adam(model.parameters(), lr=lr) return train_loop(model, ds, opt, seed, mode='adam', delay=0) def switched_run(seed, lr, delay): seed_all(seed); ds = get_dataset('dynamics', seed, n_train=400, n_test=400) model = make_model('rnn_small', ds['input_shape'], ds['out_dim']) # Two locally smooth modes: low-momentum/smaller step and high-step/low momentum. opt_l = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9) opt_r = torch.optim.SGD(model.parameters(), lr=2.0*lr, momentum=0.1) return train_loop(model, ds, (opt_l, opt_r), seed, mode='switched', delay=delay) def train_loop(model, ds, opt, seed, mode, delay): requested = device_try() # Robust CUDA fallback, including cuDNN GRU failures. devices = [requested, 'cpu'] if requested == 'cuda' else ['cpu'] last_err = None for device in devices: try: model = model.to(device) x, y = ds['xtr'].to(device), ds['ytr'].to(device) lossf = nn.MSELoss(); n = len(x) if mode == 'switched': opt_l, opt_r = opt ref = flat_params(model).clone() u = torch.zeros_like(ref); u[0] = 1.0 hist = [0.0] * (delay + 1); sections=[]; displacements=[]; modes=[] for ep in range(EPOCHS): model.train(); perm = torch.randperm(n, device=device) for j in range(0, n, BATCH): idx = perm[j:j+BATCH] if mode == 'adam': chosen = opt else: gate = hist[-1-delay] chosen = opt_r if gate > 0 else opt_l chosen.zero_grad(set_to_none=True) loss = lossf(model(x[idx]), y[idx]); loss.backward(); chosen.step() flat = flat_params(model) s = float(torch.dot(u, flat-ref).cpu()) # EMA reference defines a local switching section. ref = 0.98*ref + 0.02*flat hist.append(s); sections.append(s); modes.append(1 if mode != 'adam' and gate > 0 else 0) score = metric_eval(model, ds, device) a = np.asarray(sections[-max(20, min(100, len(sections))):], dtype=float) # NN behavior signature: empirical lagged return displacement and cycle amplitude. if len(sections) > 2: yy=np.asarray(sections, dtype=float) yprev=yy[:-1]; ynext=yy[1:] valid=np.abs(yprev) > np.quantile(np.abs(yprev), .2) slope=float(np.polyfit(yprev[valid], (ynext-yprev)[valid], 1)[0]) if valid.sum()>3 else float('nan') else: slope=float('nan') return {'metric': score, 'cycle_amplitude': float(np.std(a)), 'return_slope': slope, 'r_fraction': float(np.mean(modes)) if modes else 0.0, 'device': device} except RuntimeError as e: last_err = repr(e) if device == 'cuda': try: torch.cuda.empty_cache() except Exception: pass continue raise raise RuntimeError(last_err or 'training failed') def baseline_config(cfg): # sweep_baseline calls make_fn(cfg), then evaluates the returned seed function. return lambda seed: adam_run(seed, cfg['lr'])['metric'] def idea_config(cfg): vals=[] for seed in SEEDS: vals.append(switched_run(seed, cfg['lr'], cfg['delay'])['metric']) return {'per_seed': vals, 'mean': float(np.mean(vals)), 'config': cfg} def main(): # Baseline sweep uses same lr union and also its central Adam knob. grid=[{'lr': x} for x in LRS] base = sweep_baseline(baseline_config, grid, seeds=SEEDS) # API returns best and full; normalize by inspecting expected keys. best_cfg = base.get('best_cfg', base.get('best_config', base.get('best', grid[0]))) idea_grid=[{'lr': best_cfg['lr'], 'delay': d} for d in DELAYS] idea_runs=[idea_config(c) for c in idea_grid] best_idea=min(idea_runs, key=lambda z:z['mean']) # Full eight-seed signature on selected idea and matched baseline best lr. base_full=evaluate(baseline_config({'lr':best_cfg['lr']}), seeds=SEEDS) selected=best_idea['config'] idea_full={'per_seed': [switched_run(s, selected['lr'], selected['delay'])['metric'] for s in SEEDS], 'mean': best_idea['mean'], 'config': selected} sig=[] for s in SEEDS: z=switched_run(s, selected['lr'], selected['delay']); sig.append(z) amp=np.asarray([z['cycle_amplitude'] for z in sig]) # Test predicted amplitude scaling from actual trained behavior across delays. bydelay=[] for d in DELAYS: rr=[switched_run(s, selected['lr'], d)['cycle_amplitude'] for s in SEEDS] bydelay.append(float(np.mean(rr))) positive=np.asarray(bydelay)>1e-10 slope=float(np.polyfit(np.log(np.asarray(DELAYS)[positive]+1), np.log(np.asarray(bydelay)[positive]), 1)[0]) if positive.sum()>=2 else float('nan') signature={'prediction':'cycle amplitude should increase approximately as mu^(1/M), M=4 for n=q=1', 'predicted_exponent':0.25,'observed_delay_plus_one_loglog_slope':slope, 'amplitude_by_delay':dict(zip(map(str,DELAYS),bydelay)), 'return_slope_mean':float(np.nanmean([z['return_slope'] for z in sig])), 'confirmed':bool(np.isfinite(slope) and slope>0 and abs(slope-.25)<.20)} report=make_report('dynamics','rnn_small',{'best_config':best_cfg,'full':base_full,'sweep':base},idea_full,signature) report['idea_sweep']=idea_runs open('bench_report.json','w').write(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__=='__main__': main()