Multi-agent RL environment based on Boids, implemented with JAX using Esquilax
The environment is based on popular boids model where agents recreate flocking behaviours based on simple interaction rules. The environment implements boids as a multi-agent reinforcement problem where each boid takes individual actions and have individual localised views of the environment.
This environment has been built using Esquilax a JAX multi-agent simulation and RL library.
import flock_env
import jax
key = jax.random.PRNGKey(101)
key_reset, key_act, key_step = jax.random.split(key, 3)
# Initialise a flock environment with 10 agents
env = flock_env.SimpleFlockEnv(
reward_func=flock_env.rewards.exponential_rewards,
n_agents=10,
i_range=0.1,
)
env_params = env.default_params()
# Reset the environment and get state and observation
obs, state = env.reset(key_reset, env_params)
# Sample random action for agents
actions = jax.random.uniform(key_act, (10, 2))
# Step the environment
new_obs, new_state, rewards, dones = env.step(
key_step, env_params, state, actions
)See examples/ppo_example.ipynb for an example
of training a Proximal-Policy-Optimisation agent with this environment
using a JAX implementation of PPO
and Esquilax's built-in
RL training functionality.
The package and requirements can be installed using poetry by running
poetry install- Objects/obstacles in the environment
The previous version of this project built around Numba can be found in
/deprecated
Pre commit hooks can be installed by running
pre-commit installPre-commit checks can then be run using
task lintTests can be run with
task test