path: root/experiments/bci_run.py

#!/usr/bin/env python3
"""Run one paired continuous-BCI candidate from the frozen oral-B protocol."""
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__)))
ORAL_B_PATH = os.path.join(ROOT, "ORAL_B.md")


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_provenance():
    commit = subprocess.run(
        ["git", "rev-parse", "HEAD"], cwd=ROOT, check=True,
        capture_output=True, text=True).stdout.strip()
    dirty = bool(subprocess.run(
        ["git", "status", "--porcelain", "--untracked-files=no"], cwd=ROOT,
        check=True, capture_output=True, text=True).stdout.strip())
    runner_path = os.path.relpath(os.path.abspath(__file__), ROOT)
    protocol_path = os.path.relpath(ORAL_B_PATH, ROOT)
    tracked = all(subprocess.run(
        ["git", "ls-files", "--error-unmatch", path], cwd=ROOT,
        capture_output=True).returncode == 0 for path in (runner_path, protocol_path))
    return {"git_commit": commit, "git_dirty": dirty,
            "runner_and_protocol_tracked": tracked,
            "oral_b_protocol_sha256": sha256(ORAL_B_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 condition_settings(name):
    settings = {
        "intact": dict(control_gain=1.0, plasticity_gain=1.0,
                       learn_vectorizer=True),
        "fixed_vectorizer": dict(control_gain=1.0, plasticity_gain=1.0,
                                 learn_vectorizer=False),
        "online_lesion": dict(control_gain=0.0, plasticity_gain=1.0,
                              learn_vectorizer=True),
        "plasticity_lesion": dict(control_gain=1.0, plasticity_gain=0.0,
                                  learn_vectorizer=True),
        "both_lesion": dict(control_gain=0.0, plasticity_gain=0.0,
                            learn_vectorizer=True),
    }
    return settings[name]


def train_condition(initial, trajectories, name, *, collect):
    model = initial.clone()
    settings = condition_settings(name)
    daily_success = []
    events = []
    global_step = 0
    perturbation_events = 0
    perturbation_scalar_rewards = 0
    start = time.perf_counter()
    for day in range(model.cfg.days):
        report = run_day(
            model, trajectories, day, collect=collect, global_step=global_step,
            control_gain=settings["control_gain"],
            plasticity_gain=settings["plasticity_gain"],
            learn_vectorizer=settings["learn_vectorizer"],
            learn_predictor=True)
        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"])
            for event in report["events"]:
                if event.get("causal_target") is not None:
                    perturbation_events += 1
                    perturbation_scalar_rewards += 2 * int(event["active"].sum())
        elif settings["learn_vectorizer"]:
            # One batched event every perturb_every temporal steps.
            perturbation_events += sum(
                1 for step in range(model.cfg.steps_per_episode)
                if ((day * model.cfg.steps_per_episode + step)
                    % model.cfg.perturb_every == 0))
            # The exact active count is unavailable without collection. Report
            # the conservative full-batch count for non-primary conditions.
            perturbation_scalar_rewards += (
                2 * model.cfg.episodes_per_day
                * sum(1 for step in range(model.cfg.steps_per_episode)
                      if ((day * model.cfg.steps_per_episode + step)
                          % model.cfg.perturb_every == 0)))
    elapsed = time.perf_counter() - start
    early = sum(daily_success[:3]) / 3
    late = sum(daily_success[-3:]) / 3
    ordinary = (model.cfg.days * model.cfg.episodes_per_day
                * model.cfg.steps_per_episode)
    return model, events, {
        "daily_success": daily_success,
        "early_success": early,
        "late_success": late,
        "learning_gain": late - early,
        "training_wall_s": elapsed,
        "cost": {
            "ordinary_state_episode_steps": ordinary,
            "perturbation_events": perturbation_events,
            "causal_scalar_reward_observations": perturbation_scalar_rewards,
            "forward_equivalent_episode_steps": (
                ordinary + perturbation_scalar_rewards / model.cfg.context_dim),
            "forward_equivalent_definition": (
                "ordinary full state step plus each perturbed cursor/loss readout "
                "at 1/context_dim of a context-to-soma state step"),
        },
    }


def evaluate(model, trajectories, control_gain, *, collect):
    # Clone prevents even accidental evaluator mutation from changing another
    # paired phase condition.
    evaluated = model.clone()
    start = time.perf_counter()
    report = run_day(
        evaluated, trajectories, 0, collect=collect, global_step=0,
        control_gain=control_gain, plasticity_gain=0.0,
        learn_vectorizer=False, learn_predictor=False)
    wall = time.perf_counter() - start
    if collect:
        annotate_day_events(report["events"], 0, report["success"], 1_000_000)
    return report, wall


def build_config(args):
    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=args.eta, perturb_sigma=0.03, perturb_every=4,
        kappa=args.kappa, feedback=args.feedback)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--task-seed", type=int, required=True)
    parser.add_argument("--model-seed", type=int, required=True)
    parser.add_argument("--feedback", choices=("error", "error_velocity"), required=True)
    parser.add_argument("--kappa", type=float, choices=(0.0, 0.1, 0.3), required=True)
    parser.add_argument("--eta", type=float, choices=(0.01, 0.03), required=True)
    parser.add_argument("--tag", required=True)
    parser.add_argument("--outdir", default="results/bci_dev")
    args = parser.parse_args()
    if args.task_seed not in (0, 1, 2):
        raise ValueError("development runner permits only frozen task seeds 0,1,2")
    if args.model_seed != 0:
        raise ValueError("development runner permits only frozen model seed 0")

    provenance = source_provenance()
    if provenance["git_dirty"] or not provenance["runner_and_protocol_tracked"]:
        raise RuntimeError(
            "refusing to generate oral-B evidence from dirty or untracked source")
    cfg = build_config(args)
    train_trajectories = generate_trajectories(cfg, args.task_seed)
    evaluation_trajectories = generate_trajectories(
        cfg, args.task_seed + 100_000, days=1, episodes=256)
    initial = BCISDIL(cfg, args.model_seed)
    started = time.perf_counter()

    trained = {}
    training_events = None
    condition_reports = {}
    for name in ("intact", "fixed_vectorizer", "online_lesion",
                 "plasticity_lesion", "both_lesion"):
        model, events, report = train_condition(
            initial, train_trajectories, name, collect=name == "intact")
        trained[name] = model
        condition_reports[name] = report
        if name == "intact":
            training_events = events

    evaluation_events = None
    for name, model in trained.items():
        control_gain = condition_settings(name)["control_gain"]
        report, wall = evaluate(
            model, evaluation_trajectories, control_gain,
            collect=name == "intact")
        condition_reports[name]["final_success"] = report["success_rate"]
        condition_reports[name]["evaluation_wall_s"] = wall
        if name == "intact":
            evaluation_events = report["events"]

    acute, acute_wall = evaluate(
        trained["intact"], evaluation_trajectories, 0.0, collect=False)
    sham, sham_wall = evaluate(
        trained["intact"], evaluation_trajectories, 1.0, collect=False)
    intact_final = condition_reports["intact"]["final_success"]
    if sham["success_rate"] != intact_final:
        raise RuntimeError("paired sham evaluation changed intact success")

    signatures = signature_metrics(
        training_events, evaluation_events, cfg, trained["intact"].role)
    condition_reports["acute_online_lesion_after_intact"] = {
        "final_success": acute["success_rate"],
        "drop_from_intact": intact_final - acute["success_rate"],
        "evaluation_wall_s": acute_wall,
    }
    condition_reports["sham_after_intact"] = {
        "final_success": sham["success_rate"],
        "change_from_intact": sham["success_rate"] - intact_final,
        "evaluation_wall_s": sham_wall,
    }
    condition_reports["intact"]["gain_over_fixed_vectorizer_final"] = (
        intact_final - condition_reports["fixed_vectorizer"]["final_success"])
    condition_reports["both_lesion"]["fraction_of_intact_learning_gain"] = (
        condition_reports["both_lesion"]["learning_gain"]
        / max(1e-12, condition_reports["intact"]["learning_gain"]))

    result = {
        "schema_version": 1,
        "args": vars(args),
        "config": vars(cfg),
        "protocol": {
            "name": "oral_b_continuous_bci_development_v1",
            "selection_split": "development",
            "training_task_seed": args.task_seed,
            "evaluation_task_seed": args.task_seed + 100_000,
            "confirmation_seeds_touched": False,
            "paired_environment_across_conditions": True,
            "evaluation_plasticity": False,
        },
        "provenance": provenance,
        "conditions": condition_reports,
        "signatures": signatures,
        "finite": None,
        "wall_s": time.perf_counter() - started,
        "peak_rss_mib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0,
        "hardware": {
            "device": "cpu",
            "platform": platform.platform(),
            "processor": platform.processor(),
            "torch_version": torch.__version__,
            "threads": torch.get_num_threads(),
        },
    }
    result["finite"] = finite_tree(result)
    if not result["finite"]:
        raise RuntimeError("non-finite continuous-BCI result")
    os.makedirs(args.outdir, exist_ok=True)
    path = os.path.join(args.outdir, args.tag + ".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": intact_final,
        "intact_gain": condition_reports["intact"]["learning_gain"],
        "fixed_final": condition_reports["fixed_vectorizer"]["final_success"],
        "both_gain_fraction": condition_reports["both_lesion"][
            "fraction_of_intact_learning_gain"],
        "signatures": signatures,
    }, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()