modify the behavior for mutate_add_node and mutate_add_conn. Currently, this two mutation will just change the structure of the network, but not influence the output for the network.

tensorneat/algorithm/neat/gene/node/default.py

@@ -2,7 +2,7 @@ from typing import Tuple
 
 import jax, jax.numpy as jnp
 
-from utils import Act, Agg, act, agg, mutate_int, mutate_float
+from utils import Act, Agg, act_func, agg_func, mutate_int, mutate_float
 from . import BaseNodeGene
 
 
@@ -23,12 +23,12 @@ class DefaultNodeGene(BaseNodeGene):
         response_mutate_power: float = 0.5,
         response_mutate_rate: float = 0.7,
         response_replace_rate: float = 0.1,
-        activation_default: callable = Act.sigmoid,
-        activation_options: Tuple = (Act.sigmoid,),
-        activation_replace_rate: float = 0.1,
         aggregation_default: callable = Agg.sum,
         aggregation_options: Tuple = (Agg.sum,),
         aggregation_replace_rate: float = 0.1,
+        activation_default: callable = Act.sigmoid,
+        activation_options: Tuple = (Act.sigmoid,),
+        activation_replace_rate: float = 0.1,
     ):
         super().__init__()
         self.bias_init_mean = bias_init_mean
@@ -43,25 +43,20 @@ class DefaultNodeGene(BaseNodeGene):
         self.response_mutate_rate = response_mutate_rate
         self.response_replace_rate = response_replace_rate
 
-        self.activation_default = activation_options.index(activation_default)
-        self.activation_options = activation_options
-        self.activation_indices = jnp.arange(len(activation_options))
-        self.activation_replace_rate = activation_replace_rate
-
         self.aggregation_default = aggregation_options.index(aggregation_default)
         self.aggregation_options = aggregation_options
         self.aggregation_indices = jnp.arange(len(aggregation_options))
         self.aggregation_replace_rate = aggregation_replace_rate
 
-    def new_custom_attrs(self, state):
+        self.activation_default = activation_options.index(activation_default)
+        self.activation_options = activation_options
+        self.activation_indices = jnp.arange(len(activation_options))
+        self.activation_replace_rate = activation_replace_rate
+
+    def new_identity_attrs(self, state):
         return jnp.array(
-            [
-                self.bias_init_mean,
-                self.response_init_mean,
-                self.activation_default,
-                self.aggregation_default,
-            ]
-        )
+            [0, 1, self.aggregation_default, -1]
+        )  # activation=-1 means Act.identity
 
     def new_random_attrs(self, state, randkey):
         k1, k2, k3, k4 = jax.random.split(randkey, num=4)
@@ -69,17 +64,17 @@ class DefaultNodeGene(BaseNodeGene):
         res = (
             jax.random.normal(k2, ()) * self.response_init_std + self.response_init_mean
         )
-        act = jax.random.randint(k3, (), 0, len(self.activation_options))
-        agg = jax.random.randint(k4, (), 0, len(self.aggregation_options))
-        return jnp.array([bias, res, act, agg])
+        agg = jax.random.choice(k3, self.aggregation_indices)
+        act = jax.random.choice(k4, self.activation_indices)
 
-    def mutate(self, state, randkey, node):
+        return jnp.array([bias, res, agg, act])
+
+    def mutate(self, state, randkey, attrs):
         k1, k2, k3, k4 = jax.random.split(randkey, num=4)
-        index = node[0]
-
+        bias, res, agg, act = attrs
         bias = mutate_float(
             k1,
-            node[1],
+            bias,
             self.bias_init_mean,
             self.bias_init_std,
             self.bias_mutate_power,
@@ -89,7 +84,7 @@ class DefaultNodeGene(BaseNodeGene):
 
         res = mutate_float(
             k2,
-            node[2],
+            res,
             self.response_init_mean,
             self.response_init_std,
             self.response_mutate_power,
@@ -97,33 +92,33 @@ class DefaultNodeGene(BaseNodeGene):
             self.response_replace_rate,
         )
 
-        act = mutate_int(
-            k3, node[3], self.activation_indices, self.activation_replace_rate
-        )
-
         agg = mutate_int(
-            k4, node[4], self.aggregation_indices, self.aggregation_replace_rate
+            k4, agg, self.aggregation_indices, self.aggregation_replace_rate
         )
 
-        return jnp.array([index, bias, res, act, agg])
+        act = mutate_int(k3, act, self.activation_indices, self.activation_replace_rate)
 
-    def distance(self, state, node1, node2):
+        return jnp.array([bias, res, agg, act])
+
+    def distance(self, state, attrs1, attrs2):
+        bias1, res1, agg1, act1 = attrs1
+        bias2, res2, agg2, act2 = attrs2
         return (
-            jnp.abs(node1[1] - node2[1])  # bias
-            + jnp.abs(node1[2] - node2[2])  # response
-            + (node1[3] != node2[3])  # activation
-            + (node1[4] != node2[4])  # aggregation
+            jnp.abs(bias1 - bias2)  # bias
+            + jnp.abs(res1 - res2)  # response
+            + (agg1 != agg2)  # aggregation
+            + (act1 != act2)  # activation
         )
 
     def forward(self, state, attrs, inputs, is_output_node=False):
-        bias, res, act_idx, agg_idx = attrs
+        bias, res, agg, act = attrs
 
-        z = agg(agg_idx, inputs, self.aggregation_options)
+        z = agg_func(agg, inputs, self.aggregation_options)
         z = bias + res * z
 
         # the last output node should not be activated
         z = jax.lax.cond(
-            is_output_node, lambda: z, lambda: act(act_idx, z, self.activation_options)
+            is_output_node, lambda: z, lambda: act_func(act, z, self.activation_options)
         )
 
         return z