path: root/experiments/bci_td_run.py

#!/usr/bin/env python3
"""Run one D4-gated temporal-difference oral-B recovery candidate."""
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 BCIConfig, BCISDIL, generate_trajectories, run_day
from sdil.bci_metrics import annotate_day_events, signature_metrics


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROTOCOL_PATH = os.path.join(ROOT, "ORAL_B_RECOVERY.md")
CONDITIONS = ("intact", "fixed_vectorizer", "plasticity_lesion", "oracle_role")


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 provenance(d4_gate_path):
    def run(command):
        return subprocess.run(
            command, cwd=ROOT, check=True, capture_output=True,
            text=True).stdout.strip()
    runner = os.path.relpath(os.path.abspath(__file__), ROOT)
    protocol = os.path.relpath(PROTOCOL_PATH, ROOT)
    d4_gate = os.path.relpath(os.path.abspath(d4_gate_path), ROOT)
    return {
        "git_commit": run(["git", "rev-parse", "HEAD"]),
        "git_tracked_dirty": bool(run(
            ["git", "status", "--porcelain", "--untracked-files=no"])),
        "runner_and_protocol_tracked": all(subprocess.run(
            ["git", "ls-files", "--error-unmatch", path], cwd=ROOT,
            capture_output=True).returncode == 0
            for path in (runner, protocol)),
        "protocol_sha256": sha256(PROTOCOL_PATH),
        "d4_gate_tracked": subprocess.run(
            ["git", "ls-files", "--error-unmatch", d4_gate], cwd=ROOT,
            capture_output=True).returncode == 0,
        "d4_gate_sha256": sha256(d4_gate_path),
    }


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 build_config(eta):
    return BCIConfig(
        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=eta, perturb_sigma=0.03, perturb_every=4,
        kappa=0.0, feedback="performance_velocity")


def neutral_warmup(initial, task_seed, model_seed, condition, batches=100):
    """Fit neutral coupling and, when allowed, the cursor-role vectorizer."""
    model = initial.clone()
    generator = torch.Generator(device="cpu").manual_seed(
        400_000 + 100 * task_seed + model_seed)
    active = torch.ones(model.cfg.episodes_per_day, dtype=torch.bool)
    role_cursor_observations = 0
    for _ in range(batches):
        soma = torch.randn(
            model.cfg.episodes_per_day, model.cfg.n_neurons,
            generator=generator).tanh()
        ordinary = model.coupling * soma
        model.update_predictor(soma, ordinary, active)
        xi = torch.empty_like(soma).bernoulli_(
            0.5, generator=generator).mul_(2).sub_(1)
        sigma = model.cfg.perturb_sigma
        # The environment produces two scalar cursor observations. The role
        # target function cannot access the environment's role map.
        cursor_plus = (soma + sigma * xi) @ model.role
        cursor_minus = (soma - sigma * xi) @ model.role
        role_cursor_observations += 2 * soma.shape[0]
        target = model.causal_role_targets(cursor_plus, cursor_minus, xi)
        if condition == "intact" or condition == "plasticity_lesion":
            model.update_vectorizer(
                torch.ones(soma.shape[0], 1),
                torch.zeros_like(soma), target, active)
    if condition == "oracle_role":
        # Explicit diagnostic ceiling; never eligible as the selected learner.
        model.A[:, 0].copy_(model.role)
    predictor_error = float((model.P - model.coupling).abs().max())
    role_cosine = float(torch.nn.functional.cosine_similarity(
        model.A[:, 0], model.role, dim=0).item())
    return model, {
        "batches": batches,
        "examples": batches * model.cfg.episodes_per_day,
        "instruction_present": False,
        "role_cursor_scalar_observations": role_cursor_observations,
        "predictor_max_abs_error": predictor_error,
        "role_cosine_after_warmup": role_cosine,
        "rng_seed": 400_000 + 100 * task_seed + model_seed,
    }


def condition_settings(name):
    return {
        "intact": {"plasticity_gain": 1.0, "learn_vectorizer": True},
        "fixed_vectorizer": {
            "plasticity_gain": 1.0, "learn_vectorizer": False},
        "plasticity_lesion": {
            "plasticity_gain": 0.0, "learn_vectorizer": True},
        "oracle_role": {"plasticity_gain": 1.0, "learn_vectorizer": False},
    }[name]


