new architecture

algorithm/neat/gene/conn/__init__.py

@@ -0,0 +1,2 @@
+from .base import BaseConnGene
+from .default import DefaultConnGene

algorithm/neat/gene/conn/base.py

@@ -0,0 +1,12 @@
+from .. import BaseGene
+
+
+class BaseConnGene(BaseGene):
+    "Base class for connection genes."
+    fixed_attrs = ['input_index', 'output_index', 'enabled']
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, attrs, inputs):
+        raise NotImplementedError

algorithm/neat/gene/conn/default.py

@@ -0,0 +1,51 @@
+import jax.numpy as jnp
+
+from utils import mutate_float
+from . import BaseConnGene
+
+
+class DefaultConnGene(BaseConnGene):
+    "Default connection gene, with the same behavior as in NEAT-python."
+
+    fixed_attrs = ['input_index', 'output_index', 'enabled']
+    attrs = ['weight']
+
+    def __init__(
+            self,
+            weight_init_mean: float = 0.0,
+            weight_init_std: float = 1.0,
+            weight_mutate_power: float = 0.5,
+            weight_mutate_rate: float = 0.8,
+            weight_replace_rate: float = 0.1,
+    ):
+        super().__init__()
+        self.weight_init_mean = weight_init_mean
+        self.weight_init_std = weight_init_std
+        self.weight_mutate_power = weight_mutate_power
+        self.weight_mutate_rate = weight_mutate_rate
+        self.weight_replace_rate = weight_replace_rate
+
+    def new_custom_attrs(self):
+        return jnp.array([self.weight_init_mean])
+
+    def mutate(self, key, conn):
+        input_index = conn[0]
+        output_index = conn[1]
+        enabled = conn[2]
+        weight = mutate_float(key,
+                              conn[3],
+                              self.weight_init_mean,
+                              self.weight_init_std,
+                              self.weight_mutate_power,
+                              self.weight_mutate_rate,
+                              self.weight_replace_rate
+                              )
+
+        return jnp.array([input_index, output_index, enabled, weight])
+
+    def distance(self, attrs1, attrs2):
+        return (attrs1[2] != attrs2[2]) + jnp.abs(attrs1[3] - attrs2[3])  # enable + weight
+
+    def forward(self, attrs, inputs):
+        weight = attrs[0]
+        return inputs * weight