instadeepai / mava Goto Github PK

Mava provides simplified code for quickly iterating on ideas in multi-agent reinforcement learning (MARL) with useful implementations of MARL algorithms in JAX allowing for easy parallelisation across devices with JAX's pmap. Mava is a project originating in the Research Team at InstaDeep.

To join us in these efforts, please feel free to reach out, raise issues or read our contribution guidelines (or just star 🌟 to stay up to date with the latest developments)!

Mava currently offers the following building blocks for MARL research:

• 🥑 Implementations of MARL algorithms: Implementations of multi-agent PPO systems that follow both the Centralised Training with Decentralised Execution (CTDE) and Decentralised Training with Decentralised Execution (DTDE) MARL paradigms.
• 🍬 Environment Wrappers: Example wrappers for mapping Jumanji environments to an environment that is compatible with Mava. At the moment, we support Robotic Warehouse and Level-Based Foraging with plans to support more environments soon. We have also recently added support for the SMAX environment from JaxMARL.
• 🎓 Educational Material: Quickstart notebook to demonstrate how Mava can be used and to highlight the added value of JAX-based MARL.
• 🧪 Statistically robust evaluation: Mava natively supports logging to json files which adhere to the standard suggested by Gorsane et al. (2022). This enables easy downstream experiment plotting and aggregation using the tools found in the MARL-eval library.

For comparing Mava’s stability to other JAX-based baseline algorithms, we train Mava’s recurrent IPPO and MAPPO systems on a broad range of SMAX tasks. In all cases we do not rerun baselines but instead take results for final win rates from the JaxMARL technical report. For the full SMAX experiments results, please see the following page.

All of the experiments below were performed using an NVIDIA Quadro RTX 4000 GPU with 8GB Memory.

In order to show the utility of end-to-end JAX-based MARL systems and JAX-based environments we compare the speed of Mava against EPyMARL as measured in total training wallclock time on simple Robotic Warehouse (RWARE) tasks with 2 and 4 agents. Our aim is to illustrate the speed increases that are possible with using end-to-end Jax-based systems and we do not necessarily make an effort to achieve optimal performance. For EPyMARL, we use the hyperparameters as recommended by Papoudakis et al. (2020) and for Mava we performed a basic grid search. In both cases, systems were trained up to 20 million total environment steps using 16 vectorised environments.

In order to benefit from the wallclock speed-ups afforded by JAX-based systems it is required that environments also be written in JAX. It is for this reason that Mava does not use the exact same version of the RWARE environment as EPyMARL but instead uses a JAX-based implementation of RWARE found in Jumanji, under the name RobotWarehouse. One of the notable differences in the underlying environment logic is that RobotWarehouse will not attempt to resolve agent collisions but will instead terminate an episode when agents do collide. In our experiments, this appeared to make the environment more challenging. For this reason we show the performance of Mava on Jumanji with and without termination upon collision indicated with w/o collision in the figure legends. For a more detailed discussion, please see the following page.

Mava also supports Jumanji's LBF. We evaluate Mava's recurrent MAPPO system on LBF, against EPyMARL (we used original LBF for EPyMARL) in 2 and 4 agent settings up to 20 million timesteps. Both systems were trained using 16 vectorized environments. For the EPyMARL systems we use a NVIDIA A100 GPU and for the Mava systems we use a GeForce RTX 3050 laptop GPU with 4GB of memory. To show how Mava can generalise to different hardware, we also train the Mava systems on a TPU v3-8. We plan to publish comprehensive performance benchmarks for all Mava's algorithms across various LBF scenarios soon.

Furthermore, we illustrate the speed of Mava by showing the steps per second as the number of parallel environments is increased. These steps per second scaling plots were computed using a standard laptop GPU, specifically an RTX-3060 GPU with 6GB memory.

The current code in Mava is adapted from PureJaxRL which provides high-quality single-file implementations with research-friendly features. In turn, PureJaxRL is inspired by the code philosophy from CleanRL. Along this vein of easy-to-use and understandable RL codebases, Mava is not designed to be a modular library and is not meant to be imported. Our repository focuses on simplicity and clarity in its implementations while utilising the advantages offered by JAX such as pmap and vmap, making it an excellent resource for researchers and practitioners to build upon.

