diff options
Diffstat (limited to 'experiments/bci_v2_recovery_run.py')
| -rw-r--r-- | experiments/bci_v2_recovery_run.py | 377 |
1 files changed, 377 insertions, 0 deletions
diff --git a/experiments/bci_v2_recovery_run.py b/experiments/bci_v2_recovery_run.py new file mode 100644 index 0000000..29507b3 --- /dev/null +++ b/experiments/bci_v2_recovery_run.py @@ -0,0 +1,377 @@ +#!/usr/bin/env python3 +"""Run one frozen oral-B-v2 cold-start recovery development cell.""" +import argparse +from dataclasses import replace +import hashlib +import json +import os +import platform +import resource +import subprocess +import sys +import time + +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from experiments.bci_v2_run import ( + ACUTE_MODES, + CONDITIONS, + build_config, + evaluate_performance, + finite_tree, + neutral_warmup, + train, +) +from sdil.bci import generate_trajectories +from sdil.bci_v2 import BCIV2, run_day_v2 +from sdil.bci_v2_recovery_metrics import ( + annotate_target_events, + recovery_signature_metrics, +) + + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +PROTOCOL_PATH = os.path.join(ROOT, "ORAL_B_V2_RECOVERY.md") +D4_GATE_PATH = os.path.join( + ROOT, "results", "kp_dynamic_projection_confirmation_gate.json" +) +OLD_R2_GATE_PATH = os.path.join( + ROOT, "results", "bci_td_confirmation_gate.json" +) +FAILED_V2_GATE_PATH = os.path.join( + ROOT, "results", "bci_v2_dev_gate.json" +) +TASK_SEEDS = (23, 24, 25) +MODEL_SEEDS = (0,) +TARGETS = (1.55, 1.60, 1.65, 1.70, 1.75, 1.80) +CHALLENGE_EPISODES = 128 +FIXED_CONFIG = { + "forward_eta": 0.1, + "gamma": 0.8, + "critic_eta": 0.03, + "velocity_reward_scale": 1.0, +} + + +def sha256(path): + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def source_paths(): + return { + "runner": os.path.abspath(__file__), + "development_analyzer": os.path.join( + ROOT, + "experiments", + "analyze_bci_v2_recovery_development.py", + ), + "confirmation_runner": os.path.join( + ROOT, "experiments", "bci_v2_recovery_confirmation.py" + ), + "confirmation_analyzer": os.path.join( + ROOT, + "experiments", + "analyze_bci_v2_recovery_confirmation.py", + ), + "base_dynamics": os.path.join(ROOT, "sdil", "bci.py"), + "v2_dynamics": os.path.join(ROOT, "sdil", "bci_v2.py"), + "v2_metrics": os.path.join(ROOT, "sdil", "bci_v2_metrics.py"), + "recovery_metrics": os.path.join( + ROOT, "sdil", "bci_v2_recovery_metrics.py" + ), + "protocol": PROTOCOL_PATH, + "d4_gate": D4_GATE_PATH, + "old_r2_gate": OLD_R2_GATE_PATH, + "failed_v2_gate": FAILED_V2_GATE_PATH, + } + + +def provenance(extra_paths=None): + def run(command): + return subprocess.run( + command, + cwd=ROOT, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + paths = source_paths() + if extra_paths: + paths.update(extra_paths) + relative = { + name: os.path.relpath(path, ROOT) + for name, path in paths.items() + } + tracked = { + name: subprocess.run( + ["git", "ls-files", "--error-unmatch", path], + cwd=ROOT, + capture_output=True, + ).returncode == 0 + for name, path in relative.items() + } + return { + "git_commit": run(["git", "rev-parse", "HEAD"]), + "git_tracked_dirty": bool(run([ + "git", "status", "--porcelain", "--untracked-files=no" + ])), + "tracked_inputs": tracked, + "input_sha256": { + name: sha256(path) for name, path in paths.items() + }, + } + + +def require_parent_gates(): + with open(D4_GATE_PATH) as handle: + d4 = json.load(handle) + with open(OLD_R2_GATE_PATH) as handle: + old_r2 = json.load(handle) + with open(FAILED_V2_GATE_PATH) as handle: + failed_v2 = json.load(handle) + if not ( + d4.get("status") == "passed" + and d4.get("review_score_after") == 7 + and old_r2.get("status") == "failed" + and old_r2.get("review_score_after") == 7 + and failed_v2.get("protocol") == "oral_b_v2_development_v1" + and failed_v2.get("status") == "failed" + and failed_v2.get("complete_grid") is True + and failed_v2.get("oral_b_v2_confirmation_opened") is False + and failed_v2.get("review_score_after") == 7 + ): + raise ValueError( + "cold-start recovery requires D4 and both preserved failures" + ) + + +def build_recovery_config(): + return replace( + build_config( + FIXED_CONFIG["forward_eta"], + FIXED_CONFIG["gamma"], + FIXED_CONFIG["critic_eta"], + ), + velocity_reward_scale=FIXED_CONFIG[ + "velocity_reward_scale" + ], + ) + + +def evaluate_target_ladder(model, task_seed, seed_offset): + mode_events = {name: [] for name in ACUTE_MODES} + costs = {name: 0 for name in ACUTE_MODES} + seeds = {} + for target_index, target in enumerate(TARGETS, start=1): + challenge_cfg = replace(model.cfg, target=target) + trajectory_seed = ( + seed_offset + 1_000 * task_seed + target_index - 1 + ) + seeds[str(target)] = trajectory_seed + trajectories = generate_trajectories( + challenge_cfg, + trajectory_seed, + days=1, + episodes=CHALLENGE_EPISODES, + ) + for mode in ACUTE_MODES: + settings = { + "intact": {}, + "acute_critic_lesion": {"critic_enabled": False}, + "acute_outcome_lesion": { + "terminal_outcome_enabled": False + }, + }[mode] + assay_model = model.clone() + assay_model.cfg = challenge_cfg + report = run_day_v2( + assay_model, + trajectories, + 0, + horizon=challenge_cfg.steps_per_episode, + plasticity_gain=0.0, + learn_role=False, + probe_role=False, + learn_predictor=False, + learn_critic=False, + collect=True, + **settings, + ) + episode_offset = ( + 2_000_000 + + (target_index - 1) * CHALLENGE_EPISODES + ) + annotate_target_events( + report["events"], + 0, + report["success"], + episode_offset, + target_index, + ) + mode_events[mode].extend(report["events"]) + costs[mode] += report["active_transitions"] + return mode_events, { + "targets": list(TARGETS), + "episodes_per_target": CHALLENGE_EPISODES, + "trajectory_seeds": seeds, + "active_state_episode_steps_by_mode": costs, + "selection_over_targets": False, + "maximum_steps_per_episode": model.cfg.steps_per_episode, + } + + +def run_cell( + task_seed, + model_seed, + *, + split, + performance_seed_offset, + challenge_seed_offset, +): + cfg = build_recovery_config() + trajectories = generate_trajectories(cfg, task_seed) + performance_seed = performance_seed_offset + task_seed + performance_trajectories = generate_trajectories( + cfg, performance_seed, days=1, episodes=256 + ) + initial = BCIV2(cfg, model_seed) + conditions = {} + trained = {} + warmups = {} + training_events = None + started = time.perf_counter() + for name in CONDITIONS: + model, warmup = neutral_warmup( + initial, task_seed, model_seed, name + ) + events, report = train( + model, trajectories, name, collect=name == "intact" + ) + warmups[name] = warmup + conditions[name] = report + trained[name] = model + if name == "intact": + training_events = events + for name in CONDITIONS: + report = evaluate_performance( + trained[name], performance_trajectories + ) + conditions[name]["final_success"] = report["success_rate"] + conditions[name]["evaluation_active_state_episode_steps"] = ( + report["active_transitions"] + ) + challenge_events, challenge_cost = evaluate_target_ladder( + trained["intact"], task_seed, challenge_seed_offset + ) + signatures = recovery_signature_metrics( + training_events, + challenge_events, + cfg, + trained["intact"].role, + TARGETS, + ) + return { + "config": vars(cfg), + "warmup": warmups, + "conditions": conditions, + "signatures": signatures, + "assays": { + "performance_evaluation_seed": performance_seed, + "performance_evaluation_episodes": 256, + "challenge": challenge_cost, + }, + "wall_s": time.perf_counter() - started, + "peak_rss_mib": ( + resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 + ), + "hardware": { + "device": "cpu", + "platform": platform.platform(), + "torch_version": torch.__version__, + "threads": torch.get_num_threads(), + }, + "split": split, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--task-seed", type=int, choices=TASK_SEEDS, required=True + ) + parser.add_argument( + "--model-seed", type=int, choices=MODEL_SEEDS, default=0 + ) + parser.add_argument( + "--outdir", default="results/bci_v2_recovery_dev" + ) + args = parser.parse_args() + require_parent_gates() + source = provenance() + if ( + source["git_tracked_dirty"] + or not all(source["tracked_inputs"].values()) + ): + raise RuntimeError( + "v2 recovery R1 requires clean, tracked, frozen inputs" + ) + cell = run_cell( + args.task_seed, + args.model_seed, + split="development", + performance_seed_offset=460_000, + challenge_seed_offset=470_000, + ) + result = { + "schema_version": 3, + "protocol": { + "name": "oral_b_v2_cold_start_recovery_development_v1", + "selection_split": "development_validation", + "hyperparameter_selection": False, + "confirmation_seeds_touched": False, + "fixed_config": FIXED_CONFIG, + "protocol_sha256": source["input_sha256"]["protocol"], + "d4_gate_sha256": source["input_sha256"]["d4_gate"], + "old_r2_gate_sha256": source["input_sha256"]["old_r2_gate"], + "failed_v2_gate_sha256": source["input_sha256"][ + "failed_v2_gate" + ], + }, + "args": vars(args), + "provenance": source, + **cell, + } + result["finite"] = finite_tree(result) + if not result["finite"]: + raise RuntimeError("non-finite v2 recovery development record") + os.makedirs(args.outdir, exist_ok=True) + path = os.path.join( + args.outdir, + f"bci_v2_recovery_t{args.task_seed}_m0.json", + ) + if os.path.exists(path): + raise FileExistsError(f"refusing to overwrite {path}") + with open(path, "w") as handle: + json.dump(result, handle, indent=2, sort_keys=True) + handle.write("\n") + print(json.dumps({ + "path": path, + "intact_final": result["conditions"]["intact"]["final_success"], + "challenge_success_fraction": result["signatures"][ + "challenge_success_fraction" + ], + "terminal_outcome_accuracy": result["signatures"][ + "terminal_residual_outcome_balanced_acc" + ], + "finite": result["finite"], + }, indent=2)) + + +if __name__ == "__main__": + main() |