tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

make fully stateful in module ga.

Browse Source
This commit is contained in:
wls2002
2024-05-25 16:54:19 +08:00
parent c0a7503168
commit bc378cec7f
1 changed files with 16 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
tensorneat/algorithm/neat/ga/mutation/default.py
View File

@@ -17,15 +17,15 @@ class DefaultMutation(BaseMutation):
self.node_add = node_add self.node_add = node_add
self.node_delete = node_delete self.node_delete = node_delete
def __call__(self, state, randkey, genome, nodes, conns, new_node_key): def __call__(self, state, key, genome, nodes, conns, new_node_key):
k1, k2 = jax.random.split(randkey) k1, k2 = jax.random.split(key)
nodes, conns = self.mutate_structure(k1, genome, nodes, conns, new_node_key) nodes, conns = self.mutate_structure(k1, genome, nodes, conns, new_node_key)
nodes, conns = self.mutate_values(k2, genome, nodes, conns) nodes, conns = self.mutate_values(k2, genome, nodes, conns)
return nodes, conns return nodes, conns
def mutate_structure(self, randkey, genome, nodes, conns, new_node_key): def mutate_structure(self, key, genome, nodes, conns, new_node_key):
def mutate_add_node(key_, nodes_, conns_): def mutate_add_node(key_, nodes_, conns_):
i_key, o_key, idx = self.choice_connection_key(key_, conns_) i_key, o_key, idx = self.choice_connection_key(key_, conns_)
@@ -106,11 +106,11 @@ class DefaultMutation(BaseMutation):
is_already_exist, is_already_exist,
already_exist, already_exist,
lambda: lambda:
jax.lax.cond( jax.lax.cond(
is_cycle, is_cycle,
nothing, nothing,
successful successful
) )
) )
elif genome.network_type == 'recurrent': elif genome.network_type == 'recurrent':
@@ -136,7 +136,7 @@ class DefaultMutation(BaseMutation):
successfully_delete_connection successfully_delete_connection
) )
k1, k2, k3, k4 = jax.random.split(randkey, num=4) k1, k2, k3, k4 = jax.random.split(key, num=4)
r1, r2, r3, r4 = jax.random.uniform(k1, shape=(4,)) r1, r2, r3, r4 = jax.random.uniform(k1, shape=(4,))
def no(key_, nodes_, conns_): def no(key_, nodes_, conns_):
@@ -149,8 +149,8 @@ class DefaultMutation(BaseMutation):
return nodes, conns return nodes, conns
def mutate_values(self, randkey, genome, nodes, conns): def mutate_values(self, key, genome, nodes, conns):
k1, k2 = jax.random.split(randkey, num=2) k1, k2 = jax.random.split(key, num=2)
nodes_keys = jax.random.split(k1, num=nodes.shape[0]) nodes_keys = jax.random.split(k1, num=nodes.shape[0])
conns_keys = jax.random.split(k2, num=conns.shape[0]) conns_keys = jax.random.split(k2, num=conns.shape[0])
@@ -163,11 +163,11 @@ class DefaultMutation(BaseMutation):
return new_nodes, new_conns return new_nodes, new_conns
def choice_node_key(self, rand_key, nodes, input_idx, output_idx, def choice_node_key(self, key, nodes, input_idx, output_idx,
allow_input_keys: bool = False, allow_output_keys: bool = False): allow_input_keys: bool = False, allow_output_keys: bool = False):
""" """
Randomly choose a node key from the given nodes. It guarantees that the chosen node not be the input or output node. Randomly choose a node key from the given nodes. It guarantees that the chosen node not be the input or output node.
:param rand_key: :param key:
:param nodes: :param nodes:
:param input_idx: :param input_idx:
:param output_idx: :param output_idx:
@@ -185,17 +185,17 @@ class DefaultMutation(BaseMutation):
if not allow_output_keys: if not allow_output_keys:
mask = jnp.logical_and(mask, ~jnp.isin(node_keys, output_idx)) mask = jnp.logical_and(mask, ~jnp.isin(node_keys, output_idx))
idx = fetch_random(rand_key, mask) idx = fetch_random(key, mask)
key = jnp.where(idx != I_INT, nodes[idx, 0], jnp.nan) key = jnp.where(idx != I_INT, nodes[idx, 0], jnp.nan)
return key, idx return key, idx
def choice_connection_key(self, rand_key, conns): def choice_connection_key(self, key, conns):
""" """
Randomly choose a connection key from the given connections. Randomly choose a connection key from the given connections.
:return: i_key, o_key, idx :return: i_key, o_key, idx
""" """
idx = fetch_random(rand_key, ~jnp.isnan(conns[:, 0])) idx = fetch_random(key, ~jnp.isnan(conns[:, 0]))
i_key = jnp.where(idx != I_INT, conns[idx, 0], jnp.nan) i_key = jnp.where(idx != I_INT, conns[idx, 0], jnp.nan)
o_key = jnp.where(idx != I_INT, conns[idx, 1], jnp.nan) o_key = jnp.where(idx != I_INT, conns[idx, 1], jnp.nan)