new architecture

pipeline.py

@@ -1,7 +1,3 @@
-"""
-pipeline for jitable env like func_fit, gymnax
-"""
-
 from functools import partial
 from typing import Type
 
@@ -16,24 +12,28 @@ from core import State, Algorithm, Problem
 
 class Pipeline:
 
-    def __init__(self, config: Config, algorithm: Algorithm, problem_type: Type[Problem]):
+    def __init__(
+        self,
+        algorithm: Algorithm,
+        problem: Problem,
+        seed: int = 42,
+        fitness_target: float = 1,
+        generation_limit: int = 1000,
+        pop_size: int = 100,
+    ):
+        assert problem.jitable, "Currently, problem must be jitable"
 
-        assert problem_type.jitable, "problem must be jitable"
-
-        self.config = config
         self.algorithm = algorithm
-        self.problem = problem_type(config.problem)
+        self.problem = problem
+        self.seed = seed
+        self.fitness_target = fitness_target
+        self.generation_limit = generation_limit
+        self.pop_size = pop_size
 
         print(self.problem.input_shape, self.problem.output_shape)
 
-        if isinstance(algorithm, NEAT):
-            assert config.neat.inputs == self.problem.input_shape[-1], f"problem input shape {self.problem.input_shape}"
-
-        elif isinstance(algorithm, HyperNEAT):
-            assert config.hyperneat.inputs == self.problem.input_shape[-1], f"problem input shape {self.problem.input_shape}"
-
-        else:
-            raise NotImplementedError
+        # TODO: make each algorithm's input_num and output_num
+        assert algorithm.input_num == self.problem.input_shape[-1], f"problem input shape {self.problem.input_shape}"
 
         self.act_func = self.algorithm.act
 
@@ -45,19 +45,19 @@ class Pipeline:
         self.generation_timestamp = None
 
     def setup(self):
-        key = jax.random.PRNGKey(self.config.basic.seed)
+        key = jax.random.PRNGKey(self.seed)
         algorithm_key, evaluate_key = jax.random.split(key, 2)
-        state = State()
-        state = self.algorithm.setup(algorithm_key, state)
-        return state.update(
-            evaluate_key=evaluate_key
+
+        # TODO: Problem should has setup function to maintain state
+        return State(
+            alg=self.algorithm.setup(algorithm_key),
+            pro=self.problem.setup(evaluate_key),
         )
 
     @partial(jax.jit, static_argnums=(0,))
     def step(self, state):
-
         key, sub_key = jax.random.split(state.evaluate_key)
-        keys = jax.random.split(key, self.config.basic.pop_size)
+        keys = jax.random.split(key, self.pop_size)
 
         pop = self.algorithm.ask(state)
 
@@ -72,7 +72,7 @@ class Pipeline:
 
     def auto_run(self, ini_state):
         state = ini_state
-        for _ in range(self.config.basic.generation_limit):
+        for _ in range(self.generation_limit):
 
             self.generation_timestamp = time.time()
 
@@ -84,7 +84,7 @@ class Pipeline:
 
             self.analysis(state, previous_pop, fitnesses)
 
-            if max(fitnesses) >= self.config.basic.fitness_target:
+            if max(fitnesses) >= self.fitness_target:
                 print("Fitness limit reached!")
                 return state, self.best_genome
 
@@ -120,3 +120,4 @@ class Pipeline:
         print("start compile")
         self.step.lower(self, state).compile()
         print(f"compile finished, cost time: {time.time() - tic}s")
+