bold-lab-ai/JaxMARL

Multi-Agent Reinforcement Learning with JAX

What it solves

JaxMARL provides a unified, GPU-enabled framework for Multi-Agent Reinforcement Learning (MARL). It addresses the need for a single library that combines a wide variety of MARL environments and popular baseline algorithms, allowing researchers to evaluate methods across different tasks efficiently.

How it works

Built on JAX, the library leverages GPU acceleration to provide high-performance, vectorized environments. It follows an API inspired by PettingZoo and Gymnax, where agents' actions, observations, and rewards are handled as dictionaries. The library also includes SMAX, a simplified, vectorized version of the StarCraft Multi-Agent Challenge that removes the dependency on the StarCraft II game engine.

Who it’s for

It is designed for researchers and developers working in multi-agent reinforcement learning who need high-throughput training and a standardized way to benchmark against established baselines.

Highlights

  • Diverse Environment Suite: Includes a wide range of tasks such as MPE, Overcooked, Hanabi, and robotic control via Multi-Agent Brax.
  • GPU-Enabled Efficiency: Fully implemented in JAX for high-speed execution on GPUs and TPUs.
  • Baseline Algorithms: Provides single-file implementations of popular algorithms like IPPO, MAPPO, QMIX, and VDN.
  • SMAX Environment: A novel, engine-free version of the StarCraft micro-management challenge.
  • Sim2Real Support: Integration with JaxRobotarium for multi-robot environments with open access sim2real capabilities.

Related

  • Project
    instadeepai/Mava

    A distributed multi-agent reinforcement learning research library built in JAX that provides high-performance implementations of state-of-the-art MARL algorithms.

  • Project
    RobertTLange/gymnax

    A collection of reinforcement learning environments reimplemented in JAX to enable high-throughput, vectorized experiments on GPUs and TPUs.

  • Project
    dynamiqs/dynamiqs

    A GPU-accelerated and differentiable Python library for high-performance simulation of quantum systems using JAX.

  • Project
    extropic-ai/thrml

    A JAX library for building and sampling probabilistic graphical models, focusing on efficient block Gibbs sampling and energy-based models.

  • Project
    Farama-Foundation/PettingZoo

    A Python library for multi-agent reinforcement learning research that provides a standardized API and a diverse set of environments, similar to a multi-agent version of Gymnasium.