tweaked trace-stdp and mstdpet

Author: Alexander Ororbia

ngclearn/components/synapses/hebbian/traceSTDPSynapse.py

@@ -81,6 +81,7 @@ def __init__(
         self.Aminus = A_minus ## LTD strength
         self.Rscale = resist_scale ## post-transformation scale factor
         self.w_bound = w_bound #1. ## soft weight constraint
+        self.w_eps = 0. ## w_eps = 0.01
 
         ## Compartment setup
         preVals = jnp.zeros((self.batch_size, shape[0]))
@@ -113,6 +114,7 @@ def _compute_update(
         else:
             ## calculate post-synaptic term
             dWpost = jnp.matmul((x_pre - x_tar).T, post * Aplus)
+
             dWpre = 0.
             if Aminus > 0.:
                 ## calculate pre-synaptic term
@@ -124,16 +126,16 @@ def _compute_update(
     @transition(output_compartments=["weights", "dWeights"])
     @staticmethod
     def evolve(
-            dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights, eta
+            dt, w_bound, w_eps, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights, eta
     ):
         dWeights = TraceSTDPSynapse._compute_update(
             dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
         )
         ## do a gradient ascent update/shift
         weights = weights + dWeights * eta
         ## enforce non-negativity
-        eps = 0.01 # 0.001
-        weights = jnp.clip(weights, eps, w_bound - eps)  # jnp.abs(w_bound))
+        #w_eps = 0. # 0.01 # 0.001
+        weights = jnp.clip(weights, w_eps, w_bound - w_eps)  # jnp.abs(w_bound))
         return weights, dWeights
 
     @transition(output_compartments=["inputs", "outputs", "preSpike", "postSpike", "preTrace", "postTrace", "dWeights"])

ngclearn/components/synapses/modulated/MSTDPETSynapse.py

@@ -81,6 +81,7 @@ def __init__(
             name, shape, A_plus, A_minus, eta=eta, mu=mu, pretrace_target=pretrace_target, weight_init=weight_init,
             resist_scale=resist_scale, p_conn=p_conn, w_bound=w_bound, batch_size=batch_size, **kwargs
         )
+        self.w_eps = 0.
         ## MSTDP/MSTDP-ET meta-parameters
         self.tau_elg = tau_elg
         self.elg_decay = elg_decay
@@ -91,28 +92,23 @@ def __init__(
     @transition(output_compartments=["weights", "dWeights", "eligibility"])
     @staticmethod
     def evolve(
-            dt, w_bound, preTrace_target, mu, Aplus, Aminus, tau_elg, elg_decay, preSpike, postSpike, preTrace,
-            postTrace, weights, eta, modulator, eligibility
-    ):
-        ## compute local synaptic update (via STDP)
-        dW_dt = TraceSTDPSynapse._compute_update(
-            dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
-        ) ## produce dW/dt (ODE for synaptic change dynamics)
+            dt, w_bound, w_eps, preTrace_target, mu, Aplus, Aminus, tau_elg, elg_decay, preSpike, postSpike, preTrace,
+            postTrace, weights, dWeights, eta, modulator, eligibility
+    ):        
         if tau_elg > 0.: ## perform dynamics of M-STDP-ET
-            ## update eligibility trace given current local update
-            # dElg_dt = -eligibility * elg_decay + dW_dt * update_scale
-            # eligibility = eligibility + dElg_dt * dt/elg_tau
-            eligibility = eligibility * jnp.exp(-dt / tau_elg) * elg_decay + dW_dt
-        else: ## perform dynamics of M-STDP (no eligibility trace)
-            eligibility = dW_dt
-        ## Perform a trace/update times a modulatory signal (e.g., reward)
-        dWeights = eligibility * modulator
-
+            eligibility = eligibility * jnp.exp(-dt / tau_elg) * elg_decay + dWeights/tau_elg
+        else: ## otherwise, just do M-STDP
+            eligibility = dWeights ## dynamics of M-STDP had no eligibility tracing
         ## do a gradient ascent update/shift
-        weights = weights + dWeights * eta ## modulate update
-        ## enforce non-negativity
-        eps = 0.01
-        weights = jnp.clip(weights, eps, w_bound - eps)  # jnp.abs(w_bound))
+        weights = weights + eligibility * modulator * eta  ## do modulated update
+        dW_dt = TraceSTDPSynapse._compute_update( ## use Hebbian/STDP rule to obtain a non-modulated update
+            dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
+        )
+        dWeights = dW_dt ## can think of this as eligibility at time t
+
+        #w_eps = 0. # 0.01
+        weights = jnp.clip(weights, w_eps, w_bound - w_eps)  # jnp.abs(w_bound))
+
         return weights, dWeights, eligibility
 
     @transition(