Make network.reset_state_variables() actually reach the learning rules

bindsnet/learning/MCC_learning.py

@@ -302,8 +302,18 @@ def _connection_update(self, **kwargs) -> None:
 
         super().update()
 
-    def reset_state_variables(self):
-        return
+    def reset_state_variables(self) -> None:
+        # language=rst
+        """
+        Clear the moving-average buffers so a new episode does not average over
+        updates accumulated during the previous one.
+        """
+
+        if self.average_update > 0:
+            self.average_buffer_pre.zero_()
+            self.average_buffer_post.zero_()
+            self.average_buffer_index_pre = 0
+            self.average_buffer_index_post = 0
 
 
 class Hebbian(MCC_LearningRule):
@@ -389,8 +399,12 @@ def _connection_update(self, **kwargs) -> None:
 
         super().update()
 
-    def reset_state_variables(self):
-        return
+    def reset_state_variables(self) -> None:
+        # language=rst
+        """
+        Nothing to reset: the rule holds no state between steps, deriving each
+        update from the layers' current spikes and traces.
+        """
 
 
 class MSTDP(MCC_LearningRule):
@@ -457,6 +471,14 @@ def __init__(
         self.tc_plus = torch.tensor(kwargs.get("tc_plus", 20.0))
         self.tc_minus = torch.tensor(kwargs.get("tc_minus", 20.0))
 
+        # State the update path fills in lazily: the previous step's spikes,
+        # kept by the fast path for its rank-1 update, and the dense path's
+        # eligibility. None means "not built yet", which is also the state
+        # ``reset_state_variables`` restores.
+        self._prev_source_s = None
+        self._prev_target_s = None
+        self.eligibility = None
+
         # Initialize variables for average update and continues update
         self.average_update = kwargs.get("average_update", 0)
         self.continues_update = kwargs.get("continues_update", False)
@@ -540,7 +562,7 @@ def _connection_update(self, **kwargs) -> None:
             and not self.feature_value.is_sparse
         )
         if fast:
