audit: isolate BCI role learner from environment map

1 files changed, 8 insertions, 1 deletions

diff --git a/experiments/bci_smoke.py b/experiments/bci_smoke.py
index 4724d26..1fda89e 100644
--- a/experiments/bci_smoke.py
+++ b/experiments/bci_smoke.py
@@ -104,7 +104,14 @@ def check_temporal_difference_role_vectorizer():
         cfg.episodes_per_day, cfg.n_neurons, generator=generator)
     xi = torch.empty_like(soma).bernoulli_(
         0.5, generator=generator).mul_(2).sub_(1)
-    target = model.causal_role_targets(soma, xi)
+    cursor_plus = (soma + cfg.perturb_sigma * xi) @ model.role
+    cursor_minus = (soma - cfg.perturb_sigma * xi) @ model.role
+    target = model.causal_role_targets(cursor_plus, cursor_minus, xi)
+    saved_role = model.role.clone()
+    model.role.zero_()
+    assert torch.equal(
+        target, model.causal_role_targets(cursor_plus, cursor_minus, xi))
+    model.role.copy_(saved_role)
     mean_target = target.mean(0)
     cosine = torch.nn.functional.cosine_similarity(
         mean_target, model.role, dim=0).item()