EcoJAX

This repository aims to provide a simple and modular codebase to run experiments on large neural-agents ecological models using JAX.

Installation

Clone the repository and create a virtual environment. The repo work (at least) in python 3.10 to 3.12.

git clone [email protected]:tboulet/EcoJAX.git
cd EcoJAX
python -m venv venv
source venv/bin/activate   # on Windows, use `venv\Scripts\activate.bat`

Install JAX

On linux (note your CUDA version may vary, and I'm unsure of the JAX version, 0.4.30 works fine) :

pip install jax[cuda12_pip] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

On Windows :

pip install jax[cpu]

Install the requirements

pip install -r requirements.txt

Eventually install the package in editable mode

pip install -e .

Usage

To run any experiment, run the run.py file with the desired configuration. We use Hydra as our configuration manager.

python run.py

To use a specific configuration, use the --config-name argument.

python run.py --config-name=my_config

You can modify a specific argument thanks to Hydra's override system.

python run.py path.to.arg.in.config=new_value

In particular, the 3 main components interacting with each other in the experiments are the environment, the agents species (or agents) and the model. You can modify the configuration of each of these components. The meaning of those three components is explained below.

python run.py env=gridworld agents=ne model=cnn

Components

The environment, the agents and the model interact with each other in the context of $n$ agents living in an environment. Note that $n$ is not the current number of agents but the maximal number of agents in the simulation, and stay fixed during the whole simulation. This is a design choice to allow for efficient batching of the various arrays which is necessary for JAX parallelism.

Environment

The environment is the world in which the agents evolve. It receives as input a batch of $n$ actions and returns two things :

• a batch of $n$ observations, one for each agent
• the "eco-information" that contains which agents just died, which agents are just born (and from which parents), etc.

The only environment currently available is the gridworld environment (and variations).

Agents Species

The agents species is a class that contains the logic of the agents. It receives as input a batch of $n$ observations and returns a batch of $n$ actions. It also should manage internally the way the agents inherit their genes.

Agents species currently available are :

• ne (Neural Evolution) : the agents are neural networks that are evolved using a genetic algorithm

Model

The model is the neural network architecture that will be used by the agents species.

Models currently available are :

• mlp (Multi-Layer Perceptron) : flatten and concatenate the observations and pass them through a simple feed-forward neural network
• cnn (Convolutional Neural Network) : do the same as the MLP but image-like observations are passed through a convolutional neural network
• random : the model is a random function that returns random actions