-            if hasattr(self, "_prev_target_s"):
+            if self._prev_target_s is not None:
                 if isinstance(reward, torch.Tensor):
                     # Keep reward on-device (no host sync for tensor rewards).
                     update = (
@@ -560,7 +582,7 @@ def _connection_update(self, **kwargs) -> None:
         else:
             # Dense-eligibility path: averaging buffers, custom reductions, or
             # sparse weights.
-            if not hasattr(self, "eligibility"):
+            if self.eligibility is None:
                 self.eligibility = torch.zeros(
                     batch_size,
                     *self.feature_value.shape,
@@ -607,8 +629,25 @@ def _connection_update(self, **kwargs) -> None:
 
         super().update()
 
-    def reset_state_variables(self):
-        return
+    def reset_state_variables(self) -> None:
+        # language=rst
+        """
+        Clear every variable that carries across time steps, so a new episode
+        starts from the same state as a freshly-built rule.
+        """
+
+        if self.eligibility is not None:
+            self.eligibility.zero_()
+        self.p_plus.zero_()
+        self.p_minus.zero_()
+        if self.average_update > 0:
+            self.average_buffer.zero_()
+            self.average_buffer_index = 0
+        # The fast path keeps the previous step's spikes to build the next
+        # rank-1 update. Drop them, or the first step of a new episode pairs
+        # with the last step of the old one.
+        self._prev_source_s = None
+        self._prev_target_s = None
 
 
 class MSTDPET(MCC_LearningRule):
@@ -803,6 +842,16 @@ def _connection_update(self, **kwargs) -> None:
         super().update()
 
     def reset_state_variables(self) -> None:
+        # language=rst
+        """
+        Clear every variable that carries across time steps, so a new episode
+        starts from the same state as a freshly-built rule.
+        """
+
         self.eligibility.zero_()
         self.eligibility_trace.zero_()
-        return
+        self.p_plus.zero_()
+        self.p_minus.zero_()
+        if self.average_update > 0:
+            self.average_buffer.zero_()
+            self.average_buffer_index = 0

bindsnet/network/topology_features.py

@@ -167,15 +167,17 @@ def cast_dtype_if_needed(value, value_dtype):
         else:
             return value
 
-    @abstractmethod
     def reset_state_variables(self) -> None:
         # language=rst
         """
-        Contains resetting logic for the feature.
+        Reset the feature between samples or episodes. Features that hold no
+        state of their own inherit this, which forwards the reset to the
+        feature's learning rule; a feature with its own state overrides it and
+        calls ``super()`` first.
         """
+
         if self.learning_rule:
             self.learning_rule.reset_state_variables()
-        pass
 
     @abstractmethod
     def compute(self, s) -> Union[torch.Tensor, float, int]:
@@ -449,9 +451,6 @@ def compute(self, s) -> Union[torch.Tensor, float, int]:
             return self.sparse_bernoulli()
         return torch.bernoulli(self.value)
 
-    def reset_state_variables(self) -> None:
-        pass
-
     def prime_feature(self, connection, device, **kwargs) -> None:
         ## Initialize value ###
         if self.value is None:
@@ -528,9 +527,6 @@ def __init__(
     def compute(self, s) -> torch.Tensor:
         return self.value
 
-    def reset_state_variables(self) -> None:
-        pass
-
     def prime_feature(self, connection, device, **kwargs) -> None:
         # Check if feature is already primed
         if self.is_primed:
@@ -581,9 +577,6 @@ def __init__(self) -> None:
         """
         pass
 
-    def reset_state_variables(self) -> None:
-        pass
-
     def compute(self, s) -> Union[torch.Tensor, float, int]:
         return s.float().mean() * torch.ones(
             self.source_n, self.target_n, device=s.device
@@ -657,9 +650,6 @@ def __init__(
             batch_size=batch_size,
         )
 
-    def reset_state_variables(self) -> None:
-        pass
-
     def compute(self, s) -> Union[torch.Tensor, float, int]:
         if self.enforce_polarity:
             pos_mask = ~torch.logical_xor(self.value > 0, self.positive_mask)
@@ -734,9 +724,6 @@ def __init__(
     # Bias is additive: folds as ``B <- B + value`` in the connection's pipeline.
     op = "add"
 
-    def reset_state_variables(self) -> None:
-        pass
-
     def compute(self, s) -> Union[torch.Tensor, float, int]:
         # Additive offset added to every synapse (independent of the spikes).
         return self.value
@@ -779,9 +766,6 @@ def __init__(
             batch_size=batch_size,
         )
 
-    def reset_state_variables(self) -> None:
-        pass
-
     def compute(self, s) -> Union[torch.Tensor, float, int]:
         return self.value
 
@@ -839,9 +823,6 @@ def __init__(
     # Degradation is subtractive: folded as ``B <- B - degrade_function(value)``.
     op = "sub"
 
-    def reset_state_variables(self) -> None:
-        pass
-
     def compute(self, s) -> Union[torch.Tensor, float, int]:
         # Subtractive offset (via degrade_function) applied to every synapse.
         if self.degrade_function is not None:
@@ -959,11 +940,11 @@ def compute(self, s) -> Union[torch.Tensor, float, int]:
     def reset_state_variables(
         self,
     ):
+        super().reset_state_variables()
         self.spike_buffer = torch.zeros_like(self.spike_buffer)
         self.counter = 0
         self.start_counter = False
         self.value = self.init_value.clone().detach()  # initial mask
-        pass
 
 
 class AdaptationBaseOtherSynaps(AbstractFeature):
@@ -1076,11 +1057,11 @@ def compute(self, s) -> Union[torch.Tensor, float, int]:
     def reset_state_variables(
         self,
     ):
+        super().reset_state_variables()
         self.spike_buffer = torch.zeros_like(self.spike_buffer)
         self.counter = 0
         self.start_counter = False
         self.value = self.init_value.clone().detach()  # initial mask
-        pass
 
 
 ### Sub Features ###

test/network/test_learning.py

@@ -1,16 +1,21 @@
+import pytest
 import torch
 
 from bindsnet.learning import (
     MSTDP,
     MSTDPET,
     Hebbian,
-    PostPre,
-    Rmax,
-    WeightDependentPostPre,
 )
+from bindsnet.learning import MCC_learning as mcc
+from bindsnet.learning import PostPre, Rmax, WeightDependentPostPre
 from bindsnet.network import Network
+from bindsnet.network import topology_features as tf
 from bindsnet.network.nodes import CSRMNodes, Input, LIFNodes, SRM0Nodes
-from bindsnet.network.topology import Connection, Conv2dConnection
+from bindsnet.network.topology import (
+    Connection,
+    Conv2dConnection,
+    MulticompartmentConnection,
+)
 
 
 class TestLearningRules:
@@ -268,3 +273,140 @@ def test_rmax(self):
             time=250,
             reward=1.0,
         )
+
+
+class TestLearningRuleReset:
+    """
+    ``network.reset_state_variables()`` must clear every variable a
+    ``MulticompartmentConnection`` learning rule carries between time steps.
+
+    Regression for #777: the reset never reached the learning rules at all,
+    because every feature overrode ``reset_state_variables`` with a bare
+    ``pass``, and the rules that did get called cleared only part of their
+    state.
+    """
+
+    @staticmethod
+    def _build(rule, n=8, **rule_kwargs):
+        """A one-connection network whose single Weight uses ``rule``."""
+        network = Network(dt=1.0)
+        network.add_layer(Input(n=n, traces=True), name="input")
+        network.add_layer(LIFNodes(n=n, traces=True), name="output")
+        weight = tf.Weight(
+            name="w",
+            value=torch.rand(n, n),
+            range=[0.0, 1.0],
+            nu=(1e-2, 1e-2),
+            learning_rule=rule,
+        )
+        connection = MulticompartmentConnection(
+            source=network.layers["input"],
+            target=network.layers["output"],
+            device="cpu",
+            pipeline=[weight],
+            **rule_kwargs,
+        )
+        network.add_connection(connection, source="input", target="output")
+        return network, connection.pipeline[0].learning_rule
+
+    @staticmethod
+    def _drive(network, n=8, time=100, seed=0):
+        torch.manual_seed(seed)
+        network.run(
+            inputs={"input": torch.bernoulli(torch.rand(time, n)).byte()},
+            time=time,
+            reward=1.0,
+        )
+
+    def test_reset_reaches_the_learning_rule(self):
+        # The bug behind #777: MSTDPET's reset was never invoked, so even the
+        # two variables it did clear survived a network reset.
+        network, rule = self._build(mcc.MSTDPET)
+        self._drive(network)
+        assert rule.eligibility_trace.abs().sum() > 0  # sanity: state exists
+        network.reset_state_variables()
+        assert torch.all(rule.eligibility_trace == 0)
+
+    def test_mstdpet_reset_clears_all_state(self):
+        # Reported by @saachigoyall in #777: p_plus, p_minus and the
+        # moving-average buffer were left untouched.
+        network, rule = self._build(
+            mcc.MSTDPET, average_update=5, continues_update=True
+        )
+        self._drive(network)
+        assert rule.p_plus.abs().sum() > 0
+        assert rule.p_minus.abs().sum() > 0
+        assert rule.average_buffer.abs().sum() > 0
+
+        network.reset_state_variables()
+
+        assert torch.all(rule.eligibility == 0)
+        assert torch.all(rule.eligibility_trace == 0)
+        assert torch.all(rule.p_plus == 0)
+        assert torch.all(rule.p_minus == 0)
+        assert torch.all(rule.average_buffer == 0)
+        assert rule.average_buffer_index == 0
+
+    def test_mstdp_reset_clears_all_state(self):
+        # MSTDP's reset was a bare ``return``, clearing nothing.
+        network, rule = self._build(mcc.MSTDP, average_update=5, continues_update=True)
+        self._drive(network)
+        assert rule.p_plus.abs().sum() > 0
+        assert rule.p_minus.abs().sum() > 0
+
+        network.reset_state_variables()
+
+        assert rule.eligibility is None or torch.all(rule.eligibility == 0)
+        assert torch.all(rule.p_plus == 0)
+        assert torch.all(rule.p_minus == 0)
+        assert torch.all(rule.average_buffer == 0)
+        assert rule.average_buffer_index == 0
+
+    def test_mstdp_reset_clears_fast_path_spike_lag(self):
+        # The fast path keeps the previous step's spikes for its rank-1 update.
+        # Left in place, the first step of a new episode pairs with the last
+        # step of the old one.
+        network, rule = self._build(mcc.MSTDP)
+        self._drive(network)
+        assert rule._prev_source_s is not None
+        assert rule._prev_target_s is not None
+
+        network.reset_state_variables()
+
+        assert rule._prev_source_s is None
+        assert rule._prev_target_s is None
+
+    def test_postpre_reset_clears_average_buffers(self):
+        # PostPre's reset was a bare ``return``; both buffers survived.
+        network, rule = self._build(
+            mcc.PostPre, average_update=5, continues_update=True
+        )
+        self._drive(network)
+
+        network.reset_state_variables()
+
+        assert torch.all(rule.average_buffer_pre == 0)
+        assert torch.all(rule.average_buffer_post == 0)
+        assert rule.average_buffer_index_pre == 0
+        assert rule.average_buffer_index_post == 0
+
+    @pytest.mark.parametrize("rule", [mcc.MSTDP, mcc.MSTDPET, mcc.PostPre])
+    def test_episodes_are_independent_after_reset(self, rule):
+        # The symptom #777 reported: with a reset between them, two identical
+        # episodes must produce identical weights. Before the fix the second
+        # episode started from the first one's leftover state.
+        network, _ = self._build(rule)
+        feature = network.connections[("input", "output")].pipeline[0]
+        w0 = feature.value.clone()
+
+        self._drive(network, seed=1)
+        after_first = feature.value.clone()
+
+        network.reset_state_variables()
+        with torch.no_grad():
+            feature.value.copy_(w0)
+        self._drive(network, seed=1)
+        after_second = feature.value.clone()
+
+        assert not torch.allclose(after_first, w0)  # sanity: learning happened
+        assert torch.allclose(after_first, after_second, atol=1e-6)