generally complete, but not work well. Debug

algorithms/neat/pipeline.py

@@ -1,7 +1,13 @@
+from typing import List, Union, Tuple, Callable
+import time
+
 import jax
+import numpy as np
 
 from .species import SpeciesController
 from .genome import create_initialize_function, create_mutate_function, create_forward_function
+from .genome import batch_crossover
+from .genome import expand, expand_single
 
 
 class Pipeline:
@@ -9,9 +15,13 @@ class Pipeline:
     Neat algorithm pipeline.
     """
 
-    def __init__(self, config):
+    def __init__(self, config, seed=42):
+        self.randkey = jax.random.PRNGKey(seed)
+
         self.config = config
         self.N = config.basic.init_maximum_nodes
+        self.expand_coe = config.basic.expands_coe
+        self.pop_size = config.neat.population.pop_size
 
         self.species_controller = SpeciesController(config)
         self.initialize_func = create_initialize_function(config)
@@ -22,6 +32,11 @@ class Pipeline:
 
         self.species_controller.speciate(self.pop_nodes, self.pop_connections, self.generation)
 
+        self.new_node_keys_pool: List[int] = [max(self.output_idx) + 1]
+
+        self.generation_timestamp = time.time()
+        self.best_fitness = float('-inf')
+
     def ask(self, batch: bool):
         """
         Create a forward function for the population.
@@ -35,7 +50,120 @@ class Pipeline:
 
     def tell(self, fitnesses):
         self.generation += 1
-        print(type(fitnesses), fitnesses)
+
         self.species_controller.update_species_fitnesses(fitnesses)
 
+        crossover_pair = self.species_controller.reproduce(self.generation)
 
+        self.update_next_generation(crossover_pair)
+
+        self.species_controller.speciate(self.pop_nodes, self.pop_connections, self.generation)
+
+        self.expand()
+
+    def auto_run(self, fitness_func, analysis: Union[Callable, str] = "default"):
+        for _ in range(self.config.neat.population.generation_limit):
+            forward_func = self.ask(batch=True)
+            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)
+
+            self.tell(fitnesses)
+        print("Generation limit reached!")
+
+    def update_next_generation(self, crossover_pair: List[Union[int, Tuple[int, int]]]) -> None:
+        """
+        create the next generation
+        :param crossover_pair: created from self.reproduce()
+        """
+
+        assert self.pop_nodes.shape[0] == self.pop_size
+
+        k1, k2, self.randkey = jax.random.split(self.randkey, 3)
+
+        # crossover
+        # prepare elitism mask and crossover pair
+        elitism_mask = np.full(self.pop_size, False)
+        for i, pair in enumerate(crossover_pair):
+            if not isinstance(pair, tuple):  # elitism
+                elitism_mask[i] = True
+                crossover_pair[i] = (pair, pair)
+        crossover_pair = np.array(crossover_pair)
+
+        # batch crossover
+        wpn = self.pop_nodes[crossover_pair[:, 0]]  # winner pop nodes
+        wpc = self.pop_connections[crossover_pair[:, 0]]  # winner pop connections
+        lpn = self.pop_nodes[crossover_pair[:, 1]]  # loser pop nodes
+        lpc = self.pop_connections[crossover_pair[:, 1]]  # loser pop connections
+        crossover_rand_keys = jax.random.split(k1, self.pop_size)
+        # npn, npc = batch_crossover(crossover_rand_keys, wpn, wpc, lpn, lpc)  # new pop nodes, new pop connections
+        npn, npc = crossover_wrapper(crossover_rand_keys, wpn, wpc, lpn, lpc)
+        # mutate
+        new_node_keys = np.array(self.fetch_new_node_keys())
+        mutate_rand_keys = jax.random.split(k2, self.pop_size)
+        m_npn, m_npc = self.mutate_func(mutate_rand_keys, npn, npc, new_node_keys)
+        m_npn, m_npc = jax.device_get(m_npn), jax.device_get(m_npc)
+        # elitism don't mutate
+        # (pop_size, ) to (pop_size, 1, 1)
+        self.pop_nodes = np.where(elitism_mask[:, None, None], npn, m_npn)
+        # (pop_size, ) to (pop_size, 1, 1, 1)
+        self.pop_connections = np.where(elitism_mask[:, None, None, None], npc, m_npc)
+
+        # recycle unused node keys
+        unused = []
+        for i, nodes in enumerate(self.pop_nodes):
+            node_keys, key = nodes[:, 0], new_node_keys[i]
+            if not np.isin(key, node_keys):  # the new node key is not used
+                unused.append(key)
+        self.new_node_keys_pool = unused + self.new_node_keys_pool
+
+    def expand(self):
+        """
+        Expand the population if needed.
+        :return:
+        when the maximum node number of the population >= N
+        the population will expand
+        """
+        pop_node_keys = self.pop_nodes[:, :, 0]
+        pop_node_sizes = np.sum(~np.isnan(pop_node_keys), axis=1)
+        max_node_size = np.max(pop_node_sizes)
+        if max_node_size >= self.N:
+            print(f"expand to {self.N}!")
+            self.N = int(self.N * self.expand_coe)
+            self.pop_nodes, self.pop_connections = expand(self.pop_nodes, self.pop_connections, self.N)
+
+            # don't forget to expand representation genome in species
+            for s in self.species_controller.species:
+                s.representative = expand(*s.representative, self.N)
+
+    def fetch_new_node_keys(self):
+        # if remain unused keys are not enough, create new keys
+        if len(self.new_node_keys_pool) < self.pop_size:
+            max_unused_key = max(self.new_node_keys_pool) if self.new_node_keys_pool else -1
+            new_keys = list(range(max_unused_key + 1, max_unused_key + 1 + 10 * self.pop_size))
+            self.new_node_keys_pool.extend(new_keys)
+
+        # fetch keys from pool
+        res = self.new_node_keys_pool[:self.pop_size]
+        self.new_node_keys_pool = self.new_node_keys_pool[self.pop_size:]
+        return res
+
+    def default_analysis(self, fitnesses):
+        max_f, min_f, mean_f, std_f = max(fitnesses), min(fitnesses), np.mean(fitnesses), np.std(fitnesses)
+        species_sizes = [len(s.members) for s in self.species_controller.species.values()]
+
+        new_timestamp = time.time()
+        cost_time = new_timestamp - self.generation_timestamp
+        self.generation_timestamp = new_timestamp
+
+        print(f"Generation: {self.generation}",
+              f"fitness: {max_f}, {min_f}, {mean_f}, {std_f}, Species sizes: {species_sizes}, Cost time: {cost_time}")
+
+
+def crossover_wrapper(*args):
+    return batch_crossover(*args)