remove create_func....

algorithm/neat/__init__.py

@@ -1 +1,2 @@
 from .neat import NEAT
+from .gene import *

algorithm/neat/gene/normal.py

@@ -66,7 +66,7 @@ class NormalGene(Gene):
     node_attrs = ['bias', 'response', 'aggregation', 'activation']
     conn_attrs = ['weight']
 
-    def __init__(self, config: NormalGeneConfig):
+    def __init__(self, config: NormalGeneConfig = NormalGeneConfig()):
         self.config = config
         self.act_funcs = [Activation.name2func[name] for name in config.activation_options]
         self.agg_funcs = [Aggregation.name2func[name] for name in config.aggregation_options]
@@ -101,7 +101,7 @@ class NormalGene(Gene):
         )
 
     def update(self, state):
-        pass
+        return state
 
     def new_node_attrs(self, state):
         return jnp.array([state.bias_init_mean, state.response_init_mean,

algorithm/neat/gene/recurrent.py

@@ -19,7 +19,7 @@ class RecurrentGeneConfig(NormalGeneConfig):
 
 class RecurrentGene(NormalGene):
 
-    def __init__(self, config: RecurrentGeneConfig):
+    def __init__(self, config: RecurrentGeneConfig = RecurrentGeneConfig()):
         self.config = config
         super().__init__(config)
 

algorithm/neat/neat.py

@@ -28,9 +28,9 @@ class NEAT(Algorithm):
 
         state = state.update(
             P=self.config.basic.pop_size,
-            N=self.config.neat.maximum_nodes,
-            C=self.config.neat.maximum_conns,
-            S=self.config.neat.maximum_species,
+            N=self.config.neat.max_nodes,
+            C=self.config.neat.max_conns,
+            S=self.config.neat.max_species,
             NL=1 + len(self.gene.node_attrs),  # node length = (key) + attributes
             CL=3 + len(self.gene.conn_attrs),  # conn length = (in, out, key) + attributes
             max_stagnation=self.config.neat.max_stagnation,
@@ -80,6 +80,8 @@ class NEAT(Algorithm):
         return state.pop_genomes
 
     def tell_algorithm(self, state: State, fitness):
+        state = self.gene.update(state)
+
         k1, k2, randkey = jax.random.split(state.randkey, 3)
 
         state = state.update(