JaxMARL: Ambienti di Apprendimento Rinforzato Multi-Agente in JAX

Abstract

I benchmark svolgono un ruolo importante nello sviluppo degli algoritmi di machine learning. Ad esempio, la ricerca nel campo del reinforcement learning (RL) è stata fortemente influenzata dagli ambienti e dai benchmark disponibili. Tuttavia, gli ambienti RL sono tradizionalmente eseguiti sulla CPU, limitandone la scalabilità con le risorse computazionali tipiche del contesto accademico. I recenti progressi in JAX hanno permesso un uso più ampio dell'accelerazione hardware per superare questi ostacoli computazionali, abilitando pipeline di addestramento RL e ambienti massivamente paralleli. Ciò è particolarmente utile per la ricerca sul multi-agent reinforcement learning (MARL). Innanzitutto, è necessario considerare più agenti in ogni passo dell'ambiente, aumentando il carico computazionale, e in secondo luogo, la complessità del campionamento è incrementata a causa della non stazionarietà, dell'osservabilità parziale decentralizzata o di altre sfide tipiche del MARL. In questo articolo, presentiamo JaxMARL, la prima code base open-source che combina facilità d'uso con l'efficienza abilitata dalla GPU, supportando un ampio numero di ambienti MARL comunemente utilizzati e algoritmi baseline popolari. Considerando il tempo di esecuzione effettivo, i nostri esperimenti mostrano che la nostra pipeline di addestramento basata su JAX è fino a 12500 volte più veloce rispetto agli approcci esistenti. Ciò consente valutazioni efficienti e approfondite, con il potenziale di alleviare la crisi di valutazione del campo. Introduciamo inoltre e valutiamo SMAX, una versione vettorizzata e semplificata della popolare StarCraft Multi-Agent Challenge, che elimina la necessità di eseguire il motore di gioco StarCraft II. Ciò non solo abilita l'accelerazione GPU, ma fornisce anche un ambiente MARL più flessibile, sbloccando il potenziale per il self-play, il meta-learning e altre future applicazioni nel MARL. Forniamo il codice all'indirizzo 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.

JaxMARL: Ambienti di Apprendimento Rinforzato Multi-Agente in JAX

JaxMARL: Multi-Agent RL Environments in JAX

Abstract

Support