JaxMARL: Entornos de Aprendizaje por Refuerzo Multiagente en JAX

Resumen

Los puntos de referencia desempeñan un papel importante en el desarrollo de algoritmos de aprendizaje automático. Por ejemplo, la investigación en aprendizaje por refuerzo (RL, por sus siglas en inglés) ha estado fuertemente influenciada por los entornos y puntos de referencia disponibles. Sin embargo, los entornos de RL tradicionalmente se ejecutan en la CPU, lo que limita su escalabilidad con los recursos computacionales típicos del ámbito académico. Los avances recientes en JAX han permitido un uso más amplio de la aceleración por hardware para superar estos obstáculos computacionales, posibilitando pipelines de entrenamiento y entornos de RL masivamente paralelos. Esto es particularmente útil para la investigación en aprendizaje por refuerzo multiagente (MARL). En primer lugar, se deben considerar múltiples agentes en cada paso del entorno, lo que añade una carga computacional, y en segundo lugar, la complejidad de las muestras aumenta debido a la no estacionariedad, la observabilidad parcial descentralizada u otros desafíos propios del MARL. En este artículo, presentamos JaxMARL, la primera base de código de código abierto que combina facilidad de uso con eficiencia habilitada para GPU, y que soporta un gran número de entornos MARL comúnmente utilizados, así como algoritmos de referencia populares. Al considerar el tiempo de reloj, nuestros experimentos muestran que, por ejecución, nuestro pipeline de entrenamiento basado en JAX es hasta 12500 veces más rápido que los enfoques existentes. Esto permite evaluaciones eficientes y exhaustivas, con el potencial de aliviar la crisis de evaluación del campo. También presentamos y evaluamos SMAX, una versión vectorizada y simplificada del popular StarCraft Multi-Agent Challenge, que elimina la necesidad de ejecutar el motor del juego StarCraft II. Esto no solo permite la aceleración por GPU, sino que también proporciona un entorno MARL más flexible, desbloqueando el potencial para el autoaprendizaje, el metaaprendizaje y otras aplicaciones futuras en MARL. Proporcionamos el código en 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: Entornos de Aprendizaje por Refuerzo Multiagente en JAX

JaxMARL: Multi-Agent RL Environments in JAX

Resumen

Support