esquilax 2.1.0

JAX multi-agent simulation and ML toolset

JAX Multi-Agent RL, A-Life, and Simulation Framework

Esquilax is set of transformations and utilities intended to allow developers and researchers to quickly implement models of multi-agent systems for rl-training, evolutionary methods, and a-life.

It is intended for systems involving large number of agents, and to work alongside other JAX packages like Flax and Evosax.

Full documentation can be found here

Features

• Built on top of JAX

  This has the benefits of JAX; high-performance, built in GPU support etc., but also means Esquilax can interoperate with existing JAX ML and RL libraries.

• Interaction Algorithm Implementations

  Implements common agent interaction patterns. This allows users to concentrate on model design instead of low-level algorithm implementation details.

• Scale and Performance

  JIT compilation and GPU support enables simulations and multi-agent systems containing large numbers of agents whilst maintaining performance and training throughput.

• Functional Patterns

  Esquilax is designed around functional patterns, ensuring models can be readily parallelised, but also aiding composition and readability

• Built-in RL and Evolutionary Training

  Esquilax provides functionality for running multi-agent RL and multi-strategy neuro-evolution training, within Esquilax simulations.

Should I Use Esquilax?

Esquilax is intended for time-stepped models of large scale systems with fixed numbers of entities, where state is updated in parallel. As such you should probably not use Esquilax if:

• You want to use something other than stepped updates, e.g. continuous time, event driven models, or where agents are intended to update in sequence.
• You need variable numbers of entities or temporary entities, e.g. message passing.
• You need a high-fidelity physics/robotics simulation.

Getting Started

Esquilax can be installed from pip using

pip install esquilax

The requirements for evolutionary and rl training are not installed by default. They can be installed using the evo and rl extras respectively, e.g.:

pip install esquilax[evo]

You may need to manually install JAXlib, especially for GPU support. Installation instructions for JAX can be found here.

Examples

Example models and multi-agent policy training implemented using Esquilax can be found here. A virtual environment with additional dependencies for the examples can be setup using poetry with

poetry install --extras all --with examples

For a larger project using Esquilax see this Boid flock RL environment.

or this multi-agent rl environment, part of the Jumanji RL environment library.

Contributing

Issues

Please report any issues or feature suggestions here.

Developers

Developer notes can be found here, Esquilax is under active development and contributions are very welcome!