JaxMARL: Среды для многопользовательского обучения с подкреплением в JAX

Аннотация

Бенчмарки играют важную роль в разработке алгоритмов машинного обучения. Например, исследования в области обучения с подкреплением (RL) в значительной степени зависят от доступных сред и бенчмарков. Однако традиционно RL-среды выполняются на CPU, что ограничивает их масштабируемость при использовании типичных академических вычислительных ресурсов. Недавние достижения в JAX позволили шире использовать аппаратное ускорение для преодоления этих вычислительных ограничений, что сделало возможным создание массово параллельных конвейеров обучения и сред RL. Это особенно полезно для исследований в области многопользовательского обучения с подкреплением (MARL). Во-первых, на каждом шаге среды необходимо учитывать множество агентов, что увеличивает вычислительную нагрузку, а во-вторых, сложность выборки возрастает из-за нестационарности, децентрализованной частичной наблюдаемости или других проблем MARL. В данной статье мы представляем JaxMARL — первую открытую кодовую базу, которая сочетает простоту использования с эффективностью, поддерживаемой GPU, и включает большое количество популярных сред MARL, а также базовых алгоритмов. Наши эксперименты показывают, что с точки зрения реального времени обучения наш конвейер на основе JAX может быть до 12500 раз быстрее существующих подходов. Это позволяет проводить эффективные и тщательные оценки, что может помочь смягчить кризис оценки в данной области. Мы также представляем и тестируем SMAX — векторизованную, упрощенную версию популярного StarCraft Multi-Agent Challenge, которая устраняет необходимость запуска игрового движка StarCraft II. Это не только позволяет использовать GPU-ускорение, но и предоставляет более гибкую среду MARL, открывая возможности для самообучения, метаобучения и других будущих приложений в MARL. Код доступен по адресу 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: Среды для многопользовательского обучения с подкреплением в JAX

JaxMARL: Multi-Agent RL Environments in JAX

Аннотация

Support