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

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

@@ -10,13 +10,17 @@ from utils import (
add_conn, add_conn,
delete_node_by_pos, delete_node_by_pos,
delete_conn_by_pos, delete_conn_by_pos,
extract_node_attrs,
extract_conn_attrs,
set_node_attrs,
set_conn_attrs,
) )
class DefaultMutation(BaseMutation): class DefaultMutation(BaseMutation):
def __init__( def __init__(
self, self,
conn_add: float = 0.4, conn_add: float = 0.2,
conn_delete: float = 0, conn_delete: float = 0,
node_add: float = 0.2, node_add: float = 0.2,
node_delete: float = 0, node_delete: float = 0,
@@ -42,29 +46,38 @@ class DefaultMutation(BaseMutation):
remain_conn_space = jnp.isnan(conns[:, 0]).sum() remain_conn_space = jnp.isnan(conns[:, 0]).sum()
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_) """
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(): 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) new_conns = delete_conn_by_pos(conns_, idx)
# add a new node # add a new node with identity attrs
new_nodes = add_node( 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 # add two new connections
# first is with identity attrs
new_conns = add_conn( new_conns = add_conn(
new_conns, new_conns,
i_key, i_key,
new_node_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 = add_conn(
new_conns, new_conns,
new_node_key, new_node_key,
o_key, o_key,
genome.conn_gene.new_custom_attrs(state), original_attrs,
) )
return new_nodes, new_conns return new_nodes, new_conns
@@ -76,9 +89,12 @@ class DefaultMutation(BaseMutation):
) )
def mutate_delete_node(key_, nodes_, conns_): def mutate_delete_node(key_, nodes_, conns_):
"""
delete a node
"""
# randomly choose a node # randomly choose a node
key, idx = self.choice_node_key( key, idx = self.choose_node_key(
key_, key_,
nodes_, nodes_,
genome.input_idx, genome.input_idx,
@@ -101,17 +117,21 @@ class DefaultMutation(BaseMutation):
return new_nodes, new_conns return new_nodes, new_conns
return jax.lax.cond( return jax.lax.cond(
idx == I_INF, idx == I_INF, # no available node to delete
lambda: (nodes_, conns_), # do nothing lambda: (nodes_, conns_), # do nothing
successful_delete_node, successful_delete_node,
) )
def mutate_add_conn(key_, nodes_, conns_): def mutate_add_conn(key_, nodes_, conns_):
"""
add a connection while do not influence the output of the network
"""
# randomly choose two nodes # randomly choose two nodes
k1_, k2_ = jax.random.split(key_, num=2) k1_, k2_ = jax.random.split(key_, num=2)
# input node of the connection can be any node # 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_, k1_,
nodes_, nodes_,
genome.input_idx, genome.input_idx,
@@ -121,7 +141,7 @@ class DefaultMutation(BaseMutation):
) )
# output node of the connection can be any node except input node # 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_, k2_,
nodes_, nodes_,
genome.input_idx, genome.input_idx,
@@ -137,8 +157,9 @@ class DefaultMutation(BaseMutation):
return nodes_, conns_ return nodes_, conns_
def successful(): def successful():
# add a connection with zero attrs
return nodes_, add_conn( 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": if genome.network_type == "feedforward":
@@ -164,7 +185,7 @@ class DefaultMutation(BaseMutation):
def mutate_delete_conn(key_, nodes_, conns_): def mutate_delete_conn(key_, nodes_, conns_):
# randomly choose a connection # 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( return jax.lax.cond(
idx == I_INF, idx == I_INF,
@@ -175,42 +196,47 @@ class DefaultMutation(BaseMutation):
k1, k2, k3, k4 = jax.random.split(randkey, num=4) k1, k2, k3, k4 = jax.random.split(randkey, num=4)
r1, r2, r3, r4 = jax.random.uniform(k1, shape=(4,)) r1, r2, r3, r4 = jax.random.uniform(k1, shape=(4,))
def no(_, nodes_, conns_): def nothing(_, nodes_, conns_):
return nodes_, conns_ return nodes_, conns_
if self.node_add > 0: if self.node_add > 0:
nodes, conns = jax.lax.cond( 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: if self.node_delete > 0:
nodes, conns = jax.lax.cond( 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: if self.conn_add > 0:
nodes, conns = jax.lax.cond( 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: if self.conn_delete > 0:
nodes, conns = jax.lax.cond( 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 return nodes, conns
def mutate_values(self, state, randkey, genome, nodes, conns): def mutate_values(self, state, randkey, genome, nodes, conns):
k1, k2 = jax.random.split(randkey, num=2) k1, k2 = jax.random.split(randkey)
nodes_keys = jax.random.split(k1, num=nodes.shape[0]) nodes_randkeys = jax.random.split(k1, num=genome.max_nodes)
conns_keys = jax.random.split(k2, num=conns.shape[0]) conns_randkeys = jax.random.split(k2, num=genome.max_conns)
new_nodes = jax.vmap(genome.node_gene.mutate, in_axes=(None, 0, 0))( node_attrs = jax.vmap(extract_node_attrs)(nodes)
state, nodes_keys, 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))( new_nodes = jax.vmap(set_node_attrs)(nodes, new_node_attrs)
state, conns_keys, conns
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 # nan nodes not changed
new_nodes = jnp.where(jnp.isnan(nodes), jnp.nan, new_nodes) new_nodes = jnp.where(jnp.isnan(nodes), jnp.nan, new_nodes)
@@ -218,7 +244,7 @@ class DefaultMutation(BaseMutation):
return new_nodes, new_conns return new_nodes, new_conns
def choice_node_key( def choose_node_key(
self, self,
key, key,
nodes, nodes,
@@ -251,7 +277,7 @@ class DefaultMutation(BaseMutation):
key = jnp.where(idx != I_INF, nodes[idx, 0], jnp.nan) key = jnp.where(idx != I_INF, nodes[idx, 0], jnp.nan)
return key, idx 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. Randomly choose a connection key from the given connections.
:return: i_key, o_key, idx :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 from utils import State
@@ -12,21 +13,25 @@ class BaseGene:
def setup(self, state=State()): def setup(self, state=State()):
return state return state
def new_custom_attrs(self, state): def new_identity_attrs(self, state):
# the attrs which make the least influence on the network, used in add node or add conn in mutation # the attrs which do identity transformation, used in mutate add node
raise NotImplementedError raise NotImplementedError
def new_random_attrs(self, state, randkey): def new_random_attrs(self, state, randkey):
# random attributes of the gene. used in initialization. # random attributes of the gene. used in initialization.
raise NotImplementedError raise NotImplementedError
def mutate(self, state, randkey, gene): def mutate(self, state, randkey, attrs):
raise NotImplementedError raise NotImplementedError
def crossover(self, state, randkey, gene1, gene2): def crossover(self, state, randkey, attrs1, attrs2):
raise NotImplementedError 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 raise NotImplementedError
def forward(self, state, attrs, inputs): 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): def __init__(self):
super().__init__() super().__init__()
def crossover(self, state, randkey, gene1, gene2): def new_zero_attrs(self, state):
return jnp.where( # the attrs which make the least influence on the network, used in mutate add conn
jax.random.normal(randkey, gene1.shape) > 0, raise NotImplementedError
gene1,
gene2,
)
def forward(self, state, attrs, inputs): def forward(self, state, attrs, inputs):
raise NotImplementedError 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_mutate_rate = weight_mutate_rate
self.weight_replace_rate = weight_replace_rate self.weight_replace_rate = weight_replace_rate
def new_custom_attrs(self, state): def new_zero_attrs(self, state):
return jnp.array([self.weight_init_mean]) 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): def new_random_attrs(self, state, randkey):
weight = ( weight = (
@@ -35,12 +38,11 @@ class DefaultConnGene(BaseConnGene):
) )
return jnp.array([weight]) return jnp.array([weight])
def mutate(self, state, randkey, conn): def mutate(self, state, randkey, attrs):
input_index = conn[0] weight = attrs[0]
output_index = conn[1]
weight = mutate_float( weight = mutate_float(
randkey, randkey,
conn[2], weight,
self.weight_init_mean, self.weight_init_mean,
self.weight_init_std, self.weight_init_std,
self.weight_mutate_power, self.weight_mutate_power,
@@ -48,10 +50,12 @@ class DefaultConnGene(BaseConnGene):
self.weight_replace_rate, self.weight_replace_rate,
) )
return jnp.array([input_index, output_index, weight]) return jnp.array([weight])
def distance(self, state, attrs1, attrs2): 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): def forward(self, state, attrs, inputs):
weight = attrs[0] weight = attrs[0]

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

@@ -9,13 +9,6 @@ class BaseNodeGene(BaseGene):
def __init__(self): def __init__(self):
super().__init__() 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): def forward(self, state, attrs, inputs, is_output_node=False):
raise NotImplementedError 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 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 from . import BaseNodeGene
@@ -23,12 +23,12 @@ class DefaultNodeGene(BaseNodeGene):
response_mutate_power: float = 0.5, response_mutate_power: float = 0.5,
response_mutate_rate: float = 0.7, response_mutate_rate: float = 0.7,
response_replace_rate: float = 0.1, 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_default: callable = Agg.sum,
aggregation_options: Tuple = (Agg.sum,), aggregation_options: Tuple = (Agg.sum,),
aggregation_replace_rate: float = 0.1, aggregation_replace_rate: float = 0.1,
activation_default: callable = Act.sigmoid,
activation_options: Tuple = (Act.sigmoid,),
activation_replace_rate: float = 0.1,
): ):
super().__init__() super().__init__()
self.bias_init_mean = bias_init_mean self.bias_init_mean = bias_init_mean
@@ -43,25 +43,20 @@ class DefaultNodeGene(BaseNodeGene):
self.response_mutate_rate = response_mutate_rate self.response_mutate_rate = response_mutate_rate
self.response_replace_rate = response_replace_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_default = aggregation_options.index(aggregation_default)
self.aggregation_options = aggregation_options self.aggregation_options = aggregation_options
self.aggregation_indices = jnp.arange(len(aggregation_options)) self.aggregation_indices = jnp.arange(len(aggregation_options))
self.aggregation_replace_rate = aggregation_replace_rate 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( return jnp.array(
[ [0, 1, self.aggregation_default, -1]
self.bias_init_mean, ) # activation=-1 means Act.identity
self.response_init_mean,
self.activation_default,
self.aggregation_default,
]
)
def new_random_attrs(self, state, randkey): def new_random_attrs(self, state, randkey):
k1, k2, k3, k4 = jax.random.split(randkey, num=4) k1, k2, k3, k4 = jax.random.split(randkey, num=4)
@@ -69,17 +64,17 @@ class DefaultNodeGene(BaseNodeGene):
res = ( res = (
jax.random.normal(k2, ()) * self.response_init_std + self.response_init_mean 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.choice(k3, self.aggregation_indices)
agg = jax.random.randint(k4, (), 0, len(self.aggregation_options)) act = jax.random.choice(k4, self.activation_indices)
return jnp.array([bias, res, act, agg])
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) k1, k2, k3, k4 = jax.random.split(randkey, num=4)
index = node[0] bias, res, agg, act = attrs
bias = mutate_float( bias = mutate_float(
k1, k1,
node[1], bias,
self.bias_init_mean, self.bias_init_mean,
self.bias_init_std, self.bias_init_std,
self.bias_mutate_power, self.bias_mutate_power,
@@ -89,7 +84,7 @@ class DefaultNodeGene(BaseNodeGene):
res = mutate_float( res = mutate_float(
k2, k2,
node[2], res,
self.response_init_mean, self.response_init_mean,
self.response_init_std, self.response_init_std,
self.response_mutate_power, self.response_mutate_power,
@@ -97,33 +92,33 @@ class DefaultNodeGene(BaseNodeGene):
self.response_replace_rate, self.response_replace_rate,
) )
act = mutate_int(
k3, node[3], self.activation_indices, self.activation_replace_rate
)
agg = mutate_int( 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 ( return (
jnp.abs(node1[1] - node2[1]) # bias jnp.abs(bias1 - bias2) # bias
+ jnp.abs(node1[2] - node2[2]) # response + jnp.abs(res1 - res2) # response
+ (node1[3] != node2[3]) # activation + (agg1 != agg2) # aggregation
+ (node1[4] != node2[4]) # aggregation + (act1 != act2) # activation
) )
def forward(self, state, attrs, inputs, is_output_node=False): 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 z = bias + res * z
# the last output node should not be activated # the last output node should not be activated
z = jax.lax.cond( 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 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 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 from . import BaseNodeGene
@@ -21,12 +21,12 @@ class NodeGeneWithoutResponse(BaseNodeGene):
bias_mutate_power: float = 0.5, bias_mutate_power: float = 0.5,
bias_mutate_rate: float = 0.7, bias_mutate_rate: float = 0.7,
bias_replace_rate: float = 0.1, 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_default: callable = Agg.sum,
aggregation_options: Tuple = (Agg.sum,), aggregation_options: Tuple = (Agg.sum,),
aggregation_replace_rate: float = 0.1, aggregation_replace_rate: float = 0.1,
activation_default: callable = Act.sigmoid,
activation_options: Tuple = (Act.sigmoid,),
activation_replace_rate: float = 0.1,
): ):
super().__init__() super().__init__()
self.bias_init_mean = bias_init_mean self.bias_init_mean = bias_init_mean
@@ -35,39 +35,36 @@ class NodeGeneWithoutResponse(BaseNodeGene):
self.bias_mutate_rate = bias_mutate_rate self.bias_mutate_rate = bias_mutate_rate
self.bias_replace_rate = bias_replace_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_default = aggregation_options.index(aggregation_default)
self.aggregation_options = aggregation_options self.aggregation_options = aggregation_options
self.aggregation_indices = jnp.arange(len(aggregation_options)) self.aggregation_indices = jnp.arange(len(aggregation_options))
self.aggregation_replace_rate = aggregation_replace_rate 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( return jnp.array(
[ [0, self.aggregation_default, -1]
self.bias_init_mean, ) # activation=-1 means Act.identity
self.activation_default,
self.aggregation_default,
]
)
def new_random_attrs(self, state, randkey): 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 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.choice(k2, self.aggregation_indices)
agg = jax.random.randint(k4, (), 0, len(self.aggregation_options)) act = jax.random.choice(k3, self.activation_indices)
return jnp.array([bias, act, agg])
def mutate(self, state, randkey, node): return jnp.array([bias, agg, act])
k1, k2, k3, k4 = jax.random.split(randkey, num=4)
index = node[0] def mutate(self, state, randkey, attrs):
k1, k2, k3 = jax.random.split(randkey, num=3)
bias, agg, act = attrs
bias = mutate_float( bias = mutate_float(
k1, k1,
node[1], bias,
self.bias_init_mean, self.bias_init_mean,
self.bias_init_std, self.bias_init_std,
self.bias_mutate_power, self.bias_mutate_power,
@@ -75,32 +72,29 @@ class NodeGeneWithoutResponse(BaseNodeGene):
self.bias_replace_rate, self.bias_replace_rate,
) )
act = mutate_int(
k3, node[2], self.activation_indices, self.activation_replace_rate
)
agg = mutate_int( 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.array([bias, agg, act])
return (
jnp.abs(node1[1] - node2[1]) # bias def distance(self, state, attrs1, attrs2):
+ (node1[2] != node2[2]) # activation bias1, agg1, act1 = attrs1
+ (node1[3] != node2[3]) # aggregation bias2, agg2, act2 = attrs2
)
return jnp.abs(bias1 - bias2) + (agg1 != agg2) + (act1 != act2)
def forward(self, state, attrs, inputs, is_output_node=False): 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 z = bias + z
# the last output node should not be activated # the last output node should not be activated
z = jax.lax.cond( 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 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 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 from . import BaseNodeGene
@@ -23,12 +23,12 @@ class NormalizedNode(BaseNodeGene):
bias_mutate_power: float = 0.5, bias_mutate_power: float = 0.5,
bias_mutate_rate: float = 0.7, bias_mutate_rate: float = 0.7,
bias_replace_rate: float = 0.1, 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_default: callable = Agg.sum,
aggregation_options: Tuple = (Agg.sum,), aggregation_options: Tuple = (Agg.sum,),
aggregation_replace_rate: float = 0.1, 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_mean: float = 1.0,
alpha_init_std: float = 1.0, alpha_init_std: float = 1.0,
alpha_mutate_power: float = 0.5, alpha_mutate_power: float = 0.5,
@@ -47,16 +47,16 @@ class NormalizedNode(BaseNodeGene):
self.bias_mutate_rate = bias_mutate_rate self.bias_mutate_rate = bias_mutate_rate
self.bias_replace_rate = bias_replace_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_default = aggregation_options.index(aggregation_default)
self.aggregation_options = aggregation_options self.aggregation_options = aggregation_options
self.aggregation_indices = jnp.arange(len(aggregation_options)) self.aggregation_indices = jnp.arange(len(aggregation_options))
self.aggregation_replace_rate = aggregation_replace_rate 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_mean = alpha_init_mean
self.alpha_init_std = alpha_init_std self.alpha_init_std = alpha_init_std
self.alpha_mutate_power = alpha_mutate_power self.alpha_mutate_power = alpha_mutate_power
@@ -69,38 +69,31 @@ class NormalizedNode(BaseNodeGene):
self.beta_mutate_rate = beta_mutate_rate self.beta_mutate_rate = beta_mutate_rate
self.beta_replace_rate = beta_replace_rate self.beta_replace_rate = beta_replace_rate
def new_custom_attrs(self, state): def new_identity_attrs(self, state):
return jnp.array( return jnp.array(
[ [0, self.aggregation_default, -1, 0, 1, 1, 0]
self.bias_init_mean, ) # activation=-1 means Act.identity
self.activation_default,
self.aggregation_default,
0, # mean
1, # std
self.alpha_init_mean, # alpha
self.beta_init_mean, # beta
]
)
def new_random_attrs(self, state, randkey): 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 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)) act = jax.random.randint(k3, (), 0, len(self.activation_options))
agg = jax.random.randint(k4, (), 0, len(self.aggregation_options))
mean = 0 mean = 0
std = 1 std = 1
alpha = jax.random.normal(k5, ()) * self.alpha_init_std + self.alpha_init_mean alpha = jax.random.normal(k4, ()) * self.alpha_init_std + self.alpha_init_mean
beta = jax.random.normal(k6, ()) * self.beta_init_std + self.beta_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): def mutate(self, state, randkey, attrs):
k1, k2, k3, k4, k5, k6 = jax.random.split(randkey, num=6) k1, k2, k3, k4, k5 = jax.random.split(randkey, num=5)
index = node[0] bias, act, agg, mean, std, alpha, beta = attrs
bias = mutate_float( bias = mutate_float(
k1, k1,
node[1], bias,
self.bias_init_mean, self.bias_init_mean,
self.bias_init_std, self.bias_init_std,
self.bias_mutate_power, self.bias_mutate_power,
@@ -108,20 +101,15 @@ class NormalizedNode(BaseNodeGene):
self.bias_replace_rate, self.bias_replace_rate,
) )
act = mutate_int(
k3, node[2], self.activation_indices, self.activation_replace_rate
)
agg = mutate_int( 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] act = mutate_int(k3, act, self.activation_indices, self.activation_replace_rate)
std = node[5]
alpha = mutate_float( alpha = mutate_float(
k5, k4,
node[6], alpha,
self.alpha_init_mean, self.alpha_init_mean,
self.alpha_init_std, self.alpha_init_std,
self.alpha_mutate_power, self.alpha_mutate_power,
@@ -130,8 +118,8 @@ class NormalizedNode(BaseNodeGene):
) )
beta = mutate_float( beta = mutate_float(
k6, k5,
node[7], beta,
self.beta_init_mean, self.beta_init_mean,
self.beta_init_std, self.beta_init_std,
self.beta_mutate_power, self.beta_mutate_power,
@@ -139,40 +127,42 @@ class NormalizedNode(BaseNodeGene):
self.beta_replace_rate, 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 ( return (
jnp.abs(node1[1] - node2[1]) # bias jnp.abs(bias1 - bias2) # bias
+ (node1[2] != node2[2]) # activation + (agg1 != agg2) # aggregation
+ (node1[3] != node2[3]) # aggregation + (act1 != act2) # activation
+ (node1[6] - node2[6]) # alpha + jnp.abs(alpha1 - alpha2) # alpha
+ (node1[7] - node2[7]) # beta + jnp.abs(beta1 - beta2) # beta
) )
def forward(self, state, attrs, inputs, is_output_node=False): def forward(self, state, attrs, inputs, is_output_node=False):
""" """
post_act = (agg(inputs) + bias - mean) / std * alpha + beta 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 = bias + z
z = (z - mean) / (std + self.eps) * alpha + beta # normalization z = (z - mean) / (std + self.eps) * alpha + beta # normalization
# the last output node should not be activated # the last output node should not be activated
z = jax.lax.cond( 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 return z
def update_by_batch(self, state, attrs, batch_inputs, is_output_node=False): 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))( batch_z = jax.vmap(agg_func, in_axes=(None, 0, None))(
agg_idx, batch_inputs, self.aggregation_options agg, batch_inputs, self.aggregation_options
) )
batch_z = bias + batch_z batch_z = bias + batch_z
@@ -192,8 +182,8 @@ class NormalizedNode(BaseNodeGene):
batch_z = jax.lax.cond( batch_z = jax.lax.cond(
is_output_node, is_output_node,
lambda: batch_z, lambda: batch_z,
lambda: jax.vmap(act, in_axes=(None, 0, None))( lambda: jax.vmap(act_func, in_axes=(None, 0, None))(
act_idx, batch_z, self.activation_options 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 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 ..genome import BaseGenome
from .base import BaseSpecies from .base import BaseSpecies
@@ -557,8 +564,10 @@ class DefaultSpecies(BaseSpecies):
non_homologous_cnt = node_cnt1 + node_cnt2 - 2 * jnp.sum(intersect_mask) non_homologous_cnt = node_cnt1 + node_cnt2 - 2 * jnp.sum(intersect_mask)
# calculate the distance of homologous nodes # 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))( hnd = jax.vmap(self.genome.node_gene.distance, in_axes=(None, 0, 0))(
state, fr, sr state, fr_attrs, sr_attrs
) # homologous node distance ) # homologous node distance
hnd = jnp.where(jnp.isnan(hnd), 0, hnd) hnd = jnp.where(jnp.isnan(hnd), 0, hnd)
homologous_distance = jnp.sum(hnd * intersect_mask) 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]) 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) 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))( hcd = jax.vmap(self.genome.conn_gene.distance, in_axes=(None, 0, 0))(
state, fr, sr state, fr_attrs, sr_attrs
) # homologous connection distance ) # homologous connection distance
hcd = jnp.where(jnp.isnan(hcd), 0, hcd) hcd = jnp.where(jnp.isnan(hcd), 0, hcd)
homologous_distance = jnp.sum(hcd * intersect_mask) 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 .activation import Act, act_func, ACT_ALL
from .aggregation import Agg, agg, AGG_ALL from .aggregation import Agg, agg_func, AGG_ALL
from .tools import * from .tools import *
from .graph import * from .graph import *
from .state import State 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 calculate activation function for each node
""" """
idx = jnp.asarray(idx, dtype=jnp.int32) idx = jnp.asarray(idx, dtype=jnp.int32)
# change idx from float to int # 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 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) 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 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( return jnp.where(
jnp.isnan(unflatten[0, i, j]), jnp.isnan(unflatten[0, i, j]),
jnp.nan, 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") x, y = jnp.meshgrid(jnp.arange(N), jnp.arange(N), indexing="ij")
@@ -64,6 +66,40 @@ def flatten_conns(nodes, unflatten, C):
return conns 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 @jit
def fetch_first(mask, default=I_INF) -> Array: def fetch_first(mask, default=I_INF) -> Array:
""" """