JaxMARL : Environnements d'apprentissage par renforcement multi-agents en JAX

JaxMARL: Multi-Agent RL Environments in JAX

November 16, 2023

Auteurs: 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

cs.AI

Résumé

Les benchmarks jouent un rôle crucial dans le développement des algorithmes d'apprentissage automatique. Par exemple, la recherche en apprentissage par renforcement (RL) a été fortement influencée par les environnements et les benchmarks disponibles. Cependant, les environnements de RL sont traditionnellement exécutés sur le CPU, limitant ainsi leur évolutivité avec les ressources de calcul typiques du milieu académique. Les récentes avancées dans JAX ont permis une utilisation plus large de l'accélération matérielle pour surmonter ces obstacles computationnels, permettant des pipelines d'entraînement et des environnements de RL massivement parallèles. Cela est particulièrement utile pour la recherche en apprentissage par renforcement multi-agent (MARL). Tout d'abord, plusieurs agents doivent être pris en compte à chaque étape de l'environnement, ce qui augmente la charge computationnelle, et deuxièmement, la complexité des échantillons est accrue en raison de la non-stationnarité, de l'observabilité partielle décentralisée, ou d'autres défis du MARL. Dans cet article, nous présentons JaxMARL, la première base de code open-source qui combine la facilité d'utilisation avec l'efficacité accélérée par GPU, et qui prend en charge un grand nombre d'environnements MARL couramment utilisés ainsi que des algorithmes de référence populaires. En considérant le temps réel, nos expériences montrent que notre pipeline d'entraînement basé sur JAX est jusqu'à 12500 fois plus rapide que les approches existantes. Cela permet des évaluations efficaces et approfondies, avec le potentiel de soulager la crise d'évaluation du domaine. Nous introduisons également et évaluons SMAX, une version vectorisée et simplifiée du populaire StarCraft Multi-Agent Challenge, qui élimine la nécessité d'exécuter le moteur de jeu StarCraft II. Cela permet non seulement une accélération par GPU, mais offre également un environnement MARL plus flexible, débloquant le potentiel pour l'auto-apprentissage, le méta-apprentissage, et d'autres applications futures en MARL. Nous fournissons le code à l'adresse suivante : https://github.com/flairox/jaxmarl.

English

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.