add more rl task in examples

examples/func_fit/xor.py

@@ -1,5 +1,5 @@
 from config import *
-from pipeline import Pipeline
+from pipeline_jitable_env import Pipeline
 from algorithm import NEAT
 from algorithm.neat.gene import NormalGene, NormalGeneConfig
 from problem.func_fit import XOR, FuncFitConfig

examples/func_fit/xor_hyperneat.py

@@ -1,5 +1,5 @@
 from config import *
-from pipeline import Pipeline
+from pipeline_jitable_env import Pipeline
 from algorithm.neat import NormalGene, NormalGeneConfig
 from algorithm.hyperneat import HyperNEAT, NormalSubstrate, NormalSubstrateConfig
 from problem.func_fit import XOR3d, FuncFitConfig

examples/func_fit/xor_recurrent.py

@@ -1,5 +1,5 @@
 from config import *
-from pipeline import Pipeline
+from pipeline_jitable_env import Pipeline
 from algorithm import NEAT
 from algorithm.neat.gene import RecurrentGene, RecurrentGeneConfig
 from problem.func_fit import XOR3d, FuncFitConfig

examples/gymnax/acrobot.py

@@ -0,0 +1,39 @@
+import jax.numpy as jnp
+
+from config import *
+from pipeline_jitable_env import Pipeline
+from algorithm import NEAT
+from algorithm.neat.gene import NormalGene, NormalGeneConfig
+from problem.rl_env import GymNaxConfig, GymNaxEnv
+
+
+def example_conf():
+    return Config(
+        basic=BasicConfig(
+            seed=42,
+            fitness_target=0,
+            pop_size=10000
+        ),
+        neat=NeatConfig(
+            inputs=6,
+            outputs=3,
+        ),
+        gene=NormalGeneConfig(
+            activation_default=Act.tanh,
+            activation_options=(Act.tanh,),
+        ),
+        problem=GymNaxConfig(
+            env_name='Acrobot-v1',
+            output_transform=lambda out: jnp.argmax(out)  # the action of acrobot is {0, 1, 2}
+        )
+    )
+
+
+if __name__ == '__main__':
+    conf = example_conf()
+
+    algorithm = NEAT(conf, NormalGene)
+    pipeline = Pipeline(conf, algorithm, GymNaxEnv)
+    state = pipeline.setup()
+    pipeline.pre_compile(state)
+    state, best = pipeline.auto_run(state)

examples/gymnax/cartpole.py

@@ -1,7 +1,7 @@
 import jax.numpy as jnp
 
 from config import *
-from pipeline import Pipeline
+from pipeline_jitable_env import Pipeline
 from algorithm import NEAT
 from algorithm.neat.gene import NormalGene, NormalGeneConfig
 from problem.rl_env import GymNaxConfig, GymNaxEnv

examples/gymnax/mountain_car.py

@@ -0,0 +1,39 @@
+import jax.numpy as jnp
+
+from config import *
+from pipeline_jitable_env import Pipeline
+from algorithm import NEAT
+from algorithm.neat.gene import NormalGene, NormalGeneConfig
+from problem.rl_env import GymNaxConfig, GymNaxEnv
+
+
+def example_conf():
+    return Config(
+        basic=BasicConfig(
+            seed=42,
+            fitness_target=0,
+            pop_size=10000
+        ),
+        neat=NeatConfig(
+            inputs=2,
+            outputs=3,
+        ),
+        gene=NormalGeneConfig(
+            activation_default=Act.sigmoid,
+            activation_options=(Act.sigmoid,),
+        ),
+        problem=GymNaxConfig(
+            env_name='MountainCar-v0',
+            output_transform=lambda out: jnp.argmax(out)  # the action of cartpole is {0, 1, 2}
+        )
+    )
+
+
+if __name__ == '__main__':
+    conf = example_conf()
+
+    algorithm = NEAT(conf, NormalGene)
+    pipeline = Pipeline(conf, algorithm, GymNaxEnv)
+    state = pipeline.setup()
+    pipeline.pre_compile(state)
+    state, best = pipeline.auto_run(state)

examples/gymnax/mountain_car_continuous.py

@@ -0,0 +1,38 @@
+import jax.numpy as jnp
+
+from config import *
+from pipeline_jitable_env import Pipeline
+from algorithm import NEAT
+from algorithm.neat.gene import NormalGene, NormalGeneConfig
+from problem.rl_env import GymNaxConfig, GymNaxEnv
+
+
+def example_conf():
+    return Config(
+        basic=BasicConfig(
+            seed=42,
+            fitness_target=100,
+            pop_size=10000
+        ),
+        neat=NeatConfig(
+            inputs=2,
+            outputs=1,
+        ),
+        gene=NormalGeneConfig(
+            activation_default=Act.tanh,
+            activation_options=(Act.tanh,),
+        ),
+        problem=GymNaxConfig(
+            env_name='MountainCarContinuous-v0'
+        )
+    )
+
+
+if __name__ == '__main__':
+    conf = example_conf()
+
+    algorithm = NEAT(conf, NormalGene)
+    pipeline = Pipeline(conf, algorithm, GymNaxEnv)
+    state = pipeline.setup()
+    pipeline.pre_compile(state)
+    state, best = pipeline.auto_run(state)

examples/gymnax/pendulum.py

@@ -0,0 +1,39 @@
+import jax.numpy as jnp
+
+from config import *
+from pipeline_jitable_env import Pipeline
+from algorithm import NEAT
+from algorithm.neat.gene import NormalGene, NormalGeneConfig
+from problem.rl_env import GymNaxConfig, GymNaxEnv
+
+
+def example_conf():
+    return Config(
+        basic=BasicConfig(
+            seed=42,
+            fitness_target=0,
+            pop_size=10000
+        ),
+        neat=NeatConfig(
+            inputs=3,
+            outputs=1,
+        ),
+        gene=NormalGeneConfig(
+            activation_default=Act.tanh,
+            activation_options=(Act.tanh,),
+        ),
+        problem=GymNaxConfig(
+            env_name='Pendulum-v1',
+            output_transform=lambda out: out * 2  # the action of pendulum is [-2, 2]
+        )
+    )
+
+
+if __name__ == '__main__':
+    conf = example_conf()
+
+    algorithm = NEAT(conf, NormalGene)
+    pipeline = Pipeline(conf, algorithm, GymNaxEnv)
+    state = pipeline.setup()
+    pipeline.pre_compile(state)
+    state, best = pipeline.auto_run(state)

examples/gymnax/reacher.py

@@ -0,0 +1,36 @@
+from config import *
+from pipeline_jitable_env import Pipeline
+from algorithm import NEAT
+from algorithm.neat.gene import NormalGene, NormalGeneConfig
+from problem.rl_env import GymNaxConfig, GymNaxEnv
+
+
+def example_conf():
+    return Config(
+        basic=BasicConfig(
+            seed=42,
+            fitness_target=500,
+            pop_size=10000
+        ),
+        neat=NeatConfig(
+            inputs=8,
+            outputs=2,
+        ),
+        gene=NormalGeneConfig(
+            activation_default=Act.sigmoid,
+            activation_options=(Act.sigmoid,),
+        ),
+        problem=GymNaxConfig(
+            env_name='Reacher-misc',
+        )
+    )
+
+
+if __name__ == '__main__':
+    conf = example_conf()
+
+    algorithm = NEAT(conf, NormalGene)
+    pipeline = Pipeline(conf, algorithm, GymNaxEnv)
+    state = pipeline.setup()
+    pipeline.pre_compile(state)
+    state, best = pipeline.auto_run(state)