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

2 files changed, 703 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")
diff --git a/sdil/bci_v2.py b/sdil/bci_v2.py
new file mode 100644
index 0000000..f2fdf6c
--- /dev/null
+++ b/sdil/bci_v2.py
@@ -0,0 +1,410 @@
+"""Actor--critic somato-dendritic innovation for the oral-B-v2 branch.
+
+This module is deliberately separate from :mod:`sdil.bci`: the failed R2
+implementation and its source-bound records remain immutable.
+"""
+from dataclasses import dataclass
+
+import torch
+
+from sdil.bci import BCIConfig, causal_roles
+
+
+@dataclass(frozen=True)
+class BCIV2Config(BCIConfig):
+    gamma: float = 0.9
+    critic_eta: float = 0.03
+    eligibility_decay: float = 0.8
+    velocity_reward_scale: float = 0.25
+    terminal_reward: float = 1.0
+
+    def validate(self):
+        super().validate()
+        if self.feedback != "performance_velocity":
+            raise ValueError("oral-B-v2 requires performance_velocity")
+        if not 0 <= self.gamma < 1:
+            raise ValueError("gamma must lie in [0,1)")
+        if self.critic_eta <= 0:
+            raise ValueError("critic_eta must be positive")
+        if not 0 <= self.eligibility_decay < 1:
+            raise ValueError("eligibility_decay must lie in [0,1)")
+        if self.velocity_reward_scale < 0 or self.terminal_reward <= 0:
+            raise ValueError("reward scales are invalid")
+
+
+class BCIV2:
+    """Manual local actor--critic with a role-vectorized TD innovation."""
+
+    def __init__(self, cfg, model_seed=0, *, device="cpu",
+                 dtype=torch.float32):
+        cfg.validate()
+        self.cfg = cfg
+        self.device = device
+        self.dtype = dtype
+        self.role = causal_roles(cfg, device=device, dtype=dtype)
+        generator = torch.Generator(device="cpu").manual_seed(model_seed)
+        self.W = (
+            0.15
+            * torch.randn(
+                cfg.n_neurons,
+                cfg.context_dim,
+                generator=generator,
+                dtype=dtype,
+            )
+            / cfg.context_dim ** 0.5
+        ).to(device)
+        self.A = (
+            0.05
+            * torch.randn(cfg.n_neurons, generator=generator, dtype=dtype)
+        ).to(device)
+        self.coupling = (
+            cfg.coupling_scale
+            * torch.exp(
+                0.2
+                * torch.randn(
+                    cfg.n_neurons, generator=generator, dtype=dtype
+                )
+            )
+        ).to(device)
+        self.P = torch.zeros(cfg.n_neurons, device=device, dtype=dtype)
+        self.P_bias = torch.zeros(
+            cfg.n_neurons, device=device, dtype=dtype
+        )
+        # A scalar value circuit reads the surrounding population. Its update
+        # is a local delta rule driven by the same scalar TD innovation.
+        self.critic = torch.zeros(
+            cfg.n_neurons + 1, device=device, dtype=dtype
+        )
+        self.model_seed = model_seed
+
+    def clone(self):
+        other = BCIV2(
+            self.cfg,
+            self.model_seed,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        for name in (
+            "W",
+            "A",
+            "coupling",
+            "P",
+            "P_bias",
+            "critic",
+        ):
+            setattr(other, name, getattr(self, name).clone())
+        return other
+
+    def value_features(self, soma):
+        return torch.cat(
+            (
+                torch.ones(
+                    soma.shape[0],
+                    1,
+                    device=soma.device,
+                    dtype=soma.dtype,
+                ),
+                soma,
+            ),
+            dim=1,
+        )
+
+    def value(self, features):
+        return features @ self.critic
+
+    def causal_role_targets(self, cursor_plus, cursor_minus, perturbations):
+        sigma = self.cfg.perturb_sigma
+        directional_derivative = (
+            cursor_plus - cursor_minus
+        ) / (2.0 * sigma)
+        return directional_derivative.unsqueeze(1) * perturbations
+
+    def update_predictor(self, soma, ordinary, active):
+        if not bool(active.any()):
+            return
+        selected_soma = soma[active]
+        selected_target = ordinary[active]
+        residual = selected_target - (
+            self.P * selected_soma + self.P_bias
+        )
+        soma_centered = selected_soma - selected_soma.mean(0)
+        residual_centered = residual - residual.mean(0)
+        variance = soma_centered.square().mean(0)
+        self.P += self.cfg.predictor_eta * (
+            (residual_centered * soma_centered).mean(0)
+            / (variance + 1e-6)
+        )
+        self.P_bias += self.cfg.predictor_eta * residual.mean(0)
+
+    def update_role(self, target, active):
+        if not bool(active.any()):
+            return
+        self.A += self.cfg.vectorizer_eta * (
+            target[active].mean(0) - self.A
+        )
+
+    def update_critic(self, previous_features, delta, active):
+        if not bool(active.any()):
+            return
+        self.critic += self.cfg.critic_eta * (
+            delta[active].unsqueeze(1) * previous_features[active]
+        ).mean(0)
+
+    def initial_episode_state(self, episodes):
+        soma = torch.zeros(
+            episodes,
+            self.cfg.n_neurons,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        features = self.value_features(soma)
+        return {
+            "soma": soma,
+            "abs_error": torch.full(
+                (episodes,),
+                abs(self.cfg.target),
+                device=self.device,
+                dtype=self.dtype,
+            ),
+            "value_features": features,
+            "value": self.value(features),
+            "eligibility": torch.zeros(
+                episodes,
+                self.cfg.n_neurons,
+                self.cfg.context_dim,
+                device=self.device,
+                dtype=self.dtype,
+            ),
+            "active": torch.ones(
+                episodes, device=self.device, dtype=torch.bool
+            ),
+            "success": torch.zeros(
+                episodes, device=self.device, dtype=torch.bool
+            ),
+        }
+
+    def step(
+        self,
+        context,
+        process_noise,
+        perturbations,
+        state,
+        global_step,
+        *,
+        final_step,
+        plasticity_gain=1.0,
+        learn_role=True,
+        learn_predictor=True,
+        learn_critic=True,
+        critic_enabled=True,
+        terminal_outcome_enabled=True,
+    ):
+        """Advance one BCI transition and apply only manual local updates."""
+        with torch.no_grad():
+            active = state["active"]
+            u = (
+                context @ self.W.t()
+                + self.cfg.inertia * state["soma"]
+                + process_noise
+            )
+            proposed_soma = torch.tanh(u)
+            soma = torch.where(
+                active.unsqueeze(1), proposed_soma, state["soma"]
+            )
+            cursor = soma @ self.role
+            error = self.cfg.target - cursor
+            abs_error = error.abs()
+            velocity = state["abs_error"] - abs_error
+            newly_successful = active & (cursor >= self.cfg.target)
+            terminal = active & (newly_successful | final_step)
+            outcome = newly_successful.to(self.dtype)
+
+            features = self.value_features(soma)
+            current_value = (
+                self.value(features)
+                if critic_enabled
+                else torch.zeros_like(error)
+            )
+            previous_value = (
+                state["value"]
+                if critic_enabled
+                else torch.zeros_like(error)
+            )
+            terminal_bonus = (
+                self.cfg.terminal_reward * outcome
+                if terminal_outcome_enabled
+                else torch.zeros_like(outcome)
+            )
+            reward = (
+                self.cfg.velocity_reward_scale * velocity + terminal_bonus
+            )
+            bootstrap = torch.where(
+                terminal, torch.zeros_like(current_value), current_value
+            )
+            delta = (
+                reward
+                + self.cfg.gamma * bootstrap
+                - previous_value
+            )
+            delta = delta * active
+
+            ordinary = self.coupling * soma
+            instructional = delta.unsqueeze(1) * self.A.unsqueeze(0)
+            raw = ordinary + instructional
+            baseline = self.P * soma + self.P_bias
+            innovation = (raw - baseline) * active.unsqueeze(1)
+
+            gain = 1.0 - soma.square()
+            local_eligibility = (
+                gain.unsqueeze(2) * context.unsqueeze(1)
+            )
+            eligibility = (
+                self.cfg.eligibility_decay * state["eligibility"]
+                + local_eligibility
+            )
+            eligibility = eligibility * active[:, None, None]
+
+            did_perturb = (
+                learn_role and global_step % self.cfg.perturb_every == 0
+            )
+            causal_target = None
+            if did_perturb:
+                sigma = self.cfg.perturb_sigma
+                cursor_plus = (soma + sigma * perturbations) @ self.role
+                cursor_minus = (soma - sigma * perturbations) @ self.role
+                causal_target = self.causal_role_targets(
+                    cursor_plus, cursor_minus, perturbations
+                )
+
+            if plasticity_gain and bool(active.any()):
+                update = (
+                    innovation[:, :, None] * eligibility
+                )[active].mean(0)
+                self.W += plasticity_gain * self.cfg.forward_eta * update
+            if did_perturb:
+                self.update_role(causal_target, active)
+            if learn_critic and critic_enabled:
+                self.update_critic(
+                    state["value_features"], delta, active
+                )
+            if learn_predictor:
+                self.update_predictor(soma, ordinary, active)
+
+            next_active = active & ~terminal
+            next_success = state["success"] | newly_successful
+            next_features = torch.where(
+                next_active.unsqueeze(1),
+                features,
+                state["value_features"],
+            )
+            # Re-evaluate the continuing state's value after the local critic
+            # update.  This is the usual online TD(0) convention: the next
+            # transition starts from V(s_{t+1}; theta_{t+1}), rather than a
+            # stale prediction made with theta_t.
+            updated_current_value = (
+                self.value(features)
+                if critic_enabled
+                else torch.zeros_like(error)
+            )
+            next_value = torch.where(
+                next_active,
+                updated_current_value,
+                state["value"],
+            )
+            next_state = {
+                "soma": soma,
+                "abs_error": torch.where(
+                    next_active, abs_error, state["abs_error"]
+                ),
+                "value_features": next_features,
+                "value": next_value,
+                "eligibility": eligibility,
+                "active": next_active,
+                "success": next_success,
+            }
+        record = {
+            "soma": soma,
+            "base_soma": soma,
+            "cursor": cursor,
+            "error": error,
+            "abs_error": abs_error,
+            "previous_abs_error": state["abs_error"],
+            "performance_velocity": velocity,
+            "value_prediction": previous_value,
+            "next_value_prediction": current_value,
+            "reward": reward,
+            "td_innovation": delta,
+            "raw_apical": raw,
+            "innovation": innovation,
+            "ordinary_apical": ordinary,
+            "instructional_apical": instructional,
+            "eligibility": eligibility,
+            "context": context,
+            "causal_target": causal_target,
+            "did_perturb": did_perturb,
+            "active": active,
+            "terminal": terminal,
+            "outcome": outcome.bool(),
+        }
+        return next_state, record
+
+
+def run_day_v2(
+    model,
+    trajectories,
+    day,
+    *,
+    plasticity_gain=1.0,
+    learn_role=True,
+    learn_predictor=True,
+    learn_critic=True,
+    critic_enabled=True,
+    terminal_outcome_enabled=True,
+    collect=False,
+    global_step=0,
+    horizon=None,
+):
+    """Run a paired episode batch with an explicit terminal outcome event."""
+    episodes = trajectories.context.shape[1]
+    available_steps = trajectories.context.shape[2]
+    horizon = model.cfg.steps_per_episode if horizon is None else horizon
+    if not 1 <= horizon <= min(model.cfg.steps_per_episode, available_steps):
+        raise ValueError("v2 horizon exceeds the available trajectory")
+    state = model.initial_episode_state(episodes)
+    events = []
+    for step_index in range(horizon):
+        previous_soma = state["soma"]
+        state, record = model.step(
+            trajectories.context[day, :, step_index],
+            trajectories.process_noise[day, :, step_index],
+            trajectories.perturbations[day, :, step_index],
+            state,
+            global_step,
+            final_step=step_index == horizon - 1,
+            plasticity_gain=plasticity_gain,
+            learn_role=learn_role,
+            learn_predictor=learn_predictor,
+            learn_critic=learn_critic,
+            critic_enabled=critic_enabled,
+            terminal_outcome_enabled=terminal_outcome_enabled,
+        )
+        if collect:
+            event = {
+                key: value.detach().cpu().clone()
+                for key, value in record.items()
+                if isinstance(value, torch.Tensor)
+            }
+            event["previous_soma"] = (
+                previous_soma.detach().cpu().clone()
+            )
+            event["step"] = step_index
+            events.append(event)
+        global_step += 1
+    if bool(state["active"].any()):
+        raise RuntimeError("every v2 episode must terminate")
+    return {
+        "success": state["success"].detach().cpu(),
+        "success_rate": state["success"].float().mean().item(),
+        "events": events,
+        "global_step": global_step,
+    }