diff options
Diffstat (limited to 'sdil/bci_v2.py')
| -rw-r--r-- | sdil/bci_v2.py | 15 |
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, } |