# SocialJax: An Evaluation Suite for Multi-agent Reinforcement Learning in Sequential Social Dilemmas

arXiv: https://arxiv.org/abs/2503.14576 · HTML: https://mlsanity.com/paper/2503.14576/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 30 references that are in the map.
- First posted on arXiv 18 Mar 2025.

## Builds on (most cited first, 21 of 21)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Is Independent Learning All You Need in the StarCraft Multi-Agent Challenge?](https://mlsanity.com/paper/2011.09533/) (2020)
- [BabyAI: A Platform to Study the Sample Efficiency of Grounded Language Learning](https://mlsanity.com/paper/1810.08272/) (2018)
- [Multi-agent Reinforcement Learning in Sequential Social Dilemmas](https://mlsanity.com/paper/1702.03037/) (2017)
- [OpenSpiel: A Framework for Reinforcement Learning in Games](https://mlsanity.com/paper/1908.09453/) (2019)
- [Behaviour Suite for Reinforcement Learning](https://mlsanity.com/paper/1908.03568/) (2019)
- [Inequity aversion improves cooperation in intertemporal social dilemmas](https://mlsanity.com/paper/1803.08884/) (2018)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Scalable Evaluation of Multi-Agent Reinforcement Learning with Melting Pot](https://mlsanity.com/paper/2107.06857/) (2021)
- [V-MPO: On-Policy Maximum a Posteriori Policy Optimization for Discrete and Continuous Control](https://mlsanity.com/paper/1909.12238/) (2019)
- [Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2306.09884/) (2023)
- [Learning to Incentivize Other Learning Agents](https://mlsanity.com/paper/2006.06051/) (2020)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning](https://mlsanity.com/paper/2303.17503/) (2023)
- [Social diversity and social preferences in mixed-motive reinforcement learning](https://mlsanity.com/paper/2002.02325/) (2020)
- [Maintaining cooperation in complex social dilemmas using deep reinforcement learning](https://mlsanity.com/paper/1707.01068/) (2017)
- [Melting Pot 2.0](https://mlsanity.com/paper/2211.13746/) (2022)
- [DeepMind Lab2D](https://mlsanity.com/paper/2011.07027/) (2020)

## Built on by (most cited first, 8 of 8)

- [MEAL: A Benchmark for Continual Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2506.14990/) (2025)
- [NomicLaw: Emergent Trust and Strategic Argumentation in LLMs During Collaborative Law-Making](https://mlsanity.com/paper/2508.05344/) (2025)
- [A Recipe for Stable Offline Multi-agent Reinforcement Learning](https://mlsanity.com/paper/2603.08399/) (2026)
- [Altruism and Fair Objective in Mixed-Motive Markov games](https://mlsanity.com/paper/2602.08389/) (2026)
- [JaxAHT: A JAX-Based Library for Ad Hoc Teamwork](https://mlsanity.com/paper/2609.13716/) (2026)
- [Multi-Agent Reinforcement Learning for Task Offloading in Wireless Edge Networks](https://mlsanity.com/paper/2509.01257/) (2025)
- [α-fair heterogeneous agent reinforcement learning](https://mlsanity.com/paper/2606.13076/) (2026)
- [Factoriax: A GPU-Accelerated Factorio-Style Simulator for Reinforcement Learning](https://mlsanity.com/paper/2610.05569/) (2026)

## Abstract (arXiv)

Sequential social dilemmas pose a significant challenge in the field of multi-agent reinforcement learning (MARL), requiring environments that accurately reflect the tension between individual and collective interests. Previous benchmarks and environments, such as Melting Pot, provide an evaluation protocol that measures generalization to new social partners in various test scenarios. However, running reinforcement learning algorithms in traditional environments requires substantial computational resources. In this paper, we introduce SocialJax, a suite of sequential social dilemma environments and algorithms implemented in JAX. JAX is a high-performance numerical computing library for Python that enables significant improvements in operational efficiency. Our experiments demonstrate that the SocialJax training pipeline achieves at least 50texttimes speed-up in real-time performance compared to Melting Pot RLlib baselines. Additionally, we validate the effectiveness of baseline algorithms within SocialJax environments. Finally, we use Schelling diagrams to verify the social dilemma properties of these environments, ensuring that they accurately capture the dynamics of social dilemmas.