update reinforce synapse

Author: rxng8

ngclearn/components/synapses/modulated/REINFORCESynapse.py

@@ -1,9 +1,8 @@
 # %%
 
 from jax import random, numpy as jnp, jit
-from ngcsimlib.logger import info
-from ngcsimlib.compartment import Compartment
-from ngcsimlib.parser import compilable
+from ngclearn import compilable, Compartment
+
 from ngclearn.utils.model_utils import clip, d_clip
 import jax
 import jax.numpy as jnp
@@ -194,20 +193,19 @@ def evolve(self, dt):
     # @transition(output_compartments=["inputs", "outputs", "objective", "rewards", "dWeights", "accumulated_gradients", "step_count", "seed"])
     # @staticmethod
     @compilable
-    def reset(self, batch_size, shape):
-        preVals = jnp.zeros((batch_size, shape[0]))
-        postVals = jnp.zeros((batch_size, shape[1]))
+    def reset(self):
+        preVals = jnp.zeros((self.batch_size, self.shape[0]))
+        postVals = jnp.zeros((self.batch_size, self.shape[1]))
         inputs = preVals
         outputs = postVals
         objective = jnp.zeros(())
-        rewards = jnp.zeros((batch_size,))
-        dWeights = jnp.zeros(shape)
-        accumulated_gradients = jnp.zeros((shape[0], shape[1] * 2))
+        rewards = jnp.zeros((self.batch_size,))
+        dWeights = jnp.zeros(self.shape)
+        accumulated_gradients = jnp.zeros((self.shape[0], self.shape[1] * 2))
         step_count = jnp.zeros(())
         seed = jax.random.PRNGKey(42)
 
-
-        not self.inputs.targeted and self.inputs.set(inputs)
+        hasattr(self.inputs, 'targeted') and not self.inputs.targeted and self.inputs.set(inputs)
         self.outputs.set(outputs)
         self.objective.set(objective)
         self.rewards.set(rewards)
@@ -260,20 +258,6 @@ def help(cls): ## component help function
                 "hyperparameters": hyperparams}
         return info
 
-    def __repr__(self):
-        comps = [varname for varname in dir(self) if isinstance(getattr(self, varname), Compartment)]
-        maxlen = max(len(c) for c in comps) + 5
-        lines = f"[{self.__class__.__name__}] PATH: {self.name}\n"
-        for c in comps:
-            stats = tensorstats(getattr(self, c).get())
-            if stats is not None:
-                line = [f"{k}: {v}" for k, v in stats.items()]
-                line = ", ".join(line)
-            else:
-                line = "None"
-            lines += f"  {f'({c})'.ljust(maxlen)}{line}\n"
-        return lines
-
 
 if __name__ == '__main__':
     from ngcsimlib.context import Context

tests/components/synapses/modulated/test_REINFORCESynapse.py

@@ -2,17 +2,12 @@
 
 import jax
 from jax import numpy as jnp, random, jit
-from ngcsimlib.context import Context
 import numpy as np
 np.random.seed(42)
 from ngclearn.components.synapses.modulated.REINFORCESynapse import REINFORCESynapse, gaussian_logpdf
-from ngcsimlib.compilers import compile_command, wrap_command
 from numpy.testing import assert_array_equal
 
-from ngcsimlib.compilers.process import Process, transition
-from ngcsimlib.component import Component
-from ngcsimlib.compartment import Compartment
-from ngcsimlib.context import Context
+from ngclearn import Context, MethodProcess
 
 import jax
 import jax.numpy as jnp
@@ -39,22 +34,16 @@ def test_REINFORCESynapse1():
             scalar_stddev=-1.0
         )
 
-        evolve_process = (Process("evolve_proc") >> a.evolve)
-        ctx.wrap_and_add_command(jit(evolve_process.pure), name="adapt")
+        evolve_process = (MethodProcess("evolve_proc") >> a.evolve)
+        reset_process = (MethodProcess("reset_proc") >> a.reset)
 
-        reset_process = (Process("reset_proc") >> a.reset)
-        ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
-
-        @Context.dynamicCommand
         def clamp_inputs(x):
             a.inputs.set(x)
 
-        @Context.dynamicCommand
         def clamp_rewards(x):
             assert x.ndim == 1, "Rewards must be a 1D array"
             a.rewards.set(x)
 
-        @Context.dynamicCommand
         def clamp_weights(x):
             a.weights.set(x)
 
@@ -80,7 +69,7 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
     expected_weights = jnp.concatenate([expected_weights_mu, expected_weights_logstd], axis=-1)
     initial_ngclearn_weights = jnp.concatenate([expected_weights_mu, expected_weights_logstd], axis=-1)[None]
     expected_gradient_list = []
-    ctx.reset()
+    reset_process.run()
 
     # Loop through 3 steps
     for step in range(10):
@@ -94,12 +83,12 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
         clamp_weights(initial_ngclearn_weights)
         clamp_rewards(outputs)
         clamp_inputs(inputs)
