update some examples

examples/gymnax/cartpole.py

@@ -1,41 +1,46 @@
 import jax.numpy as jnp
 
-from pipeline import Pipeline
-from algorithm.neat import *
+from tensorneat.pipeline import Pipeline
+from tensorneat.algorithm.neat import NEAT
+from tensorneat.genome import DefaultGenome, BiasNode
 
-from problem.rl_env import GymNaxEnv
+from tensorneat.problem.rl import GymNaxEnv
+from tensorneat.common import Act, Agg
 
 
-def action_policy(randkey, forward_func, obs):
-    return jnp.argmax(forward_func(obs))
-
 
 if __name__ == "__main__":
+    # the network has 2 outputs, the max one will be the action
+    # as the action of cartpole is {0, 1}
+
     pipeline = Pipeline(
         algorithm=NEAT(
-            species=DefaultSpecies(
-                genome=DefaultGenome(
-                    num_inputs=4,
-                    num_outputs=2,
-                    max_nodes=50,
-                    max_conns=100,
-                    # output_transform=lambda out: jnp.argmax(
-                    #     out
-                    # ),  # the action of cartpole is {0, 1}
+            pop_size=1000,
+            species_size=20,
+            survival_threshold=0.1,
+            compatibility_threshold=1.0,
+            genome=DefaultGenome(
+                num_inputs=4,
+                num_outputs=2,
+                init_hidden_layers=(),
+                node_gene=BiasNode(
+                    activation_options=Act.tanh,
+                    aggregation_options=Agg.sum,
                 ),
-                pop_size=10000,
-                species_size=10,
+                output_transform=jnp.argmax,
             ),
         ),
         problem=GymNaxEnv(
-            env_name="CartPole-v1", repeat_times=5, action_policy=action_policy
+            env_name="CartPole-v1",
+            repeat_times=5,
         ),
-        generation_limit=10000,
+        seed=42,
+        generation_limit=100,
         fitness_target=500,
     )
 
     # initialize state
     state = pipeline.setup()
-    # print(state)
+
     # run until terminate
     state, best = pipeline.auto_run(state)