fixes (#192)

courtois-neuromod · Jan 19, 2024 · de59662 · de59662
1 parent f5d7c94
commit de59662
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/cneuromax/fitting/neuroevolution/net/cpu/static/rnnfc.py b/cneuromax/fitting/neuroevolution/net/cpu/static/rnnfc.py
@@ -41,6 +41,9 @@ def __init__(self: "CPUStaticRNNFC", config: CPUStaticRNNFCConfig) -> None:
         self.h: Float32[Tensor, " hidden_size"] = torch.zeros(
             size=(config.hidden_size,),
         )
+        for param in self.parameters():
+            param.requires_grad = False
+            param.data = torch.zeros_like(param.data)
 
     def reset(self: "CPUStaticRNNFC") -> None:
         """Resets the hidden state of the RNN."""

diff --git a/cneuromax/fitting/neuroevolution/space/base.py b/cneuromax/fitting/neuroevolution/space/base.py
@@ -18,7 +18,7 @@ class BaseSpaceConfig:
         eval_num_steps: See :attr:`~.BaseSpaceConfig.eval_num_steps`.
     """
 
-    eval_num_steps: An[int, ge(0)] = 0
+    eval_num_steps: An[int, ge(0)] = "${config.eval_num_steps}"  # type: ignore[assignment]
 
 
 class BaseSpace(metaclass=ABCMeta):

diff --git a/cneuromax/fitting/neuroevolution/space/reinforcement.py b/cneuromax/fitting/neuroevolution/space/reinforcement.py
@@ -142,7 +142,7 @@ def evaluate(
                     curr_gen=curr_gen,
                 )
             if agent.curr_eval_num_steps == self.config.eval_num_steps:
-                out["done"] = True
+                break
         self.run_post_eval(agent=agent, out=out, curr_gen=curr_gen)
         return np.array(
             (