| """Paper Algorithm 2, 300 oracle queries x 10 independent trials. |
| |
| Run only inside a case-owned SLURM job. --resume resumes at a query boundary; |
| it never warm-starts a new posterior draw from a previous iteration's model. |
| """ |
| import json |
| import time |
| from pathlib import Path |
|
|
| from common import (case_guard, config_hash, hardware_record, load_config, |
| load_data, model_state_on_cpu, negative_total_flow_rates, parse_args, predict, |
| restore_rng, rng_state, save_json, save_torch, seed_all, |
| write_csv) |
| import torch |
| from model.NOTS import NOTS |
|
|
|
|
| def fit_posterior(cfg, xs, ys, seed, iteration, loss_file): |
| |
| model = NOTS(**cfg['model']).to(xs.device) |
| tc = cfg['training'] |
| if tc['normalization'] != 'none': |
| raise ValueError('Audited paper-priority protocol requires physical-space MSE') |
| optimizer = torch.optim.AdamW(model.parameters(), lr=tc['lr'], weight_decay=0) |
| scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=tc['epochs']) |
| model.train() |
| for epoch in range(tc['epochs']): |
| mse_total = loss_total = 0.0 |
| count = 0 |
| for first in range(0, len(xs), tc['batch_size']): |
| bx, by = xs[first:first + tc['batch_size']], ys[first:first + tc['batch_size']] |
| optimizer.zero_grad(set_to_none=True) |
| pressure = model(bx) |
| mse = (pressure - by).square().mean() |
| penalty = model.l2_penalty() |
| loss = mse + tc['lambda_reg'] * penalty |
| if not torch.isfinite(loss): |
| raise FloatingPointError(f'Nonfinite training loss seed={seed} query={iteration}') |
| loss.backward() |
| optimizer.step() |
| mse_total += mse.item() * len(bx) |
| loss_total += loss.item() * len(bx) |
| count += len(bx) |
| lr = scheduler.get_last_lr()[0] |
| scheduler.step() |
| record = {'seed': seed, 'iteration': iteration, 'epoch': epoch + 1, |
| 'n_observations': len(xs), 'mse': mse_total / count, |
| 'loss': loss_total / count, 'lr': lr} |
| with open(loss_file, 'a') as f: |
| f.write(json.dumps(record) + '\n') |
| print('TRAIN ' + json.dumps(record), flush=True) |
| return model |
|
|
|
|
| @torch.no_grad() |
| def observed_mse(model, xs, ys, batch_size): |
| err = 0.0 |
| for first in range(0, len(xs), batch_size): |
| err += (model(xs[first:first + batch_size]) - ys[first:first + batch_size]).square().sum().item() |
| return err / ys.numel() |
|
|
|
|
| def trial(cfg, seed, x, oracle_y, truth, noise_std, data_meta, resume): |
| out = Path(cfg['paths']['output_dir']) |
| rd = out / 'results' / f'trial_{seed}' |
| rd.mkdir(parents=True, exist_ok=True) |
| state_file = rd / 'resume.pt' |
| fingerprint = config_hash(cfg) |
| seed_all(seed) |
| records, indices, noisy_y = [], [], [] |
| model = None |
| if resume and state_file.exists(): |
| state = torch.load(state_file, map_location='cpu', weights_only=False) |
| if state['config_hash'] != fingerprint or state['data_sha256'] != data_meta['sha256']: |
| raise RuntimeError('Resume config/data differs from original experiment') |
| records, indices = state['records'], state['indices'] |
| noisy_y = list(state['noisy_y'].to(x.device).unbind()) |
| restore_rng(state['rng']) |
| if len(records) == cfg['experiment']['iterations'] and 'final_model_state' in state: |
| model = NOTS(**cfg['model']).to(x.device) |
| model.load_state_dict(state['final_model_state'], strict=True) |
| print(f'RESUME seed={seed} completed_queries={len(records)}', flush=True) |
| elif state_file.exists(): |
| raise FileExistsError(f'{state_file} exists; use --resume to avoid overwriting an experiment') |
| start = time.monotonic() |
| for t in range(len(records), cfg['experiment']['iterations']): |
| tick = time.monotonic() |
| acquisition = None |
| if t < cfg['experiment']['n_init']: |
| chosen = torch.randint(len(x), ()).item() |
| else: |
| fit_x = x[torch.tensor(indices, device=x.device)] |
| fit_y = torch.stack(noisy_y) |
| model = fit_posterior(cfg, fit_x, fit_y, seed, t + 1, rd / 'training_loss.jsonl') |
| values = negative_total_flow_rates(predict(model, x, cfg['experiment']['prediction_batch_size']), x) |
| chosen = values.argmax().item() |
| acquisition = values[chosen].item() |
| |
| observation = oracle_y[chosen] + torch.randn_like(oracle_y[chosen]) * noise_std |
| indices.append(chosen) |
| noisy_y.append(observation) |
| regret = truth.max().item() - truth[chosen].item() |
| cumulative = sum(row['regret'] for row in records) + regret |
| row = {'iteration': t + 1, 'candidate_index': chosen, |
| 'outcome': truth[chosen].item(), |
| 'observation': negative_total_flow_rates(observation.unsqueeze(0), x[chosen:chosen + 1]).item(), |
| 'acquisition': acquisition, 'regret': regret, |
| 'cumulative_regret': cumulative, 'average_regret': cumulative / (t + 1), |
| 'n_observations_used_for_fit': t, |
| 'elapsed_seconds': time.monotonic() - tick} |
| records.append(row) |
| write_csv(rd / 'trajectory.csv', records) |
| boundary_state = {'config_hash': fingerprint, 'data_sha256': data_meta['sha256'], |
| 'records': records, 'indices': indices, |
| 'noisy_y': torch.stack(noisy_y).cpu(), 'rng': rng_state()} |
| if t + 1 == cfg['experiment']['iterations']: |
| boundary_state['final_model_state'] = model_state_on_cpu(model) |
| save_torch(state_file, boundary_state) |
| print(f'QUERY seed={seed} ' + json.dumps(row), flush=True) |
| save_json(rd / 'progress.json', {'seed': seed, 'completed_queries': len(records), |
| 'status': 'running', 'last_metrics': row}) |
| checkpoint_path = out / 'weight' / f'trial_{seed}.pt' |
| if model is None: |
| |
| if not checkpoint_path.exists(): |
| raise RuntimeError('Completed trajectory lacks final trained checkpoint') |
| return checkpoint_path |
| |
| used_indices = indices[:-1] |
| final_train_mse = observed_mse(model, x[torch.tensor(used_indices, device=x.device)], |
| torch.stack(noisy_y[:-1]), cfg['training']['batch_size']) |
| checkpoint = {'model_state': model_state_on_cpu(model), |
| 'model_config': cfg['model'], 'normalization': {'kind': 'identity'}, |
| 'seed': seed, 'iteration': len(records), 'fit_indices': used_indices, |
| 'n_fit_observations': len(used_indices), 'train_mse': final_train_mse, |
| 'config_hash': fingerprint, 'data_sha256': data_meta['sha256'], |
| 'selection_rule': 'minimum observed training MSE among final trial models; test never used'} |
| save_torch(checkpoint_path, checkpoint) |
| |
| best_path = out / 'weight/best_model.pt' |
| old_mse = float('inf') |
| if best_path.exists(): |
| old_mse = torch.load(best_path, map_location='cpu', weights_only=False)['train_mse'] |
| if final_train_mse < old_mse: |
| save_torch(best_path, checkpoint) |
| save_json(rd / 'progress.json', {'seed': seed, 'completed_queries': len(records), |
| 'status': 'completed', 'last_metrics': records[-1], |
| 'checkpoint': str(checkpoint_path), 'train_mse': final_train_mse, |
| 'seconds_this_invocation': time.monotonic() - start}) |
| print(f'TRIAL_COMPLETE seed={seed} average_regret={records[-1]["average_regret"]:.9g} train_mse={final_train_mse:.9g}', flush=True) |
| return checkpoint_path |
|
|
|
|
| def random_baseline(cfg, truth): |
| out = Path(cfg['paths']['output_dir']) / 'results' |
| for seed in cfg['experiment']['seeds']: |
| |
| generator = torch.Generator().manual_seed(seed) |
| choices = torch.randint(len(truth), (cfg['experiment']['iterations'],), generator=generator) |
| regrets = truth.cpu().max() - truth.cpu()[choices] |
| records = [] |
| cumulative = 0.0 |
| for t, (idx, regret) in enumerate(zip(choices.tolist(), regrets.tolist()), 1): |
| cumulative += regret |
| records.append({'iteration': t, 'candidate_index': idx, 'regret': regret, |
| 'cumulative_regret': cumulative, 'average_regret': cumulative / t}) |
| write_csv(out / f'rs_{seed}.csv', records) |
| print(f'RS seed={seed} average_regret={records[-1]["average_regret"]:.9g}', flush=True) |
|
|
|
|
| def main(): |
| parser = parse_args(__doc__) |
| parser.add_argument('--resume', action='store_true') |
| parser.add_argument('--seed', type=int, help='Optional single configured trial for recovery') |
| args = parser.parse_args() |
| cfg = load_config(args.config) |
| device = case_guard(cfg) |
| if cfg['experiment']['iterations'] != 300 or len(cfg['experiment']['seeds']) != 10: |
| raise ValueError('Formal experiment requires 300 queries and 10 independent seeds') |
| x, y, meta = load_data(cfg) |
| out = Path(cfg['paths']['output_dir']) |
| save_json(out / 'results/data_manifest.json', meta) |
| save_json(out / 'results/hardware.json', hardware_record()) |
| x, y = x.to(device), y.to(device) |
| truth = negative_total_flow_rates(y, x) |
| |
| noise_std = (cfg['experiment']['noise_variance_fraction'] * y.var(dim=0, correction=1)).sqrt() |
| save_torch(out / 'results/oracle.pt', {'objectives': truth.cpu(), 'noise_std': noise_std.cpu(), |
| 'best_index': truth.argmax().item(), 'worst_index': truth.argmin().item(), |
| 'data_sha256': meta['sha256']}) |
| print(f'ORACLE max={truth.max().item():.9g} min={truth.min().item():.9g} ' + json.dumps(meta), flush=True) |
| seeds = [args.seed] if args.seed is not None else cfg['experiment']['seeds'] |
| if not set(seeds).issubset(cfg['experiment']['seeds']): |
| raise ValueError('Seed is outside configured independent trials') |
| if cfg['experiment']['run_random_baseline']: |
| random_baseline(cfg, truth) |
| for seed in seeds: |
| trial(cfg, seed, x, y, truth, noise_std, meta, args.resume) |
| print('TRAINING_REQUEST_FINISHED', flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
