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

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.

Browse Source
This commit is contained in:
wls2002
2024-06-01 20:42:42 +08:00
parent 4ad9f0a85a
commit e65200a94e
14 changed files with 281 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
tensorneat/algorithm/neat/ga/crossover/default.py
View File

@@ -1,6 +1,12 @@
import jax, jax.numpy as jnp
from .base import BaseCrossover
from utils.tools import (
extract_node_attrs,
extract_conn_attrs,
set_node_attrs,
set_conn_attrs,
)
class DefaultCrossover(BaseCrossover):
@@ -20,21 +26,33 @@ class DefaultCrossover(BaseCrossover):
# For not homologous genes, use the value of nodes1(winner)
# For homologous genes, use the crossover result between nodes1 and nodes2
new_nodes = jnp.where(
jnp.isnan(nodes1) | jnp.isnan(nodes2),
nodes1,
jax.vmap(genome.node_gene.crossover, in_axes=(None, 0, 0, 0))(state, randkeys1, nodes1, nodes2),
node_attrs1 = jax.vmap(extract_node_attrs)(nodes1)
node_attrs2 = jax.vmap(extract_node_attrs)(nodes2)
new_node_attrs = jnp.where(
jnp.isnan(node_attrs1) | jnp.isnan(node_attrs2), # one of them is nan
node_attrs1, # not homologous genes or both nan, use the value of nodes1(winner)
jax.vmap(genome.node_gene.crossover, in_axes=(None, 0, 0, 0))(
state, randkeys1, node_attrs1, node_attrs2
), # homologous or both nan
)
new_nodes = jax.vmap(set_node_attrs)(nodes1, new_node_attrs)
# crossover connections
con_keys1, con_keys2 = conns1[:, :2], conns2[:, :2]
conns2 = self.align_array(con_keys1, con_keys2, conns2, is_conn=True)
new_conns = jnp.where(
jnp.isnan(conns1) | jnp.isnan(conns2),
conns1,
jax.vmap(genome.conn_gene.crossover, in_axes=(None, 0, 0, 0))(state, randkeys2, conns1, conns2),
conns_attrs1 = jax.vmap(extract_conn_attrs)(conns1)
conns_attrs2 = jax.vmap(extract_conn_attrs)(conns2)
new_conn_attrs = jnp.where(
jnp.isnan(conns_attrs1) | jnp.isnan(conns_attrs2),
conns_attrs1, # not homologous genes or both nan, use the value of conns1(winner)
jax.vmap(genome.conn_gene.crossover, in_axes=(None, 0, 0, 0))(
state, randkeys2, conns_attrs1, conns_attrs2
), # homologous or both nan
)
new_conns = jax.vmap(set_conn_attrs)(conns1, new_conn_attrs)
return new_nodes, new_conns

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

@@ -10,13 +10,17 @@ from utils import (
add_conn,
delete_node_by_pos,
delete_conn_by_pos,
extract_node_attrs,
extract_conn_attrs,
set_node_attrs,
set_conn_attrs,
)
class DefaultMutation(BaseMutation):
def __init__(
self,
conn_add: float = 0.4,
conn_add: float = 0.2,
conn_delete: float = 0,
node_add: float = 0.2,
node_delete: float = 0,
@@ -42,29 +46,38 @@ class DefaultMutation(BaseMutation):
remain_conn_space = jnp.isnan(conns[:, 0]).sum()
def mutate_add_node(key_, nodes_, conns_):
i_key, o_key, idx = self.choice_connection_key(key_, conns_)
"""
add a node while do not influence the output of the network
"""
i_key, o_key, idx = self.choose_connection_key(
key_, conns_
) # choose a connection
def successful_add_node():
# remove the original connection
# remove the original connection and record its attrs
original_attrs = extract_conn_attrs(conns_[idx])
new_conns = delete_conn_by_pos(conns_, idx)
# add a new node
# add a new node with identity attrs
new_nodes = add_node(
nodes_, new_node_key, genome.node_gene.new_custom_attrs(state)
nodes_, new_node_key, genome.node_gene.new_identity_attrs(state)
)
# add two new connections
# first is with identity attrs
new_conns = add_conn(
new_conns,
i_key,
new_node_key,
genome.conn_gene.new_custom_attrs(state),
genome.conn_gene.new_identity_attrs(state),
)
# second is with the origin attrs
new_conns = add_conn(
new_conns,
new_node_key,
o_key,
genome.conn_gene.new_custom_attrs(state),
original_attrs,
)
return new_nodes, new_conns
@@ -76,9 +89,12 @@ class DefaultMutation(BaseMutation):
)
def mutate_delete_node(key_, nodes_, conns_):
"""
delete a node
"""
# randomly choose a node
key, idx = self.choice_node_key(
key, idx = self.choose_node_key(
key_,
nodes_,
genome.input_idx,
@@ -101,17 +117,21 @@ class DefaultMutation(BaseMutation):
return new_nodes, new_conns
return jax.lax.cond(
idx == I_INF,
idx == I_INF, # no available node to delete
lambda: (nodes_, conns_), # do nothing
successful_delete_node,
)
def mutate_add_conn(key_, nodes_, conns_):
"""
add a connection while do not influence the output of the network
"""
# randomly choose two nodes
k1_, k2_ = jax.random.split(key_, num=2)
# input node of the connection can be any node
i_key, from_idx = self.choice_node_key(
i_key, from_idx = self.choose_node_key(
k1_,
nodes_,
genome.input_idx,
@@ -121,7 +141,7 @@ class DefaultMutation(BaseMutation):
)
# output node of the connection can be any node except input node
o_key, to_idx = self.choice_node_key(
o_key, to_idx = self.choose_node_key(
k2_,
nodes_,
genome.input_idx,
@@ -137,8 +157,9 @@ class DefaultMutation(BaseMutation):
return nodes_, conns_
def successful():
# add a connection with zero attrs
return nodes_, add_conn(
conns_, i_key, o_key, genome.conn_gene.new_custom_attrs(state)
conns_, i_key, o_key, genome.conn_gene.new_zero_attrs(state)
)
if genome.network_type == "feedforward":
@@ -164,7 +185,7 @@ class DefaultMutation(BaseMutation):
def mutate_delete_conn(key_, nodes_, conns_):
# randomly choose a connection
i_key, o_key, idx = self.choice_connection_key(key_, conns_)
i_key, o_key, idx = self.choose_connection_key(key_, conns_)
return jax.lax.cond(
idx == I_INF,
@@ -175,42 +196,47 @@ class DefaultMutation(BaseMutation):
k1, k2, k3, k4 = jax.random.split(randkey, num=4)
r1, r2, r3, r4 = jax.random.uniform(k1, shape=(4,))
def no(_, nodes_, conns_):
def nothing(_, nodes_, conns_):
return nodes_, conns_
if self.node_add > 0:
nodes, conns = jax.lax.cond(
r1 < self.node_add, mutate_add_node, no, k1, nodes, conns
r1 < self.node_add, mutate_add_node, nothing, k1, nodes, conns
)
if self.node_delete > 0:
nodes, conns = jax.lax.cond(
r2 < self.node_delete, mutate_delete_node, no, k2, nodes, conns
r2 < self.node_delete, mutate_delete_node, nothing, k2, nodes, conns
)
if self.conn_add > 0:
nodes, conns = jax.lax.cond(
r3 < self.conn_add, mutate_add_conn, no, k3, nodes, conns
r3 < self.conn_add, mutate_add_conn, nothing, k3, nodes, conns
)
if self.conn_delete > 0:
nodes, conns = jax.lax.cond(
r4 < self.conn_delete, mutate_delete_conn, no, k4, nodes, conns
r4 < self.conn_delete, mutate_delete_conn, nothing, k4, nodes, conns
)
return nodes, conns
def mutate_values(self, state, randkey, genome, nodes, conns):
k1, k2 = jax.random.split(randkey, num=2)
nodes_keys = jax.random.split(k1, num=nodes.shape[0])
conns_keys = jax.random.split(k2, num=conns.shape[0])
k1, k2 = jax.random.split(randkey)
nodes_randkeys = jax.random.split(k1, num=genome.max_nodes)
conns_randkeys = jax.random.split(k2, num=genome.max_conns)
new_nodes = jax.vmap(genome.node_gene.mutate, in_axes=(None, 0, 0))(
state, nodes_keys, nodes
node_attrs = jax.vmap(extract_node_attrs)(nodes)
new_node_attrs = jax.vmap(genome.node_gene.mutate, in_axes=(None, 0, 0))(
state, nodes_randkeys, node_attrs
)
new_conns = jax.vmap(genome.conn_gene.mutate, in_axes=(None, 0, 0))(
state, conns_keys, conns
new_nodes = jax.vmap(set_node_attrs)(nodes, new_node_attrs)
conn_attrs = jax.vmap(extract_conn_attrs)(conns)
new_conn_attrs = jax.vmap(genome.conn_gene.mutate, in_axes=(None, 0, 0))(
state, conns_randkeys, conn_attrs
)
new_conns = jax.vmap(set_conn_attrs)(conns, new_conn_attrs)
# nan nodes not changed
new_nodes = jnp.where(jnp.isnan(nodes), jnp.nan, new_nodes)
@@ -218,7 +244,7 @@ class DefaultMutation(BaseMutation):
return new_nodes, new_conns
def choice_node_key(
def choose_node_key(
self,
key,
nodes,
@@ -251,7 +277,7 @@ class DefaultMutation(BaseMutation):
key = jnp.where(idx != I_INF, nodes[idx, 0], jnp.nan)
return key, idx
def choice_connection_key(self, key, conns):
def choose_connection_key(self, key, conns):
"""
Randomly choose a connection key from the given connections.
:return: i_key, o_key, idx

17
tensorneat/algorithm/neat/gene/base.py
View File

@@ -1,3 +1,4 @@
import jax, jax.numpy as jnp
from utils import State
@@ -12,21 +13,25 @@ class BaseGene:
def setup(self, state=State()):
return state
def new_custom_attrs(self, state):
# the attrs which make the least influence on the network, used in add node or add conn in mutation
def new_identity_attrs(self, state):
# the attrs which do identity transformation, used in mutate add node
raise NotImplementedError
def new_random_attrs(self, state, randkey):
# random attributes of the gene. used in initialization.
raise NotImplementedError
def mutate(self, state, randkey, gene):
def mutate(self, state, randkey, attrs):
raise NotImplementedError
def crossover(self, state, randkey, gene1, gene2):
raise NotImplementedError
def crossover(self, state, randkey, attrs1, attrs2):
return jnp.where(
jax.random.normal(randkey, attrs1.shape) > 0,
attrs1,
attrs2,
)
def distance(self, state, gene1, gene2):
def distance(self, state, attrs1, attrs2):
raise NotImplementedError
def forward(self, state, attrs, inputs):

9
tensorneat/algorithm/neat/gene/conn/base.py
View File

@@ -9,12 +9,9 @@ class BaseConnGene(BaseGene):
def __init__(self):
super().__init__()
def crossover(self, state, randkey, gene1, gene2):
return jnp.where(
jax.random.normal(randkey, gene1.shape) > 0,
gene1,
gene2,
)
def new_zero_attrs(self, state):
# the attrs which make the least influence on the network, used in mutate add conn
raise NotImplementedError
def forward(self, state, attrs, inputs):
raise NotImplementedError

20
tensorneat/algorithm/neat/gene/conn/default.py
View File

@@ -25,8 +25,11 @@ class DefaultConnGene(BaseConnGene):
self.weight_mutate_rate = weight_mutate_rate
self.weight_replace_rate = weight_replace_rate
def new_custom_attrs(self, state):
return jnp.array([self.weight_init_mean])
def new_zero_attrs(self, state):
return jnp.array([0.0]) # weight = 0
def new_identity_attrs(self, state):
return jnp.array([1.0]) # weight = 1
def new_random_attrs(self, state, randkey):
weight = (
@@ -35,12 +38,11 @@ class DefaultConnGene(BaseConnGene):
)
return jnp.array([weight])
def mutate(self, state, randkey, conn):
input_index = conn[0]
output_index = conn[1]
def mutate(self, state, randkey, attrs):
weight = attrs[0]
weight = mutate_float(
randkey,
conn[2],
weight,
self.weight_init_mean,
self.weight_init_std,
self.weight_mutate_power,
@@ -48,10 +50,12 @@ class DefaultConnGene(BaseConnGene):
self.weight_replace_rate,
)
return jnp.array([input_index, output_index, weight])
return jnp.array([weight])
def distance(self, state, attrs1, attrs2):
return jnp.abs(attrs1[0] - attrs2[0])
weight1 = attrs1[0]
weight2 = attrs2[0]
return jnp.abs(weight1 - weight2)
def forward(self, state, attrs, inputs):
weight = attrs[0]

7
tensorneat/algorithm/neat/gene/node/base.py
View File

@@ -9,13 +9,6 @@ class BaseNodeGene(BaseGene):
def __init__(self):
super().__init__()
def crossover(self, state, randkey, gene1, gene2):
return jnp.where(
jax.random.normal(randkey, gene1.shape) > 0,
gene1,
gene2,
)
def forward(self, state, attrs, inputs, is_output_node=False):
raise NotImplementedError

73
tensorneat/algorithm/neat/gene/node/default.py
View File

@@ -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

72
tensorneat/algorithm/neat/gene/node/default_without_response.py
View File

@@ -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
@@ -21,12 +21,12 @@ class NodeGeneWithoutResponse(BaseNodeGene):
bias_mutate_power: float = 0.5,
bias_mutate_rate: float = 0.7,
bias_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
@@ -35,39 +35,36 @@ class NodeGeneWithoutResponse(BaseNodeGene):
self.bias_mutate_rate = bias_mutate_rate
self.bias_replace_rate = bias_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.activation_default,
self.aggregation_default,
]
)
[0, 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)
k1, k2, k3 = jax.random.split(randkey, num=3)
bias = jax.random.normal(k1, ()) * self.bias_init_std + self.bias_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, act, agg])
agg = jax.random.choice(k2, self.aggregation_indices)
act = jax.random.choice(k3, self.activation_indices)
def mutate(self, state, randkey, node):
k1, k2, k3, k4 = jax.random.split(randkey, num=4)
index = node[0]
return jnp.array([bias, agg, act])
def mutate(self, state, randkey, attrs):
k1, k2, k3 = jax.random.split(randkey, num=3)
bias, agg, act = attrs
bias = mutate_float(
k1,
node[1],
bias,
self.bias_init_mean,
self.bias_init_std,
self.bias_mutate_power,
@@ -75,32 +72,29 @@ class NodeGeneWithoutResponse(BaseNodeGene):
self.bias_replace_rate,
)
act = mutate_int(
k3, node[2], self.activation_indices, self.activation_replace_rate
)
agg = mutate_int(
k4, node[3], self.aggregation_indices, self.aggregation_replace_rate
k2, agg, self.aggregation_indices, self.aggregation_replace_rate
)
return jnp.array([index, bias, act, agg])
act = mutate_int(k3, act, self.activation_indices, self.activation_replace_rate)
def distance(self, state, node1, node2):
return (
jnp.abs(node1[1] - node2[1]) # bias
+ (node1[2] != node2[2]) # activation
+ (node1[3] != node2[3]) # aggregation
)
return jnp.array([bias, agg, act])
def distance(self, state, attrs1, attrs2):
bias1, agg1, act1 = attrs1
bias2, agg2, act2 = attrs2
return jnp.abs(bias1 - bias2) + (agg1 != agg2) + (act1 != act2)
def forward(self, state, attrs, inputs, is_output_node=False):
bias, act_idx, agg_idx = attrs
bias, agg, act = attrs
z = agg(agg_idx, inputs, self.aggregation_options)
z = agg_func(agg, inputs, self.aggregation_options)
z = bias + 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

100
tensorneat/algorithm/neat/gene/node/normalized.py
View File

@@ -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 NormalizedNode(BaseNodeGene):
bias_mutate_power: float = 0.5,
bias_mutate_rate: float = 0.7,
bias_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,
alpha_init_mean: float = 1.0,
alpha_init_std: float = 1.0,
alpha_mutate_power: float = 0.5,
@@ -47,16 +47,16 @@ class NormalizedNode(BaseNodeGene):
self.bias_mutate_rate = bias_mutate_rate
self.bias_replace_rate = bias_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
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.alpha_init_mean = alpha_init_mean
self.alpha_init_std = alpha_init_std
self.alpha_mutate_power = alpha_mutate_power
@@ -69,38 +69,31 @@ class NormalizedNode(BaseNodeGene):
self.beta_mutate_rate = beta_mutate_rate
self.beta_replace_rate = beta_replace_rate
def new_custom_attrs(self, state):
def new_identity_attrs(self, state):
return jnp.array(
[
self.bias_init_mean,
self.activation_default,
self.aggregation_default,
0, # mean
1, # std
self.alpha_init_mean, # alpha
self.beta_init_mean, # beta
]
)
[0, self.aggregation_default, -1, 0, 1, 1, 0]
) # activation=-1 means Act.identity
def new_random_attrs(self, state, randkey):
k1, k2, k3, k4, k5, k6 = jax.random.split(randkey, num=6)
k1, k2, k3, k4, k5 = jax.random.split(randkey, num=5)
bias = jax.random.normal(k1, ()) * self.bias_init_std + self.bias_init_mean
agg = jax.random.randint(k2, (), 0, len(self.aggregation_options))
act = jax.random.randint(k3, (), 0, len(self.activation_options))
agg = jax.random.randint(k4, (), 0, len(self.aggregation_options))
mean = 0
std = 1
alpha = jax.random.normal(k5, ()) * self.alpha_init_std + self.alpha_init_mean
beta = jax.random.normal(k6, ()) * self.beta_init_std + self.beta_init_mean
alpha = jax.random.normal(k4, ()) * self.alpha_init_std + self.alpha_init_mean
beta = jax.random.normal(k5, ()) * self.beta_init_std + self.beta_init_mean
return jnp.array([bias, act, agg, mean, std, alpha, beta])
return jnp.array([bias, agg, act, mean, std, alpha, beta])
def mutate(self, state, randkey, node):
k1, k2, k3, k4, k5, k6 = jax.random.split(randkey, num=6)
index = node[0]
def mutate(self, state, randkey, attrs):
k1, k2, k3, k4, k5 = jax.random.split(randkey, num=5)
bias, act, agg, mean, std, alpha, beta = attrs
bias = mutate_float(
k1,
node[1],
bias,
self.bias_init_mean,
self.bias_init_std,
self.bias_mutate_power,
@@ -108,20 +101,15 @@ class NormalizedNode(BaseNodeGene):
self.bias_replace_rate,
)
act = mutate_int(
k3, node[2], self.activation_indices, self.activation_replace_rate
)
agg = mutate_int(
k4, node[3], self.aggregation_indices, self.aggregation_replace_rate
k2, agg, self.aggregation_indices, self.aggregation_replace_rate
)
mean = node[4]
std = node[5]
act = mutate_int(k3, act, self.activation_indices, self.activation_replace_rate)
alpha = mutate_float(
k5,
node[6],
k4,
alpha,
self.alpha_init_mean,
self.alpha_init_std,
self.alpha_mutate_power,
@@ -130,8 +118,8 @@ class NormalizedNode(BaseNodeGene):
)
beta = mutate_float(
k6,
node[7],
k5,
beta,
self.beta_init_mean,
self.beta_init_std,
self.beta_mutate_power,
@@ -139,40 +127,42 @@ class NormalizedNode(BaseNodeGene):
self.beta_replace_rate,
)
return jnp.array([index, bias, act, agg, mean, std, alpha, beta])
return jnp.array([bias, agg, act, mean, std, alpha, beta])
def distance(self, state, node1, node2):
def distance(self, state, attrs1, attrs2):
bias1, agg1, act1, mean1, std1, alpha1, beta1 = attrs1
bias2, agg2, act2, mean2, std2, alpha2, beta2 = attrs2
return (
jnp.abs(node1[1] - node2[1]) # bias
+ (node1[2] != node2[2]) # activation
+ (node1[3] != node2[3]) # aggregation
+ (node1[6] - node2[6]) # alpha
+ (node1[7] - node2[7]) # beta
jnp.abs(bias1 - bias2) # bias
+ (agg1 != agg2) # aggregation
+ (act1 != act2) # activation
+ jnp.abs(alpha1 - alpha2) # alpha
+ jnp.abs(beta1 - beta2) # beta
)
def forward(self, state, attrs, inputs, is_output_node=False):
"""
post_act = (agg(inputs) + bias - mean) / std * alpha + beta
"""
bias, act_idx, agg_idx, mean, std, alpha, beta = attrs
bias, agg, act, mean, std, alpha, beta = attrs
z = agg(agg_idx, inputs, self.aggregation_options)
z = agg_func(agg, inputs, self.aggregation_options)
z = bias + z
z = (z - mean) / (std + self.eps) * alpha + beta # normalization
# 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
def update_by_batch(self, state, attrs, batch_inputs, is_output_node=False):
bias, act_idx, agg_idx, mean, std, alpha, beta = attrs
bias, agg, act, mean, std, alpha, beta = attrs
batch_z = jax.vmap(agg, in_axes=(None, 0, None))(
agg_idx, batch_inputs, self.aggregation_options
batch_z = jax.vmap(agg_func, in_axes=(None, 0, None))(
agg, batch_inputs, self.aggregation_options
)
batch_z = bias + batch_z
@@ -192,8 +182,8 @@ class NormalizedNode(BaseNodeGene):
batch_z = jax.lax.cond(
is_output_node,
lambda: batch_z,
lambda: jax.vmap(act, in_axes=(None, 0, None))(
act_idx, batch_z, self.activation_options
lambda: jax.vmap(act_func, in_axes=(None, 0, None))(
act, batch_z, self.activation_options
),
)

18
tensorneat/algorithm/neat/species/default.py
View File

@@ -1,5 +1,12 @@
import jax, jax.numpy as jnp
from utils import State, rank_elements, argmin_with_mask, fetch_first
from utils import (
State,
rank_elements,
argmin_with_mask,
fetch_first,
extract_conn_attrs,
extract_node_attrs,
)
from ..genome import BaseGenome
from .base import BaseSpecies
@@ -557,8 +564,10 @@ class DefaultSpecies(BaseSpecies):
non_homologous_cnt = node_cnt1 + node_cnt2 - 2 * jnp.sum(intersect_mask)
# calculate the distance of homologous nodes
fr_attrs = jax.vmap(extract_node_attrs)(fr)
sr_attrs = jax.vmap(extract_node_attrs)(sr)
hnd = jax.vmap(self.genome.node_gene.distance, in_axes=(None, 0, 0))(
state, fr, sr
state, fr_attrs, sr_attrs
) # homologous node distance
hnd = jnp.where(jnp.isnan(hnd), 0, hnd)
homologous_distance = jnp.sum(hnd * intersect_mask)
@@ -593,8 +602,11 @@ class DefaultSpecies(BaseSpecies):
intersect_mask = jnp.all(fr[:, :2] == sr[:, :2], axis=1) & ~jnp.isnan(fr[:, 0])
non_homologous_cnt = con_cnt1 + con_cnt2 - 2 * jnp.sum(intersect_mask)
fr_attrs = jax.vmap(extract_conn_attrs)(fr)
sr_attrs = jax.vmap(extract_conn_attrs)(sr)
hcd = jax.vmap(self.genome.conn_gene.distance, in_axes=(None, 0, 0))(
state, fr, sr
state, fr_attrs, sr_attrs
) # homologous connection distance
hcd = jnp.where(jnp.isnan(hcd), 0, hcd)
homologous_distance = jnp.sum(hcd * intersect_mask)

4
tensorneat/utils/__init__.py
View File

@@ -1,5 +1,5 @@
from .activation import Act, act, ACT_ALL
from .aggregation import Agg, agg, AGG_ALL
from .activation import Act, act_func, ACT_ALL
from .aggregation import Agg, agg_func, AGG_ALL
from .tools import *
from .graph import *
from .state import State

11
tensorneat/utils/activation.py
View File

@@ -68,11 +68,18 @@ ACT_ALL = (
)
def act(idx, z, act_funcs):
def act_func(idx, z, act_funcs):
"""
calculate activation function for each node
"""
idx = jnp.asarray(idx, dtype=jnp.int32)
# change idx from float to int
res = jax.lax.switch(idx, act_funcs, z)
# -1 means identity activation
res = jax.lax.cond(
idx == -1,
lambda: z,
lambda: jax.lax.switch(idx, act_funcs, z),
)
return res

2
tensorneat/utils/aggregation.py
View File

@@ -53,7 +53,7 @@ class Agg:
AGG_ALL = (Agg.sum, Agg.product, Agg.max, Agg.min, Agg.maxabs, Agg.median, Agg.mean)
def agg(idx, z, agg_funcs):
def agg_func(idx, z, agg_funcs):
"""
calculate activation function for inputs of node
"""

38
tensorneat/utils/tools.py
View File

@@ -47,7 +47,9 @@ def flatten_conns(nodes, unflatten, C):
return jnp.where(
jnp.isnan(unflatten[0, i, j]),
jnp.nan,
jnp.concatenate([jnp.array([node_keys[i], node_keys[j]]), unflatten[:, i, j]]),
jnp.concatenate(
[jnp.array([node_keys[i], node_keys[j]]), unflatten[:, i, j]]
),
)
x, y = jnp.meshgrid(jnp.arange(N), jnp.arange(N), indexing="ij")
@@ -64,6 +66,40 @@ def flatten_conns(nodes, unflatten, C):
return conns
def extract_node_attrs(node):
"""
node: Array(NL, )
extract the attributes of a node
"""
return node[1:] # 0 is for idx
def set_node_attrs(node, attrs):
"""
node: Array(NL, )
attrs: Array(NL-1, )
set the attributes of a node
"""
return node.at[1:].set(attrs) # 0 is for idx
def extract_conn_attrs(conn):
"""
conn: Array(CL, )
extract the attributes of a connection
"""
return conn[2:] # 0, 1 is for in-idx and out-idx
def set_conn_attrs(conn, attrs):
"""
conn: Array(CL, )
attrs: Array(CL-2, )
set the attributes of a connection
"""
return conn.at[2:].set(attrs) # 0, 1 is for in-idx and out-idx
@jit
def fetch_first(mask, default=I_INF) -> Array:
"""