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

1 files changed, 410 insertions, 0 deletions

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,
+    }