diff --git a/README.md b/README.md
index 4b5d31b..eb45a79 100644
--- a/README.md
+++ b/README.md
@@ -1,88 +1,106 @@
-# NEATax: Tensorized NEAT implementation in JAX
+
+
+
+
+
+
-TensorNEAT is a powerful tool that utilizes JAX to implement the NEAT (NeuroEvolution of Augmenting Topologies)
-algorithm. It provides support for parallel execution of tasks such as network forward computation, mutation,
-and crossover at the population level.
+# TensorNEAT: Tensorized NEAT implementation in JAX
+
+
+
+
+
+
+## Introduction
+🚀TensorNEAT, a part of EvoX project, aims to enhance the NEAT (NeuroEvolution of Augmenting Topologies) algorithm by incorporating GPU acceleration. Utilizing JAX for parallel computations, it extends NEAT's capabilities to modern computational environments, making advanced neuroevolution accessible and fast.
## Requirements
-* available [JAX](https://github.com/google/jax#installation) environment;
-* [gymnax](https://github.com/RobertTLange/gymnax) (optional).
-
+TensorNEAT requires:
+- jax (version >= 0.4.16)
+- jaxlib (version >= 0.3.0)
+- brax [optional]
+- gymnax [optional]
+
## Example
Simple Example for XOR problem:
```python
-from config import *
from pipeline import Pipeline
-from algorithm import NEAT
-from algorithm.neat.gene import NormalGene, NormalGeneConfig
-from problem.func_fit import XOR, FuncFitConfig
+from algorithm.neat import *
+
+from problem.func_fit import XOR3d
if __name__ == '__main__':
- # running config
- config = Config(
- basic=BasicConfig(
- seed=42,
- fitness_target=-1e-2,
- pop_size=10000
+ pipeline = Pipeline(
+ algorithm=NEAT(
+ species=DefaultSpecies(
+ genome=DefaultGenome(
+ num_inputs=3,
+ num_outputs=1,
+ max_nodes=50,
+ max_conns=100,
+ ),
+ pop_size=10000,
+ species_size=10,
+ compatibility_threshold=3.5,
+ ),
),
- neat=NeatConfig(
- inputs=2,
- outputs=1
- ),
- gene=NormalGeneConfig(),
- problem=FuncFitConfig(
- error_method='rmse'
- )
+ problem=XOR3d(),
+ generation_limit=10000,
+ fitness_target=-1e-8
)
- # define algorithm: NEAT with NormalGene
- algorithm = NEAT(config, NormalGene)
- # full pipeline
- pipeline = Pipeline(config, algorithm, XOR)
+
# initialize state
state = pipeline.setup()
+ # print(state)
# run until terminate
state, best = pipeline.auto_run(state)
# show result
pipeline.show(state, best)
```
-Simple Example for RL envs in gymnax(CartPole-v0):
+Simple Example for RL envs in Brax (Ant):
```python
-import jax.numpy as jnp
-
-from config import *
from pipeline import Pipeline
-from algorithm import NEAT
-from algorithm.neat.gene import NormalGene, NormalGeneConfig
-from problem.rl_env import GymNaxConfig, GymNaxEnv
+from algorithm.neat import *
+
+from problem.rl_env import BraxEnv
+from utils import Act
if __name__ == '__main__':
- conf = Config(
- basic=BasicConfig(
- seed=42,
- fitness_target=500,
- pop_size=10000
+ pipeline = Pipeline(
+ algorithm=NEAT(
+ species=DefaultSpecies(
+ genome=DefaultGenome(
+ num_inputs=27,
+ num_outputs=8,
+ max_nodes=50,
+ max_conns=100,
+ node_gene=DefaultNodeGene(
+ activation_options=(Act.tanh,),
+ activation_default=Act.tanh,
+ )
+ ),
+ pop_size=1000,
+ species_size=10,
+ ),
),
- neat=NeatConfig(
- inputs=4,
- outputs=1,
+ problem=BraxEnv(
+ env_name='ant',
),
- gene=NormalGeneConfig(
- activation_default=Act.sigmoid,
- activation_options=(Act.sigmoid,),
- ),
- problem=GymNaxConfig(
- env_name='CartPole-v1',
- output_transform=lambda out: jnp.where(out[0] > 0.5, 1, 0) # the action of cartpole is {0, 1}
- )
+ generation_limit=10000,
+ fitness_target=5000
)
- algorithm = NEAT(conf, NormalGene)
- pipeline = Pipeline(conf, algorithm, GymNaxEnv)
+ # initialize state
state = pipeline.setup()
+ # print(state)
+ # run until terminate
state, best = pipeline.auto_run(state)
```
-`/examples` folder contains more examples.
+more examples are in `tensorneat/examples`.
-## TO BE COMPLETE...
diff --git a/imgs/evox_logo_dark.png b/imgs/evox_logo_dark.png
new file mode 100644
index 0000000..325e89c
Binary files /dev/null and b/imgs/evox_logo_dark.png differ
diff --git a/imgs/evox_logo_light.png b/imgs/evox_logo_light.png
new file mode 100644
index 0000000..70230a3
Binary files /dev/null and b/imgs/evox_logo_light.png differ
diff --git a/t.py b/t.py
deleted file mode 100644
index 5335a7a..0000000
--- a/t.py
+++ /dev/null
@@ -1,4 +0,0 @@
-import jax.numpy as jnp
-
-a = jnp.zeros((0, 9, 9))
-print(a)
\ No newline at end of file
diff --git a/algorithm/__init__.py b/tensorneat/algorithm/__init__.py
similarity index 100%
rename from algorithm/__init__.py
rename to tensorneat/algorithm/__init__.py
diff --git a/algorithm/base.py b/tensorneat/algorithm/base.py
similarity index 100%
rename from algorithm/base.py
rename to tensorneat/algorithm/base.py
diff --git a/algorithm/hyperneat/__init__.py b/tensorneat/algorithm/hyperneat/__init__.py
similarity index 100%
rename from algorithm/hyperneat/__init__.py
rename to tensorneat/algorithm/hyperneat/__init__.py
diff --git a/algorithm/hyperneat/hyperneat.py b/tensorneat/algorithm/hyperneat/hyperneat.py
similarity index 100%
rename from algorithm/hyperneat/hyperneat.py
rename to tensorneat/algorithm/hyperneat/hyperneat.py
diff --git a/algorithm/hyperneat/substrate/__init__.py b/tensorneat/algorithm/hyperneat/substrate/__init__.py
similarity index 100%
rename from algorithm/hyperneat/substrate/__init__.py
rename to tensorneat/algorithm/hyperneat/substrate/__init__.py
diff --git a/algorithm/hyperneat/substrate/base.py b/tensorneat/algorithm/hyperneat/substrate/base.py
similarity index 100%
rename from algorithm/hyperneat/substrate/base.py
rename to tensorneat/algorithm/hyperneat/substrate/base.py
diff --git a/algorithm/hyperneat/substrate/default.py b/tensorneat/algorithm/hyperneat/substrate/default.py
similarity index 100%
rename from algorithm/hyperneat/substrate/default.py
rename to tensorneat/algorithm/hyperneat/substrate/default.py
diff --git a/algorithm/hyperneat/substrate/full.py b/tensorneat/algorithm/hyperneat/substrate/full.py
similarity index 100%
rename from algorithm/hyperneat/substrate/full.py
rename to tensorneat/algorithm/hyperneat/substrate/full.py
diff --git a/algorithm/neat/__init__.py b/tensorneat/algorithm/neat/__init__.py
similarity index 100%
rename from algorithm/neat/__init__.py
rename to tensorneat/algorithm/neat/__init__.py
diff --git a/algorithm/neat/ga/__init__.py b/tensorneat/algorithm/neat/ga/__init__.py
similarity index 100%
rename from algorithm/neat/ga/__init__.py
rename to tensorneat/algorithm/neat/ga/__init__.py
diff --git a/algorithm/neat/ga/crossover/__init__.py b/tensorneat/algorithm/neat/ga/crossover/__init__.py
similarity index 100%
rename from algorithm/neat/ga/crossover/__init__.py
rename to tensorneat/algorithm/neat/ga/crossover/__init__.py
diff --git a/algorithm/neat/ga/crossover/base.py b/tensorneat/algorithm/neat/ga/crossover/base.py
similarity index 100%
rename from algorithm/neat/ga/crossover/base.py
rename to tensorneat/algorithm/neat/ga/crossover/base.py
diff --git a/algorithm/neat/ga/crossover/default.py b/tensorneat/algorithm/neat/ga/crossover/default.py
similarity index 100%
rename from algorithm/neat/ga/crossover/default.py
rename to tensorneat/algorithm/neat/ga/crossover/default.py
diff --git a/algorithm/neat/ga/mutation/__init__.py b/tensorneat/algorithm/neat/ga/mutation/__init__.py
similarity index 100%
rename from algorithm/neat/ga/mutation/__init__.py
rename to tensorneat/algorithm/neat/ga/mutation/__init__.py
diff --git a/algorithm/neat/ga/mutation/base.py b/tensorneat/algorithm/neat/ga/mutation/base.py
similarity index 100%
rename from algorithm/neat/ga/mutation/base.py
rename to tensorneat/algorithm/neat/ga/mutation/base.py
diff --git a/algorithm/neat/ga/mutation/default.py b/tensorneat/algorithm/neat/ga/mutation/default.py
similarity index 100%
rename from algorithm/neat/ga/mutation/default.py
rename to tensorneat/algorithm/neat/ga/mutation/default.py
diff --git a/algorithm/neat/gene/__init__.py b/tensorneat/algorithm/neat/gene/__init__.py
similarity index 100%
rename from algorithm/neat/gene/__init__.py
rename to tensorneat/algorithm/neat/gene/__init__.py
diff --git a/algorithm/neat/gene/base.py b/tensorneat/algorithm/neat/gene/base.py
similarity index 100%
rename from algorithm/neat/gene/base.py
rename to tensorneat/algorithm/neat/gene/base.py
diff --git a/algorithm/neat/gene/conn/__init__.py b/tensorneat/algorithm/neat/gene/conn/__init__.py
similarity index 100%
rename from algorithm/neat/gene/conn/__init__.py
rename to tensorneat/algorithm/neat/gene/conn/__init__.py
diff --git a/algorithm/neat/gene/conn/base.py b/tensorneat/algorithm/neat/gene/conn/base.py
similarity index 100%
rename from algorithm/neat/gene/conn/base.py
rename to tensorneat/algorithm/neat/gene/conn/base.py
diff --git a/algorithm/neat/gene/conn/default.py b/tensorneat/algorithm/neat/gene/conn/default.py
similarity index 100%
rename from algorithm/neat/gene/conn/default.py
rename to tensorneat/algorithm/neat/gene/conn/default.py
diff --git a/algorithm/neat/gene/node/__init__.py b/tensorneat/algorithm/neat/gene/node/__init__.py
similarity index 100%
rename from algorithm/neat/gene/node/__init__.py
rename to tensorneat/algorithm/neat/gene/node/__init__.py
diff --git a/algorithm/neat/gene/node/base.py b/tensorneat/algorithm/neat/gene/node/base.py
similarity index 100%
rename from algorithm/neat/gene/node/base.py
rename to tensorneat/algorithm/neat/gene/node/base.py
diff --git a/algorithm/neat/gene/node/default.py b/tensorneat/algorithm/neat/gene/node/default.py
similarity index 100%
rename from algorithm/neat/gene/node/default.py
rename to tensorneat/algorithm/neat/gene/node/default.py
diff --git a/algorithm/neat/genome/__init__.py b/tensorneat/algorithm/neat/genome/__init__.py
similarity index 100%
rename from algorithm/neat/genome/__init__.py
rename to tensorneat/algorithm/neat/genome/__init__.py
diff --git a/algorithm/neat/genome/base.py b/tensorneat/algorithm/neat/genome/base.py
similarity index 100%
rename from algorithm/neat/genome/base.py
rename to tensorneat/algorithm/neat/genome/base.py
diff --git a/algorithm/neat/genome/default.py b/tensorneat/algorithm/neat/genome/default.py
similarity index 100%
rename from algorithm/neat/genome/default.py
rename to tensorneat/algorithm/neat/genome/default.py
diff --git a/algorithm/neat/genome/recurrent.py b/tensorneat/algorithm/neat/genome/recurrent.py
similarity index 100%
rename from algorithm/neat/genome/recurrent.py
rename to tensorneat/algorithm/neat/genome/recurrent.py
diff --git a/algorithm/neat/neat.py b/tensorneat/algorithm/neat/neat.py
similarity index 100%
rename from algorithm/neat/neat.py
rename to tensorneat/algorithm/neat/neat.py
diff --git a/algorithm/neat/species/__init__.py b/tensorneat/algorithm/neat/species/__init__.py
similarity index 100%
rename from algorithm/neat/species/__init__.py
rename to tensorneat/algorithm/neat/species/__init__.py
diff --git a/algorithm/neat/species/base.py b/tensorneat/algorithm/neat/species/base.py
similarity index 100%
rename from algorithm/neat/species/base.py
rename to tensorneat/algorithm/neat/species/base.py
diff --git a/algorithm/neat/species/default.py b/tensorneat/algorithm/neat/species/default.py
similarity index 100%
rename from algorithm/neat/species/default.py
rename to tensorneat/algorithm/neat/species/default.py
diff --git a/examples/brax/ant.py b/tensorneat/examples/brax/ant.py
similarity index 100%
rename from examples/brax/ant.py
rename to tensorneat/examples/brax/ant.py
diff --git a/examples/brax/half_cheetah.gif b/tensorneat/examples/brax/half_cheetah.gif
similarity index 100%
rename from examples/brax/half_cheetah.gif
rename to tensorneat/examples/brax/half_cheetah.gif
diff --git a/examples/brax/half_cheetah.py b/tensorneat/examples/brax/half_cheetah.py
similarity index 100%
rename from examples/brax/half_cheetah.py
rename to tensorneat/examples/brax/half_cheetah.py
diff --git a/examples/brax/reacher.py b/tensorneat/examples/brax/reacher.py
similarity index 100%
rename from examples/brax/reacher.py
rename to tensorneat/examples/brax/reacher.py
diff --git a/examples/func_fit/xor.py b/tensorneat/examples/func_fit/xor.py
similarity index 100%
rename from examples/func_fit/xor.py
rename to tensorneat/examples/func_fit/xor.py
diff --git a/examples/func_fit/xor3d_hyperneat.py b/tensorneat/examples/func_fit/xor3d_hyperneat.py
similarity index 100%
rename from examples/func_fit/xor3d_hyperneat.py
rename to tensorneat/examples/func_fit/xor3d_hyperneat.py
diff --git a/examples/func_fit/xor_recurrent.py b/tensorneat/examples/func_fit/xor_recurrent.py
similarity index 100%
rename from examples/func_fit/xor_recurrent.py
rename to tensorneat/examples/func_fit/xor_recurrent.py
diff --git a/examples/gymnax/arcbot.py b/tensorneat/examples/gymnax/arcbot.py
similarity index 100%
rename from examples/gymnax/arcbot.py
rename to tensorneat/examples/gymnax/arcbot.py
diff --git a/examples/gymnax/cartpole.py b/tensorneat/examples/gymnax/cartpole.py
similarity index 100%
rename from examples/gymnax/cartpole.py
rename to tensorneat/examples/gymnax/cartpole.py
diff --git a/examples/gymnax/cartpole_hyperneat.py b/tensorneat/examples/gymnax/cartpole_hyperneat.py
similarity index 100%
rename from examples/gymnax/cartpole_hyperneat.py
rename to tensorneat/examples/gymnax/cartpole_hyperneat.py
diff --git a/examples/gymnax/mountain_car.py b/tensorneat/examples/gymnax/mountain_car.py
similarity index 100%
rename from examples/gymnax/mountain_car.py
rename to tensorneat/examples/gymnax/mountain_car.py
diff --git a/examples/gymnax/mountain_car_continuous.py b/tensorneat/examples/gymnax/mountain_car_continuous.py
similarity index 100%
rename from examples/gymnax/mountain_car_continuous.py
rename to tensorneat/examples/gymnax/mountain_car_continuous.py
diff --git a/examples/gymnax/pendulum.py b/tensorneat/examples/gymnax/pendulum.py
similarity index 100%
rename from examples/gymnax/pendulum.py
rename to tensorneat/examples/gymnax/pendulum.py
diff --git a/examples/gymnax/reacher.py b/tensorneat/examples/gymnax/reacher.py
similarity index 100%
rename from examples/gymnax/reacher.py
rename to tensorneat/examples/gymnax/reacher.py
diff --git a/pipeline.py b/tensorneat/pipeline.py
similarity index 100%
rename from pipeline.py
rename to tensorneat/pipeline.py
diff --git a/problem/__init__.py b/tensorneat/problem/__init__.py
similarity index 100%
rename from problem/__init__.py
rename to tensorneat/problem/__init__.py
diff --git a/problem/base.py b/tensorneat/problem/base.py
similarity index 100%
rename from problem/base.py
rename to tensorneat/problem/base.py
diff --git a/problem/func_fit/__init__.py b/tensorneat/problem/func_fit/__init__.py
similarity index 100%
rename from problem/func_fit/__init__.py
rename to tensorneat/problem/func_fit/__init__.py
diff --git a/problem/func_fit/func_fit.py b/tensorneat/problem/func_fit/func_fit.py
similarity index 100%
rename from problem/func_fit/func_fit.py
rename to tensorneat/problem/func_fit/func_fit.py
diff --git a/problem/func_fit/xor.py b/tensorneat/problem/func_fit/xor.py
similarity index 100%
rename from problem/func_fit/xor.py
rename to tensorneat/problem/func_fit/xor.py
diff --git a/problem/func_fit/xor3d.py b/tensorneat/problem/func_fit/xor3d.py
similarity index 100%
rename from problem/func_fit/xor3d.py
rename to tensorneat/problem/func_fit/xor3d.py
diff --git a/problem/rl_env/__init__.py b/tensorneat/problem/rl_env/__init__.py
similarity index 100%
rename from problem/rl_env/__init__.py
rename to tensorneat/problem/rl_env/__init__.py
diff --git a/problem/rl_env/brax_env.py b/tensorneat/problem/rl_env/brax_env.py
similarity index 100%
rename from problem/rl_env/brax_env.py
rename to tensorneat/problem/rl_env/brax_env.py
diff --git a/problem/rl_env/gymnax_env.py b/tensorneat/problem/rl_env/gymnax_env.py
similarity index 100%
rename from problem/rl_env/gymnax_env.py
rename to tensorneat/problem/rl_env/gymnax_env.py
diff --git a/problem/rl_env/rl_jit.py b/tensorneat/problem/rl_env/rl_jit.py
similarity index 100%
rename from problem/rl_env/rl_jit.py
rename to tensorneat/problem/rl_env/rl_jit.py
diff --git a/test/__init__.py b/tensorneat/test/__init__.py
similarity index 100%
rename from test/__init__.py
rename to tensorneat/test/__init__.py
diff --git a/test/test_genome.py b/tensorneat/test/test_genome.py
similarity index 100%
rename from test/test_genome.py
rename to tensorneat/test/test_genome.py
diff --git a/utils/__init__.py b/tensorneat/utils/__init__.py
similarity index 100%
rename from utils/__init__.py
rename to tensorneat/utils/__init__.py
diff --git a/utils/activation.py b/tensorneat/utils/activation.py
similarity index 100%
rename from utils/activation.py
rename to tensorneat/utils/activation.py
diff --git a/utils/aggregation.py b/tensorneat/utils/aggregation.py
similarity index 100%
rename from utils/aggregation.py
rename to tensorneat/utils/aggregation.py
diff --git a/utils/graph.py b/tensorneat/utils/graph.py
similarity index 100%
rename from utils/graph.py
rename to tensorneat/utils/graph.py
diff --git a/utils/state.py b/tensorneat/utils/state.py
similarity index 100%
rename from utils/state.py
rename to tensorneat/utils/state.py
diff --git a/utils/tools.py b/tensorneat/utils/tools.py
similarity index 100%
rename from utils/tools.py
rename to tensorneat/utils/tools.py
+ 