At the moment Mava is not meant to be installed as a library, but rather to be used as a research tool.

You can use Mava by cloning the repo and pip installing as follows:

git clone https://github.com/instadeepai/mava.git
cd mava
pip install -e .

We have tested Mava on Python 3.9. Note that because the installation of JAX differs depending on your hardware accelerator, we advise users to explicitly install the correct JAX version (see the official installation guide). For more in-depth installation guides including Docker builds and virtual environments, please see our detailed installation guide.

To get started with training your first Mava system, simply run one of the system files. e.g.,

python mava/systems/ff_ippo.py

Mava makes use of Hydra for config management. In order to see our default system configs please see the mava/configs/ directory. A benefit of Hydra is that configs can either be set in config yaml files or overwritten from the terminal on the fly. For an example of running a system on the LBF environment, the above code can simply be adapted as follows:

python mava/systems/ff_ippo.py env=lbf

Different scenarios can also be run by making the following config updates from the terminal:

python mava/systems/ff_ippo.py env=rware env/scenario=tiny-4ag

Additionally, we also have a Quickstart notebook that can be used to quickly create and train your first Multi-agent system.

Mava can be used in a wide array of advanced systems. As an example, we demonstrate recording experience data from one of our PPO systems into a Flashbax Vault. This vault can then easily be integrated into offline MARL systems, such as those found in OG-MARL. See the Advanced README for more information.

Please read our contributing docs for details on how to submit pull requests, our Contributor License Agreement and community guidelines.

We plan to iteratively expand Mava in the following increments:

• 🌴 Support for more environments.
• 🔁 More robust recurrent systems.
• 🌳 Support for non JAX-based environments.
• 🦾 Support for off-policy algorithms.
• 🎛 Continuous action space environments and algorithms.

Please do follow along as we develop this next phase!

Originally Mava was written in Tensorflow 2. Support for the TF2-based framework and systems has now been fully deprecated. If you would still like to use it, please install v0.1.3 of Mava (i.e. pip install id-mava==0.1.3).

InstaDeep's MARL ecosystem in JAX. In particular, we suggest users check out the following sister repositories:

• 🔌 OG-MARL: datasets with baselines for offline MARL in JAX.
• 🌴 Jumanji: a diverse suite of scalable reinforcement learning environments in JAX.
• 😎 Matrax: a collection of matrix games in JAX.
• ⚡ Flashbax: accelerated replay buffers in JAX.
• 📈 MARL-eval: standardised experiment data aggregation and visualisation for MARL.

Related. Other libraries related to accelerated MARL in JAX.

• 🦊 JaxMARL: accelerated MARL environments with baselines in JAX.
• 🌀 DeepMind Anakin for the Anakin podracer architecture to train RL agents at scale.
• ♟️ Pgx: JAX implementations of classic board games, such as Chess, Go and Shogi.
• 🔼 Minimax: JAX implementations of autocurricula baselines for RL.

If you use Mava in your work, please cite the accompanying technical report:

@article{dekock2023mava,
    title={Mava: a research library for distributed multi-agent reinforcement learning in JAX},
    author={Ruan de Kock and Omayma Mahjoub and Sasha Abramowitz and Wiem Khlifi and Callum Rhys Tilbury
    and Claude Formanek and Andries P. Smit and Arnu Pretorius},
    year={2023},
    journal={arXiv preprint arXiv:2107.01460},
    url={https://arxiv.org/pdf/2107.01460.pdf},
}

We would like to thank all the authors who contributed to the previous TF version of Mava: Kale-ab Tessera, St John Grimbly, Kevin Eloff, Siphelele Danisa, Lawrence Francis, Jonathan Shock, Herman Kamper, Willie Brink, Herman Engelbrecht, Alexandre Laterre, Karim Beguir. Their contributions can be found in our TF technical report.

The development of Mava was supported with Cloud TPUs from Google's TPU Research Cloud (TRC) 🌤.