floxs 0.1.0

Multi-agent RL flock and swarm environments implemented in JAX

Floxs

Swarm and Flock Multi-agent RL environments, implemented with JAX using Esquilax

Multi-agent RL environments with a focus on larger scale flocks and swarms on continuos spaces. Implemented using Esquilax a JAX multi-agent simulation and RL library, and the Jumanji RL environment API.

from floxs.flock.env import Flock
import jax


env = Flock()

key = jax.random.PRNGKey(101)
state, ts = env.reset(key)
states = [state]

for _ in range(100):
    key, k = jax.random.split(key)
    actions = jax.random.uniform(k, (env.generator.num_boids, 2), minval=0.5, maxval=1.0)
    state, ts = env.step(state, actions)
    states.append(state)

# Save an animation of the environment
env.animate(states, interval=100, save_path="animation.gif")

See the Jumanji docs for more usage information.

Usage

The package and requirements can be installed using poetry by running

poetry install

Developers

Pre-Commit Hooks

Pre commit hooks can be installed by running

pre-commit install

Pre-commit checks can then be run using

task lint

Tests

Tests can be run with

task test

Build Documentation

Docs can be built using Sphinx by running

task docs

Built docs will be generated in the docs/build folder.