DOCSEARCHbyAlgolia
Submit Docs

JaxMARL

repository·main·Indexed 21 days ago

https://github.com/flairox/jaxmarl

A Multi-Agent Reinforcement Learning (MARL) library built in JAX, providing GPU-efficient environments and baseline algorithms. It includes pure JAX implementations of IPPO, MAPPO, and Q-Learning algorithms such as PQN-VDN, IQL, VDN, QMIX, TransfQMix, and SHAQ. The library supports environments including MPE, SMAX, Hanabi, Overcooked, and JaxNav, and features a CTRolloutManager wrapper for centralized training and parallel execution.

Tokens
21.3K
Snippets
55
Records
95
Agent score
72%

What's inside jaxmarl

  1. Overview of OvercookedV2 features

    main

    OvercookedV2 is an extended version of the original Overcooked environment designed for complex coordination challenges. Key features include:

    • Configurable agent view radius for partial observability.
    • Multiple ingredients and recipes.
    • Asymmetric information via recipe indicators.
    • Randomized starting positions and directions.
    • Grounded communication using button recipe indicators.
    • Flexible layout creation via ASCII strings.

  2. Overview of QLearning Baselines

    main

    JaxMARL provides pure JAX implementations of several Multi-Agent Reinforcement Learning (MARL) Q-Learning algorithms. These baselines are designed for high performance and training speed.

    Supported Algorithms:

    • PQN-VDN (Parallelised Q-Network): Most performant baseline for returns and training speed. Supports MPE, SMAX, Hanabi, and Overcooked.
    • IQL (Independent Q-Learners): Supports MPE, SMAX, and Overcooked.
    • VDN (Value Decomposition Network): Supports MPE, SMAX, Hanabi, and Overcooked.
    • QMIX: Supports MPE, SMAX, and Hanabi (Note: Not supported for Overcooked).
    • TransfQMix: Supports MPE_Spread and SMAX (requires observation matrices).
    • SHAQ (Shapley Value Theory): Implementation follows the original paper code.

    Algorithm Compatibility Matrix:

    AlgorithmMPESMAXOvercookedHanabi
    PQN-VDNYesYesYesYes
    IQLYesYesYesNo
    VDNYesYesYesYes
    QMIXYesYesNoYes
    TransfQMixYes (Spread)YesNoNo

  3. Overview of SMAX environments

    main
    SMAX is a purely JAX-based implementation of SMAC-like environments. It is designed for studying decentralized unit micromanagement. Each scenario involves fixed teams of units competing against each other. Because it is implemented in JAX, it is suitable for high-performance reinforcement learning workflows.

  4. Overview of the GridWorld environment

    main
    The GridWorld in JaxMARL is a procedurally-generated environment modeled after MiniGrid. It follows the gymnax interface, making it compatible with standard JAX-based reinforcement learning workflows. It is intended to serve as a foundation for building more complex gridworld environments, such as Overcooked or Cultural Learning GridWorlds.

  5. Overview of Multi-Agent Brax Environments

    main

    This directory contains multi-agent factorisations of MuJoCo tasks based on the FACMAC paper. Each agent controls a specific subset of joints and observes only a local state.

    Note: This environment is deprecated because Brax itself is deprecated; it is intended to be migrated to MJX.

    Available Environments

    EnvironmentDescription
    ant_4x24 agents, 2 joints each. One agent controls each leg.
    halfcheetah_6x16 agents, 1 joint each. One agent controls each joint.
    hopper_3x13 agents, 1 joint each. One agent controls each joint.
    humanoid_9|82 agents, 9 and 8 joints. One agent controls the upper body, the other the lower body.
    walker2d_2x32 agents, 3 joints each. Factored into right and left leg.

  6. Available Q-Learning Algorithms in JaxMARL

    main

    JaxMARL provides pure JAX implementations of several Multi-Agent Reinforcement Learning (MARL) Q-Learning algorithms.

    Supported Algorithms:

    • PQN-VDN (Parallelised Q-Network): Currently the most performant baseline for Q-Learning in terms of returns and training speed.
    • IQL (Independent Q-Learners)
    • VDN (Value Decomposition Network)
    • QMIX
    • TransfQMix (Transformers for Leveraging the Graph Structure of MARL Problems): Note that this currently only supports MPE_Spread and SMAX. For other environments, you must wrap observation vectors into matrices using jaxmarl.wrappers.transformers.
    • SHAQ (Incorporating Shapley Value Theory into Multi-Agent Q-Learning)

    Algorithm Compatibility:

    • Standard (IQL, VDN, QMIX): Supports MPE, SMAX, and Overcooked (Note: QMIX is not supported on Overcooked).
    • PQN-VDN: Supports MPE, SMAX, Hanabi, and Overcooked.

  7. What is SMAX

    main
    SMAX is a purely JAX-based implementation of the SMAC (StarCraft Multi-Agent Challenge) environment. It focuses on decentralized unit micromanagement across various scenarios where each scenario features fixed teams of units. Because it is implemented in JAX, it is designed for high-performance reinforcement learning.

  8. Overview of the Coin Game environment

    main

    The Coin Game is a multi-agent grid-world environment implemented in JaxMARL, designed to simulate social dilemmas with high-dimensional dynamic states. It is based on the implementation used in research such as Model-Free Opponent Shaping (Lu et al.) and Learning with Opponent-Learning Awareness (Foerster et al.).

    Environment Mechanics:

    • Setup: Two players (labeled red and blue) compete in a 3x3 grid.
    • Objective: Players attempt to pick up coins (also labeled red and blue) by moving into the same position as the coin.
    • Rewards:
      • Picking up a coin of your own color: +1 reward.
      • Picking up a coin of the opponent's color: The opponent receives a -2 reward.
    • Social Dilemma: If both agents act purely greedily to pick up every coin, the expected reward for both agents is 0, illustrating the tension between individual greed and collective cooperation.

  9. How the CTRolloutManager wrapper works

    main

    All Q-Learning algorithms utilize the CTRolloutManager environment wrapper (located in jaxmarl.wrappers.baselines). This wrapper is essential for centralized training and parallel execution. It performs the following tasks:

    1. Batchification: Batchifies the step and reset functions to enable running multiple environments in parallel.
    2. Centralized Training Support: Adds a global observation (obs["__all__"]) and a global reward (rewards["__all__"]) to the returns of env.step.
    3. Preprocessing: Preprocesses and uniforms observation vectors (e.g., flattening, padding, and adding features like ID one-hot encoding).

    You can modify this wrapper to suit specific research needs.

  10. Understand Overcooked action and observation spaces

    main

    Action Space

    There are 6 possible actions:

    • right
    • down
    • left
    • up
    • interact
    • no-op

    Observation Space

    Observations are designed for ConvNets. They are sparse, mostly binary encodings with the shape layout_height x layout_width x n_channels, where n_channels = 26. For the specific mapping of each channel, refer to the get_obs(...) method in overcooked.py.

  11. How CTRolloutManager works in QLearning baselines

    main

    All Q-Learning algorithms utilize the CTRolloutManager environment wrapper (located in jaxmarl.wrappers.baselines). This wrapper is essential for centralized training and parallel execution. It performs the following tasks:

    • Batchification: Batches the step and reset functions to enable running multiple environments in parallel.
    • Centralized Training Support: Injects a global observation (obs["__all__"]) and a global reward (rewards["__all__"]) into the returns of env.step.
    • Preprocessing: Uniforms and preprocesses observation vectors (e.g., flattening, padding, and adding features like ID one-hot encoding).

    Developers can modify this wrapper to suit specific research needs.

  12. Understand Overcooked-JAX Action and Observation spaces

    main

    Action Space

    The environment provides 6 possible actions:

    • right
    • down
    • left
    • up
    • interact
    • no-op

    Observation Space

    Observations are designed for ConvNets and follow the original Overcooked-AI featurization. Each observation is a sparse, mostly binary encoding with the shape layout_height x layout_width x 26.

    For a detailed mapping of the 26 channels, refer to the get_obs(...) method in overcooked.py.