JaxMARL: Ambientes de Aprendizado por Reforço Multiagente em JAX

Resumo

Os benchmarks desempenham um papel importante no desenvolvimento de algoritmos de aprendizado de máquina. Por exemplo, a pesquisa em aprendizado por reforço (RL, do inglês Reinforcement Learning) tem sido fortemente influenciada pelos ambientes e benchmarks disponíveis. No entanto, os ambientes de RL tradicionalmente são executados na CPU, limitando sua escalabilidade com os recursos computacionais típicos da academia. Avanços recentes no JAX permitiram o uso mais amplo de aceleração de hardware para superar esses obstáculos computacionais, possibilitando pipelines de treinamento e ambientes de RL massivamente paralelos. Isso é particularmente útil para pesquisas em aprendizado por reforço multiagente (MARL, do inglês Multi-Agent Reinforcement Learning). Primeiramente, múltiplos agentes devem ser considerados a cada passo do ambiente, aumentando a carga computacional, e, em segundo lugar, a complexidade de amostragem é maior devido à não estacionariedade, observabilidade parcial descentralizada ou outros desafios do MARL. Neste artigo, apresentamos o JaxMARL, a primeira base de código de código aberto que combina facilidade de uso com eficiência habilitada por GPU, e suporta um grande número de ambientes de MARL comumente utilizados, bem como algoritmos de referência populares. Ao considerar o tempo de execução, nossos experimentos mostram que, por execução, nosso pipeline de treinamento baseado em JAX é até 12500 vezes mais rápido que as abordagens existentes. Isso permite avaliações eficientes e abrangentes, com o potencial de aliviar a crise de avaliação do campo. Também introduzimos e avaliamos o SMAX, uma versão vetorizada e simplificada do popular StarCraft Multi-Agent Challenge, que elimina a necessidade de executar o motor do jogo StarCraft II. Isso não apenas permite a aceleração por GPU, mas também fornece um ambiente de MARL mais flexível, desbloqueando o potencial para auto-jogo, meta-aprendizado e outras aplicações futuras no MARL. Disponibilizamos o código em 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: Ambientes de Aprendizado por Reforço Multiagente em JAX

JaxMARL: Multi-Agent RL Environments in JAX

Resumo

Support