"Neuro Evolution of Augmented Topologies" JAX implementation.

Usage:

python -m examples.train_<example>

# Common flags:
# --visualize-training Send real-time training data to the TUI.
# --population=<int>
# --generations=<int>
# --max-episode-steps=<int>
# --feedforward-only Disable recurrent connections.
# --max-nds=<int> Max nodes per network (JAX relies on static shapes for performance).

Visualization TUI

training_visualizer/generate_grpc.sh
cd training_visualizer
go run .

See training_visualizer/generate_grpc.sh for grpc requirements.

File Structure:

src/data.py, src/hyperparams.py: Main files defining the data for the algo and the policy, and the initialization.

src/solver.py : High level logic of the neat algorithm. Abides by evojax's ask-tell and best_params api; in this way evojax handles the algo-policy-task loop.

src/algo/ : The different chunks of the neat algorithm.

src/policy.py : The functions used for the forward pass, both for the recurrent and feedforward-only settings.

src/topological_sorting.py : For preparing the forward pass data when using the feedfoward only setting.

tasks/: Environments for the individuals to run against.

TUI

training_visualizer/server.go : Go program that renders the visualizations in the terminal using a library called tcell. Abides by training_visualizer/grpc_api/schema.proto.

training_visualizer/sugiyama.go : Exposes a function, used in server.go, that renders networks through the Sugiyama method. Unfinished prototype.

training_visualizer/grpc_api/schema.proto: Defines the data to send. Alludes to src/data.py, src/hyperparams.py.

src/extra_modules/training_visualizer/client.py : Prepares and sends the data to training_visualizer/server.go.

**.pb.go, **pb2.py, **pb2_grpc.py: Autogenerated files.

Acknowledgments:

• EvoJAX: Orchestrates the rollout loop, and provides utils.

• tcell: The TUI was made with this library.

• Mujoco Playground: The physics tasks.

Sources of ideas:

• Stanley, K. O., & Miikkulainen, R. (2002). Evolving Neural Networks Through Augmenting Topologies Evolutionary Computation, 10(2), 99-127.

• Buckland, M. AI Techniques for Game Programming.

• Sugiyama, K. Graph Drawing and its Applications for Software and Knowledge Engineers.

• Kindermann, P. Lectures on Visualization of Graphs.

• neat-c

• NEAT-Python

• TensorNeat

License

Apache-2.0 pbvrl

I reserve the right to make it private in the future. See LICENSE.md or CONTRIBUTING.md

Continuity

I will maintain it for as long as I find manageable.

Roadmap

• Add more examples
• Benchmark against top alternatives and showcase results
• Finish TUI: Sugiyama view
• Finish TUI: Handle both feedforward and recurrent networks
• Support jax.shard_map (Replace EvoJAX (archived) with own rollout loop)

Maybe

• Finish pytest tests (They might not be as useful as the visualization in surfacing bugs)
• Expose it as a package
• Finish multiagent example