"""Optional plasticity layer for Brain 5D. Sprint 2C extends the Sprint 2B nearest-neighbour STDP/eligibility layer with reward-modulated three-factor learning. The engine remains outside ``src.core`` and observes completed core steps through the generic post-step hook. This module provides the LearningEngine class, which implements: - Pair-based STDP with nearest-neighbour pairing - Eligibility traces for reward-modulated learning - Three-factor (reward-modulated) plasticity - Configurable parameters via LearningParameters Design Principles: 1. The engine is optional and can be enabled/disabled via configuration. 2. It observes the network via post-step hooks, never owning the runtime loop. 3. All plasticity is bounded (weights clamped to [min_weight, max_weight]). 4. Rewards can be delayed and are applied when due. Example: >>> from src.learning import LearningEngine >>> engine = LearningEngine(network, config) >>> engine.attach() >>> # ... run simulation ... >>> stats = engine.stats >>> print(f"STDP updates: {stats.stdp_weight_updates}") >>> engine.detach() """ from __future__ import annotations import math import time from collections.abc import Mapping from dataclasses import asdict, dataclass, field from typing import TYPE_CHECKING, Any, cast from .eligibility import EligibilityTrace from .reward import RewardSignal if TYPE_CHECKING: from src.core.network import NeuralNetwork, StepResult from src.core.synapse import Synapse Config = Mapping[str, Any] def _as_mapping(value: Any, name: str) -> Mapping[str, Any]: """Validate and cast a config subsection to a typed mapping.""" if not isinstance(value, Mapping): raise TypeError(f"{name} config must be a mapping") return cast("Mapping[str, Any]", value) # ============================================================================ # Learning Parameters # ============================================================================ @dataclass(frozen=True, slots=True) class LearningParameters: """Validated plasticity parameters loaded from configuration. This parameter bundle controls all aspects of the learning engine: - STDP: amplitudes, time constants, weight bounds - Eligibility: trace decay time constant - Reward: learning rate, delay, clamping behavior Attributes: stdp_enabled: Whether STDP is enabled. a_plus: LTP amplitude. a_minus: LTD amplitude. tau_plus: LTP time constant (ticks). tau_minus: LTD time constant (ticks). min_weight: Minimum weight (clamping lower bound). max_weight: Maximum weight (clamping upper bound). eligibility_enabled: Whether eligibility traces are enabled. eligibility_tau_ticks: Eligibility trace decay time constant (ticks). reward_enabled: Whether reward-modulated plasticity is enabled. reward_learning_rate: Learning rate for reward-modulated updates. reward_delay_ticks: Delay before rewards are applied. reward_clamp_weights: Whether to clamp weights after reward updates. reward_reset_trace: Whether to reset eligibility after reward application. reward_trace_epsilon: Minimum trace value to consider for updates. """ stdp_enabled: bool = False a_plus: float = 0.1 a_minus: float = 0.12 tau_plus: float = 20.0 tau_minus: float = 20.0 min_weight: float = 0.0 max_weight: float = 1.0 eligibility_enabled: bool = False eligibility_tau_ticks: float = 200.0 reward_enabled: bool = False reward_learning_rate: float = 0.01 reward_delay_ticks: int = 0 reward_clamp_weights: bool = True reward_reset_trace: bool = False reward_trace_epsilon: float = 1e-12 @classmethod def from_config(cls, config: Config) -> LearningParameters: """Build and validate learning parameters from a config mapping. Args: config: Configuration dictionary containing 'stdp', 'eligibility', and 'reward' sections. Returns: Validated LearningParameters instance. Raises: TypeError: If any section is not a mapping. ValueError: If any parameter is invalid. """ stdp = _as_mapping(config.get("stdp", {}), "stdp") eligibility = _as_mapping(config.get("eligibility", {}), "eligibility") reward = _as_mapping(config.get("reward", {}), "reward") params = cls( stdp_enabled=bool(stdp.get("enabled", False)), a_plus=float(stdp.get("a_plus", 0.1)), a_minus=float(stdp.get("a_minus", 0.12)), tau_plus=float(stdp.get("tau_plus", 20.0)), tau_minus=float(stdp.get("tau_minus", 20.0)), min_weight=float(stdp.get("min_weight", 0.0)), max_weight=float(stdp.get("max_weight", 1.0)), eligibility_enabled=bool(eligibility.get("enabled", False)), eligibility_tau_ticks=float(eligibility.get("tau_ticks", 200.0)), reward_enabled=bool(reward.get("enabled", False)), reward_learning_rate=float(reward.get("learning_rate", 0.01)), reward_delay_ticks=int(reward.get("delay_ticks", 0)), reward_clamp_weights=bool(reward.get("clamp_weights", True)), reward_reset_trace=bool(reward.get("reset_trace_after_reward", False)), reward_trace_epsilon=float(reward.get("trace_epsilon", 1e-12)), ) params.validate() return params def validate(self) -> None: """Raise ValueError for invalid plasticity parameters.""" # Non-negative values non_negative = { "stdp.a_plus": self.a_plus, "stdp.a_minus": self.a_minus, "reward.learning_rate": self.reward_learning_rate, "reward.trace_epsilon": self.reward_trace_epsilon, } for name, value in non_negative.items(): if value < 0.0 or not math.isfinite(value): raise ValueError(f"{name} must be finite and >= 0") # Positive time constants if self.tau_plus <= 0.0 or not math.isfinite(self.tau_plus): raise ValueError("stdp.tau_plus must be finite and > 0") if self.tau_minus <= 0.0 or not math.isfinite(self.tau_minus): raise ValueError("stdp.tau_minus must be finite and > 0") if self.eligibility_tau_ticks <= 0.0 or not math.isfinite( self.eligibility_tau_ticks ): raise ValueError("eligibility.tau_ticks must be finite and > 0") # Weight bounds if self.min_weight > self.max_weight: raise ValueError("stdp.min_weight must be <= stdp.max_weight") # Reward delay if self.reward_delay_ticks < 0: raise ValueError("reward.delay_ticks must be >= 0") # Reward requires eligibility if self.reward_enabled and not self.eligibility_enabled: raise ValueError("reward learning requires eligibility.enabled=true") def to_dict(self) -> dict[str, Any]: """Convert to dictionary for serialization.""" return asdict(self) # ============================================================================ # Learning State # ============================================================================ @dataclass(slots=True) class _SynapseLearningState: """State for a single synapse tracked by the learning engine. Attributes: pre_id: ID of the presynaptic neuron. synapse: Reference to the Synapse object. last_pre_tick: Tick of the last presynaptic spike, or None. last_post_tick: Tick of the last postsynaptic spike, or None. eligibility: Eligibility trace for reward-modulated learning. """ pre_id: int synapse: Synapse last_pre_tick: int | None = None last_post_tick: int | None = None eligibility: EligibilityTrace = field(default_factory=EligibilityTrace) @dataclass(slots=True) class _SynapseTickEvent: """Spike event for a synapse in a single tick. Attributes: pre_id: ID of the presynaptic neuron. synapse: Reference to the Synapse object. pre_spiked: Whether the presynaptic neuron spiked this tick. post_spiked: Whether the postsynaptic neuron spiked this tick. """ pre_id: int synapse: Synapse pre_spiked: bool = False post_spiked: bool = False # ============================================================================ # Learning Statistics # ============================================================================ @dataclass(frozen=True, slots=True) class LearningStats: """Runtime statistics for the optional learning layer. Attributes: updates: Total number of update calls. stdp_weight_updates: Number of STDP weight updates applied. reward_weight_updates: Number of reward-modulated weight updates applied. rewards_received: Number of rewards received. rewards_applied: Number of rewards applied. pending_rewards: Number of rewards currently pending (delayed). last_update_ms: Time taken for the last update in milliseconds. total_update_ms: Total time spent on updates in milliseconds. """ updates: int stdp_weight_updates: int reward_weight_updates: int rewards_received: int rewards_applied: int pending_rewards: int last_update_ms: float total_update_ms: float def to_dict(self) -> dict[str, Any]: """Convert to dictionary for serialization.""" return { "updates": self.updates, "stdp_weight_updates": self.stdp_weight_updates, "reward_weight_updates": self.reward_weight_updates, "rewards_received": self.rewards_received, "rewards_applied": self.rewards_applied, "pending_rewards": self.pending_rewards, "last_update_ms": self.last_update_ms, "total_update_ms": self.total_update_ms, } # ============================================================================ # Learning Engine # ============================================================================ class LearningEngine: """Nearest-neighbour STDP, eligibility and reward-modulated plasticity. This engine implements the complete learning pipeline: 1. STDP: Pair-based, nearest-neighbour with configurable amplitudes and time constants 2. Eligibility: Trace accumulation for reward-modulated learning 3. Reward-modulated plasticity: Three-factor learning with delayed rewards The engine observes the network via post-step hooks and applies updates after each completed tick. It maintains state for each synapse and automatically refreshes its topology when the network changes. Example: >>> engine = LearningEngine(network, config) >>> engine.attach() >>> # Run simulation... >>> stats = engine.stats >>> print(f"STDP updates: {stats.stdp_weight_updates}") >>> engine.detach() """ def __init__(self, network: NeuralNetwork, config: Config) -> None: """Initialize the learning engine. Args: network: The neural network to observe and modify. config: Configuration dictionary with learning parameters. Raises: TypeError: If configuration sections are invalid. ValueError: If parameters are invalid. """ self.network = network self.params = LearningParameters.from_config(config) # Stable synapse identity: (pre_id, target_id) tuple. # This is deterministic across process restarts, unlike id(synapse) # which depends on Python object memory addresses (ASLR). self._states: dict[tuple[int, int], _SynapseLearningState] = {} self._incoming: dict[int, list[tuple[int, Synapse]]] = {} self._known_synapse_count = -1 self._pending_rewards: list[RewardSignal] = [] self._attached = False self._updates = 0 self._stdp_weight_updates = 0 self._reward_weight_updates = 0 self._rewards_received = 0 self._rewards_applied = 0 self._last_update_ms = 0.0 self._total_update_ms = 0.0 self.refresh_topology() # ======================================================================== # Properties # ======================================================================== @property def enabled(self) -> bool: """Return whether any plasticity component is active.""" return ( self.params.stdp_enabled or self.params.eligibility_enabled or self.params.reward_enabled ) @property def stats(self) -> LearningStats: """Return an immutable snapshot of learning runtime statistics.""" return LearningStats( updates=self._updates, stdp_weight_updates=self._stdp_weight_updates, reward_weight_updates=self._reward_weight_updates, rewards_received=self._rewards_received, rewards_applied=self._rewards_applied, pending_rewards=len(self._pending_rewards), last_update_ms=self._last_update_ms, total_update_ms=self._total_update_ms, ) @property def is_attached(self) -> bool: """Return whether the engine is attached to the network.""" return self._attached # ======================================================================== # Lifecycle Management # ======================================================================== def attach(self) -> None: """Register this engine on the network's generic post-step hook.""" if not self._attached: self.network.add_post_step_hook(self.update) self._attached = True def detach(self) -> None: """Detach the engine from the network hook.""" if self._attached: self.network.remove_post_step_hook(self.update) self._attached = False def refresh_topology(self) -> None: """Rebuild indexes while preserving state of still-live synapses. Uses stable (pre_id, target_id) keys instead of id(synapse) so that learning state survives process-restart restore. Parallel synapses (disabled in production config) are disambiguated by index. """ incoming: dict[int, list[tuple[int, Synapse]]] = {} live_keys: set[tuple[int, int]] = set() for pre_id, synapses in self.network.synapses.items(): for synapse in synapses: key = (pre_id, synapse.target_id) live_keys.add(key) incoming.setdefault(synapse.target_id, []).append((pre_id, synapse)) if key not in self._states: self._states[key] = _SynapseLearningState( pre_id=pre_id, synapse=synapse, eligibility=EligibilityTrace(self.params.eligibility_tau_ticks), ) # Remove states for synapses that no longer exist self._states = { key: state for key, state in self._states.items() if key in live_keys } self._incoming = incoming self._known_synapse_count = self.network.synapse_count # ======================================================================== # Core Update # ======================================================================== def update(self, step_result: StepResult) -> None: """Observe one completed core tick and apply plasticity updates. This is called automatically by the network's post-step hook. Args: step_result: The result of the completed network step. """ if not self.enabled: return start = time.perf_counter() # Refresh topology if synapses have changed if self.network.synapse_count != self._known_synapse_count: self.refresh_topology() tick = int(step_result.tick) # Use sorted() for deterministic iteration order — set iteration # is hash-based and non-deterministic across process restarts. spike_ids = sorted(set(step_result.spike_ids)) if spike_ids: events: dict[tuple[int, int], _SynapseTickEvent] = {} # Collect presynaptic spikes for pre_id in spike_ids: for synapse in self.network.synapses.get(pre_id, ()): key = (pre_id, synapse.target_id) event = events.setdefault( key, _SynapseTickEvent(pre_id=pre_id, synapse=synapse), ) event.pre_spiked = True # Collect postsynaptic spikes for post_id in spike_ids: for pre_id, synapse in self._incoming.get(post_id, ()): key = (pre_id, synapse.target_id) event = events.setdefault( key, _SynapseTickEvent(pre_id=pre_id, synapse=synapse), ) event.post_spiked = True # Process each synapse event in deterministic order for key in sorted(events): event = events[key] self._process_synapse_event(event, tick) # Apply due rewards self._apply_due_rewards(tick) # Update statistics self._updates += 1 self._last_update_ms = (time.perf_counter() - start) * 1000.0 self._total_update_ms += self._last_update_ms # ======================================================================== # Reward Management # ======================================================================== def set_reward(self, value: float, tick: int) -> None: """Submit an external scalar reward. A zero-delay reward is applied immediately at ``tick`` so callers can reward the just-completed step. Delayed rewards are queued and applied by subsequent ``update`` calls when their due tick is reached. Args: value: The reward value (positive = reinforcement, negative = punishment). tick: The tick at which the reward is emitted. """ if not self.params.reward_enabled: return reward = RewardSignal(value=float(value), tick=int(tick)) self._rewards_received += 1 if self.params.reward_delay_ticks == 0: self._apply_reward(reward, tick) return self._pending_rewards.append(reward) def reset_state(self) -> None: """Forget timing, eligibility and pending rewards without changing weights. This is useful for resetting the learning state between episodes while preserving the learned weights. """ for key in sorted(self._states): state = self._states[key] state.last_pre_tick = None state.last_post_tick = None state.eligibility.reset() self._pending_rewards.clear() # ======================================================================== # Query Methods # ======================================================================== def get_eligibility( self, pre_id: int, post_id: int, tick: int | None = None, ) -> float: """Read the eligibility trace for one non-parallel connection. Args: pre_id: ID of the presynaptic neuron. post_id: ID of the postsynaptic neuron. tick: Tick to read the trace at (default: current tick). Returns: The eligibility trace value. Raises: KeyError: If no synapse exists between pre_id and post_id. ValueError: If multiple parallel synapses match the query. """ matches = [ synapse for synapse in self.network.synapses.get(pre_id, ()) if synapse.target_id == post_id ] if not matches: raise KeyError(f"No synapse {pre_id}->{post_id}") if len(matches) > 1: raise ValueError("Multiple parallel synapses match; query is ambiguous") key = (pre_id, post_id) state = self._states[key] return state.eligibility.read(tick) # ======================================================================== # Internal Methods # ======================================================================== def _process_synapse_event(self, event: _SynapseTickEvent, tick: int) -> None: """Process a single synapse event with STDP and eligibility. Uses stable (pre_id, target_id) key instead of id(synapse) to ensure deterministic behaviour across process restarts. """ key = (event.pre_id, event.synapse.target_id) state = self._states[key] raw_delta = 0.0 # LTD: POST before PRE if event.pre_spiked and state.last_post_tick is not None: dt = state.last_post_tick - tick if dt < 0: raw_delta -= self.params.a_minus * math.exp(dt / self.params.tau_minus) # LTP: PRE before POST if event.post_spiked and state.last_pre_tick is not None: dt = tick - state.last_pre_tick if dt > 0: raw_delta += self.params.a_plus * math.exp(-dt / self.params.tau_plus) # Update eligibility if self.params.eligibility_enabled and raw_delta != 0.0: state.eligibility.add(raw_delta, tick) # Apply STDP weight change if self.params.stdp_enabled and raw_delta != 0.0: old_weight = event.synapse.weight event.synapse.weight = self._bounded_weight(old_weight + raw_delta) if event.synapse.weight != old_weight: event.synapse.mark_dirty() self._stdp_weight_updates += 1 # Update spike timing if event.pre_spiked: state.last_pre_tick = tick if event.post_spiked: state.last_post_tick = tick def _apply_due_rewards(self, tick: int) -> None: """Apply all rewards that are due at the current tick.""" if not self.params.reward_enabled or not self._pending_rewards: return pending: list[RewardSignal] = [] for reward in self._pending_rewards: if reward.is_due(tick, self.params.reward_delay_ticks): due_tick = reward.due_tick(self.params.reward_delay_ticks) self._apply_reward(reward, due_tick) else: pending.append(reward) self._pending_rewards = pending def _apply_reward(self, reward: RewardSignal, effective_tick: int) -> None: """Apply a reward to all synapses with non-zero eligibility.""" changed = False for key in sorted(self._states): state = self._states[key] eligibility = state.eligibility.read(effective_tick) # Skip if eligibility is too small if abs(eligibility) <= self.params.reward_trace_epsilon: continue # Compute weight delta delta = self.params.reward_learning_rate * reward.value * eligibility # Apply weight change old_weight = state.synapse.weight candidate = old_weight + delta if self.params.reward_clamp_weights: candidate = self._bounded_weight(candidate) state.synapse.weight = candidate if candidate != old_weight: state.synapse.mark_dirty() self._reward_weight_updates += 1 changed = True # Reset trace if configured if self.params.reward_reset_trace: state.eligibility.reset() self._rewards_applied += 1 _ = changed # Keep linter happy def _bounded_weight(self, weight: float) -> float: """Clamp a weight to [min_weight, max_weight].""" return max(self.params.min_weight, min(self.params.max_weight, weight)) # ======================================================================== # String Representation # ======================================================================== def __repr__(self) -> str: """Return a string representation of the learning engine.""" return ( f"LearningEngine(enabled={self.enabled}, " f"attached={self._attached}, " f"updates={self._updates}, " f"stdp_updates={self._stdp_weight_updates}, " f"reward_updates={self._reward_weight_updates})" ) # ============================================================================ # Module Exports # ============================================================================ __all__ = [ "LearningEngine", "LearningParameters", "LearningStats", "RewardSignal", ]