JaxMARL: Multi-Agent RL Environments in JAX
- URL: http://arxiv.org/abs/2311.10090v4
- Date: Tue, 19 Dec 2023 14:55:15 GMT
- Title: JaxMARL: Multi-Agent RL Environments in JAX
- Authors: Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook,
Andrei Lupu, Gardar Ingvarsson, Timon Willi, Akbir Khan, Christian Schroeder
de Witt, Alexandra Souly, Saptarashmi Bandyopadhyay, Mikayel Samvelyan, Minqi
Jiang, Robert Tjarko Lange, Shimon Whiteson, Bruno Lacerda, Nick Hawes, Tim
Rocktaschel, Chris Lu, Jakob Nicolaus Foerster
- Abstract summary: We present JaxMARL, the first open-source code base that combines ease-of-use with GPU enabled efficiency.
Our experiments show that per-run our JAX-based training pipeline is up to 12500x faster than existing approaches.
We also introduce and benchmark SMAX, a vectorised, simplified version of the popular StarCraft Multi-Agent Challenge.
- Score: 107.7560737385902
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Benchmarks play an important role in the development of machine learning
algorithms. For example, research in reinforcement learning (RL) has been
heavily influenced by available environments and benchmarks. However, RL
environments are traditionally run on the CPU, limiting their scalability with
typical academic compute. Recent advancements in JAX have enabled the wider use
of hardware acceleration to overcome these computational hurdles, enabling
massively parallel RL training pipelines and environments. This is particularly
useful for multi-agent reinforcement learning (MARL) research. First of all,
multiple agents must be considered at each environment step, adding
computational burden, and secondly, the sample complexity is increased due to
non-stationarity, decentralised partial observability, or other MARL
challenges. In this paper, we present JaxMARL, the first open-source code base
that combines ease-of-use with GPU enabled efficiency, and supports a large
number of commonly used MARL environments as well as popular baseline
algorithms. When considering wall clock time, our experiments show that per-run
our JAX-based training pipeline is up to 12500x faster than existing
approaches. This enables efficient and thorough evaluations, with the potential
to alleviate the evaluation crisis of the field. We also introduce and
benchmark SMAX, a vectorised, simplified version 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. We provide code at
https://github.com/flairox/jaxmarl.
Related papers
- Benchmarking Predictive Coding Networks -- Made Simple [48.652114040426625]
We first propose a library called PCX, whose focus lies on performance and simplicity.
We use PCX to implement a large set of benchmarks for the community to use for their experiments.
arXiv Detail & Related papers (2024-07-01T10:33:44Z) - XuanCe: A Comprehensive and Unified Deep Reinforcement Learning Library [18.603206638756056]
XuanCe is a comprehensive and unified deep reinforcement learning (DRL) library.
XuanCe offers a wide range of functionalities, including over 40 classical DRL and multi-agent DRL algorithms.
XuanCe is open-source and can be accessed at https://agi-brain.com/agi-brain/xuance.git.
arXiv Detail & Related papers (2023-12-25T14:45:39Z) - BenchMARL: Benchmarking Multi-Agent Reinforcement Learning [8.130948896195878]
BenchMARL is the first training library created to enable standardized benchmarking across different algorithms, models, and environments.
BenchMARL uses TorchRL as its backend, granting it high performance and maintained state-of-the-art implementations.
arXiv Detail & Related papers (2023-12-03T18:15:58Z) - FusionAI: Decentralized Training and Deploying LLMs with Massive
Consumer-Level GPUs [57.12856172329322]
We envision a decentralized system unlocking the potential vast untapped consumer-level GPU.
This system faces critical challenges, including limited CPU and GPU memory, low network bandwidth, the variability of peer and device heterogeneity.
arXiv Detail & Related papers (2023-09-03T13:27:56Z) - SequeL: A Continual Learning Library in PyTorch and JAX [50.33956216274694]
SequeL is a library for Continual Learning that supports both PyTorch and JAX frameworks.
It provides a unified interface for a wide range of Continual Learning algorithms, including regularization-based approaches, replay-based approaches, and hybrid approaches.
We release SequeL as an open-source library, enabling researchers and developers to easily experiment and extend the library for their own purposes.
arXiv Detail & Related papers (2023-04-21T10:00:22Z) - marl-jax: Multi-Agent Reinforcement Leaning Framework [7.064383217512461]
We present marl-jax, a multi-agent reinforcement learning software package for training and evaluating social generalization of the agents.
The package is designed for training a population of agents in multi-agent environments and evaluating their ability to generalize to diverse background agents.
arXiv Detail & Related papers (2023-03-24T05:05:01Z) - EnvPool: A Highly Parallel Reinforcement Learning Environment Execution
Engine [69.47822647770542]
parallel environment execution is often the slowest part of the whole system but receives little attention.
With a curated design for paralleling RL environments, we have improved the RL environment simulation speed across different hardware setups.
On a high-end machine, EnvPool achieves 1 million frames per second for the environment execution on Atari environments and 3 million frames per second on MuJoCo environments.
arXiv Detail & Related papers (2022-06-21T17:36:15Z) - ElegantRL-Podracer: Scalable and Elastic Library for Cloud-Native Deep
Reinforcement Learning [141.58588761593955]
We present a library ElegantRL-podracer for cloud-native deep reinforcement learning.
It efficiently supports millions of cores to carry out massively parallel training at multiple levels.
At a low-level, each pod simulates agent-environment interactions in parallel by fully utilizing nearly 7,000 GPU cores in a single GPU.
arXiv Detail & Related papers (2021-12-11T06:31:21Z) - TLeague: A Framework for Competitive Self-Play based Distributed
Multi-Agent Reinforcement Learning [28.795986840557475]
TLeague aims at large-scale training and implements several main-stream-MARL algorithms.
We present experiments over StarCraft II, ViZDoom and Pommerman to show the efficiency and effectiveness of TLeague.
arXiv Detail & Related papers (2020-11-25T17:24:20Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.