REVIEW 16 cited by
JaxMARL: Multi-Agent RL Environments and Algorithms in JAX
Not yet reviewed by Pith; the record is open.
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
Signed reviews
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.
Forward citations
Cited by 16 Pith papers
-
Expected Return Symmetries
Expected return symmetries, transformations that preserve the expected return of optimal policies, contain environment symmetries as a subgroup and improve zero-shot coordination in Hanabi and Overcooked V2.
-
Recurrent Structural Policy Gradient for Partially Observable Mean Field Games
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.
-
Benchmarking Partial Observability in Reinforcement Learning with a Suite of Memory-Improvable Domains
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.
-
Benchmarking Massively Parallelized Multi-Task Reinforcement Learning for Robotics Tasks
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...
-
Chargax: A JAX Accelerated EV Charging Simulator
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.
-
Ad-Hoc Human-AI Coordination Challenge
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.
-
Multi-Agent Reinforcement Learning for Inverse Design in Photonic Integrated Circuits
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.
-
Hadamax Encoding: Elevating Performance in Model-Free Atari
Hadamax, a Hadamard-product and max-pooling encoder, improves PQN's median human-normalized Atari-57 score by about 80% with no algorithmic changes.
-
Scaling Multi Agent Reinforcement Learning for Underwater Acoustic Tracking via Autonomous Vehicles
A JAX-based GPU environment and a transformer-plus-curriculum MARL method train policies that transfer to the Gazebo LRAUV simulator and track up to 5 targets with around 5 m average error.
-
Cross-environment Cooperation Enables Zero-shot Multi-agent Coordination
Training a self-play agent across many procedurally generated cooperative tasks yields better zero-shot coordination with novel partners and novel layouts than training on one task with many partners.
-
LEMUR: Learning to Align with Multi-Objective Reinforcement Learning from Preference Feedback
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.
-
Preventing Learning Stagnation in PPO by Scaling to 1 Million Parallel Environments
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.
-
Self-Supervised Goal-Reaching Results in Multi-Agent Cooperation and Exploration
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.
-
Remembering the Markov Property in Cooperative MARL
Cooperative MARL agents solve many benchmarks with brittle conventions that bypass observations and memory, so the field needs tasks that genuinely require grounded, history-based reasoning.
-
Towards Autonomous Micromobility through Scalable Urban Simulation
URBAN-SIM generates diverse interactive city scenes at high training speed and URBAN-BENCH measures four robot types on eight micromobility tasks, with scale-up training lifting navigation success from 5% to 83%.
-
A Research Agenda for Usability and Generalisation in Reinforcement Learning
RL environments should be described in user-friendly domain-specific languages or natural language, so non-engineers can define tasks and agents can generalize to new tasks.
Discussion (0). Continue with ORCID to comment.