diff --git a/tensorneat/algorithm/neat/ga/crossover/default.py b/tensorneat/algorithm/neat/ga/crossover/default.py index 073f8ef..f55203c 100644 --- a/tensorneat/algorithm/neat/ga/crossover/default.py +++ b/tensorneat/algorithm/neat/ga/crossover/default.py @@ -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 diff --git a/tensorneat/algorithm/neat/ga/mutation/default.py b/tensorneat/algorithm/neat/ga/mutation/default.py index 823e0dd..db3d64c 100644 --- a/tensorneat/algorithm/neat/ga/mutation/default.py +++ b/tensorneat/algorithm/neat/ga/mutation/default.py @@ -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 diff --git a/tensorneat/algorithm/neat/gene/base.py b/tensorneat/algorithm/neat/gene/base.py index 99a57da..8ef4f71 100644 --- a/tensorneat/algorithm/neat/gene/base.py +++ b/tensorneat/algorithm/neat/gene/base.py @@ -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): diff --git a/tensorneat/algorithm/neat/gene/conn/base.py b/tensorneat/algorithm/neat/gene/conn/base.py index 017b671..66de56e 100644 --- a/tensorneat/algorithm/neat/gene/conn/base.py +++ b/tensorneat/algorithm/neat/gene/conn/base.py @@ -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 diff --git a/tensorneat/algorithm/neat/gene/conn/default.py b/tensorneat/algorithm/neat/gene/conn/default.py index e645d36..00520cc 100644 --- a/tensorneat/algorithm/neat/gene/conn/default.py +++ b/tensorneat/algorithm/neat/gene/conn/default.py @@ -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] diff --git a/tensorneat/algorithm/neat/gene/node/base.py b/tensorneat/algorithm/neat/gene/node/base.py index d694a44..64d242e 100644 --- a/tensorneat/algorithm/neat/gene/node/base.py +++ b/tensorneat/algorithm/neat/gene/node/base.py @@ -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 diff --git a/tensorneat/algorithm/neat/gene/node/default.py b/tensorneat/algorithm/neat/gene/node/default.py index bb90780..cd89eab 100644 --- a/tensorneat/algorithm/neat/gene/node/default.py +++ b/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 diff --git a/tensorneat/algorithm/neat/gene/node/default_without_response.py b/tensorneat/algorithm/neat/gene/node/default_without_response.py index e29ae4f..5f52ea8 100644 --- a/tensorneat/algorithm/neat/gene/node/default_without_response.py +++ b/tensorneat/algorithm/neat/gene/node/default_without_response.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 @@ -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 diff --git a/tensorneat/algorithm/neat/gene/node/normalized.py b/tensorneat/algorithm/neat/gene/node/normalized.py index 7dfb045..9e5e4b0 100644 --- a/tensorneat/algorithm/neat/gene/node/normalized.py +++ b/tensorneat/algorithm/neat/gene/node/normalized.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 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 ), ) diff --git a/tensorneat/algorithm/neat/species/default.py b/tensorneat/algorithm/neat/species/default.py index 5dcdd5a..64c3d9f 100644 --- a/tensorneat/algorithm/neat/species/default.py +++ b/tensorneat/algorithm/neat/species/default.py @@ -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) diff --git a/tensorneat/utils/__init__.py b/tensorneat/utils/__init__.py index 2b1711d..c46db75 100644 --- a/tensorneat/utils/__init__.py +++ b/tensorneat/utils/__init__.py @@ -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 diff --git a/tensorneat/utils/activation.py b/tensorneat/utils/activation.py index c4bb494..a8f192e 100644 --- a/tensorneat/utils/activation.py +++ b/tensorneat/utils/activation.py @@ -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 diff --git a/tensorneat/utils/aggregation.py b/tensorneat/utils/aggregation.py index 114abc2..3ada04a 100644 --- a/tensorneat/utils/aggregation.py +++ b/tensorneat/utils/aggregation.py @@ -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 """ diff --git a/tensorneat/utils/tools.py b/tensorneat/utils/tools.py index d0030ce..b33e48f 100644 --- a/tensorneat/utils/tools.py +++ b/tensorneat/utils/tools.py @@ -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: """