JaxMARL: Multi-Agent RL Environments and Algorithms in JAX
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2023
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866916465231265792 |
|---|---|
| author | Rutherford, Alexander Ellis, Benjamin Gallici, Matteo Cook, Jonathan Lupu, Andrei Ingvarsson, Gardar Willi, Timon Hammond, Ravi Khan, Akbir de Witt, Christian Schroeder Souly, Alexandra Bandyopadhyay, Saptarashmi Samvelyan, Mikayel Jiang, Minqi Lange, Robert Tjarko Whiteson, Shimon Lacerda, Bruno Hawes, Nick Rocktaschel, Tim Lu, Chris Foerster, Jakob Nicolaus |
| author_facet | Rutherford, Alexander Ellis, Benjamin Gallici, Matteo Cook, Jonathan Lupu, Andrei Ingvarsson, Gardar Willi, Timon Hammond, Ravi Khan, Akbir de Witt, Christian Schroeder Souly, Alexandra Bandyopadhyay, Saptarashmi Samvelyan, Mikayel Jiang, Minqi Lange, Robert Tjarko Whiteson, Shimon Lacerda, Bruno Hawes, Nick Rocktaschel, Tim Lu, Chris Foerster, Jakob Nicolaus |
| contents | Benchmarks are crucial in the development of machine learning algorithms, with available environments significantly influencing reinforcement learning (RL) research. Traditionally, RL environments run on the CPU, which limits their scalability with typical academic compute. However, recent advancements in JAX have enabled the wider use of hardware acceleration, enabling massively parallel RL training pipelines and environments. While this has been successfully applied to single-agent RL, it has not yet been widely adopted for multi-agent scenarios. In this paper, we present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments and popular baseline algorithms. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches, and up to 12500x when multiple training runs are vectorized. This enables efficient and thorough evaluations, potentially alleviating the evaluation crisis in the field. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge, which removes the need to run the StarCraft II game engine. This not only enables GPU acceleration, but also provides a more flexible MARL environment, unlocking the potential for self-play, meta-learning, and other future applications in MARL. The code is available at https://github.com/flairox/jaxmarl. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2311_10090 |
| institution | arXiv |
| publishDate | 2023 |
| record_format | arxiv |
| spellingShingle | JaxMARL: Multi-Agent RL Environments and Algorithms in JAX Rutherford, Alexander Ellis, Benjamin Gallici, Matteo Cook, Jonathan Lupu, Andrei Ingvarsson, Gardar Willi, Timon Hammond, Ravi Khan, Akbir de Witt, Christian Schroeder Souly, Alexandra Bandyopadhyay, Saptarashmi Samvelyan, Mikayel Jiang, Minqi Lange, Robert Tjarko Whiteson, Shimon Lacerda, Bruno Hawes, Nick Rocktaschel, Tim Lu, Chris Foerster, Jakob Nicolaus Machine Learning Artificial Intelligence Multiagent Systems Benchmarks are crucial in the development of machine learning algorithms, with available environments significantly influencing reinforcement learning (RL) research. Traditionally, RL environments run on the CPU, which limits their scalability with typical academic compute. However, recent advancements in JAX have enabled the wider use of hardware acceleration, enabling massively parallel RL training pipelines and environments. While this has been successfully applied to single-agent RL, it has not yet been widely adopted for multi-agent scenarios. In this paper, we present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments and popular baseline algorithms. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches, and up to 12500x when multiple training runs are vectorized. This enables efficient and thorough evaluations, potentially alleviating the evaluation crisis in the field. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge, which removes the need to run the StarCraft II game engine. This not only enables GPU acceleration, but also provides a more flexible MARL environment, unlocking the potential for self-play, meta-learning, and other future applications in MARL. The code is available at https://github.com/flairox/jaxmarl. |
| title | JaxMARL: Multi-Agent RL Environments and Algorithms in JAX |
| topic | Machine Learning Artificial Intelligence Multiagent Systems |
| url | https://arxiv.org/abs/2311.10090 |