import sys, json, math, 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, evaluate, sweep_baseline, make_report TRACK = 'tabular' MODEL = 'mlp_tiny' EPOCHS = 18 BATCH = 64 PROPOSALS = 4 DEPTHS = (1, 2, 4, 8) LR_GRID = (1e-3, 3e-3, 1e-2) 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 def fold_escape(rho, M=PROPOSALS): rho = float(np.clip(rho, 0.0, 1.0)) return 1.0 - (1.0 - rho) ** M def math_check(): rng = np.random.default_rng(1268) rows = [] rho = 0.173 for M in (1, 2, 4, 8, 16, 32): p = fold_escape(rho, M) observed = rng.binomial(120000, p) / 120000.0 rows.append({'M': M, 'predicted': p, 'observed': float(observed), 'abs_error': float(abs(p - observed))}) optima = [] for stale in (0.0, 0.005, 0.015, 0.03, 0.05): vals = [] for m in range(1, 17): rho_m = 0.04 + 0.28 * (1.0 - math.exp(-m / 2.0)) - stale * max(0, m - 3) vals.append(fold_escape(rho_m, 8)) optima.append({'stale': stale, 'optimal_depth': int(np.argmax(vals) + 1)}) return {'formula_check': rows, 'max_abs_error': max(r['abs_error'] for r in rows), 'stale_curve_optima': optima} def make_pools(ds, seed): rng = np.random.default_rng(seed) n = len(ds['xtr']) order = rng.permutation(n) # Ordered pools emulate successive elite batches; each pool has a level # based on its target quantile, while preserving the same data for both arms. pools = [] for start in range(0, n, BATCH): ix = order[start:start + BATCH] pools.append((ds['xtr'][ix], ds['ytr'][ix])) return pools def run_system(seed, lr, fixed_depth=None, adaptive=False, return_signature=False): seed_all(seed) ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=200) pools = make_pools(ds, seed + 991) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') try: net = make_model(MODEL, ds['input_shape'], ds['out_dim']).to(device) opt = torch.optim.Adam(net.parameters(), lr=lr) lossf = nn.MSELoss() history, chosen = [], [] current = 1 if adaptive else int(fixed_depth) tau = 0.015 for ep in range(EPOCHS): # Recent pools are the replay window. A fixed arm uses the same # construction but never changes depth. if adaptive and ep % 3 == 0: candidates = [current] if current > 1: candidates.append(max(1, current // 2)) if current < max(DEPTHS): candidates.append(min(max(DEPTHS), current * 2)) net.eval() estimates = {} with torch.no_grad(): for dep in candidates: recent = pools[-dep:] xx = torch.cat([p[0] for p in recent]).to(device) yy = torch.cat([p[1] for p in recent]).to(device) # A level-k escape is prediction above the current # pool's upper target quartile; this is independent of # the training loss used for final evaluation. threshold = float(torch.quantile(yy.flatten(), 0.75)) rho = float((net(xx).flatten() > threshold).float().mean()) estimates[dep] = fold_escape(rho) if current * 2 in estimates and estimates[current * 2] - estimates[current] > tau: current = current * 2 elif current // 2 in estimates and estimates[current] - estimates[current // 2] < -tau: current = max(1, current // 2) chosen.append(current) recent = pools[-current:] xx = torch.cat([p[0] for p in recent]).to(device) yy = torch.cat([p[1] for p in recent]).to(device) perm = torch.randperm(len(xx), device=device) net.train(); total = 0.0 for i in range(0, len(xx), BATCH): ix = perm[i:i + BATCH] loss = lossf(net(xx[ix]), yy[ix]) opt.zero_grad(); loss.backward(); opt.step() total += float(loss.detach()) * len(ix) history.append(total / len(xx)) net.eval() with torch.no_grad(): pred = net(ds['xte'].to(device)) metric = float(((pred - ds['yte'].to(device)) ** 2).mean().cpu()) # Trained-model signature: measured rho and observed M-fold # escapes on independent test proposals, not an analytic toy. q = float(torch.quantile(ds['yte'].flatten(), 0.75)) p = pred.flatten() rho_obs = float((p > q).float().mean().cpu()) groups = p[:(len(p) // PROPOSALS) * PROPOSALS].reshape(-1, PROPOSALS) observed_escape = float((groups > q).any(dim=1).float().mean().cpu()) predicted_escape = fold_escape(rho_obs, PROPOSALS) sig = {'rho_hat': rho_obs, 'predicted_escape': predicted_escape, 'observed_group_escape': observed_escape, 'absolute_error': abs(predicted_escape - observed_escape), 'confirmed': abs(predicted_escape - observed_escape) < 0.10, 'mean_depth': float(np.mean(chosen)), 'depth_trajectory': chosen} return (metric, sig) if return_signature else metric except Exception: # Explicit CPU fallback for shared/limited CUDA environments. if device.type == 'cuda': torch.cuda.empty_cache() torch.set_default_device('cpu') return run_system(seed, lr, fixed_depth, adaptive, return_signature) raise def baseline_factory(cfg): return lambda seed: run_system(seed, cfg['lr'], fixed_depth=cfg['depth']) def idea_factory(cfg): return lambda seed: run_system(seed, cfg['lr'], adaptive=True) def main(): math_result = math_check() grid = [{'lr': lr, 'depth': depth} for lr in LR_GRID for depth in DEPTHS] base = sweep_baseline(baseline_factory, grid) best_lr = float(base['best_cfg']['lr']) # The idea grid is exactly the selected baseline lr plus two nearby lr # values, all of which are present in the baseline union grid. nearby = [lr for lr in LR_GRID if lr != best_lr] idea_cfgs = [{'lr': x} for x in [best_lr] + nearby[:2]] idea_trials = [] for cfg in idea_cfgs: r = evaluate(idea_factory(cfg)) idea_trials.append({'cfg': cfg, 'result': r}) best_trial = min(idea_trials, key=lambda z: z['result']['mean']) idea = best_trial['result'] sig_metric, signature = run_system(0, best_trial['cfg']['lr'], adaptive=True, return_signature=True) report = make_report(TRACK, MODEL, base, idea, {'math_sanity': math_result, 'trained_model_signature': signature, 'idea_lr_trials': idea_trials, 'track_choice': 'tabular: replay-memory training intervention; same MLP and regression task in both arms'}) report['idea']['selected_cfg'] = best_trial['cfg'] Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()