1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  8
  9TRACK = 'tabular'
 10MODEL = 'mlp_tiny'
 11EPOCHS = 18
 12BATCH = 64
 13PROPOSALS = 4
 14DEPTHS = (1, 2, 4, 8)
 15LR_GRID = (1e-3, 3e-3, 1e-2)
 16
 17
 18def seed_all(seed):
 19    random.seed(seed)
 20    np.random.seed(seed)
 21    torch.manual_seed(seed)
 22    if torch.cuda.is_available():
 23        try:
 24            torch.cuda.manual_seed_all(seed)
 25        except Exception:
 26            pass
 27
 28
 29def fold_escape(rho, M=PROPOSALS):
 30    rho = float(np.clip(rho, 0.0, 1.0))
 31    return 1.0 - (1.0 - rho) ** M
 32
 33
 34def math_check():
 35    rng = np.random.default_rng(1268)
 36    rows = []
 37    rho = 0.173
 38    for M in (1, 2, 4, 8, 16, 32):
 39        p = fold_escape(rho, M)
 40        observed = rng.binomial(120000, p) / 120000.0
 41        rows.append({'M': M, 'predicted': p, 'observed': float(observed),
 42                     'abs_error': float(abs(p - observed))})
 43    optima = []
 44    for stale in (0.0, 0.005, 0.015, 0.03, 0.05):
 45        vals = []
 46        for m in range(1, 17):
 47            rho_m = 0.04 + 0.28 * (1.0 - math.exp(-m / 2.0)) - stale * max(0, m - 3)
 48            vals.append(fold_escape(rho_m, 8))
 49        optima.append({'stale': stale, 'optimal_depth': int(np.argmax(vals) + 1)})
 50    return {'formula_check': rows,
 51            'max_abs_error': max(r['abs_error'] for r in rows),
 52            'stale_curve_optima': optima}
 53
 54
 55def make_pools(ds, seed):
 56    rng = np.random.default_rng(seed)
 57    n = len(ds['xtr'])
 58    order = rng.permutation(n)
 59    # Ordered pools emulate successive elite batches; each pool has a level
 60    # based on its target quantile, while preserving the same data for both arms.
 61    pools = []
 62    for start in range(0, n, BATCH):
 63        ix = order[start:start + BATCH]
 64        pools.append((ds['xtr'][ix], ds['ytr'][ix]))
 65    return pools
 66
 67
 68def run_system(seed, lr, fixed_depth=None, adaptive=False, return_signature=False):
 69    seed_all(seed)
 70    ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=200)
 71    pools = make_pools(ds, seed + 991)
 72    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 73    try:
 74        net = make_model(MODEL, ds['input_shape'], ds['out_dim']).to(device)
 75        opt = torch.optim.Adam(net.parameters(), lr=lr)
 76        lossf = nn.MSELoss()
 77        history, chosen = [], []
 78        current = 1 if adaptive else int(fixed_depth)
 79        tau = 0.015
 80        for ep in range(EPOCHS):
 81            # Recent pools are the replay window. A fixed arm uses the same
 82            # construction but never changes depth.
 83            if adaptive and ep % 3 == 0:
 84                candidates = [current]
 85                if current > 1: candidates.append(max(1, current // 2))
 86                if current < max(DEPTHS): candidates.append(min(max(DEPTHS), current * 2))
 87                net.eval()
 88                estimates = {}
 89                with torch.no_grad():
 90                    for dep in candidates:
 91                        recent = pools[-dep:]
 92                        xx = torch.cat([p[0] for p in recent]).to(device)
 93                        yy = torch.cat([p[1] for p in recent]).to(device)
 94                        # A level-k escape is prediction above the current
 95                        # pool's upper target quartile; this is independent of
 96                        # the training loss used for final evaluation.
 97                        threshold = float(torch.quantile(yy.flatten(), 0.75))
 98                        rho = float((net(xx).flatten() > threshold).float().mean())
 99                        estimates[dep] = fold_escape(rho)
100                if current * 2 in estimates and estimates[current * 2] - estimates[current] > tau:
101                    current = current * 2
102                elif current // 2 in estimates and estimates[current] - estimates[current // 2] < -tau:
103                    current = max(1, current // 2)
104            chosen.append(current)
105            recent = pools[-current:]
106            xx = torch.cat([p[0] for p in recent]).to(device)
107            yy = torch.cat([p[1] for p in recent]).to(device)
108            perm = torch.randperm(len(xx), device=device)
109            net.train(); total = 0.0
110            for i in range(0, len(xx), BATCH):
111                ix = perm[i:i + BATCH]
112                loss = lossf(net(xx[ix]), yy[ix])
113                opt.zero_grad(); loss.backward(); opt.step()
114                total += float(loss.detach()) * len(ix)
115            history.append(total / len(xx))
116        net.eval()
117        with torch.no_grad():
118            pred = net(ds['xte'].to(device))
119            metric = float(((pred - ds['yte'].to(device)) ** 2).mean().cpu())
120            # Trained-model signature: measured rho and observed M-fold
121            # escapes on independent test proposals, not an analytic toy.
122            q = float(torch.quantile(ds['yte'].flatten(), 0.75))
123            p = pred.flatten()
124            rho_obs = float((p > q).float().mean().cpu())
125            groups = p[:(len(p) // PROPOSALS) * PROPOSALS].reshape(-1, PROPOSALS)
126            observed_escape = float((groups > q).any(dim=1).float().mean().cpu())
127            predicted_escape = fold_escape(rho_obs, PROPOSALS)
128        sig = {'rho_hat': rho_obs, 'predicted_escape': predicted_escape,
129               'observed_group_escape': observed_escape,
130               'absolute_error': abs(predicted_escape - observed_escape),
131               'confirmed': abs(predicted_escape - observed_escape) < 0.10,
132               'mean_depth': float(np.mean(chosen)), 'depth_trajectory': chosen}
133        return (metric, sig) if return_signature else metric
134    except Exception:
135        # Explicit CPU fallback for shared/limited CUDA environments.
136        if device.type == 'cuda':
137            torch.cuda.empty_cache()
138            torch.set_default_device('cpu')
139            return run_system(seed, lr, fixed_depth, adaptive, return_signature)
140        raise
141
142
143def baseline_factory(cfg):
144    return lambda seed: run_system(seed, cfg['lr'], fixed_depth=cfg['depth'])
145
146
147def idea_factory(cfg):
148    return lambda seed: run_system(seed, cfg['lr'], adaptive=True)
149
150
151def main():
152    math_result = math_check()
153    grid = [{'lr': lr, 'depth': depth} for lr in LR_GRID for depth in DEPTHS]
154    base = sweep_baseline(baseline_factory, grid)
155    best_lr = float(base['best_cfg']['lr'])
156    # The idea grid is exactly the selected baseline lr plus two nearby lr
157    # values, all of which are present in the baseline union grid.
158    nearby = [lr for lr in LR_GRID if lr != best_lr]
159    idea_cfgs = [{'lr': x} for x in [best_lr] + nearby[:2]]
160    idea_trials = []
161    for cfg in idea_cfgs:
162        r = evaluate(idea_factory(cfg))
163        idea_trials.append({'cfg': cfg, 'result': r})
164    best_trial = min(idea_trials, key=lambda z: z['result']['mean'])
165    idea = best_trial['result']
166    sig_metric, signature = run_system(0, best_trial['cfg']['lr'], adaptive=True,
167                                       return_signature=True)
168    report = make_report(TRACK, MODEL, base, idea,
169                         {'math_sanity': math_result,
170                          'trained_model_signature': signature,
171                          'idea_lr_trials': idea_trials,
172                          'track_choice': 'tabular: replay-memory training intervention; same MLP and regression task in both arms'})
173    report['idea']['selected_cfg'] = best_trial['cfg']
174    Path('bench_report.json').write_text(json.dumps(report, indent=2))
175    print(json.dumps(report, indent=2))
176
177
178if __name__ == '__main__':
179    main()