protocol: freeze independent oral-B-v2 gates

1 files changed, 546 insertions, 0 deletions

diff --git a/experiments/bci_v2_run.py b/experiments/bci_v2_run.py
new file mode 100644
index 0000000..5edebad
--- /dev/null
+++ b/experiments/bci_v2_run.py
@@ -0,0 +1,546 @@
+#!/usr/bin/env python3
+"""Run one frozen oral-B-v2 development cell."""
+import argparse
+import hashlib
+import json
+import math
+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 sdil.bci import generate_trajectories
+from sdil.bci_v2 import BCIV2, BCIV2Config, run_day_v2
+from sdil.bci_v2_metrics import annotate_events, signature_metrics_v2
+
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+PROTOCOL_PATH = os.path.join(ROOT, "ORAL_B_V2.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"
+)
+TASK_SEEDS = (20, 21, 22)
+MODEL_SEEDS = (0,)
+FORWARD_ETAS = (0.03, 0.1)
+GAMMAS = (0.8, 0.95)
+CRITIC_ETAS = (0.01, 0.03)
+HORIZONS = (4, 8, 12, 16, 20, 24, 28)
+CHALLENGE_EPISODES = 128
+CONDITIONS = (
+    "intact",
+    "fixed_role",
+    "plasticity_lesion",
+    "critic_training_lesion",
+    "outcome_training_lesion",
+    "oracle_role",
+)
+ACUTE_MODES = (
+    "intact",
+    "acute_critic_lesion",
+    "acute_outcome_lesion",
+)
+
+
+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 finite_tree(value):
+    if isinstance(value, dict):
+        return all(finite_tree(item) for item in value.values())
+    if isinstance(value, (list, tuple)):
+        return all(finite_tree(item) for item in value)
+    if isinstance(value, (int, float)):
+        return math.isfinite(value)
+    return True
+
+
+def source_paths():
+    return {
+        "runner": os.path.abspath(__file__),
+        "development_analyzer": os.path.join(
+            ROOT, "experiments", "analyze_bci_v2_development.py"
+        ),
+        "confirmation_runner": os.path.join(
+            ROOT, "experiments", "bci_v2_confirmation.py"
+        ),
+        "confirmation_analyzer": os.path.join(
+            ROOT, "experiments", "analyze_bci_v2_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"),
+        "protocol": PROTOCOL_PATH,
+        "d4_gate": D4_GATE_PATH,
+        "old_r2_gate": OLD_R2_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)
+    if not (
+        d4.get("protocol")
+        == "kp_dynamic_neutral_projection_confirmation_v1"
+        and d4.get("status") == "passed"
+        and d4.get("review_score_after") == 7
+    ):
+        raise ValueError("oral-B-v2 requires the complete D4 accept gate")
+    if not (
+        old_r2.get("protocol") == "oral_b_td_confirmation_v1"
+        and old_r2.get("status") == "failed"
+        and old_r2.get("review_score_after") == 7
+    ):
+        raise ValueError("oral-B-v2 preserves and requires the failed old R2")
+    return d4, old_r2
+
+
+def build_config(forward_eta, gamma, critic_eta):
+    return BCIV2Config(
+        n_plus=5,
+        n_minus=5,
+        n_background=30,
+        context_dim=16,
+        steps_per_episode=28,
+        episodes_per_day=64,
+        days=14,
+        target=0.8,
+        inertia=0.65,
+        process_noise=0.12,
+        context_ar=0.8,
+        coupling_scale=1.0,
+        predictor_eta=0.2,
+        vectorizer_eta=0.03,
+        forward_eta=forward_eta,
+        perturb_sigma=0.03,
+        perturb_every=4,
+        kappa=0.0,
+        feedback="performance_velocity",
+        gamma=gamma,
+        critic_eta=critic_eta,
+        eligibility_decay=0.8,
+        velocity_reward_scale=0.25,
+        terminal_reward=1.0,
+    )
+
+
+def condition_settings(name):
+    return {
+        "intact": {},
+        "fixed_role": {"learn_role": False},
+        "plasticity_lesion": {"plasticity_gain": 0.0},
+        "critic_training_lesion": {
+            "learn_critic": False,
+            "critic_enabled": False,
+        },
+        "outcome_training_lesion": {
+            "terminal_outcome_enabled": False,
+        },
+        "oracle_role": {"learn_role": False},
+    }[name]
+
+
+def neutral_warmup(initial, task_seed, model_seed, condition, batches=100):
+    """Fit neutral traffic and expose every condition to paired role probes."""
+    model = initial.clone()
+    generator = torch.Generator(device="cpu").manual_seed(
+        600_000 + 100 * task_seed + model_seed
+    )
+    active = torch.ones(model.cfg.episodes_per_day, dtype=torch.bool)
+    update_role = condition not in ("fixed_role", "oracle_role")
+    scalar_observations = 0
+    for _ in range(batches):
+        soma = torch.randn(
+            model.cfg.episodes_per_day,
+            model.cfg.n_neurons,
+            generator=generator,
+        ).tanh()
+        model.update_predictor(
+            soma, model.coupling * soma, active
+        )
+        perturbations = torch.empty_like(soma).bernoulli_(
+            0.5, generator=generator
+        ).mul_(2).sub_(1)
+        sigma = model.cfg.perturb_sigma
+        cursor_plus = (soma + sigma * perturbations) @ model.role
+        cursor_minus = (soma - sigma * perturbations) @ model.role
+        scalar_observations += 2 * soma.shape[0]
+        target = model.causal_role_targets(
+            cursor_plus, cursor_minus, perturbations
+        )
+        if update_role:
+            model.update_role(target, active)
+    if condition == "oracle_role":
+        model.A.copy_(model.role)
+    return model, {
+        "batches": batches,
+        "examples": batches * model.cfg.episodes_per_day,
+        "instruction_present": False,
+        "role_cursor_scalar_observations": scalar_observations,
+        "role_update_enabled": update_role,
+        "predictor_max_abs_error": float(
+            (model.P - model.coupling).abs().max()
+        ),
+        "role_cosine_after_warmup": float(
+            torch.nn.functional.cosine_similarity(
+                model.A, model.role, dim=0
+            ).item()
+        ),
+        "rng_seed": 600_000 + 100 * task_seed + model_seed,
+    }
+
+
+def train(model, trajectories, condition, collect):
+    settings = condition_settings(condition)
+    daily_success = []
+    events = []
+    global_step = 0
+    active_transitions = 0
+    role_probe_examples = 0
+    started = time.perf_counter()
+    for day in range(model.cfg.days):
+        report = run_day_v2(
+            model,
+            trajectories,
+            day,
+            global_step=global_step,
+            collect=collect,
+            learn_predictor=True,
+            probe_role=True,
+            **settings,
+        )
+        global_step = report["global_step"]
+        active_transitions += report["active_transitions"]
+        role_probe_examples += report["role_probe_examples"]
+        daily_success.append(report["success_rate"])
+        if collect:
+            annotate_events(
+                report["events"],
+                day,
+                report["success"],
+                day * model.cfg.episodes_per_day,
+                model.cfg.steps_per_episode,
+            )
+            events.extend(report["events"])
+    return events, {
+        "daily_success": daily_success,
+        "early_success": sum(daily_success[:3]) / 3,
+        "late_success": sum(daily_success[-3:]) / 3,
+        "learning_gain": (
+            sum(daily_success[-3:])
+            - sum(daily_success[:3])
+        ) / 3,
+        "training_wall_s": time.perf_counter() - started,
+        "role_cosine_after_training": float(
+            torch.nn.functional.cosine_similarity(
+                model.A, model.role, dim=0
+            ).item()
+        ),
+        "critic_l2_after_training": float(model.critic.norm().item()),
+        "cost": {
+            "maximum_state_episode_steps": (
+                model.cfg.days
+                * model.cfg.episodes_per_day
+                * model.cfg.steps_per_episode
+            ),
+            "active_state_episode_steps": active_transitions,
+            "role_probe_examples": role_probe_examples,
+            "cursor_scalar_observations": 2 * role_probe_examples,
+            "terminal_outcome_observations": (
+                model.cfg.days * model.cfg.episodes_per_day
+            ),
+            "task_loss_queries": 0,
+            "reverse_mode_calls": 0,
+        },
+    }
+
+
+def evaluate_performance(model, trajectories):
+    return run_day_v2(
+        model.clone(),
+        trajectories,
+        0,
+        horizon=model.cfg.steps_per_episode,
+        plasticity_gain=0.0,
+        learn_role=False,
+        probe_role=False,
+        learn_predictor=False,
+        learn_critic=False,
+        collect=False,
+    )
+
+
+def evaluate_challenge_ladder(model, task_seed, seed_offset):
+    """Paired acute readouts over a fixed, non-selected horizon ladder."""
+    mode_events = {name: [] for name in ACUTE_MODES}
+    costs = {name: 0 for name in ACUTE_MODES}
+    seeds = {}
+    for horizon_index, horizon in enumerate(HORIZONS):
+        trajectory_seed = (
+            seed_offset + 1_000 * task_seed + horizon_index
+        )
+        seeds[str(horizon)] = trajectory_seed
+        trajectories = generate_trajectories(
+            model.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]
+            report = run_day_v2(
+                model.clone(),
+                trajectories,
+                0,
+                horizon=horizon,
+                plasticity_gain=0.0,
+                learn_role=False,
+                probe_role=False,
+                learn_predictor=False,
+                learn_critic=False,
+                collect=True,
+                **settings,
+            )
+            episode_offset = (
+                1_000_000 + horizon_index * CHALLENGE_EPISODES
+            )
+            annotate_events(
+                report["events"],
+                0,
+                report["success"],
+                episode_offset,
+                horizon,
+            )
+            mode_events[mode].extend(report["events"])
+            costs[mode] += report["active_transitions"]
+    return mode_events, {
+        "horizons": list(HORIZONS),
+        "episodes_per_horizon": CHALLENGE_EPISODES,
+        "trajectory_seeds": seeds,
+        "active_state_episode_steps_by_mode": costs,
+        "selection_over_horizons": False,
+    }
+
+
+def run_cell(
+    task_seed,
+    model_seed,
+    forward_eta,
+    gamma,
+    critic_eta,
+    *,
+    split,
+    performance_seed_offset,
+    challenge_seed_offset,
+):
+    cfg = build_config(forward_eta, gamma, critic_eta)
+    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_challenge_ladder(
+        trained["intact"], task_seed, challenge_seed_offset
+    )
+    signatures = signature_metrics_v2(
+        training_events,
+        challenge_events,
+        cfg,
+        trained["intact"].role,
+    )
+    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(
+        "--forward-eta", type=float, choices=FORWARD_ETAS, required=True
+    )
+    parser.add_argument(
+        "--gamma", type=float, choices=GAMMAS, required=True
+    )
+    parser.add_argument(
+        "--critic-eta", type=float, choices=CRITIC_ETAS, required=True
+    )
+    parser.add_argument("--outdir", default="results/bci_v2_dev")
+    args = parser.parse_args()
+    require_parent_gates()
+    source = provenance()
+    if (
+        source["git_tracked_dirty"]
+        or not all(source["tracked_inputs"].values())
+    ):
+        raise RuntimeError(
+            "oral-B-v2 R1 requires clean, tracked, frozen inputs"
+        )
+    cell = run_cell(
+        args.task_seed,
+        args.model_seed,
+        args.forward_eta,
+        args.gamma,
+        args.critic_eta,
+        split="development",
+        performance_seed_offset=400_000,
+        challenge_seed_offset=410_000,
+    )
+    result = {
+        "schema_version": 2,
+        "protocol": {
+            "name": "oral_b_v2_development_v1",
+            "selection_split": "development",
+            "confirmation_seeds_touched": False,
+            "grid_size": (
+                len(TASK_SEEDS)
+                * len(FORWARD_ETAS)
+                * len(GAMMAS)
+                * len(CRITIC_ETAS)
+            ),
+            "protocol_sha256": source["input_sha256"]["protocol"],
+            "d4_gate_sha256": source["input_sha256"]["d4_gate"],
+            "old_r2_gate_sha256": source["input_sha256"]["old_r2_gate"],
+        },
+        "args": vars(args),
+        "provenance": source,
+        **cell,
+    }
+    result["finite"] = finite_tree(result)
+    if not result["finite"]:
+        raise RuntimeError("non-finite oral-B-v2 development record")
+    os.makedirs(args.outdir, exist_ok=True)
+    eta = str(args.forward_eta).replace(".", "p")
+    gamma = str(args.gamma).replace(".", "p")
+    critic = str(args.critic_eta).replace(".", "p")
+    path = os.path.join(
+        args.outdir,
+        (
+            f"bci_v2_e{eta}_g{gamma}_c{critic}"
+            f"_t{args.task_seed}_m{args.model_seed}.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"],
+        "terminal_outcome_accuracy": result["signatures"][
+            "terminal_residual_outcome_balanced_acc"
+        ],
+        "finite": result["finite"],
+    }, indent=2))
+
+
+if __name__ == "__main__":
+    main()