def train(model, trajectories, condition, collect):
    settings = condition_settings(condition)
    daily_success = []
    events = []
    global_step = 0
    start = time.perf_counter()
    for day in range(model.cfg.days):
        report = run_day(
            model, trajectories, day, global_step=global_step,
            control_gain=0.0, plasticity_gain=settings["plasticity_gain"],
            learn_vectorizer=settings["learn_vectorizer"],
            learn_predictor=True, collect=collect)
        global_step = report["global_step"]
        daily_success.append(report["success_rate"])
        if collect:
            annotate_day_events(
                report["events"], day, report["success"],
                episode_offset=day * model.cfg.episodes_per_day)
            events.extend(report["events"])
    perturbation_events = sum(
        1 for step in range(model.cfg.days * model.cfg.steps_per_episode)
        if step % model.cfg.perturb_every == 0)
    conservative_cursor_observations = (
        2 * model.cfg.episodes_per_day * perturbation_events
        if settings["learn_vectorizer"] else 0)
    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() - start,
        "role_cosine_after_training": float(
            torch.nn.functional.cosine_similarity(
                model.A[:, 0], model.role, dim=0).item()),
        "cost": {
            "ordinary_state_episode_steps": (
                model.cfg.days * model.cfg.episodes_per_day
                * model.cfg.steps_per_episode),
            "online_role_perturbation_events": perturbation_events,
            "conservative_online_cursor_scalar_observations": (
                conservative_cursor_observations),
            "ordinary_performance_change_observations": (
                model.cfg.days * model.cfg.episodes_per_day
                * model.cfg.steps_per_episode),
        },
    }


def evaluate(model, trajectories, collect):
    report = run_day(
        model.clone(), trajectories, 0, control_gain=0.0,
        plasticity_gain=0.0, learn_vectorizer=False, learn_predictor=False,
        collect=collect, global_step=0)
    if collect:
        annotate_day_events(report["events"], 0, report["success"], 1_000_000)
    return report


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--d4_gate", default="results/kp_dynamic_projection_confirmation_gate.json")
    parser.add_argument("--task-seed", type=int, choices=(0, 1, 2), required=True)
    parser.add_argument("--model-seed", type=int, choices=(0,), default=0)
    parser.add_argument("--eta", type=float, choices=(0.03, 0.1), required=True)
    parser.add_argument("--outdir", default="results/bci_td_dev")
    args = parser.parse_args()
    with open(args.d4_gate) as handle:
        d4 = json.load(handle)
    if (d4.get("protocol") !=
            "kp_dynamic_neutral_projection_confirmation_v1"
            or d4.get("status") != "passed"
            or d4.get("review_score_after") != 7):
        raise ValueError("oral-B R1 requires the complete audited D4 pass")
    source = provenance(args.d4_gate)
    if (source["git_tracked_dirty"]
            or not source["runner_and_protocol_tracked"]
            or not source["d4_gate_tracked"]):
        raise RuntimeError(
            "R1 requires clean tracked source, protocol, and D4 gate")

    cfg = build_config(args.eta)
    trajectories = generate_trajectories(cfg, args.task_seed)
    evaluation = generate_trajectories(
        cfg, args.task_seed + 200_000, days=1, episodes=256)
    initial = BCISDIL(cfg, args.model_seed)
    conditions = {}
    trained = {}
    training_events = None
    warmups = {}
    started = time.perf_counter()
    for name in CONDITIONS:
        model, warmup = neutral_warmup(
            initial, args.task_seed, args.model_seed, name)
        events, report = train(
            model, trajectories, name, collect=name == "intact")
        warmups[name] = warmup
        trained[name] = model
        conditions[name] = report
        if name == "intact":
            training_events = events
    evaluation_events = None
    for name in CONDITIONS:
        report = evaluate(
            trained[name], evaluation, collect=name == "intact")
        conditions[name]["final_success"] = report["success_rate"]
        if name == "intact":
            evaluation_events = report["events"]
    signatures = signature_metrics(
        training_events, evaluation_events, cfg, trained["intact"].role)
    result = {
        "schema_version": 1,
        "protocol": {
            "name": "oral_b_td_development_v1",
            "selection_split": "development",
            "confirmation_seeds_touched": False,
            "evaluation_task_seed": args.task_seed + 200_000,
            "d4_gate_source": args.d4_gate,
            "d4_gate_sha256": source["d4_gate_sha256"],
        },
        "args": vars(args), "config": vars(cfg), "provenance": source,
        "warmup": warmups, "conditions": conditions,
        "signatures": signatures,
        "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(),
        },
    }
    result["finite"] = finite_tree(result)
    if not result["finite"]:
        raise RuntimeError("non-finite oral-B R1 record")
    os.makedirs(args.outdir, exist_ok=True)
    eta_tag = str(args.eta).replace(".", "p")
    path = os.path.join(
        args.outdir, f"bci_td_v1_eta{eta_tag}_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, "finite": result["finite"],
        "intact_final": conditions["intact"]["final_success"],
        "sign_inversion": signatures["causal_role_sign_inversion_index"],
    }, indent=2))


if __name__ == "__main__":
    main()