"""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): # New model, all fresh weights and optimizer state on every invocation. 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() # Only the selected pressure becomes an observed training example. 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: # A completed trial needs no extra posterior draw after resumption. if not checkpoint_path.exists(): raise RuntimeError('Completed trajectory lacks final trained checkpoint') return checkpoint_path # Final model was fit before the 300th oracle query (299 observations). 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) # This is sequential within one allocation, so no concurrent writers. 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']: # Isolated stream, same fixed candidate pool, repeated queries permitted. 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) # Appendix D.3 specifies a variance ratio, hence the square root here. 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()