-        ctx.adapt(t=1., dt=dt)
-        print(f"[ngclearn] objective: {a.objective.value}")
-        print(f"[ngclearn] weights: {a.weights.value}")
-        print(f"[ngclearn] dWeights: {a.dWeights.value}")
-        print(f"[ngclearn] step_count: {a.step_count.value}")
-        print(f"[ngclearn] accumulated_gradients: {a.accumulated_gradients.value}")
+        evolve_process.run(t=1., dt=dt)
+        print(f"[ngclearn] objective: {a.objective.get()}")
+        print(f"[ngclearn] weights: {a.weights.get()}")
+        print(f"[ngclearn] dWeights: {a.dWeights.get()}")
+        print(f"[ngclearn] step_count: {a.step_count.get()}")
+        print(f"[ngclearn] accumulated_gradients: {a.accumulated_gradients.get()}")
         # -------- Expectation ---------
         print("--------------")
         expected_objective, expected_grads = grad_fn(
@@ -116,12 +105,12 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
         print(f"[Expectation] dWeights: {expected_grads}")
         print(f"[Expectation] objective: {expected_objective}")
         np.testing.assert_allclose(
-            a.dWeights.value[0],
+            a.dWeights.get()[0],
             expected_grads,
             atol=1e-8
         )
         np.testing.assert_allclose(
-            a.objective.value,
+            a.objective.get(),
             expected_objective,
             atol=1e-8
         )
@@ -131,13 +120,13 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
     decay_list = jnp.asarray([decay**i for i in range(len(expected_gradient_list))])[::-1]
     expected_accumulated_gradients = jnp.mean(jnp.stack(expected_gradient_list, 0) * decay_list[:, None, None], axis=0)
     np.testing.assert_allclose(
-        a.accumulated_gradients.value[0],
+        a.accumulated_gradients.get()[0],
         expected_accumulated_gradients,
         atol=1e-9
     )
 
 
-# test_REINFORCESynapse1()
+test_REINFORCESynapse1()
 
 
 def test_REINFORCESynapse2():
@@ -163,22 +152,16 @@ def test_REINFORCESynapse2():
             scalar_stddev=scalar_stddev
         )
 
-        evolve_process = (Process("evolve_proc") >> a.evolve)
-        ctx.wrap_and_add_command(jit(evolve_process.pure), name="adapt")
-
-        reset_process = (Process("reset_proc") >> a.reset)
-        ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
+        evolve_process = (MethodProcess("evolve_proc") >> a.evolve)
+        reset_process = (MethodProcess("reset_proc") >> a.reset)
 
-        @Context.dynamicCommand
         def clamp_inputs(x):
             a.inputs.set(x)
 
-        @Context.dynamicCommand
         def clamp_rewards(x):
             assert x.ndim == 1, "Rewards must be a 1D array"
             a.rewards.set(x)
 
-        @Context.dynamicCommand
         def clamp_weights(x):
             a.weights.set(x)
 
@@ -205,7 +188,7 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
     expected_weights = jnp.concatenate([expected_weights_mu, expected_weights_logstd], axis=-1)
     initial_ngclearn_weights = jnp.concatenate([expected_weights_mu, expected_weights_logstd], axis=-1)[None]
     expected_gradient_list = []
-    ctx.reset()
+    reset_process.run()
 
     # Loop through 3 steps
     for step in range(10):
@@ -219,12 +202,12 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
         clamp_weights(initial_ngclearn_weights)
         clamp_rewards(outputs)
         clamp_inputs(inputs)
-        ctx.adapt(t=1., dt=dt)
-        print(f"[ngclearn] objective: {a.objective.value}")
-        print(f"[ngclearn] weights: {a.weights.value}")
-        print(f"[ngclearn] dWeights: {a.dWeights.value}")
-        print(f"[ngclearn] step_count: {a.step_count.value}")
-        print(f"[ngclearn] accumulated_gradients: {a.accumulated_gradients.value}")
+        evolve_process.run(t=1., dt=dt)
+        print(f"[ngclearn] objective: {a.objective.get()}")
+        print(f"[ngclearn] weights: {a.weights.get()}")
+        print(f"[ngclearn] dWeights: {a.dWeights.get()}")
+        print(f"[ngclearn] step_count: {a.step_count.get()}")
+        print(f"[ngclearn] accumulated_gradients: {a.accumulated_gradients.get()}")
         # -------- Expectation ---------
         print("--------------")
         expected_objective, expected_grads = grad_fn(
@@ -241,12 +224,12 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
         print(f"[Expectation] dWeights: {expected_grads}")
         print(f"[Expectation] objective: {expected_objective}")
         np.testing.assert_allclose(
-            a.dWeights.value[0],
+            a.dWeights.get()[0],
             expected_grads,
             atol=1e-8
         )
         np.testing.assert_allclose(
-            a.objective.value,
+            a.objective.get(),
             expected_objective,
             atol=1e-8
         )
@@ -256,10 +239,10 @@ def fn(params: dict, inputs: jax.Array, outputs: jax.Array, seed: jax.Array):
     decay_list = jnp.asarray([decay**i for i in range(len(expected_gradient_list))])[::-1]
     expected_accumulated_gradients = jnp.mean(jnp.stack(expected_gradient_list, 0) * decay_list[:, None, None], axis=0)
     np.testing.assert_allclose(
-        a.accumulated_gradients.value[0],
+        a.accumulated_gradients.get()[0],
         expected_accumulated_gradients,
         atol=1e-9
     )
 
-# test_REINFORCESynapse2()
+test_REINFORCESynapse2()