complete HyperNEAT!

algorithm/neat/__init__.py

@@ -1,3 +1,2 @@
 from .neat import NEAT
-from .gene import NormalGene, RecurrentGene
-from .pipeline import Pipeline
+from .gene import BaseGene, NormalGene, RecurrentGene

algorithm/neat/gene/base.py

@@ -33,12 +33,10 @@ class BaseGene:
     def distance_conn(state, conn1: Array, conn2: Array):
         return conn1
 
-
     @staticmethod
-    def forward_transform(nodes, conns):
+    def forward_transform(state, nodes, conns):
         return nodes, conns
 
-
     @staticmethod
     def create_forward(config):
-        return None
+        return None

algorithm/neat/gene/normal.py

@@ -4,7 +4,7 @@ from jax import Array, numpy as jnp
 from .base import BaseGene
 from .activation import Activation
 from .aggregation import Aggregation
-from ..utils import unflatten_connections, I_INT
+from algorithm.utils import unflatten_connections, I_INT
 from ..genome import topological_sort
 
 
@@ -84,7 +84,7 @@ class NormalGene(BaseGene):
         return (con1[2] != con2[2]) + jnp.abs(con1[3] - con2[3])  # enable + weight
 
     @staticmethod
-    def forward_transform(nodes, conns):
+    def forward_transform(state, nodes, conns):
         u_conns = unflatten_connections(nodes, conns)
         conn_enable = jnp.where(~jnp.isnan(u_conns[0]), True, False)
 

algorithm/neat/gene/recurrent.py

@@ -4,13 +4,13 @@ from jax import Array, numpy as jnp, vmap
 from .normal import NormalGene
 from .activation import Activation
 from .aggregation import Aggregation
-from ..utils import unflatten_connections, I_INT
+from algorithm.utils import unflatten_connections
 
 
 class RecurrentGene(NormalGene):
 
     @staticmethod
-    def forward_transform(nodes, conns):
+    def forward_transform(state, nodes, conns):
         u_conns = unflatten_connections(nodes, conns)
 
         # remove un-enable connections and remove enable attr

algorithm/neat/genome/basic.py

@@ -6,7 +6,7 @@ from jax import Array, numpy as jnp
 
 from algorithm import State
 from ..gene import BaseGene
-from ..utils import fetch_first
+from algorithm.utils import fetch_first
 
 
 def initialize_genomes(state: State, gene_type: Type[BaseGene]):
@@ -48,6 +48,7 @@ def count(nodes: Array, cons: Array):
     cons_cnt = jnp.sum(~jnp.isnan(cons[:, 0]))
     return node_cnt, cons_cnt
 
