From ccb39d63ed5c744a05f6618691c9458d18de7ced Mon Sep 17 00:00:00 2001 From: Hananel Hazan Date: Sun, 6 Sep 2026 22:05:15 -0400 Subject: [PATCH] fix: make network.reset_state_variables() actually reach the learning rules Builds on @saachigoyall's fix. Her diagnosis was right: MSTDPET cleared only 2 of its 6 state variables. But that change alone had no observable effect, because the reset never reached any learning rule in the first place. AbstractFeature.reset_state_variables forwards to self.learning_rule, and it is the only place that does. Every concrete feature overrode it with a bare 'pass' and none called super(), so the forwarding line was unreachable. After network.reset_state_variables() nothing was cleared, not even the two variables MSTDPET already handled. Features (topology_features.py): - Drop the seven bare-'pass' overrides (Probability, Mask, MeanField, Weight, Bias, Intensity, Degradation) so they inherit the base implementation, and drop @abstractmethod from it, which is why those overrides existed at all. - The two adaptation features keep their own reset logic and now chain to super() first. Rules (MCC_learning.py): - MSTDPET now also clears p_plus, p_minus and the moving-average buffer (@saachigoyall's change). - MSTDP cleared nothing. It now clears eligibility, p_plus, p_minus, the moving-average buffer, and the fast path's one-step spike lag. That lag is newer than this PR: without clearing it, the first step of an episode pairs with the last step of the previous one, which is the contamination this PR set out to fix. - PostPre cleared nothing; it now clears both averaging buffers. - Hebbian holds no state; its no-op is now documented as deliberate. - MSTDP's lazily-built state (_prev_source_s, _prev_target_s, eligibility) is declared None in __init__ and guarded with 'is None' rather than hasattr, so reset has something defined to restore. Tests: replaces the original test, which passed tc_plus/average_update to Weight (whose signature never accepted them) and so raised TypeError rather than running. Eight cases now, seven of which fail without the source change, including an end-to-end check that two identical episodes separated by a reset produce identical weights. Full suite 91 passed. Per-step update cost unchanged over three interleaved A/B rounds at n=64 and n=256, all differences under 1% with the sign varying between rounds. Co-Authored-By: Saachi Goyal <156711741+saachigoyall@users.noreply.github.com> Co-Authored-By: Claude Opus 5 --- bindsnet/learning/MCC_learning.py | 67 ++++++++++-- bindsnet/network/topology_features.py | 33 ++---- test/network/test_learning.py | 150 +++++++++++++++++++++++++- 3 files changed, 211 insertions(+), 39 deletions(-) diff --git a/bindsnet/learning/MCC_learning.py b/bindsnet/learning/MCC_learning.py index be067e31f..77bdd78b4 100644 --- a/bindsnet/learning/MCC_learning.py +++ b/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 diff --git a/bindsnet/network/topology_features.py b/bindsnet/network/topology_features.py index 7abcd6368..6b7713a77 100644 --- a/bindsnet/network/topology_features.py +++ b/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 ### diff --git a/test/network/test_learning.py b/test/network/test_learning.py index 926b921a4..16903781e 100644 --- a/test/network/test_learning.py +++ b/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)