bci: audit v2 outcome-surprise signatures

1 files changed, 13 insertions, 2 deletions

diff --git a/sdil/bci_v2.py b/sdil/bci_v2.py
index f2fdf6c..303a272 100644
--- a/sdil/bci_v2.py
+++ b/sdil/bci_v2.py
@@ -194,6 +194,7 @@ class BCIV2:
         final_step,
         plasticity_gain=1.0,
         learn_role=True,
+        probe_role=True,
         learn_predictor=True,
         learn_critic=True,
         critic_enabled=True,
@@ -265,7 +266,7 @@ class BCIV2:
             eligibility = eligibility * active[:, None, None]
 
             did_perturb = (
-                learn_role and global_step % self.cfg.perturb_every == 0
+                probe_role and global_step % self.cfg.perturb_every == 0
             )
             causal_target = None
             if did_perturb:
@@ -281,7 +282,7 @@ class BCIV2:
                     innovation[:, :, None] * eligibility
                 )[active].mean(0)
                 self.W += plasticity_gain * self.cfg.forward_eta * update
-            if did_perturb:
+            if did_perturb and learn_role:
                 self.update_role(causal_target, active)
             if learn_critic and critic_enabled:
                 self.update_critic(
@@ -356,6 +357,7 @@ def run_day_v2(
     *,
     plasticity_gain=1.0,
     learn_role=True,
+    probe_role=True,
     learn_predictor=True,
     learn_critic=True,
     critic_enabled=True,
@@ -372,6 +374,8 @@ def run_day_v2(
         raise ValueError("v2 horizon exceeds the available trajectory")
     state = model.initial_episode_state(episodes)
     events = []
+    active_transitions = 0
+    role_probe_examples = 0
     for step_index in range(horizon):
         previous_soma = state["soma"]
         state, record = model.step(
@@ -383,6 +387,7 @@ def run_day_v2(
             final_step=step_index == horizon - 1,
             plasticity_gain=plasticity_gain,
             learn_role=learn_role,
+            probe_role=probe_role,
             learn_predictor=learn_predictor,
             learn_critic=learn_critic,
             critic_enabled=critic_enabled,
@@ -399,6 +404,10 @@ def run_day_v2(
             )
             event["step"] = step_index
             events.append(event)
+        active_count = int(record["active"].sum().item())
+        active_transitions += active_count
+        if record["did_perturb"]:
+            role_probe_examples += active_count
         global_step += 1
     if bool(state["active"].any()):
         raise RuntimeError("every v2 episode must terminate")
@@ -407,4 +416,6 @@ def run_day_v2(
         "success_rate": state["success"].float().mean().item(),
         "events": events,
         "global_step": global_step,
+        "active_transitions": active_transitions,
+        "role_probe_examples": role_probe_examples,
     }