JaxMARL: Multi-Agent RL-Umgebungen in JAX

Zusammenfassung

Benchmarks spielen eine wichtige Rolle bei der Entwicklung von Algorithmen für maschinelles Lernen. Beispielsweise wurde die Forschung im Bereich des bestärkenden Lernens (Reinforcement Learning, RL) stark von verfügbaren Umgebungen und Benchmarks beeinflusst. Traditionell werden RL-Umgebungen jedoch auf der CPU ausgeführt, was ihre Skalierbarkeit mit typischer akademischer Rechenleistung einschränkt. Jüngste Fortschritte in JAX haben die breitere Nutzung von Hardware-Beschleunigung ermöglicht, um diese rechenintensiven Hürden zu überwinden, was massiv parallele RL-Trainingspipelines und Umgebungen ermöglicht. Dies ist besonders nützlich für die Forschung im Bereich des Multi-Agenten-Reinforcement-Learnings (MARL). Zum einen müssen mehrere Agenten in jedem Umgebungsschritt berücksichtigt werden, was die Rechenlast erhöht, und zum anderen steigt die Probenkomplexität aufgrund von Nicht-Stationarität, dezentraler partieller Beobachtbarkeit oder anderen MARL-Herausforderungen. In diesem Artikel stellen wir JaxMARL vor, die erste Open-Source-Codebasis, die Benutzerfreundlichkeit mit GPU-optimierter Effizienz kombiniert und eine Vielzahl häufig verwendeter MARL-Umgebungen sowie beliebte Baseline-Algorithmen unterstützt. In Bezug auf die Echtzeit zeigen unsere Experimente, dass unsere JAX-basierte Trainingspipeline pro Durchlauf bis zu 12500-mal schneller ist als bestehende Ansätze. Dies ermöglicht effiziente und gründliche Bewertungen und hat das Potenzial, die Evaluationskrise in diesem Bereich zu lindern. Wir stellen außerdem SMAX vor, eine vektorisierte, vereinfachte Version der beliebten StarCraft Multi-Agent Challenge, die den Betrieb der StarCraft II-Spiel-Engine überflüssig macht. Dies ermöglicht nicht nur GPU-Beschleunigung, sondern bietet auch eine flexiblere MARL-Umgebung, die das Potenzial für Selbstspiel, Meta-Lernen und andere zukünftige Anwendungen in MARL freisetzt. Den Code stellen wir unter https://github.com/flairox/jaxmarl zur Verfügung.

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.

JaxMARL: Multi-Agent RL-Umgebungen in JAX

JaxMARL: Multi-Agent RL Environments in JAX

Zusammenfassung

Support