+
 def add_node(nodes: Array, cons: Array, new_key: int, attrs: Array) -> Tuple[Array, Array]:
     """
     Add a new node to the genome.

algorithm/neat/genome/graph.py

@@ -6,7 +6,7 @@ Only used in feed-forward networks.
 import jax
 from jax import jit, Array, numpy as jnp
 
-from ..utils import fetch_first, I_INT
+from algorithm.utils import fetch_first, I_INT
 
 
 @jit

algorithm/neat/genome/mutate.py

@@ -4,9 +4,9 @@ import jax
 from jax import Array, numpy as jnp, vmap
 
 from algorithm import State
-from .basic import add_node, add_connection, delete_node_by_idx, delete_connection_by_idx, count
+from .basic import add_node, add_connection, delete_node_by_idx, delete_connection_by_idx
 from .graph import check_cycles
-from ..utils import fetch_random, fetch_first, I_INT, unflatten_connections
+from algorithm.utils import fetch_random, fetch_first, I_INT, unflatten_connections
 from ..gene import BaseGene
 
 

algorithm/neat/neat.py

@@ -3,22 +3,25 @@ from typing import Type
 import jax
 import jax.numpy as jnp
 
-from algorithm.state import State
+from algorithm import Algorithm, State
 from .gene import BaseGene
 from .genome import initialize_genomes
 from .population import create_tell
 
 
-class NEAT:
+class NEAT(Algorithm):
     def __init__(self, config, gene_type: Type[BaseGene]):
+        super().__init__()
         self.config = config
         self.gene_type = gene_type
 
-        self.tell_func = jax.jit(create_tell(config, self.gene_type))
+        self.tell = create_tell(config, self.gene_type)
+        self.ask = None
+        self.forward = self.gene_type.create_forward(config)
+        self.forward_transform = self.gene_type.forward_transform
 
-    def setup(self, randkey):
-
-        state = State(
+    def setup(self, randkey, state=State()):
+        state = state.update(
             P=self.config['pop_size'],
             N=self.config['maximum_nodes'],
             C=self.config['maximum_conns'],
@@ -69,7 +72,4 @@ class NEAT:
         # move to device
         state = jax.device_put(state)
 
-        return state
-
-    def step(self, state, fitness):
-        return self.tell_func(state, fitness)
+        return state

algorithm/neat/pipeline.py

@@ -1,77 +0,0 @@
-import time
-from typing import Union, Callable
-
-import jax
-from jax import vmap, jit
-import numpy as np
-
-
-class Pipeline:
-    """
-    Neat algorithm pipeline.
-    """
-
-    def __init__(self, config, algorithm):
-        self.config = config
-        self.algorithm = algorithm
-        randkey = jax.random.PRNGKey(config['random_seed'])
-        self.state = algorithm.setup(randkey)
-
-        self.best_genome = None
-        self.best_fitness = float('-inf')
-        self.generation_timestamp = time.time()
-
-        self.evaluate_time = 0
-
-        self.forward_func = algorithm.gene_type.create_forward(config)
-        self.batch_forward_func = jit(vmap(self.forward_func, in_axes=(0, None)))
-        self.pop_batch_forward_func = jit(vmap(self.batch_forward_func, in_axes=(None, 0)))
-
-        self.pop_transform_func = jit(vmap(algorithm.gene_type.forward_transform))
-
-    def ask(self):
-        pop_transforms = self.pop_transform_func(self.state.pop_nodes, self.state.pop_conns)
-        return lambda inputs: self.pop_batch_forward_func(inputs, pop_transforms)
-
-    def tell(self, fitness):
-        self.state = self.algorithm.step(self.state, fitness)
-
-    def auto_run(self, fitness_func, analysis: Union[Callable, str] = "default"):
-        for _ in range(self.config['generation_limit']):
-            forward_func = self.ask()
-
-            fitnesses = fitness_func(forward_func)
-
-            if analysis is not None:
-                if analysis == "default":
-                    self.default_analysis(fitnesses)
-                else:
-                    assert callable(analysis), f"What the fuck you passed in? A {analysis}?"
-                    analysis(fitnesses)
-
-            if max(fitnesses) >= self.config['fitness_threshold']:
-                print("Fitness limit reached!")
-                return self.best_genome
-
-            self.tell(fitnesses)
-        print("Generation limit reached!")
-        return self.best_genome
-
-    def default_analysis(self, fitnesses):
-        max_f, min_f, mean_f, std_f = max(fitnesses), min(fitnesses), np.mean(fitnesses), np.std(fitnesses)
-
-        new_timestamp = time.time()
-        cost_time = new_timestamp - self.generation_timestamp
-        self.generation_timestamp = new_timestamp
-
-        max_idx = np.argmax(fitnesses)
-        if fitnesses[max_idx] > self.best_fitness:
-            self.best_fitness = fitnesses[max_idx]
-            self.best_genome = (self.state.pop_nodes[max_idx], self.state.pop_conns[max_idx])
-
-        member_count = jax.device_get(self.state.species_info[:, 3])
-        species_sizes = [int(i) for i in member_count if i > 0]
-
-        print(f"Generation: {self.state.generation}",
-              f"species: {len(species_sizes)}, {species_sizes}",
-              f"fitness: {max_f}, {min_f}, {mean_f}, {std_f}, Cost time: {cost_time}")

algorithm/neat/population.py

@@ -3,7 +3,7 @@ from typing import Type
 import jax
 from jax import numpy as jnp, vmap
 
-from .utils import rank_elements, fetch_first
+from algorithm.utils import rank_elements, fetch_first
 from .genome import create_mutate, create_distance, crossover
 from .gene import BaseGene
 

algorithm/neat/utils.py

@@ -1,72 +0,0 @@
-from functools import partial
-
-import numpy as np
-import jax
-from jax import numpy as jnp, Array, jit, vmap
-
-I_INT = np.iinfo(jnp.int32).max  # infinite int
-EMPTY_NODE = np.full((1, 5), jnp.nan)
-EMPTY_CON = np.full((1, 4), jnp.nan)
-
-
-@jit
-def unflatten_connections(nodes: Array, conns: Array):
-    """
-    transform the (C, CL) connections to (CL-2, N, N)
-    :param nodes: (N, NL)
-    :param cons: (C, CL)
-    :return:
-    """
-    N = nodes.shape[0]
-    CL = conns.shape[1]
-    node_keys = nodes[:, 0]
-    i_keys, o_keys = conns[:, 0], conns[:, 1]
-    i_idxs = vmap(key_to_indices, in_axes=(0, None))(i_keys, node_keys)
-    o_idxs = vmap(key_to_indices, in_axes=(0, None))(o_keys, node_keys)
-    res = jnp.full((CL - 2, N, N), jnp.nan)
-
-    # Is interesting that jax use clip when attach data in array
-    # however, it will do nothing set values in an array
-    # put all attributes include enable in res
-    res = res.at[:, i_idxs, o_idxs].set(conns[:, 2:].T)
-
-    return res
-
-
-def key_to_indices(key, keys):
-    return fetch_first(key == keys)
-
-
-@jit
-def fetch_first(mask, default=I_INT) -> Array:
-    """
-    fetch the first True index
-    :param mask: array of bool
-    :param default: the default value if no element satisfying the condition
-    :return: the index of the first element satisfying the condition. if no element satisfying the condition, return default value
-    """
-    idx = jnp.argmax(mask)
-    return jnp.where(mask[idx], idx, default)
-
-
-@jit
-def fetch_random(rand_key, mask, default=I_INT) -> Array:
-    """
-    similar to fetch_first, but fetch a random True index
-    """
-    true_cnt = jnp.sum(mask)
-    cumsum = jnp.cumsum(mask)
-    target = jax.random.randint(rand_key, shape=(), minval=1, maxval=true_cnt + 1)
-    mask = jnp.where(true_cnt == 0, False, cumsum >= target)
-    return fetch_first(mask, default)
-
-
-@partial(jit, static_argnames=['reverse'])
-def rank_elements(array, reverse=False):
-    """
-    rank the element in the array.
-    if reverse is True, the rank is from small to large. default large to small
-    """
-    if not reverse:
-        array = -array
-    return jnp.argsort(jnp.argsort(array))