1import os, sys, json, math, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, sweep_baseline, evaluate, make_report
  7
  8SEEDS = tuple(range(8))
  9# Union is shared: baseline is evaluated at every idea lr.
 10LRS = [1e-3, 3e-3, 6e-3]
 11DELAYS = [0, 2, 4]
 12EPOCHS = 12
 13BATCH = 128
 14
 15
 16def seed_all(seed):
 17    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 18    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 19
 20
 21def device_try():
 22    return 'cuda' if torch.cuda.is_available() else 'cpu'
 23
 24
 25def flat_params(model):
 26    return torch.cat([p.detach().reshape(-1) for p in model.parameters()])
 27
 28
 29def metric_eval(model, ds, device):
 30    model.eval()
 31    with torch.no_grad():
 32        pred = model(ds['xte'].to(device))
 33        return float(((pred - ds['yte'].to(device)) ** 2).mean().cpu())
 34
 35
 36def adam_run(seed, lr):
 37    seed_all(seed); ds = get_dataset('dynamics', seed, n_train=400, n_test=400)
 38    model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 39    # Explicitly use the canonical baseline optimization, with the same budget.
 40    opt = torch.optim.Adam(model.parameters(), lr=lr)
 41    return train_loop(model, ds, opt, seed, mode='adam', delay=0)
 42
 43
 44def switched_run(seed, lr, delay):
 45    seed_all(seed); ds = get_dataset('dynamics', seed, n_train=400, n_test=400)
 46    model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 47    # Two locally smooth modes: low-momentum/smaller step and high-step/low momentum.
 48    opt_l = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
 49    opt_r = torch.optim.SGD(model.parameters(), lr=2.0*lr, momentum=0.1)
 50    return train_loop(model, ds, (opt_l, opt_r), seed, mode='switched', delay=delay)
 51
 52
 53def train_loop(model, ds, opt, seed, mode, delay):
 54    requested = device_try()
 55    # Robust CUDA fallback, including cuDNN GRU failures.
 56    devices = [requested, 'cpu'] if requested == 'cuda' else ['cpu']
 57    last_err = None
 58    for device in devices:
 59        try:
 60            model = model.to(device)
 61            x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 62            lossf = nn.MSELoss(); n = len(x)
 63            if mode == 'switched':
 64                opt_l, opt_r = opt
 65            ref = flat_params(model).clone()
 66            u = torch.zeros_like(ref); u[0] = 1.0
 67            hist = [0.0] * (delay + 1); sections=[]; displacements=[]; modes=[]
 68            for ep in range(EPOCHS):
 69                model.train(); perm = torch.randperm(n, device=device)
 70                for j in range(0, n, BATCH):
 71                    idx = perm[j:j+BATCH]
 72                    if mode == 'adam':
 73                        chosen = opt
 74                    else:
 75                        gate = hist[-1-delay]
 76                        chosen = opt_r if gate > 0 else opt_l
 77                    chosen.zero_grad(set_to_none=True)
 78                    loss = lossf(model(x[idx]), y[idx]); loss.backward(); chosen.step()
 79                    flat = flat_params(model)
 80                    s = float(torch.dot(u, flat-ref).cpu())
 81                    # EMA reference defines a local switching section.
 82                    ref = 0.98*ref + 0.02*flat
 83                    hist.append(s); sections.append(s); modes.append(1 if mode != 'adam' and gate > 0 else 0)
 84            score = metric_eval(model, ds, device)
 85            a = np.asarray(sections[-max(20, min(100, len(sections))):], dtype=float)
 86            # NN behavior signature: empirical lagged return displacement and cycle amplitude.
 87            if len(sections) > 2:
 88                yy=np.asarray(sections, dtype=float)
 89                yprev=yy[:-1]; ynext=yy[1:]
 90                valid=np.abs(yprev) > np.quantile(np.abs(yprev), .2)
 91                slope=float(np.polyfit(yprev[valid], (ynext-yprev)[valid], 1)[0]) if valid.sum()>3 else float('nan')
 92            else: slope=float('nan')
 93            return {'metric': score, 'cycle_amplitude': float(np.std(a)),
 94                    'return_slope': slope, 'r_fraction': float(np.mean(modes)) if modes else 0.0,
 95                    'device': device}
 96        except RuntimeError as e:
 97            last_err = repr(e)
 98            if device == 'cuda':
 99                try: torch.cuda.empty_cache()
100                except Exception: pass
101                continue
102            raise
103    raise RuntimeError(last_err or 'training failed')
104
105
106def baseline_config(cfg):
107    # sweep_baseline calls make_fn(cfg), then evaluates the returned seed function.
108    return lambda seed: adam_run(seed, cfg['lr'])['metric']
109
110
111def idea_config(cfg):
112    vals=[]
113    for seed in SEEDS:
114        vals.append(switched_run(seed, cfg['lr'], cfg['delay'])['metric'])
115    return {'per_seed': vals, 'mean': float(np.mean(vals)), 'config': cfg}
116
117
118def main():
119    # Baseline sweep uses same lr union and also its central Adam knob.
120    grid=[{'lr': x} for x in LRS]
121    base = sweep_baseline(baseline_config, grid, seeds=SEEDS)
122    # API returns best and full; normalize by inspecting expected keys.
123    best_cfg = base.get('best_cfg', base.get('best_config', base.get('best', grid[0])))
124    idea_grid=[{'lr': best_cfg['lr'], 'delay': d} for d in DELAYS]
125    idea_runs=[idea_config(c) for c in idea_grid]
126    best_idea=min(idea_runs, key=lambda z:z['mean'])
127    # Full eight-seed signature on selected idea and matched baseline best lr.
128    base_full=evaluate(baseline_config({'lr':best_cfg['lr']}), seeds=SEEDS)
129    selected=best_idea['config']
130    idea_full={'per_seed': [switched_run(s, selected['lr'], selected['delay'])['metric'] for s in SEEDS],
131               'mean': best_idea['mean'], 'config': selected}
132    sig=[]
133    for s in SEEDS:
134        z=switched_run(s, selected['lr'], selected['delay']); sig.append(z)
135    amp=np.asarray([z['cycle_amplitude'] for z in sig])
136    # Test predicted amplitude scaling from actual trained behavior across delays.
137    bydelay=[]
138    for d in DELAYS:
139        rr=[switched_run(s, selected['lr'], d)['cycle_amplitude'] for s in SEEDS]
140        bydelay.append(float(np.mean(rr)))
141    positive=np.asarray(bydelay)>1e-10
142    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')
143    signature={'prediction':'cycle amplitude should increase approximately as mu^(1/M), M=4 for n=q=1',
144      'predicted_exponent':0.25,'observed_delay_plus_one_loglog_slope':slope,
145      'amplitude_by_delay':dict(zip(map(str,DELAYS),bydelay)),
146      'return_slope_mean':float(np.nanmean([z['return_slope'] for z in sig])),
147      'confirmed':bool(np.isfinite(slope) and slope>0 and abs(slope-.25)<.20)}
148    report=make_report('dynamics','rnn_small',{'best_config':best_cfg,'full':base_full,'sweep':base},idea_full,signature)
149    report['idea_sweep']=idea_runs
150    open('bench_report.json','w').write(json.dumps(report,indent=2))
151    print(json.dumps(report,indent=2))
152
153if __name__=='__main__': main()