Pith. sign in

REVIEW 9 cited by

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

Not yet reviewed by Pith; the record is open.

Re-run · record.json Download PDF Read on arXiv ↗

This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.

SPECIMEN: schema-true, not a live event

T0 review · schema-true

One-sentence machine reading of the paper's core claim.

pith:XXXXXXXX · record.json · timestamp

arxiv 2311.10090 v6 pith:B67RZONZ submitted 2023-11-16 cs.LG cs.AIcs.MA

Alexander Rutherford , Benjamin Ellis , Matteo Gallici , Jonathan Cook , Andrei Lupu , Gardar Ingvarsson , Timon Willi , Ravi Hammond
show 13 more authors
Akbir Khan Christian Schroeder de Witt Alexandra Souly Saptarashmi Bandyopadhyay Mikayel Samvelyan Minqi Jiang Robert Tjarko Lange Shimon Whiteson Bruno Lacerda Nick Hawes Tim Rocktaschel Chris Lu Jakob Nicolaus Foerster
This is my paper
classification cs.LGcs.AIcs.MA
keywords environmentsalgorithmsjaxmarlmarlmulti-agenttrainingaccelerationavailable
verification ladder T0 review T1 audit T2 compute T3 formal
0 comments
read the original abstract

Benchmarks are crucial in the development of machine learning algorithms, with available environments significantly influencing reinforcement learning (RL) research. Traditionally, RL environments run on the CPU, which limits their scalability with typical academic compute. However, recent advancements in JAX have enabled the wider use of hardware acceleration, enabling massively parallel RL training pipelines and environments. While this has been successfully applied to single-agent RL, it has not yet been widely adopted for multi-agent scenarios. In this paper, we present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments and popular baseline algorithms. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches, and up to 12500x when multiple training runs are vectorized. This enables efficient and thorough evaluations, potentially alleviating the evaluation crisis in the field. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation 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. The code is available at https://github.com/flairox/jaxmarl.

Discussion (0). Sign in to comment.

Forward citations

Cited by 9 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

  1. Recurrent Structural Policy Gradient for Partially Observable Mean Field Games

    cs.AI 2026-02 conditional novelty 6.0 of 10

    RSPG extends structural policy gradients to shared-observation partially observable mean-field games, learning recurrent policies with order-of-magnitude faster convergence than model-free RL.

  2. Benchmarking Partial Observability in Reinforcement Learning with a Suite of Memory-Improvable Domains

    cs.LG 2025-07 conditional novelty 6.0 of 10

    A JAX benchmark suite of nine memory-improvable partially observable RL environments, with evidence that recurrent and transformer agents beat memoryless baselines but fall short of state-augmented ceilings.

  3. Benchmarking Massively Parallelized Multi-Task Reinforcement Learning for Robotics Tasks

    cs.RO 2025-07 conditional novelty 6.0 of 10

    The paper introduces MTBench, a GPU-accelerated benchmark for massively parallel multi-task RL, and reports experiments suggesting on-policy methods outperform off-policy baselines while value learning limits MTRL per...

  4. Chargax: A JAX Accelerated EV Charging Simulator

    cs.LG 2025-07 conditional novelty 6.0 of 10

    Chargax is a JAX-based EV charging simulator that accelerates reinforcement learning training by 100x to 1000x compared to existing environments, with modular real-world scenarios.

  5. Ad-Hoc Human-AI Coordination Challenge

    cs.AI 2025-06 conditional novelty 6.0 of 10

    AH2AC2 provides an open, reproducible Hanabi benchmark for human-AI ad-hoc coordination, with withheld proxy agents and baselines showing a large gap to human-level play.

  6. Multi-Agent Reinforcement Learning for Inverse Design in Photonic Integrated Circuits

    cs.LG 2025-06 conditional novelty 6.0 of 10

    A multi-agent bandit RL formulation for binary photonic topology optimization produces designs that outperform gradient-based optimization on eight simulated 2D and 3D tasks.

  7. LEMUR: Learning to Align with Multi-Objective Reinforcement Learning from Preference Feedback

    cs.AI 2026-07 conditional novelty 5.0 of 10

    LEMUR jointly learns a separate reward model for each teacher's preferences and uses them to train a population of multi-objective policies, beating baselines that merge feedback into one reward.

  8. Preventing Learning Stagnation in PPO by Scaling to 1 Million Parallel Environments

    cs.LG 2026-03 conditional novelty 5.0 of 10

    PPO plateaus can be avoided by increasing the number of parallel environments, which reduces both the outer-loop step size and update noise; scaling to 1M environments sustained improvement to 1T transitions.

  9. Self-Supervised Goal-Reaching Results in Multi-Agent Cooperation and Exploration

    cs.LG 2025-09 conditional novelty 5.0 of 10

    Self-supervised multi-agent goal-reaching, where each agent independently learns a contrastive critic of its own observations, achieves cooperation and exploration in sparse-reward MARL tasks where standard baselines fail.

Pith tools