Metadata-Version: 2.4
Name: esquilax
Version: 3.0.0
Summary: JAX multi-agent simulation and ML toolset
Keywords: agent-based modelling,simulation,reinforcement-learning
Author: zombie-einstein
Author-email: zombie-einstein <zombie-einstein@proton.me>
License-Expression: MIT
License-File: LICENSE
Classifier: Development Status :: 4 - Beta
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Programming Language :: Python :: Implementation :: CPython
Classifier: Programming Language :: Python :: Implementation :: PyPy
Classifier: Environment :: GPU
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Artificial Life
Requires-Dist: chex>=0.1.92
Requires-Dist: jax>=0.12.0.1
Requires-Dist: jax-tqdm>=0.4.0
Requires-Dist: matplotlib>=3.11.2 ; extra == 'examples'
Requires-Dist: pillow>=13.0.0 ; extra == 'examples'
Maintainer: zombie-einstein
Maintainer-email: zombie-einstein <zombie-einstein@proton.me>
Requires-Python: >=3.12, <4.0
Provides-Extra: examples
Description-Content-Type: text/markdown

<div align="center">
  <img src="https://github.com/zombie-einstein/esquilax/raw/main/.github/images/text_logo.png" />
  <br>
  <em>JAX Multi-Agent RL, A-Life, and Simulation Framework</em>
</div>
<br>

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](https://github.com/google/flax) and
[Evosax](https://github.com/RobertTLange/evosax).

**Full documentation can be found
[here](https://zombie-einstein.github.io/esquilax/)**

## 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

## 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

``` bash
pip install esquilax
```

You may need to manually install JAXlib, especially for GPU support.
Installation instructions for JAX can be found
[here](https://github.com/google/jax?tab=readme-ov-file#installation).

## Examples

Example models implemented using Esquilax
can be found [here](https://github.com/zombie-einstein/esquilax/tree/main/examples).

Examples can be run using [uv](https://docs.astral.sh/uv/)

```bash
uv run --extras examples ./examples/...
```

For a project using Esquilax see
[Floxs](https://github.com/zombie-einstein/floxs) a collection of multi-agent
RL flock/swarm environments or
[this](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/swarms/search_and_rescue)
multi-agent rl environment, part of the [Jumanji](https://github.com/instadeepai/jumanji)
RL environment library.

## Contributing

### Issues

Please report any issues or feature suggestions
[here](https://github.com/zombie-einstein/esquilax/issues).

### Developers

Developer notes can be found
[here](https://github.com/zombie-einstein/esquilax/blob/main/.github/docs/developers.md),
Esquilax is under active development and contributions are very welcome!
