bci: add endpoint-free oral-B-v2 mechanics

1 files changed, 293 insertions, 0 deletions

diff --git a/experiments/bci_v2_smoke.py b/experiments/bci_v2_smoke.py
new file mode 100644
index 0000000..fc6509b
--- /dev/null
+++ b/experiments/bci_v2_smoke.py
@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+"""Endpoint-free mechanics checks for the independent oral-B-v2 branch."""
+import os
+import sys
+
+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
+
+
+def check_role_estimator():
+    cfg = BCIV2Config(
+        days=2, episodes_per_day=8192, steps_per_episode=2,
+        feedback="performance_velocity",
+    )
+    model = BCIV2(cfg, model_seed=1)
+    generator = torch.Generator().manual_seed(2)
+    soma = 0.2 * torch.randn(
+        cfg.episodes_per_day, cfg.n_neurons, generator=generator
+    )
+    perturbations = torch.empty_like(soma).bernoulli_(
+        0.5, generator=generator
+    ).mul_(2).sub_(1)
+    sigma = cfg.perturb_sigma
+    cursor_plus = (soma + sigma * perturbations) @ model.role
+    cursor_minus = (soma - sigma * perturbations) @ model.role
+    saved_role = model.role.clone()
+    target = model.causal_role_targets(
+        cursor_plus, cursor_minus, perturbations
+    )
+    model.role.zero_()
+    assert torch.equal(
+        target,
+        model.causal_role_targets(
+            cursor_plus, cursor_minus, perturbations
+        ),
+    )
+    model.role.copy_(saved_role)
+    estimate = target.mean(0)
+    cosine = torch.nn.functional.cosine_similarity(
+        estimate, model.role, dim=0
+    ).item()
+    assert cosine > 0.99
+    active = torch.ones(cfg.episodes_per_day, dtype=torch.bool)
+    model.A.zero_()
+    model.update_role(target, active)
+    learned_cosine = torch.nn.functional.cosine_similarity(
+        model.A, model.role, dim=0
+    ).item()
+    assert learned_cosine > 0.99
+    print(
+        "v2 antithetic causal-role estimator: "
+        f"target_cos={cosine:.6f}, update_cos={learned_cosine:.6f}"
+    )
+
+
+def check_predictor():
+    cfg = BCIV2Config(
+        days=2, episodes_per_day=256, steps_per_episode=2,
+        feedback="performance_velocity",
+    )
+    model = BCIV2(cfg, model_seed=3)
+    generator = torch.Generator().manual_seed(4)
+    active = torch.ones(cfg.episodes_per_day, dtype=torch.bool)
+    for _ in range(100):
+        soma = torch.randn(
+            cfg.episodes_per_day,
+            cfg.n_neurons,
+            generator=generator,
+        ).tanh()
+        model.update_predictor(soma, model.coupling * soma, active)
+    assert (model.P - model.coupling).abs().max() < 1e-4
+    assert model.P_bias.abs().max() < 1e-5
+    print("v2 neutral predictor identifies soma-predictable traffic: exact")
+
+
+def check_td_and_eligibility_updates():
+    cfg = BCIV2Config(
+        days=2, episodes_per_day=4, steps_per_episode=3,
+        target=10.0, process_noise=0.0, inertia=0.0,
+        predictor_eta=0.2, vectorizer_eta=0.03,
+        forward_eta=0.1, perturb_every=99,
+        gamma=0.8, critic_eta=0.05, eligibility_decay=0.6,
+        velocity_reward_scale=0.25, terminal_reward=1.0,
+        feedback="performance_velocity",
+    )
+    initial = BCIV2(cfg, model_seed=5)
+    initial.P.copy_(initial.coupling)
+    initial.A.copy_(initial.role)
+    initial.critic.copy_(torch.linspace(
+        -0.1, 0.1, cfg.n_neurons + 1
+    ))
+    model = initial.clone()
+    state = model.initial_episode_state(cfg.episodes_per_day)
+    context = torch.arange(
+        cfg.episodes_per_day * cfg.context_dim, dtype=torch.float32
+    ).reshape(cfg.episodes_per_day, cfg.context_dim) / 100.0
+    noise = torch.zeros(cfg.episodes_per_day, cfg.n_neurons)
+    perturbations = torch.ones_like(noise)
+
+    previous_critic = model.critic.clone()
+    previous_w = model.W.clone()
+    previous_features = state["value_features"].clone()
+    previous_value = state["value"].clone()
+    next_state, record = model.step(
+        context,
+        noise,
+        perturbations,
+        state,
+        global_step=1,
+        final_step=False,
+        learn_role=False,
+        learn_predictor=False,
+    )
+    expected_local = (
+        (1.0 - record["soma"].square()).unsqueeze(2)
+        * context.unsqueeze(1)
+    )
+    assert torch.allclose(record["eligibility"], expected_local)
+    expected_delta = (
+        record["reward"]
+        + cfg.gamma * record["next_value_prediction"]
+        - previous_value
+    )
+    assert torch.allclose(record["td_innovation"], expected_delta)
+    expected_critic = previous_critic + cfg.critic_eta * (
+        expected_delta.unsqueeze(1) * previous_features
+    ).mean(0)
+    assert torch.allclose(model.critic, expected_critic)
+    expected_w = previous_w + cfg.forward_eta * (
+        record["innovation"][:, :, None] * expected_local
+    ).mean(0)
+    assert torch.allclose(model.W, expected_w)
+    assert torch.allclose(
+        next_state["value"], model.value(next_state["value_features"])
+    )
+
+    second_state, second = model.step(
+        context,
+        noise,
+        perturbations,
+        next_state,
+        global_step=2,
+        final_step=False,
+        plasticity_gain=0.0,
+        learn_role=False,
+        learn_predictor=False,
+        learn_critic=False,
+    )
+    second_local = (
+        (1.0 - second["soma"].square()).unsqueeze(2)
+        * context.unsqueeze(1)
+    )
+    assert torch.allclose(
+        second["eligibility"],
+        cfg.eligibility_decay * expected_local + second_local,
+    )
+    assert torch.equal(second_state["active"], torch.ones(4, dtype=torch.bool))
+    print("v2 TD(0), local eligibility, critic, and forward updates: exact")
+
+
+def check_terminal_outcomes_and_lesions():
+    cfg = BCIV2Config(
+        days=2, episodes_per_day=2, steps_per_episode=2,
+        target=0.2, inertia=0.0, process_noise=0.0,
+        forward_eta=0.1, gamma=0.9, critic_eta=0.03,
+        velocity_reward_scale=0.0, terminal_reward=1.0,
+        feedback="performance_velocity",
+    )
+    initial = BCIV2(cfg, model_seed=6)
+    initial.P.copy_(initial.coupling)
+    initial.A.copy_(initial.role)
+    initial.critic.zero_()
+    initial.critic[0] = 0.6
+    context = torch.zeros(2, cfg.context_dim)
+    context[:, 0] = 1.0
+    noise = torch.zeros(2, cfg.n_neurons)
+    noise[0, :cfg.n_plus] = 1.0
+    perturbations = torch.ones_like(noise)
+    state = initial.initial_episode_state(2)
+
+    intact = initial.clone()
+    next_state, record = intact.step(
+        context,
+        noise,
+        perturbations,
+        state,
+        global_step=1,
+        final_step=True,
+        plasticity_gain=0.0,
+        learn_role=False,
+        learn_predictor=False,
+        learn_critic=False,
+    )
+    assert record["terminal"].all()
+    assert record["outcome"].tolist() == [True, False]
+    assert torch.allclose(
+        record["td_innovation"], torch.tensor([0.4, -0.6])
+    )
+    assert not bool(next_state["active"].any())
+
+    no_critic = initial.clone()
+    _, critic_lesion = no_critic.step(
+        context,
+        noise,
+        perturbations,
+        state,
+        global_step=1,
+        final_step=True,
+        plasticity_gain=0.0,
+        learn_role=False,
+        learn_predictor=False,
+        learn_critic=False,
+        critic_enabled=False,
+    )
+    assert torch.allclose(
+        critic_lesion["td_innovation"], torch.tensor([1.0, 0.0])
+    )
+
+    no_outcome = initial.clone()
+    _, outcome_lesion = no_outcome.step(
+        context,
+        noise,
+        perturbations,
+        state,
+        global_step=1,
+        final_step=True,
+        plasticity_gain=0.0,
+        learn_role=False,
+        learn_predictor=False,
+        learn_critic=False,
+        terminal_outcome_enabled=False,
+    )
+    assert torch.allclose(
+        outcome_lesion["td_innovation"], torch.tensor([-0.6, -0.6])
+    )
+    print("v2 terminal reward, timeout surprise, and acute lesions: exact")
+
+
+def check_horizon_and_pairing():
+    cfg = BCIV2Config(
+        days=2, episodes_per_day=16, steps_per_episode=8,
+        target=10.0, feedback="performance_velocity",
+    )
+    trajectories = generate_trajectories(cfg, 7)
+    first = BCIV2(cfg, model_seed=8)
+    second = first.clone()
+    report = run_day_v2(
+        first,
+        trajectories,
+        0,
+        collect=True,
+        horizon=4,
+        plasticity_gain=0.0,
+        learn_role=False,
+        learn_predictor=False,
+        learn_critic=False,
+    )
+    repeated = run_day_v2(
+        second,
+        trajectories,
+        0,
+        collect=True,
+        horizon=4,
+        plasticity_gain=0.0,
+        learn_role=False,
+        learn_predictor=False,
+        learn_critic=False,
+    )
+    assert report["global_step"] == 4
+    assert len(report["events"]) == 4
+    terminal_count = torch.stack([
+        event["terminal"] for event in report["events"]
+    ]).sum(0)
+    assert torch.equal(terminal_count, torch.ones(16, dtype=torch.long))
+    for left, right in zip(report["events"], repeated["events"]):
+        assert all(
+            torch.equal(left[key], right[key])
+            for key in left
+            if isinstance(left[key], torch.Tensor)
+        )
+    print("v2 fixed-horizon terminal accounting and paired replay: exact")
+
+
+if __name__ == "__main__":
+    check_role_estimator()
+    check_predictor()
+    check_td_and_eligibility_updates()
+    check_terminal_outcomes_and_lesions()
+    check_horizon_and_pairing()
+    print("ALL ORAL-B-V2 MECHANICS CHECKS PASSED")