62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
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."
|
|
|
|
custom_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 new_random_attrs(self, key):
|
|
return jnp.array([mutate_float(key,
|
|
self.weight_init_mean,
|
|
self.weight_init_mean,
|
|
1.0,
|
|
0,
|
|
0,
|
|
1.0,
|
|
)
|
|
])
|
|
|
|
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
|