SocialJax: An Evaluation Suite for Multi-agent Reinforcement Learning in Sequential Social Dilemmas
Authors: Zihao Guo
Shuqing Shi
Richard Willis
Tristan Tomilin
Joel Z Leibo
Yali Du
Venue: ICLR 2026 Poster (2026)
Keywords: MARL
Sequential Social Dilemmas
Primary area: datasets and benchmarks
TL;DR: We introduce SocialJax, a high-performance JAX-based suite of sequential social dilemma environments and algorithms.

Abstract: 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 50\texttimes{} speed-up in real-time performance compared to Melting Pot’s 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.

Published as a conference paper at ICLR 2026
SOCIALJAX: AN EVALUATION SUITE FOR MULTI- AGENT REMNETTCEMENT LEARNING IN SEQUENTIAL SOCIAL DILEMMAS
Zihao Guo Shuqing Shi"* Richard Willis! Tristan Tomilin
Joel Z. Leibo Yali Du
1 KKin's College London, EEinhovven University of Technologg, "The Alan Turing Institute,
4 *Google DeepMind
(z Taao...... guo shuqing . shi, richard.willin. yali du) @kcl .ac. t. to 1 iiiituee.nnnn nl, jzl@google.com
ABSTRACT
Sequential social dilemmas pose 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 Melting Pot, provide 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 sequen- tial social dilemma environments and algorithms implemented in JAX. JAX is 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 50x speed-up in real-time performance compared to Melting Pot's RLlib baselines. Additionally, we validate the effec- tiveness 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. Our code is available https:::: / tuu.....ub....... com/ cooper r at vex Socci alJax
INTRODUCTION
Solving sequential social dilemmas remains pivotal challenge in multi-agent reinforcement learning (MARL). Efficient and rich environments and benchmarks are crucial for enabling rigorous evaluation and meaningful comparison of algorithms. Traditional MARL environments (Lowe et al. 2017] Samvelyan et al.) 2019 Bard et al. 2020 Carroll et al. 2019) have overwhelmingly focused on fully cooperative competitive tasks. Only handful of benchmarks, such as Prisoner's ' Dilemma of OpenSpiel (Lanctot et al. 2019). Stag Hunt and Chicken in PettingZoo (Terry et al.) [202I). address social dilemmas. However, these environments rely on CPU-based parallelism, which limits their simulation throughput and scalability, and they offer only narrow variety of scenarios without standardized evaluation protocols for sequential social dilemmas.
Although there are environments (Leibo et al. 2021 Agapiou et al. 2022) specifically designed for sequential social dilemmas, conducting experiments these environments highly challenging due to the substantial number of environment time steps required for training. Training in such environments demands extremely high hardware specifications (over 1000 CPUs for Melting Pot) or significant amount of computing time. Moreover, while the Melting Pot framework (Agapiou) ctt al. 2022. Leibo et al. (22T)) includes several baseline algorithms, such as A3C (Mnih et al. 2016), V-MPO (Song al. 2020), and OPRE (Vezhnevets et al., 2020), these implementations publicly available, which poses challenges for reproducibility and further development.
not
Recent libraries such JAX (Bradbury al. 2018) facilitate efficient execution of NumPy-style Python code on hardware accelerators(GGPs and TPUs) by automatically translating high-level
"Equal contribution.

---

Published as a conference paper at ICLR 2026
G Coins
(b) Gifl
(c) Communs: Hhrves Open
(6) Prisomer's Dilemma
Coop Mining
(1) Clean Up
(g) Commons: Harvest Closed
(b) Commons: Harvest Partnership
(1 Mushrooms
Figure The SocialJax suite contains nine multi-agent reinforcement learning environments designed to evaluate social dilemmas. The agents of all environments are restricted to partial observability of their surroundings through designated observation window.
functions into optimized, parallelizable operations. In reinforcement learning (RL) research, several approaches like PureJaxRL (Lu et al. 2022) have been proposed to leverage JAX (Bradbury et al. 2018) to run environments GPUs. Parallelizing environments accelerate computation. However, PureJAXRL is limited to single-agen -agett environments and does not support multi-agent scenarios. JaxMARL (Rutherford et al. (024) extends JAX's efficiency gains to MARL by providing a suite of JAX-based environments and algorithms. While JaxMARL incorporates limited set of social dilemma environments (e.g., Coins and STORM), it remains primarily geared toward traditional cooperative tasks and lacks dedicated benchmarks specifically
GPUs could significantly
designed for sequential social dilemmas.
In this paper, we develop a suite of environments for sequential social dilemmas based on JAX to improve computational efficiency. We also provide JAX-based implementations multiple MARL algorithms and construct complete training pipeline, significantly accelerating the training process. Solving sequential social dilemma problems often requires large number of time steps and diverse evaluation environments. SocialJax addresses these challenges by offering both computationally efficient and diverse environments. Furthermore, we validate the social dilemma properties of each environment using Schelling diagrams. In addition, environment, enabling a more precise quantification of the cooperative and competitive behaviors of work is the first evaluation framework leveraging JAX for
design metrics beyond return for each
agents. To the best of our knowledge, sequential social dilemmas. Our contributions are as follows:
Implementations of Social Dilemma Environments: We implement diverse suite of sequential social dilemma environments in JAX (see Figure D each exhibiting mixed-incentive dynamics, and achieve substantial improvements in simulation speed.
Implementation of Algorithms in JAX: We implement series of JAX --aaeeed MARL algorithms, including Independent PPO (IPPO) (De Witt et al. 2020 Rutherford
2024] under both individual
and common reward settings, Multi-Agent PPO (MAPPO) (Yu et al. 2022). Value Decomposition Networks (VDN) Sunnhhag et al. 2018). Social Value Orientation (SV (SVO) (ccKee et al. 2020),, and PPO using reward exchange (IPPO-RE) (Willis et al. 2025). We also provide training results with and without parameter sharing. We show that training algorithms in SocialJax is at least 50 times faster than Melting Pot 2.0 on machine with 14 CPUs and Nvidia A100 GPU.
We benchmark the performance of various environments and algorithms in terms of simulation throughput and training speed. Additionally, we evaluate the performance of the algorithms we include across various environments. Finally, we design specific evaluation metrics for each environment to measure agent behaviors either prosocial or antisocial.
Performance Benchmark and Analysis:
We further analyze the behavior of selfish and cooperative agents to understand the underlying incentive structures using Schelling diagrams (Pérolat et al.] 2017] Schelling 1973). Finally, analyze the performance of the implemented algorithmss in our environments, including cases where agent behavior varies with different SVO (McKee al. 2020 reward angles.
2

---

Published as a conference paper at ICLR 2026
2
RELATED WORK
In this section, review related work on sequential social dilemmas (Du al. 2023) and JAX-based reinforcement learning environments, algorithms, and benchmarks. Sequential Social Dilemma (Vinitsky al. 2019) introduces testbed for evaluating social dilemmas. However, its scope is restricted to only two environments on Clean Up and Harvest. OpenSpiel (Lanctot et al. 2019) offers collection of environments, some of which feature both cooperative and competitive dynamics, allowing for the exploration of social dilemmas. Melting Pot (Leibo et al.] (202T) was the first to establish a standardized benchmark for MARL specifically designed evaluate agent performance in sequential social dilemmas. By providing a diverse set of carefully designed multi-agent scenarios, it enables systematic testing of agent behaviors in cooperative, competitive, and mixed- motivee settings. Melting Pot 2.0 (Agapiou et al., (022) introduces more comprehensive and refined suite of testing environments that capture a wide range of interdependent relationships and incentive structures. The existing social dilemma environments are CPU-based, and their execution efficiency is constrained by CPU performance. Recent work also explores mechanism design approaches social dilemmas (Zhang al. 2025).
JAX (Bradbury et al. 2018) is an open-source library developed by Google, designed for high- performance numerical computing and machine learning tasks. Based on the GPU and TPU ac- celeration capabilities introduced by JAX, along with its support for parallelism and vectorization, JAX can effectively accelerate both environments and algorithms hardware accelerators (GPUs, TPUS). Functorch (He & Zou 2021) provides JAX-like composable function transforms for PyTorch, making it easier accelerate RL algorithms hardware accelerators. (Lu al. 2022) introduces parallelized framework leveraging JAX, enabling both environment and model training to run on GPUs. Jumanji (Bonnet al.) 2024) has developed a diverse range of single-agent environments in JAX, ranging from simple games to NP-hard combinatorial problems. Gymmax (Lange 2022) provides JAX versions of different environments, including classic control, Bsuite (Osband al.
2020), MinAtar (Young & Tian 2019), and collection of classic RL tasks. Pgx (Koyamada et al.
2023) implements a variety of board game simulation environments using JAX, such as Chess, Shogi, and Go. Brax (Freeman al.) [202I) is a physics engine written in JAX for RL, which re-implements MuJoCo. KLand-MiniGri (Nikulin et al. 2024) is a grid-world environment and benchmark fully implemented in JAX. Similarly, work is also demonstrated using grid-worlh environments. Mava (Mahjoub et al. 2024) and JaxMARL (Rutherford et al. 2024) both provide JAX-based MARL environments and baseline algorithms. Among the JAX environments and algorithms, only the Storm and Coins environments in JaxMARL are related to social dilemmas, while none of them specifically address sequential social dilemmas. We construct multiple environments focused social dilemma challenges.
SOCIALJAX
In this section, we define the concept of social dilemmas, introduce the environments we implement, describe the algorithms used, and present the metrics we design measure cooperative and com- petitive behaviors. Our work presents an efficient testing framework for sequential social dilemmas, complete with evaluation environments and benchmark baselines. Our SocialJax encompasses diverse suite of environments, which collectively capture a wide range of sequential social dilemma scenarios, including public good dilemmas and common pool resource problems. We evaluate the environments using IPPO individual reward (both with and without parameter sharing), IPPO common reward, MAPPO, SVO and IPPO-RE.
3.1 SEQUENTIAL SOCIAL DILEMMAS
In N N-player partially observable Markov games (Littman) 1994), agents possess partial observation of the environment. At each state, players choose actions from their respective action sets A² , A-V State transitions T S x A joint actions taken by all agents, where A ( represents the set of discrete probability distributions over S. We define the observation space of player agent has its corresponding individual reward defined as r" : S x A' x ... x A + R. In an N- N-player sequential social dilemma, cooperators and defectors each have their disjoint
A + A(S) are determined by the current state and the
O, {o'ls E S,
0 i)). Each
3

---

Published as a conference paper at ICLR 2026
policies(= c...... Te *...... Rd ..... TT Il's HhT with ++ m N. l the number of cooperators. (I) and Ra(I) are denoted as the average payoff for cooperating policies and defecting policies. When have sequential social dilemma, if it satisfies the following conditions (Leibo et al. 2017 Hughes et al. 2018:
1. Mutual cooperation is greater than mutual defection: (e > Ra(0).
2. Mutual cooperation is greater than being exploited by R. > (O).).
3. Either the fear property the greed property (or both) holds:
Fear: (i( (i) for sufficiently small
Greed: Ra(i) > Re(i) for sufficiently large i.
We use Schelling diagrams (Pérolat et al. 2017: Schelling 1973) depict the payoff's of different numbers of cooperators and defectors. The lines Re(1 + 1) and (2) depict the average reward for an agent choosing to either cooperate or defect, as a function of the number of co-players that are cooperating (see Figure3..
3.2 ENVIRONMENTS OF SOCIALJAX
Figure 1 shows the layout of our environments. SocialJax environments are mainly derived from Melting Pot 2.0 (Agapiou al. 2022). However, the layout of SocialJax environments is not the same as Melting Pot, but rather more similar to MiniGrid (Chevalier-Boisver et al. 2018). Agents observe grid matrix, with different objects represented by distinet numerical values. In all our environments, agents have the same 11 11 field of view implemented at the grid level. For further details, please refer to the Appendix [ Next, we detail the environment setup, including the specific environments in our evaluation suite.
Coins This environment was first introduced by (Lerer & Peysakhovich, 2017) to investigate how agents behave in situations where they can cooperate with or harm others by eating other agent's type of coin, making it a classic example for social dilemma.
Commons Harvest Agents need consume apples obtain rewards; however, the probability of apple regrowth is determined by the number of apples in the local neighborhood. This creates tension between selfish agents and cooperative agents. In order to discuss different social behaviors of agents, in Commons Harvest, we established three different scenarios: Commons Harvest: Closed, Commons Harvest: Open, and Commons Harvest: Partnership.
In Commons Harvest: Closed, there are two rooms in this setup PPérolat et al. 2017), each containing multiple patches of apples and a single entry. These rooms can be defended by specific players to prevent others from accessing the apples by zapping them with a beam, causing them to respawn.
In Commons Harvest: Open, all individuals can harvest apples, but must restrict their actions not collect the last apple and ensure that the apples can regrow.
The Commons Harvest: Partnership is similar Commons Harvest: Closed, specific players are required to defend the room. However, in this case, each room has two entry points, necessitating two players to cooperate in defending the room.
Clean Up This environment is public good game where players earn rewards by collecting apples. But the spawn rate of apples depends on the cleanliness of nearby river. For continuous apple growth, the agents must keep river pollution levels consistently low over time.
Coop Mining In this environment, two types ore spawn randomly in empty spaces. Iron ore (gray) can be mined individually and provides reward of +1 upon extraction, whereas gold ore requires group coordination and yields higher payoff per miner. Selfish agents tend to mine the iron ore more, while cooperative agents will try to cooperate and mine gold ore.
This environment includes multiple types of mushrooms that, when consumed by agents, deliver different rewards. Some mushrooms benefit only the individual agent, while others beneficial to the group
Mushrooms
Gift Refinement In this environment, agents can give collected tokens to other agents on the map. which upgrade to yield greater rewards. Therefore, maximizing returns requires agents to trust one another and cooperate effectively.
4

---

Published as a conference paper at ICLR 2026
Prisoner's Dilemma: Arena Agents inside this environment collect "defect" or "cooperate" tokens and then interact with each other to compare inventories. The resulting payoffs follow the classic Prisoner's Dilemma matrix, emphasizing the tension between individual incentives and collective welfare.
3.3 ALGORITHMM IN JAX
Due IPPO's widespread adoption and stability (Schulman et al. 2017 De Witt et al. (020). we implement IPPO under both parameter-sharing and non-parameeter-sharnng settings. The parameter-sharing variant of IPPO leads to faster training by using a shared policy neural net- work across all agents. The non-parameter-shrring variant, although slower, preserves agent-specific policies that prevent convergence to uniform conventions. The IPPO training curves shown in Figure 2 are based on the parameter sharing variant. Additional comparisons between parameter sharing and non-parameter sharing IPPO can be found in the Appendix C.T To encourage prosocial behavior, we introduce common rewards, while individual rewards are used to incentivize selfishness in agents. Additionally, we include MAPPO (Yu et al. 22022) as a representative centralized learning algorithm for MARL. Finally, implement SVO (McKee et al. 2020) and reward exchange (Willis al.. 2025) as representative methods specifically designed to handle social dilemmas, where agents can balance individual and collective interests.
IPPO Common Rewards refers to scenario where all agents in multi-agent system share a single, unified reward signal derived from the environment. This approach ensures that all agents aligned towards achieving the same objective, promoting collaboration and coordination among them. By sharing common reward, agents are incentivized to work together and make decisions that benefit the collective, rather than focusing solely on individual gains. This can help prevent conflicts and encourage cooperative behavior.
IPPO Individual Rewards indicates that each agent is assigned its own reward, inherently encourag- ing selfish behavior. This approach allows agents to prioritize maximizing their personal success over collective goals.
MAPPO We implement Multi-Agent PPO (MAPPO) Yu al. 2022 using both the PureJaxRL (Lu
al. 2022) and JaxMARL (Rutherford et al. 2024 frameworks. As centralized learning approach, MAPPO uses the centralized value function to help incentivize agents to perform prosocial behaviors.
VDN (Sunehag et al. 2018) decomposes the team's joint Q-value into a sum of individual agents' local Q-values, and achieves centralized training with decentralized execution by maximizing the aggregated value.
SVO We implement SVO (McKee et al. 2020) as an intrinsic motivation framework to encourage cooperative behavior in environments. To capture the relative reward distribution, reward angle 0 ((R) arctan f= is defined by SVO, where is the reward received by agent i, and F-i --1=- - Ty is the average reward of all other agents. Each agent is assigned target SVO angle ØSVO that reflects its desired distribution of reward. The agent's utility function combines extrinsic and intrinsic rewards as:
U2 o, a.) = Ti - w 10500 0(R)),
where controls the strength of the Social Value Orientation. We restrict SVO preferences to the non-negative quadrant, i.c., 0 [0",90% 90%], where 90" represents an Altruistic setting and = 0"
represents Individualistic setting.
IPPO-RE Reward exchange (Willis et al. 2025) a mechanism whereby agents exchange portions of their per-timess per-timestep rewards. Each agent retains proportion of its own reward (its self-interest. and distributes the remainder equally among co-players. Agent i's utility under reward exchange is:
1 1 Σ T3
U,(F,)) sr,
Under reward exchange, the ratio of an agent's 's utility gain when co-player receives reward to its gain when it receives the same reward itself ranges from 0 (individual reward, s = 1) to (common reward, )). Reward exchange aims find balance between these extremes, with agents
5

---

Published as a conference paper at ICLR 2026
exchanging sufficient reward to induce good collective outcomes while maintaining individual reward signals. We interpolate between the individual and common reward by training IPPO at reward ratios of 1. 5. and }. For example, in Coins (a two-player environment), use t. and
3.4 METRICS FOR SOCIALJAX ENVIRONMENTS
Since agent's behavior in the environments cannot be fully captured by the return alone, devised a tailored metric for each environment to quantify whether agents tend toward cooperation or competition. We define the semantics of the metrics for each environment follows:
Coins: Since each agent collecting their own-color coin does not harm the other, whereas collecting the other agent's coin imposes penalty on them, we evaluate cooperation by counting how many coins each agent collects that match their assigned color.
Commons Harvest: In this environment, agents must maintain sustainable number of apples ensure they continue to respawn. Therefore, we measure the number of apples remaining on the map assess whether agents are capable of preserving the long-tery collective interest.
Clean Up:
Agents need to clean polluted river tiles to enable apple regrowth. We evaluate their cooperative selfish behavior by counting the number cleaned water tiles.
Coop Mining: In this environment, agents can obtain higher rewards by cooperating to mine gold, while those who choose not cooperate can only mine for smaller reward. Therefore, evaluate agents' cooperative behavior by measuring the amount of gold mined.
The environment contains multiple types of mushrooms. Consuming blue mushrooms does not benefit the agent itself but provides rewards to the other agents. We thus measure the number of blue mushrooms consumed to evaluate the agent's willingness to sacrifice for the common good.
Mushrooms:
Gift Refinement: In this environment, agents choose to gift collected tokens to others, which results receiving a larger reward in return. thus count the number received tokens to evaluate whether agents willing trust each other.
Prisoner's Dilemma: Arena: The environment contains two types of resources, cooperative and competitive, and we introduce the quantity of cooperative resources collected by agents as a metric.
EXPERIMENTS
4.1
SETUP
To validate our library, we perform thorough tests of Common Reward IPPO, Individual Reward IPPO, MAPPO, and SVO algorithms across diverse range of Sequential Social Dilemma environments and analyze the returns and designed cooperation metrics. We assess the performance of our JAX-based environment GPU and compare its speed with the implementations provided by Melting Pot 2.0 (Agapiou et al. 2022). addition to these experiments, we plot Schelling diagrams to further verify that our environments effectively represent social dilemmas.
Table 1: Results on environment steps per second across various SocialJax environments under random actions. Different environment configurations are tested at scales from to 4096 JAX Env.
ENYS
ORIGINAL ENV
JAX ENV
128 JAX ENV
1024 JAX ENV
4096 JAX ENV
COINS
1.2 x 10"
2.0 10"
2.6 10
1.4 x
3.4 x
HARVEST: OPEN
3.7 x 10a
12x 10
1.1 10a
5.0 x
7.9 x
CLOSED
3.7 10a
10
1.1 10a
5.0 10
7.9 10
HARVEST: PARTNERSHIP
3.7 10a
10x
11x
5.0 10a
7.9 10
CLEAN UP
2.7 10a
L1X
10 10
4.3 10a
6.1 X 10
COOP MINING
3,6 10a
19x
19x 10a
1.0 10f
1.5 x 10
MUSHROOMS
42 10a
14x
14x
6.3 10a
9.8 x 10a
GIFT REFINEMENT
4.1 10a
15x 10
1.x
7.0
1.2 x
PRISONER'S DILEMMA: ARENA
4.5 10a
102
2.3
2.7 x 10

---

Published as a conference paper at ICLR 2026
120
42D
420
Return
320
120
-20
320
20
120
0.5
1/3
22
0.5
1.2
Envifomment Simslaps
Enviromment Smoslaps
Envifomment fmoslaps
(a) Coins
(b) Commons Harvest: Open
(c) Commons Harvest: Closed
600
Relures
220
00
5.5
1/5
15
2.5
0.0
0.5
14
1.5
2-0
00
55
1.5
25
Envitomment Smestops
Envrrommnt timestops
Smoslops
160
(d) Commons Harvest: Partnership
(c) Clean Up
(f) Coop Mining
220
Retume
220
1/2 1.5
15 Enviromment fimesopss
Enviroment fmmetops
140
Environment timsteps
(g) Mushrooms
(h) Gift Refinement
(i) Prisoners Dilemma: Arena
IPPO: Individual Reward
IPPO: Reward
IPPO RE
Figure 2: Training curves for range of SocialJax environments. IPPO (shared parameters) with Common Reward encourages collective interests, leading to higher overall returns, Individual Reward primarily drives selfish behavior, often resulting in lower returns. MAPPO included a centralized baseline. The SVO curve represents the case with a social angle 0 = 90%.
4.2 RESULTS
Speed Benchmarks Since both the environments and algorithms are vectorized using JAX, they are well suited for GPU acceleration. In Table [ we test up to 4096 environments in parallel and observe a significant increase in the speed of the training pipeline with the addition of more parallel environments. We compare the performance of SocialJax and Melting Pot 2.0 (Agapiou et al. 2022) by measuring steps per second on the same hardware, which consists of an NVIDIA A100 GPU and 14 CPU cores. As shown in Table []] we evaluate the speed of SocialJax across various environments and environment counts using random actions. We observe a significant speedup when running environments in parallel parallel performance of SocialJax, the original environments would need to scale up across hundreds of CPUs. Regarding wall clock time, SocialJax is capable of completing 1e9 timesteps within hours for relatively simple environments such as Coins, whereas training with Melting Pot 2.0 using Saable-Baseliness, and RLIib takes around 1, 1,300 hours and 1.50 hours, respectively. Consequently, conducting social dilemma experiments with Melting Pot 2.0 is highly challenging without large amount of CPU resources. Training Coins in SocialJax is about 50-400 times faster compared to Coins in Melting Pot 2.0. For relatively complex environments such as Clean Up. SocialJax is approximately 50-140 times faster than the original environments. There are more details about the wall clock time comparison in the Appendix
GPU compared to running single environment. To match the
IPPO Common reward and IPPO Individual reward We compare the performance of common reward IPPO (IPPO-CR) and individual reward IPPO (IPPO-IR) across different environments. In the default individual reward setting, each agent receives a reward based on the joint-action, independent of the payoffs of other agents. This creates incentives for agents to over exploit common pool resources, because only they benefit from their actions, while the social costs are shared by all. In contrast, common reward, or team reward means that all agents share single reward signal, typically the total payoff of all agents in the environment. The common reward encourages cooperation, the objective of each agent is maximize the total reward, making individual agents more inclined to adopt collaborative behaviors. In Figure 2 common reward typically performs better than individual rewards, the agents attempt maximize the total payoff.
7

---

Published as a conference paper at ICLR 2026
Table 2: Returns Across Different Environments as a Function of 0 under the SVO Algorithm.
ENVIRONMENT
0 (INDIVIDUALISTIC)
22.5"
45"
67.5
90" (ALTRUISTIC)
COINS
11.81
60.38
160.43
162.25
162.46
HARVEST: OPEN
79.71
68.96
68.28
238.40
254.54
HARVEST: CLOSED
81.05
74.74
74.68
298.97
309.94
HARVEST: PARTNERSHIP
77.44
72.45
71.58
225.94
234.30
CLEAN
0.02
0.06
50.58
1060.25
1410.53
COOP MINING
210.26
207.96
415.99
630.32
647.61
MUSHROOMS
5.94
51.26
291.55
321.22
400.85
GIFT REFINEMENT
104.29
105.09
105.22
107.38
227.95
PRISONER'S 's DILEMMA: ARENA
22.67
23.27
24.13
34.96
53.36
MAPPO As centralized learning method, MAPPO has access to information from all agents, which leads to strong performance in Clean Up (Figure 2c). and Mushrooms (Figure However, MAPPO exhibits less stable training dynamics, its centralized critic aggregates observations from all agents, increasing the learning difficulty of the network. In particular, it struggles with Commons Harvest: Partnership (Figure 2d) and Gift Refinement (Figure 2h) where it under performs IPPO-CR.
VDN outperforms the IPPO: Individual Reward baseline in several environments, including Coins, Commons Harvest variants, Mushrooms, and Prisoner's Dilemma: Arena (see Figures 2a 2d but its linear decomposition of the joint Q-value restricts its ability to model complex inter-agent interactions, leading suboptimal performance harder social dilemmas.
SVO provides powerful framework for exploring agents' cooperative and competitive tendencies. We apply SVO across all SocialJax environments investigate how the preference angle 0 influences agents' returns. We first fix A = 90% and sweep over different weight values w to identify the optimal w. Then, holding the weight fixed, we sweep over examine how agent behavior and returns vary with the angle. As summarized in Table 2 the collective return increases with larger values of 0, as agents become more inclined to cooperate at higher angles. When 0 (Individualistic orientation), agents achieve the lowest returns, whereas at 0 90 (Altruistic orientation), they attain the highest returns.
IPPO-RE results are shown with the best-performing reward ratio in Figure 2 Reward exchange substantially improves over IPPO-CR Clean Up (Figure 2e), with modest improvement in Prisoner's s Dilemma (Figure 2ip, suggesting that performance in these environments benefits from retaining some individual incentives. However, in Gift Refinement (Figure[hi reward exchange performs substantially worse than IPPO-CR, suggesting that individual incentives to free-ride particularly strong in this environment.
4.3 SOCIAL DILEMMA ENVIRONMENTS ATTRIBUTES
Schelling Diagram Analysis To validate the social dilemma properties SocialJax, we depict Schelling diagrams for each environment. The cooperative policies sampled from the agents that used a common reward, while the defector policies originate from agents trained with independent rewards. We evaluate the environments over 30 episodes and compute the average rewards for the cooperative and defecting agents.
Coins (Figure Bab. Harvest: Open (Figure 3b). Harvest: Closed (Figure 3c), Harvest: Partnership (Figure 3d Gift Refinement (Figure 3h). Mushrooms (Figure 3g] and exhibit fear: namely that agents would prefer to defect some of their co-players are defecting. An agent actually does better cooperating if all their co-players are cooperating. In Harvest: Open, we observe that it only takes a single defector agent have severely negative impact on the returns of other players. This is because defector agents typically harvest the last apple in patch, permanently depleting it. The cooperator policies have learned to leave a greater amount of apples in each patch, which causes higher regrowth rate, so they generally exhibit a greater tendency to abstain from harvesting in the presence of a defector. In Harvest: Closed, we find that more than one defector is necessary to inflict severe negative impact the other agents. This occurs because Harvest: Closed partitions resources into two distinct rooms. After lone defector depletes all apples in one room, re- re-entering the other requires visually locating its entrance and traversing through it.
8

---

Published as a conference paper at ICLR 2026
Individual payet
cooperator defeccor
cooperator
Individdin Bayoft
cooperator defeccor
defeector everage
everage
everage
cooperators
cooperators
cooperators
(a) Coins
(b) Commons Harvest: Open
(c) Commons Harvest: Closed
Individuss payott
Indiiding payert
30.
233
defectar
average
231
average
average
12
number other cooperators
cooperators
number other cooperators
(d) Commons Harvest: Partnership
(c) Clean Up
(f) Coop Mining
individual payaf
Soniino mam soyort 2
cooperator
selector
avorage
avzrag2
Svoogg
number cooperators
cooperstors
cooperators
(g) Mushrooms
(h) Gift Refinement
Prisoner Dilemma: Arena
Figure 3: Schelling diagrams of the environments, visualizing the relationship between individual incentives and collective outcomes. The dotted line represents the overall average return of cooperators and defectors.
In Clean Up (Figure 3el. instead observe greed, there is incentive defect if sufficient number of co-players are cooperating. This is because cooperators will clean the pollution, whereas the defector will not, but it requires a minimum of three cooperators clean the river for apples to grow. Coop Mining (Figure 31) and Prisoner's Dilemma: Arena (Figure 31) also primarily exhibit greed. In these environments, if sufficient number of agents cooperate, other agents can obtain higher payoff by defecting. In Coop Mining, this is because cooperators tend to neglect the iron ore, allowing the defector collect more iron ore. In Prisoner's Dilemma: Arena, when most agents collect cooperative resources, defectors can gain higher returns through interactions with cooperators.
SVO Metrics Analysis The proposed metrics quantify changes in agent behavior or environment states triggered by cooperative actions. A larger metric (Table 3) value reflects higher level of cooperative behavior exhibited by the agents. In general, a higher degree of cooperative behavior among agents leads to larger collective returns. However, there exceptions. In the three Harvest environments, the SVO algorithm encourages altruistic actions, which can reduce the agents' effi- ciency in collecting apples. Specifically, agents may intentionally avoid picking apples to allow others access to them. As result, more apples remain on the map, leading to a decrease in harvest efficiency and ultimately reduction in collective returns. In more extreme cases, agents may collect any apples at all, causing the Harvest metric to remain at high value. Under typical conditions where agents receive rewards based on apple collection, this metric performs well. As expected, IPPO-IR and SVO-0", which encourage selfish behavior, produce significantly lower metric values than more cooperative approaches such as IPPO-CR, MAPPO and SVO-90".
Reward Exchange Analysis While common reward fully aligns agent incentives, it weakens the link between an agent's 's actions and its utility, making difficult for agents to identify which of their actions contributed to reward (the credit assignment problem) and allowing agents to benefit from others' efforts without contributing themselves (the lazy agent problem). For each environment, we identify the minimum reward exchange needed resolve the conffict between individual and collective incentives by selecting the smallest ratio achieving at least 95% of the best- best-performing ratio's 's return, to account for stochastic variation across seeds. Table 4 reports the results, with the selected ratio bold.
For all environments, requires the least exchange (ratio 0.25). Harvest: Closed and Harvest: Partnership achieve good returns at ratio 0.75, but Harvest: Open requires common reward, restraining harvesting harder
reward exchange is necessary to achieve good outcomes. Coop Mining

---

Published as a conference paper at ICLR 2026
Table 3: Metries for Assessing Cooperative and Competitive Tendencies Across Environments.
IPPO-RE is implemented using the best- reward exchange ratio.
ENVS
IPPO-IR
IPPO-CR
MAPPO
VDN
IPPO-RE
SVO-90"
COINS
100.83
164.07
126.23
15.06
164.06
164.56
HARVEST: OPEN
0.91
23.68
18.45
21.41
28.57
37.67
HARVEST: CLOSED
1.71
17.57
18.57
22.47
24.22
22.70
HARVEST: PARTNERSHIP
1.42
17.75
15.56
27.34
24.73
25.12
CLEAN UP
9.28
109.54
133.99
13.71
106.71
99.08
COOP MINING
298.21
521.56
571.25
285.16
204.19
483.24
MUSHROOMS
4.07
19.55
22.96
10.02
15.57
12.33
GIFT REFINEMENT
0.03
28.37
20.24
0.01
12.20
3.83
PRISONER'S S DILEMMA
58.65
73.67
119.63
78.45
50.13
76.24
Table IPPO-RE Mean return (# std) across environments under different reward exchange ratios.
In bold, we highlight the smallest ratio that achieves at least 95% of the returns of the best ratio.
ENVIRONMENT
INDIVIDUAL
RATIO 0.25
RATIO 0.5
RATIO 0.75
COMMON REWARD
COINS
39
93 + 1
162 + 0
163 :
164 1 0
HARVEST: OPEN
67 1
70
283 42
393 + 21
420 = 35
HARVEST: CLOSED
+1+1
191
311 +
323 18
338 + 29
HARVEST: PARTNERSHIP
6
136 +
292
328
+ 17
CLEAN UP
0+:
1780
1310
COOP MINING
446
727
741
729 17
726
MUSHROOMS
15
342
415
410
398
GIFT REFINEMENT
105 10
184
65
505
PRISONER'S DILEMMA
144 40
233 48
260
226 12
+ 17
without spatial partitioning. Gift Refinement also requires common reward and shows the largest performance drop from any reduction in reward sharing, suggesting that individual incentives to free- ride are particularly strong. Clean Up achieves its best returns at ratio 0.5 0.5-notably outperforming the common reward-consistee credit assignment difficulties. The reward ratio which cooperation emerges is related to the self- interest level of the game (Willis et al. 2024); we leave formal estimation in these environments to future work.
with the hypothesis that retaining some individual incentive alleviates
5 CONCLUSION
The of hardware acceleration has revolutionized MARL research, enabling researchers over- come computational limitations and accelerate the iteration of ideas. SocialJax is a library designed leverage these advancements in the context of social dilemmas. It provides collection of social dilemma environments and baseline algorithms implemented in JAX, offering user-frient user-friendly interfaces combined with the efficiency of hardware acceleration. By achieving remarkable speed-ups over traditional CPU-hased implementations, SocialJax significantly reduces the experiment runtime. In addition, unified codebase consolidates diverse range of social dilemma environments into standardized framework, facilitating consistent and efficient experimentation and greatly benefiting future research on social dilemmas.
ACKNOWLEDGEMENT
This work was supported by the Engincering and Physical Sciences Research Council [grant number EP/Y003187/1 and UKRI849).
REFERENCES
John P Agapiou, Alexander Sasha Vezhnevets, Edgar A Duéñez-Guzmán, Jayd Matyas, Yiran Mao, Peter Sunchag, Raphael Köster, Udari Madhushani, Kavya Kopparapu, Ramona COmanescu, et al. Melting pot 2.0. arXiv preprint 2xXi::2211113744 2022.
10

---

Published as a conference paper at ICLR 2026
Nolan Bard, Jakob N. Foerster, Sarath Chandar, Neil Burch, Mare Lanctot, H. Francis Song, Emilio Parisotto, Vincent Dumoulin, Subhodeep Moitra, Edward Hughes, lain Dunning, Shibl Mourad, Hugo Larochelle, Mare G. Bellemare, and Michael Bowling. The hanabi challenge: A new frontier for Al research. Artif. Intell., 280: 103216, 2020. doi: ARTINNTTT 103216.. URL httts: .rgg/ 16. 1016/ art 2019 000.9..0332010 103216
Clément Bonnet, Daniel Luo, Donal Byrne Shikha Surana, Sasha Abramowitz, Paul Duckworth, Vincent Coyette, Laurence Illing Midgley, Elshadai Tegegn, Tristan Kalloniatis, Omayma Mahjoub, Matthew Macfarlane, Andries Smit, Nathan Grinsztajn, Raphaël Boige, Cemlyn N. Waters, Mohamed A. Mimouni, Ulrich A. Mbou Sob, Ruan de Kock, Siddarth Singh, Daniel Furelos- Blanco, Victor Le, Arnu Pretorius, and Alexandre Laterre. Jumanji: a diverse suite of scalable reinforcement learning environments in JAX. In The Twelfth International Conference Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL net forum 1===c 9wC
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http:/ /ttp:////hthbb githu......gggggggg com/ jax-
Micah Carroll, Rohin Shah, Mark K. Ho, Tom Griffiths, Sanjit A. Seshia, Pieter Abbeel, and Anca D. Dragan. On the utility learning about humans for human-ai coordination. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alcch--Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8- 14, 2019, Vancouver, BC, Canada, pp. 5175- 2019. URL /ttp::t:pptcceeeoiggggooo/o///oo neurips. catpssccppppae cc/paper/ 201 9tw.9/hahhh... f 5lb 1b8 9d 98w.a8988888888 728662286655 28a wtt.a5f111112bbbbb Abst
Maxime Chevalier-Boisvert,, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: A platform to study the sample efficiency of grounded language learning. In International Conference on Learning Representations, 2018.
Christian Schroeder De Witt, Tarun Gupta, Denys Makoviichuk, Viktor Makoviychuk, Philip HS Torr, Mingfei Sun, and Shimon Whiteson. Is independent learning all you need in the starcraft multi-agent challenge? arXiv preprint arXiv:201 1.005333, 2020.
Yali Du, Joel Z. Leibo, Usman Islam, Richard Willis, and Peter Sunehag. A review of cooperation in multi-agent learning. CoRR, abs/2312.051622 2023. doi: 10.4850/ARXXX.2312.06666 URL
https: httts://doi..rro/
www.oorg/......o.oooooooo 48550/ www.rrww..........
Lasse Espeholt, Hubert Soyer. Rémi Munos, Karen Simonyan, Volodymyr Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, Shane Legg, and Koray Kavukeuoglu. IMPALA:
scalable distributed deep-rl with importance weighted actor-learner architectures. In Jennifer G. Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüsssnn, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Muchine Learning Research, pp. 1406-1415. PMLR, 2018. URL http
7 /proceedings mlr press v80 //0///000//eeeeee httpeeppphhhataaamm html
C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax A differentiable physics engine for large scale rigid body simulation. In Joaquin Vanschoren and Sai-Kit Yeung (eds.), Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks NeurIPS Datasets and Benchmarks 2021. December 2021, virtual, 2021. URL
Horace He and Richard Zou. functorch: JAX-like composable function transforms for PyTorch.
hattps://gitie
http://ytoorcc/ffnctoooooo
2021.
Edward Hughes, Joel Z. Leibo, Matthew Phillips, Karl Tuyls, Edgar A. Duéñez-Guzmán Anto- nio Gareía Castañeda, lain Dunning, Tina Zhu, Kevin R. McKee, Raphael Koster, Heather Roff, and Thore Graepel. Inequity aversion improves cooperation in intertemporal social dilemmas.
11

---

Published as a conference paper at ICLR 2026
In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Biande Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Con- ference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 1330--344, 2018. URL
neurips. cc/ .t.m]
Marco A. Janssen, Robert Holahan, Allen Lee, and Elimor Ostrom. Lab experiments for the study of social-ecological systems. Science, 328(5978))613 201, doi: 10.1126/science. 183532. 126/ sci ence. 1183532)
doi
Sotetsu Koyamada, Shinri Okano, Soichiro Nishimori, Yuta Murata, Keigo Habara, Haruka Kita, and parallel game simulators for reinforcement learning. In Neural Information Processing Systems, 2023.
Shin Ishii. Pgx:
Marc Lanctot, Edward Lockhart, lean-Baptiste Lespiau, Vinícius Flores Zambaldi, Satyaki Upadhyay, Julien Pérolat, Sriram Srinivasan, Finbarr Timbers, Karl Tuyls, Shayegan Omidshafici, Daniel Hennes, Dustin Morrill, Paul Muller, Timo Ewalds, Ryan Faulkner, János Kramár, Bart De Vylder, Brennan Saeta, James Bradbury, David Ding, Sebastian Borgeaud, Matthew Lai, Julian Schrittwieser, Thomas W. Anthony, Edward Hughes, Ivo Danihelka, and Jonah Ryan-Davit Openspiel: A framework for reinforcement learning in games. CoRR, abs/1908.09453, 2019. URL http:///axix>:
abs 908 09453
Robert Tjarko Lange. gymmax: A JAX-based reinforcemmett reifoosseeennnnoomrrmmmaal learning environment library, 2022. URL http:// com wtt.:////ooeeeetaaeggggegg/////
Joel 7. Leibo, Vinícius Flores Zambaldi, Mare Lanetot, Janusz Marecki, and Thore Graepel. Multi- agent reinforcement learning in sequential social dilemmas. In Kate Larson, Michael Winikoff, Sanmay Das, and Edmund H. Durfee (eds.), Proceedings of the 16th Conference on Autonomous Agents and MultiAgent Systems, AAMAS 2017, São Paulo, Brazil, May 8- 12, 2017, pp. 464 464-473. wwtp/citt/iin..o/////////////oot (tt.taaiioonccmmmmd311111111
ACM, 2017. URL http:
Joel Z. Leibo, Edgar A. Duéfiez-Gzzmán, Alexander Sasha Vezhnevets, John P. Agapiou, Peter Sunehag, R. Koster, Jayd Matyas, Charlic Beattie, Igor Mordatch, and Thore Graepel. Scalable evaluation of multi-agent reinforcement learning with melting pot. In International Conference on Machine Learning, 2021.
Adam Lerer and Alexander Peysakhovich. Maintaining cooperation in complex social dilemmas using deep reinforcement learning. ArXiv, abs/1707.00166, 2017.
Michael L. Littman. Markov games as framework for multi-agent reinforcement learning. In International Conference Machine Learning, 1994.
Ryan Lowe, Yi Wu, Aviv Tamar, Jean Harb, Pieter Abbeel, and Igor Mordatch. Multi-agent actor-critic for mixed cooperative-compettive environments. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 6379-6390, 2017. URL httpps://prrcceeiinnggs///// .nw.seeeeeddinnnn.eeeeeeee. neurips http//paper/211/aapaeep html
Chris Lu, Jakub Grudzien Kuba, Alistair Letcher, Luke Metz, Christian Schröder de Witt, and Jakob N. Foerster. Discovered policy optimisation. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL hattp://paperee nips. http::/ss......aaaaaaaaaaaaaa/ap 188c78882331653572255666666 htt.://aeeeeeeatttaaatoororrrree
Omayma Mahjoub, Sasha Abramowitz, Ruan de Kock, Wiem Khlifi, Simon du Toit, Jemma Daniel, Louay Ben Nessir, Louise Beyers, Claude Formanek, Liam Clark, and Arnu Pretorius. Performant, memory efficient and scalable multi-agent reinforcement learning. CoRR, abs/2410.01706,, 2024. doi: 0044550/ARXXV.241000777 URL lattps://dat 10//10.000 48550/4885////0 212.550//rXi..211.1.
706
12

---

Published as a conference paper at ICLR 2026
Kevin R. McKee, Ian M. Gemp, Brian McWilliams, Edgar A. Dučñez-Guzmán, Edward Hughes, and Joel Z. Leibo. Social diversity and social preferences in mixed- motive reinforcement learning. In Adaptive Agents and Multi-Agent Systems, 2020.
Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Maria-Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 1928-1937, JMLR.org, 2016. httt::t::pppceeeeeee
mlr www.s/eessse...... 1w1448/mnihall htmml
Alexander Nikulin, Vladislav Kurenkov, llya Zisman, Artem Agarkov, Viacheslav Sinii, and Sergey Kolesnikov. Xland-minigrid: Scalable meta reifforeenn learning environments in JAX. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang (eds.), Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
nips. CC cc/paperpppppppeeeeeeeeeeeee... files/paper -ww-2277772A--attaaacca Datasett.
2024 http::/eeaaaaeee1113222eaaaaaaaa
and Benchhancchheeeeee Track ht.ml
Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney. Tor Lattimore, Csaba Szepesvári, Satinder Singh, Benjamin Van Roy, Richard S. Sutton, David Silver, and Hado van Hasselt. Behaviour suite for reinforcement learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia. April 26-30, 2020. OpenReview.net, 2020. TRLLittps:///peeeeeeee rygî - kSYwHl
net foormmidd
Julien Pérolat, Joel Z. Leibo, Vinícius Flores Zambaldi, Charlie Beattie, Karl Tuyls, and Thore Graepel. A multi-agent reinforcement learning model of common-pool resource appropriation. In Neural Information Processing Systems, 2017.
Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook, Andrei Lupu, Gardar Ingvars- son, Timon Willi, Ravi Hammond, Akbir Khan, Christian Schroeder de Witt, Alexandra Souly, Saptarashmi Bandyopadhyay. Mikayel Samvelyan, Minqi Jiang, Robert Tjarko Lange, Shimon Whiteson, Bruno Lacerda, Nick Hawes, Tim Rocktäschel, Chris Lu, and Jakob Nicolaus Foerster. Jaxmarl: Multi-agent rl environments and algorithms in jax. In The Thirty-iighth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024.
Mikayel Samvelyan, Tabish Rashid, C. S. D. Witt, Gregory Farquhar, Nantas Nardelli, Tim G. J. Rudner, Chia-Man Hung, Philip H. S. Torr, Jakob N. Foerster, and Shimon Whiteson. The starcraft multi-agent challenge. In Adaptive Agents and Multi-Agent Systems, 2019.
Thomas C Schelling. Hockey helmets, concealed weapons, and daylight saving: A study of binary choices with externalities. Journal of Conflict resolution, 17(3):381-428, 1973.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1rXiv:1707.063:3 2017.
H. Francis Song, Abbas Abdolmaleki, Jost Tobias Springenberg. Aidan Clark, Hubert Soyer, Jack W. Rae, Seb Noury, Arun Ahuja, Siqi Liu, Dhruva Tirumala, Nicolas Heess, Dan Belov, Martin A. Riedmiller, and Matthew M. Botvinick. V-MPO: on-policy maximum posteriori policy op- timization for discrete and continuous control. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL htllhttts///ppeeeeeett
net/ formm//oormmmii www.ooormmm..11100010000
Peter Sunehag, Guy Lever, Audrunas Gruslys, Wojciech Marian Czarnecki, Vinícius Flores Zambaldi, Max Jaderberg, Mare Lanctot, Nicolas Somnerat, Joel Z. Leibo, Karl Tuyls, and Thore Graepel. Value-decomposiiton networks for cooperative multi-agent learning based on team reward. In Elisabeth André, Sven Koenig, Mehdi Dastani, and Gita Sukthankar (eds.). Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems, AAMAS 2018, Stockholm, Sweden, July 10-15, 2018. pp. 0085-2087. International Foundation for Autonomous Agents and Multiagent Systems Richland, SC, USA/ / ACM, 2018 URL http: /1. acm. org/ ci tat on. cfm d=2388080
13

---

Published as a conference paper at ICLR 2026
J. K. Terry, Benjamin Black, Nathaniel Grammel, Mario Jayakumar, Ananth Hari, Ryan Sullivan, Luis S. Santos, Clemens Dieffendahl, Caroline Horsch, Rodrigo Perez- Vicente, Niall L. Williams, Yashas Lokesh, and Praveen Ravi. Pettingzoo: Gym for multi-agent reinforcement learning. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NewrIPS 2021, December 6-14, 2021, virtual,
pp. 15032-15043, 2021. URL / www.ppooew....eo.e.uuuuuu.o..... www.reeuuip... wttp:/s.cccppppppreeeee 80 http:///02222..........00..0.000
html
Rakshit Trivedi, Akbir Khan, Jesse Clifton, Lewis Hammond, Edgar A Duéñez-CGuzmán, Dipam Chakraborty, John P Agapiou, Jayd Matyas, Sasha Vezhnevets, Barna Pásztor, et al. Melting pot contest: Charting the future of generalized cooperative intelligence. In The Thirty-eighth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024.
Alexander Sasha Vezhnevets, Yuhuai Wu, Maria K. Eckstein, Rémi Leblond, and Joel Z. Leibo. Options as responses: Grounding behavioural hierarchies in multi-agent reinforcement learning. In International Conference on Machine Learning, 2020.
Eugene Vinitsky, Natasha Jaques, Joel Leibo, Antonio Castenada, and Edward Hughes. An open source implementation of sequential social dilemma games. / com/ euugenvvinttan www.ssk../eesettteeeeeee laspsttiiilliisllalaaaait _di 1illemma wtw.ge//ssseeeeeee/... GitHub repository.
2019.
Richard Willis, Yali Du, Joel Z. Leibo, and Michael I Luck. Resolving social dilemmas with minimal re- ward transfer. Auton. Agents Multi Agent Syst. 38(2):49, 2024. doi: 00000//1445-004-096666 URL httpp:://ol.....
1007 020000027200000000000000-0---000
Richard Willis, Yali Du, Joel Z Leibo, and Michael Luck. Quantifying the Self-Interest Level of Markov Social Dilemmas. In Proceedings of the 34th International Joint Conference on Artificial Intelligence, Montreal, Canada, August 2025. ijeai.org.
Kenny Young and Tian Tian. Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments. arXiv preprint arXiv:1903.03119 2019.
Chao Yu, Akash Velu, Eugene Vinitsky, Yu Wang, Alexandre M. Bayen, and Yi Wu. The surprising effectiveness of ppo in cooperative multi-agent games. In Neural Information Processing Systems, 2022.
Xiaoyuan Zhang, Yizhe Huang, Chengdong Ma, Zhixun Chen, Long Ma, Yali Du, Song-Chum Zhu, Yaodong Yang, and Xue Feng. Social world model-augmented mechanism design policy learning. CoRR, abs/2510. 19270, 2025. doi: 1044850/ARXX..255001122 URL https 7/doi //oo0..org//00.8555/a.x........0 20w.:885000ax1222111.111990
14

---

Published as a conference paper at ICLR 2026
A STATEMENTS OF ETHICS, REPRODUCIBILITY, AND LLM USAGE
This work complies with the ICLR Code of Ethics. No human animal subjects were involved, and all datasets, including SocialJax environments and training outputs, were used according to relevant guidelines with privacy violations. We took transparency, fairness, and integrity throughout. To support reproducibility, all code, datasets, and detailed experimental settings, including training steps, model configurations and hyperparameters, are publicly available in SocialJax repository. Large Language Models (LLMs) were used solely for linguistic refinement, such improving clarity, grammar, and coherence; they did not contribute to research design, analysis, or conceptualization. The authors take full responsibility for all scientific content.
mitigate potential biases and ensured
B ADDITIONAL DETAILS ON ENVIRONMENTS
B.1 API AND EXAMPLES
Our environment interfaces (Figure are inspired by those of mature MARL frameworks, adopting the similar API conventions as Jax.ARR which itself is inspired by PettingZoo and Gymnax, to provide intuitive and user-friendly experience for researchers.
import jax
import socialjax
from socialjax import make
nun_ageats
env make ( d'cleannup', num Lagents-nun_agents.
rng = roxdorndom.o random,,RNNKee PRNGKey (259)
rng, rrng jax. random. spl (rng)
for range (100):
ing, +rngs = jax. random, sanndom.split: (rng, num agents+1)
actions = [jax. waandom.chocce..
rngs [a],
www.acction_spaaa (0) n, arrjy.aarra((C... 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.11)
for in range (rum_agenta)
obs, state, reward, done, info env. nav.step_eny rng, state for in actions]
Figure An example of using the SocialJax API in the Clean Up environment.
Because we provide an open-source codebase, anyone can reproduce our main experimental results by configuring the environment following the repository's instructions and running the code. Figure 5 provides command line examples for executing various algorithms.
Execution speed of SocialJax environments under random actions can also be evaluated by running the code shown in Figure 6 Additionally, we provide a Colab notebook in our repository for quick experimentation and environment speed testing.
B.2 ENVIRONMENT SETTINGS AND PARAMETERS
The data format observed by agents in SocialJax environments differ from those in Melting Pot 2.0 (Agapiou et al. 2022). Melting Pot 2.0 provides observations to agents in the form of image pixels. Melting Pot 2.0 was originally designed to use 8 8 pixels per cell, resulting observations of 88 x 88 pixels in most environments. this setup, the agents have a partially observable window of 11 11 cells. In the Melting Pot Contest 2023 (Trivedi et al. 2024), all academic participants down-sampled
15

---

Published as a conference paper at ICLR 2026
Train IPPO on the Clean Up environment python algortttmm//PP//pppp__ccc.ceeeed
Train MAPPO on the Clean Up environment pythen
Train svo on the Clean Up environment python sagorithmm/S/o/svoooooooo
Figure Command-ling example for training different different algorithms (IPPO and MAPPO) and generating training curves with SocialJax.
python aped__estt/speddteet__t
Figure 6: Command-line example for testing SocialJax environment execution speed with random actions.
the cells to 1 x 1, also producing observations of 11 x 11. The agents in our environments have the same 11 x 11 field of view, but they are implemented at the grid level.
Coins Two players collect coins in a shared room, where each coin is assigned to one player's color. Each coin, upon appearing, has an equal 50% probability being assigned the first player's color or the second player's color. An agent receives reward of for collecting any coin, regardless of its color. If one player collects a coin assigned to the other player's color, the other player receives a reward of -2. Each empty tile on the map has respawn probability of p 0.0005 for each type of coin.
Commons Harvest The core concept of Commons Harvest was first defined by (Janssen et al. 2010). (Pérolat et al. 2017) introduced it in multi-agent scenarios. Melting Pot 2.0 (Agapiou et al. 2022) has made it a more comprehensive environment. There are several patches apples in the room. The players can receive reward of I, when they collect one apple. The apples will regrow with probability that is determined by the number of neighborhood apples in radius 2, when the apples are collected by the players. The probability of apple regrowth decreases as the nearby apples diminish. Apples will not regrow, if there are no apples in the neighborhood. Formally, the regrowth probability is set 0.025 if there are at least three apples nearby, 0.005 if there are exactly two, 0.001 if there is exactly one, and () when there are no apples in the neighborhood. Once all apples within patch are collected, that patch ceases to respawn apples.
Commons Harvest: Open: This experiment is conducted in an open environment with no internal obstacles, and walls are placed only along the periphery the map.
the players' roles are divided into two groups: one focuses collecting apples, while the other is tasked with preventing over-harvesting by other players (Pérolat al.. 2017). There are two rooms this setup, each containing multiple patches of apples and a single entry. These rooms can be defended by specific players to prevent others from accessing the apples by zapping them with with beam, causing them to respawn.
Commons Harvest: Closed:
Commons Harvest: Partnership In this environment, two players are required to cooperate in defending apple patches. The agents can be tested on whether they can learn to trust their partners to jointly defend their shared territory from intrusion and act sustainably when managing shared resources.
+1 reward by collecting each apple in Clean Up environment. Apples grow an orchard and their regrowth depends on the cleanliness of nearby river. Pollution accumulates in the river at constant rate and once pollution surpasses certain threshold, the apple
Clean Up The agent
16

---

Published as a conference paper at ICLR 2026
growth rate drops zero. Players have the option perform cleaning action that removes small amounts pollution from the river.
During the first 50 time steps of the game, the river remains uncontami- nated. After 50 time steps, there a 0.5 probability that one tile the river will become polluted. Apple Respawning: In the environment, let dirtCount represent the number of currently polluted river tiles, and riverCount represent the total number of river tiles, including both polluted and unpolluted ones. Then, the proportion of polluted river tiles is given by diction divtCount The probability of apple regrowth P is given by:
River Pollution Generation:
rsuerCount
ddrtFFaction 0s 0s .0, ..))
P == a min max
0r Od
where:
Og is the depletion threshold, the lower bound of the pollution range.
0, is the restoration threshold, the upper bound of the pollution range.
or is the maximum apple growth rate.
This formula calculates the probability P of apple regrowth based on the current pollution level f. with the value clipped between 0 and to ensure valid probabilities.
Coop Mining In this environment, two types of ore spawn randomly in empty spaces. Players can extract the ore to get reward. Iron ore (gray) can be mined individually and provides reward of +1 upon extraction. In contrast, gold ore (yellow) requires coordinated mining by two to four players within a 3-step window, granting a reward of +8 to each participant. When a player begins to mine gold ore, the state of the changes to "partially mined" to indicate readiness for another player to assist. Visually, this represented by brighter shade of yellow. no other player cooperates or too many players attempt mine simultaneously, the reverts to its original state, and no reward is granted.
We set the respawn rate of iron and gold propriate rates, promoting agents to make both selfish and cooperative choices, which in turn guarantees that the environment represents social dilemma. The respawn rates defined as:P.ron 0.0004 and Pgold = 0.00016. These rates are chosen to balance resource regener- ation and agent behavior, encouraging both competition and collaboration.
ensure that the regenerate at ap-
Mushrooms There are four types of mushrooms spread across the map, each offering different rewards when consumed. Eating red mushroom gives reward of 1 to the player who consumed it, while cating green mushroom gives a reward of 2, which is divided equally among all players. Consuming blue mushroom grants a reward 3, but the reward is shared among all players except the one who ate it. Additionally, eating orange mushroom causes every agent to get a -0.2 reward.
Mushroom regrowth depends on the type of mushroom eaten by players. Red mushrooms regrow with probability of 0.25 when any type of mushroom is eaten. Green mushrooms regrow with probability of 0.4 when either green or blue mushroom is eaten. Blue mushrooms regrow with probability of 0.6 when blue mushroom is eaten. Orange mushrooms always regrow with probability of when caten. Each mushroom has a digestion time, and a agent who consumes a mushroom becomes frozen during the digestion process. Red mushrooms take 10 steps to digest, while green mushrooms take 15 steps to digest, and blue mushrooms take steps to digest.
Gift Refinement In this environment, tokens spawn on the map with probability p = 0.0002. When tokens are collected, they are stored the agent's 's inventory. These stored tokens can be gifted to other agents, in which case the tokens are refined to higher level and their quantity is tripled. Tokens exist in three levels, with only the lowest level spawning naturally in the environment. Each token can be refined most twice to reach the highest level. Each agent can hold up to 15 tokens of each level. When agents execute the consunt sume action, all tokens the agent's 's inventory are converted into reward, with each token providing +1 reward regardless of refinement level.
17

---

Published as a conference paper at ICLR 2026
402
Ryt.ms
03
193
302
Enalmoment Emestaps
Ensrrnmmnt Emmsseps
Enafmommnt Limesipps
(a) Coins
(b) Commons Harvest: Open
(c) Commons Harvest: Closed
103
603
Retums
Eratrosment Emrascaa
Enviorment Jimsoless
Eratrosment Emessapa
(d) Commons Harvest: Partnership
Clean Up
Coop Mining
333
102
201
0.2
322
33
20
Enstronment Emertaps
Emessaps
Emerstaa
(g) Mushrooms
(h) Gift Refinement
(i) Prisoners Dilemma: Arena
N-PPO:
Individual Reward
N-IPO: Common Reward
PJPPO: Individual Reward
P-IPPO: Common Reward
Figure 7: Training curves of IPPO Common Reward and IPPO Individual Reward. Common Reward encourages collective interests, leading to higher overall returns, while Individual Reward primarily drives selfish behavior, often resulting in lower returns.
Prisoner's Dilemma: Arena Four individuals collect resources that represent 'defect' (red) or 'cooperate" (green) and compare inventories in encounter, which is first introduced by (Vezhnevets) et al. 2020). The inventory is represented by p = (pi..... PK). In our case, have two types of resources. When agent zaps another, agents' inventories begin tracking anew, and each agent immediately receives its corresponding reward. Consequences of the inventory comparison congruent with the classic Prisoner's Dilemma matrix game. This game exposes tension between reward for the group and reward for the individual. The matrix for the interaction is
because
Azow AT = 5
-
Each agent has mixed strategy weight, which is defined as
(Ui....,
Pi
E1=1 Pj
When one agent (the "row player") zaps and another (the "column player") is targeted, their respective payoff's are given by bilinear forms over these mixed strategies:
F Frow Vrow Arow Vcols
Tool
Yoow Acol Vcol-
After two agents complete an interaction, they respawn at designated locations within the environment and are prohibited from moving for random duration of 10 to 100 time steps.
B.3 DIVERSITY OF SOCIALJAX ENVIRONMENTS
Table mapps classical social dilemmas to the environments studied in our benchmark. Each paradigm captures distinct cooperation deffertiionddeecction tradeoff, which is instantiated in specific multi-agent
18

---

Published as a conference paper at ICLR 2026
environments. For instance, in the Prisoner's Dilemma, mutual cooperation strietly dominates mutual defection, yet unilateral defection remains individually tempting (c.g. Arena, Coins). In the Snowdrift/Chicken Game, cooperation yields benefits even if the partner defects, though mutual cooperation is preferable (e.g., Mushrooms). The Stag Hunt requires mutual cooperation to unlock high payoffs, while unilateral action results in low returns (e.g., Coop Mining, Gift Refinement). Finally, the Tragedy of the Commons illustrates the overuse of shared resources under selfish behavior, leading to system degradation (e.g., Commons Harvest variants, Clean Up).
Table Mapping classical social dilemmas to SocialJax environments.
Paradigm
Mapped Environments
Prisoner's Dilemma
Prisoner's Dilemma: Arena; Coins Mushrooms
Snowdrift Chicken Game
Stag Hunt
Coop Mining; Gift Refinement
Tragedy of the Commons
Commons Harvest (Open Closed / Partnership); Clean Up
B.4 TASK DIFFICULTY CATEGORIZATION
we include qualitative categorization T the task environments used our study along three axes: long-term impact of decisions, task multi-moaality (i.e., whether agents need to solve multiple goals), and number of agents. These axes provide intuitive guidance on the sources complexity across different social dilemmas. The final "Difficulty" column provides a coarse summary based on the combination these factors.
Table 6: Environments and their characteristics.
Environment
Long- Long-Term Impacts
Multi-Task
Agents Number
Difficulty
Clean Up
High
Yes
Very Hard
Gift Refinement
High
Yes
5
Very Hard
Harvest: Partnership
High
No
Hard
Harvest: Closed
High
No
7
Hard
Harvest: Open
High
No
Medium
Coop Mining
Medium
No
6
Medium
Prisoner's Dilemma
Medium
Yes
5
Medium
Mushrooms
High
No
5
Easy
Coins
Low
No
2
Easy
B.5 SCALABILITY OF SOCIALJAX
our framework allows users to freely adjust the number of agents via configuration files and environ- ment settings, enabling flexible scaling experiments. We conducted additional tests to evaluate the environment speed with 20 agents under random action execution, excluding Coins and Prisoner's Dilemma due their fixed design constraints, seeing Table 7 We observed only a modest slowdown in FPS compared to the performance reported in the main paper with original agent numbers. This is because our JAX-based environments rely heavily on parallel computation, where increasing the num- ber of agents primarily affects GPU memory usage rather than significantly impacting computational speed.
C ADDITIONAL RESULTS
C.1 PARAMETER SHARING AND NON-PARAMETER SHARING IPPO
We compare the training time and learning performance of shared parameter IPPO (P-IPPO) versus non-shared parameter IPPO (N-IPPO). According to Figure & the training times of shared parameter
19

---

Published as a conference paper at ICLR 2026
Table 7: Evaluate the environment speed with 20 agents under random action execution.
Envs (20 Agents)
1 JAX Env
128 JAX Env
1024 JAX Env
4096 JAX Env
Harvest: Open
1.2 X 10a
1.2 10a
5.1 10a
7.8 10a
Harvest: Closed
1.2 X 108
1.2 10S
5.1 10
7.8 10
Harvest: Partnership
1.2 103
1.2 105
5.1 10S
7.8 105
Clean Up
1.9 10a
1.9 10a
7.7 10a
1.3 106
Coop Mining
1.9 103
2.0 105
9.8 10
1.5 106
Mushrooms
1.2 103
1.2 105
5.8 105
9.5 105
Gift Refinement
1.3 103
1.4 109
6.6 10
1.2 106
Share: hummee Hondtann Natance
Environments
Figure 8: Training time to reach 1 timesteps: Shared parameter IPPO versus Non-shared parameter IPPO.
IPPO and non-shared parameter IPPO are virtually indistinguishable, with the shared parameter variant completing slightly faster.
We also present the training curves Figure 7 for shared parameter and non-shared parameter IPPO, which demonstrate essentially identical convergence behavior.
C.2 WALL CLOCK TIME SPEED COMPARISON
We compare the performance of our training pipeline with the training speed of the Stable-Baseliness and RLlib implementations with TensorFlow, based on the publicly available code from Melting Pot 2.0.
Note that the experimental results in Melting Pot 2.0 (Agapiou et al. 2022) cannot be directly compared, as the code related to their reported results has not been released. The open-source versions of the algorithms in Melting Pot 2.0 are based on RI.lib and Stable Baselines3. However, the versions of Actor-Critic Baseline (ACB) (Espeholt et al. 2018), V-MPO (Song et al.) (020), and OPRE (Vezhnevets al.) 2020) used in their reports have not been made available open source. Additionally, since Melting Pot 2.0 is pixel- grid-based, it is not feasible to align the environment and neural network settings directly. Therefore, we only compare wall-clock training times for scenarios that achieve convergence, evaluating the open-source frameworks Stable Baselines3 and RLlib alongside
environments while our SocialJax environments are
training pipeline.
As shown in Figure 9 the SocialJax training pipelines achieve roughly a 30-800x speedup over the open-source Melting Pot code, thereby substantially enhancing training efficiency. It should be noted
20

---

Published as a conference paper at ICLR 2026
PaSTi 22 120
Well 1 (taent sade
Cop Wrag
Praraen Damane ferr
Figure 9: Training time comparison across different frameworks: Stable Baselines3, RLlib, IPPO, and MAPPO.
that we did not actually execute 10 timesteps under the Stable Baselines3 and Rllib framework: rather, we estimated the time required complete 10" timesteps by extrapolating from the measured frames-per-second (FPS) performance.
D HYPERPARAMETERS FOR TRAINING
In this section, provide the hyperparameters and configuration details for different algorithms used our environments. We first provide the hyperparameter settings for IPPO, along with configuration options for enabling disabling parameter sharing, and for choosing between individual and common reward schemes. Table g shows the configuration of the IPPO algorithm. When shared rewards is set to true, it corresponds the common reward setting; when set to false, it corresponds to the individual reward setting. When PARAMETER SHARING is true, it denotes shared parameter IPPO, and when false, non-shared parameter IPPO.
21

---

Published as a conference paper at ICLR 2026
Parameter
Value
Learning Rate
LR 0.0005
Number of Environments
NUM ENVS 256
Number of Steps per Update
NUM STEPS 1000
Total Timesteps
TOTAL TIMESTEPS 3 10
Update Epochs
Number of Minibatches
NUM MINIBATCHES 500
Discount Factor
0=0.99
GAE Lambda
GAE BIE_LAMBAA=
Clip Epsilon
CLIP_EPS 0.2
Entropy Coefficient
ONT_COEF= =
Value Function Coefficient
VF = 0.5
Max Gradient Norm
MAX GRAD NORM 0.5
Activation Function
ACTIVATION relu
Environment Name
ENV_NAME clean_up
Reward Shaping Horizon
SWW_SHHAINGGHORI/AN 2.5 x 10
Shaping Begin
SHAPING BEGIN 10
Shared Rewards
shared rewards False (set True for common rewards)
CNN Enabled
can = True
JIT Compilation
jit = True
Anneal Learning Rate
ANNEAL_RR True
Seed
SEED = 30
Number of Sceds
WOM_SEEDS=1
Hyperparameter Tuning Enabled GIF Number of Frames
TUNE True
GIF NUM FRAMES = 250
Parameter Sharing
PARAMETER SHARRSSGA------ True
Number of Agents
num Lam_agents-
Number of Inner Steps
num inner steps 1000
Table 8: Hyperparameter settings and environment configuration for IPPO in the Clean Up environ- ment.
Description
Value
Learning Rate
LR 0.001
Number of Environments
NUM ENVS
Number of Steps per Update
NUM_STEPS
Total Timesteps
TIMESSES = 3 x 10
Update Epochs
UPDATE EPOCHS 2
Number of Minibatches
NUM MINIBATCHES
Discount Factor
7099
GAE Lambda
0.95
Clip Epsilon
CLIP_EPS 0.2
Entropy Coefficient
ENT WWT_COFF=000
Value Function Coefficient
VF COEF= 0.5
Max Gradient Norm
MAX GRAD TRAD_NORM- 0.5
Scale Clip Epsilon
True
Activation Function
relu
Environment Name
"lean up"
Reward Shaping Horizon
2.5 x 10
Number of Agents
7
Number of Inner Steps
1000
Shared Rewards
True
set to False individual rewards
CNN Enabled
True
JIT Compilation
True
Anneal Learning Rate
True
Table 9: Hyperparameter settings and environment configuration for MAPPO in the Clean Up environment.
22

---

Published as a conference paper at ICLR 2026
Table presents the hyperparameters for the MAPPO algorithm. For the SVO algorithm, hyperparameters are shown in Table [0] Note that during SVO training, we first perform a sweep over the parameter to find its optimal value. and then conduct sweep over the angle A. Therefore, the training script needs to be configured to sequentially perform sweep over the parameter wo (Table the angle (Table [T]. Further details environment configurations
10 followed by sweep and hyperparameter settings can be found in our open-source implementation.
sxeep_conte
"name" "cleanup_*", "method"; "grid",
"metric":
"name": "returned_ "rctrnod.pplooo..otiggiaaaaaaua "goal": "maximize",
"parameters":
"ENV wWWARGS..vo... ("values"; 10.0, 0.2, 0.4, 0.6, 0.8, 1.01
Figure 10: Sweep configuration for SVO weight in the Clean Up environment.
swccp_coafig =
"name": "elcanup_angle", "method": "grid", "metric";
"name":
"goal": "maximize",
p
"ENV_KWARGE svo_vvo_ooeo
www_ssvvvl_lsoeeeeeeeeeeeeeeeeee weiideal_ppt @eeeal_aggl__dereeeeee ": "values"; [C, 22.5, 45, 67.5, 90]
Figure 11: Sweep configuration for SVO angle in the Clean Up environment.
23

---

Published as a conference paper at ICLR 2026
Parameter
Value
Learning Rate
LR=0,0005
Number of Environments
NUM_ 2MM
Steps per Update
NUM_STEPS=IONS
Total Timesteps
TOTAL_ 3 x 10
Update Epochs
UPDATE EPOCHS-2
Minibatches
NUM MINIBATCHES=500
Discount Factor
GAMMA=0.99
GAE Lambda
GAF_LAMBDA=0.95
Clip Epsilon
CLIP_ LUP_EPS=0.2
Entropy Coefficient
ENT_COEF=DAM
Value Function Coefficient
VF COEF=00
Max Gradient Norm
MAX GRAD NORM=005
Activation Function
ACTIVATION=relu
Environment Name
ENV JNV_NAME-clean_up
Reward Shaping Horizon
REW REW_SHA WWW_SHAPINGGLRR
WWW.SHAINNG.OOZZON22555555 SIWVGGHHRZOOO.5-5- 10
Number Agents
ENV KWARGS nwmmmmmmmaaammeemm sum_aggnts=7
Inner Steps
ENV KWWRRG..umm 1mner_sttps=1000
Shared Rewards
ENV Bee.KKAARRSSSSs.. shared..dddd__ceeeee. saarredd_wwads==aaa1..
CNN Enahled
ENV WWW.KS.W.n=/// Cen.....nnTrreeeee
JIT Compilation
ENV KWARGS.jit=22
SVO Enabled
ENV SWWRRGS.vv=TTS
SVO Target Agents
ENV KWARGS.svO target agenns=(0...2.344....)
SVO Weight
WWW_KWARGS.NO SWWARGGSSvo__=O... w=0.
SVO Ideal Angle (degrees)
ENV K.WKKS.RAS. wwv__vvccl__aglleeeeeeeeeeea
Anneal Learning Rate
ANNEAL_LR==aalse
Random Seed
SEED=30
Number of Seeds
NUM SEEDS=1
Hyperparameter Tuning Reward Type
TUNE=False
REWARD=individual
GIF Frame Count
GIF BWW_NNU_FRMEES255
Table 10: Hyperparameter settings and environment configuration for SVO training in the Clean Up environment.
24


Reviews:
[{"review_id": "NEPrw5wVL0", "reviewer": "Reviewer_HwDj", "summary": "SocialJax presents a suite of Melting Pot–derived tasks implemented in a lightweight minigrid framework to probe sequential social dilemmas among multiple reinforcement-learning agents. Built on JAX, the benchmark emphasizes end-to-end vectorization to significantly speed up both training and evaluation. To validate each environment’s social incentives, if they encourage individualistic or altruistic behavior, the authors use Schelling diagrams produced via the SVO algorithm. They also evaluate several multi-agent variants of PPO, establishing standardized baselines and illustrating performance differences across tasks. The environments were evaluated to be different from each other as the number of cooperators-individual payoff curves differ between environments. The learning curves against PPO models are also quite different and thought provoking how these variations in PPO impact altruistic behaviour so much.", "questions": "I enjoyed reading the paper and following the results. That said, I was left wanting a deeper understanding of what’s new beyond Melting Pot, since the suite is derived from it. SocialJax’s focus on mixed motives (individual vs. altruistic behavior) is valuable precisely because it avoids the extremes of full competition or full cooperation. With that in mind:\n\n* One classic mixed competition–cooperation task is Level-Based Foraging (LBF, [https://github.com/semitable/lb-foraging](https://github.com/semitable/lb-foraging)), where agents compete for reward but sometimes benefit from cooperating. Would you consider expanding the suite to include LBF or similar non–Melting Pot tasks to mitigate intra-suite bias and strengthen generality?\n\n* The authors compared purely individual vs. purely shared rewards, but not mixtures. Prior work (Wang et al., *Individual Reward Assisted Multi-Agent Reinforcement Learning*, : [https://proceedings.mlr.press/v162/wang22ao.html](https://proceedings.mlr.press/v162/wang22ao.html)) shows that adding individual rewards to shared objectives in SMAC can hurt performance; recent directions in a similar vain explore higher-gradient adjustments ([https://arxiv.org/abs/2505.20579](https://arxiv.org/abs/2505.20579)) and value-shaping heuristics ([https://arxiv.org/abs/2508.17696](https://arxiv.org/abs/2508.17696)). SocialJax seems well-suited to probe this mixed-motive effect. Could the authors evaluate the IRAT method from Wang et al. on their suite since it is PPO-based, and report how well it performs?\n\n* Could the authors add a few baselines with modified versions of their environments that introduce individually rewarded “distractor” objects (e.g., apples) that disappear and pay out when stepped on? This should induce local individualistic behavior; measuring how agents balance these against common goals could deepen the robustness analysis of altruism for the different methods presented.\n\n* A key missing baseline is from the family of opponent-shaping algorithms. Proximal Advantage Alignment (“adalign”) from *Advantage Alignment Algorithms* ([https://arxiv.org/abs/2406.14662](https://arxiv.org/abs/2406.14662)) is a PPO-style method that steers policies via advantage signals and reportedly succeeds on Commons Harvest in Melting Pot. Would you include adalign to broaden algorithmic diversity and test how it's coordination compares in your tasks?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* Clear, well-structured writing that’s easy to follow.\n* Learning-curve plots are clean and legible, making comparisons straightforward.\n* Variations among PPO architectures and their returns provide informative contrasts.\n* Effective use of JAX for MARL: aligns with non-stationarity in transitions, rewards, and observations, and the performance benefits are evident.\n* Insightful comparison of individually rewarded PPO vs. common-rewarded PPO, plus PPO-RE’s exchange schedule.\n* Well-chosen environments that focus on mixed-motive cooperation which is an important aspect of multi-agent interaction and underexplored.\n* Performance descriptions and the accompanying Schelling analyses are reasonable and generally supported.\n* The Melting-Pot suite is theoretically grounded in game theory and the inclusion to JAX is welcomed", "weaknesses": "* Most environments are direct ports from Melting Pot or prior work; Coins and STORM already have JAX implementations, so the contribution feels more like technical consolidation than novel insight into mixed-motive RL.\n* All evaluated algorithms are PPO variants; including other cooperative MARL approaches would better test generality and strengthen the conclusions.\n* The Schelling diagrams (via SVO) are under-explained and lack context; it’s unclear why they’re preferable to simple metrics like cumulative individual rewards (for individualistic agents) and cumulative common rewards (for altruistic agents) or maybe even behavioural metrics. Related discussion to Schelling and why the Schelling diagrams are better is important\n* Agent partial-observation windows are not visualized in the environment images, limiting clarity about information constraints.\n* The benchmark is restricted to Melting Pot–based environments, which may introduce intra-suite bias and limit external validity.\n* The study does not examine mixed individual and common reward settings, omitting an important class of real-world mixed-motive scenarios.\n* Lack of explanation for why PPO-RE struggles in coop mining leaves the analysis incomplete.\n* Possible exploration deficits are not investigated; lack of exploration may explain lower returns for some PPO variants.\n* No off-policy or value mixer baselines", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xItlCCxToN", "reviewer": "Reviewer_RUpe", "summary": "The authors described an evaluation suite for multi-agent RL in social dilemmas.\nThis suite consisted of 9 RL environments exploring cooperation, partnership and trust\nThe authors also implemented SOTA algorithms to test the evaluation suite \nThe key contribution is implementing the environment in JAX, which has a substantial improvement in rollout performance, directly speeding up training and evaluation of RL algorithms", "questions": "* Why was the observation space changed? Is it possible to have a similar observation space for comparisons (possibly as an option)", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper outlined the problem clearly\n2. The authors presented a strong speed up improvement compared to the existing social dilemma environments\n3. There is a diverse set of environments exploring many aspects of social dilemma\n4. The authors provided a strong benchmark of RL algorithms for future comparisons", "weaknesses": "1. The authors clam to have a 50x+ speed up compared against the Melting pot 2.0. However, there were many differences in the observation space. It doesn't feel like a fair apples to apples comparison.  \n2. The effects of changing the layout grid-like observation was not discussed\n3. It might be helpful to understand how the same RL algo compares between the SocialJax vs Melting Pot 2.0", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xvi4azVvEW", "reviewer": "Reviewer_A4zv", "summary": "This paper introduces SocialJax, a JAX-based suite of nine sequential social-dilemma environments plus baseline MARL implementations (IPPO under individual/common reward, MAPPO, SVO, PPO-RE). The environments are grid-based abstractions inspired by Melting Pot (e.g., Coins, Commons Harvest variants, Clean Up, Coop Mining, Mushrooms, Gift Refinement, Prisoner’s Dilemma: Arena), with a fixed 11×11 partially observable window. The authors emphasize throughput and reproducibility: by having thousands of vectorized environments on GPU this paper shows significant speedups by using SocialJax that allow reproducibility with significantly less computational resources. They validate “social dilemma” structure via Schelling diagrams built from policy classes: “cooperators” sampled from common-reward training; “defectors” from individual-reward training, and propose simple, environment-specific cooperation metrics beyond return.", "questions": "How well do the different algorithm implementation performances align with those in the original environments (e.g. IPPO in meltinpot’s common harvest open vs socialjax's common harvest open)?\n\nSome results are unintuitive to me. In the Schelling diagrams, why is the reward of defectors lower than that of cooperators in Gift refinement as the number of other cooperators increases? What are the policies that the cooperators are learning?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well written and introduces a valuable benchmark for the MARL community. A major bottleneck of MARL research is associated with the large costs of training often available only to big industrial labs. In such a way, SocialJax is a step towards the democratization of MARL research by allowing fast training of important benchmarks with significantly lower computational constraints. I believe the main strength of this paper lies on the opportunity it might enable to push the frontiers of MARL research by making it more affordable and thus it should be accepted to the conference.", "weaknesses": "A core weakness is that SocialJax reduces the observation resolution and modality relative to the original environments, shifting from pixel observations (e.g., Melting Pot’s ~88×88 image crops for an 11×11 field of view) to discrete 11×11 grid values rendered at the grid level. This removes the perception-learning burden that makes the originals attractive for testing MARL scalability and representation learning, and it changes coordination difficulty and partial observability in ways that can inflate apparent sample efficiency. In the authors’ own description, SocialJax agents “observe a grid matrix” with a fixed 11×11 view, whereas Melting Pot provides pixel images; even Melting Pot contest entrants who downsampled still operated in a pixel pipeline, not a hand-engineered grid state. Consequently, speedups and learning curves are not apples-to-apples comparisons.\n\n**Minor errors/typos:**\n\n**Line 271:** The authors state “we evaluate cooperation by counting how many coins each agent collects that match their assigned color”. This is not a valid metric of cooperation as the agents may be learning to collect all the coins. A good metric of cooperation in this environment is “total coins of their assigned color”/ ”total coins collected” i.e. which fraction of the coins collected are of their color.\n\n**Line 192, 768:** Duplicated word “with with a beam”.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZqAzN89NJA", "reviewer": "Reviewer_CDVN", "summary": "This paper presents SocialJax, a new suite of sequential social dilemma environments and algorithms built with JAX to advance research in multi-agent reinforcement learning (MARL).  In the second part of the paper, the authors validate the performance of baseline algorithms and use Schelling diagrams to study social dilemma outcomes - this second part appears rather disjointed from the first part.", "questions": "- Can you provide some technical description about the design choices exploiting JAX characteristics in order to support GPU parallelism (also in comparison to the characteristics of the “standard” environments discussed in your paper)?\n- Did you the same GPU cores as Agapiou et al? Are the performance figures taken from that paper?", "limitations": "", "rating": 2, "confidence": 5, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "+ Standardised benchmarks for studying social dilemmas using MARL agents are useful for the community. This paper goes in that direction.", "weaknesses": "- The discussion about CPU-based parallelism is difficult to understand. In fact, it seems that the problem is not at the level of the environment for GPU parallelism, but at the level of the agent implementation. It is possible to run multiple agents in a distributed fashion as well.\n\n- In the introduction, the authors also claim that the social dilemmas are highly challenging environments. This appears to be a rather stretched claim. There are several environments that are actually much more complex than the simple dilemmas supported by SocialJAX.\n\n- The authors appear to mix the problem of model implementation and environment when they discuss MAPPO, etc. (see paragraph entitled “Implementation of Algorithms in JAX” in the Introduction). In fact, you could release a very fast library implementing MAPPO, PPO, etc. instead.\n\n- The authors say that the existing social dilemmas are restricted by CPU performance. This is not true. In fact, the models themselves can currently be run using GPUs in theory, but there is no inherent speed-up for them given their nature.  It seems that there is a mixing between agent implementation and interaction between agents in an environment. You can use standard PyTorch and use the GPU for implementing the agents. JAX allows you to do some clever synchronization in distributed training, not sure if it is useful/essential for the models that are described in this paper. \n\n- In any case, the authors do not provide a discussion of the actual implementation details that lead to a speedup with respect to a “CPU implementation”. Details about the implementation, parallelization, and synchronization are missing.\n\n- The structure of the paper is unconvincing in itself. There is no reason to introduce social dilemmas (and the mathematical details) in Section 3. The paper is supposed to be about a simulator, not about the models that are well known and widely used by the community. The set of games implemented are also quite standard (perhaps that section could be moved to an Appendix). \n\n- In Table 1, the authors should have reported information regarding training, not random actions. The authors should have reported results for an increased number of CPUs. \n\n- It is unclear if the authors used the same GPU cores in the simulations as Agapiou et al. Also, these do not appear as meaningful comparisons given the fact that these are different training libraries (for example, it is not due to the implementation of Melting Pot, etc.).\n\n- The discussion about the IPPO common reward is not relevant to this paper. It is a generic result orthogonal to the problem of the design of a MARL simulator.\n\n- The results in Table 2 and 3 and Figure 3 appear unrelated to the goal of this work. These are unrelated to the evaluation of the simulator in itself. The analysis of the metric is also quite disjointed from the core goal of the paper.\n\n- In the Appendix, the wall clock comparison is not meaningful since they are based on different implementations (some of them are also pixel-based).\n\nMinor weaknesses:\n\n- There are missing citations for certain games in Section 4.2.\n- Arena is in bold in Section 4.2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EPHf0anmAW", "reviewer": "Reviewer_m7Ji", "summary": "This paper implements a series of sequential social dilemma environments in JAX, each exhibiting mixed incentive dynamics and significantly improving the speed of environment simulation. A range of multi-agent reinforcement learning (MARL) algorithms are implemented, and training these algorithms in SocialJAX is at least 50 times faster than in MeltingPot2.0. The performance of various environments and algorithms is benchmarked in terms of both simulation throughput and training speed.", "questions": "this paper may address its issues of innovation rather than providing more engineering descriptions. If it is an engineering contribution, more detailed information should be provided during the experimental phase.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main contribution of this article is the implementation of a series of reinforcement learning algorithms using JAX, along with experiments and a comprehensive analysis.\nIf these works can be open-sourced, they would have practical significance for scenarios in the industry where there is a high demand for intelligent agents.", "weaknesses": "1. The work presented in this paper involves implementing multi-agent algorithms in the field of sequential social dilemma environments using the JAX library, which has accelerated the execution speed of the algorithms. The focus seems to be more on engineering aspects rather than innovation. The paper does not claim to have made any specific innovative engineering contributions that would have sped up the algorithms during implementation. It appears that the performance improvements are largely due to the advanced features of the JAX library.\n\n2. This work is not declared open-source, and the links provided in the abstract cannot be accessed. This might be due to time constraints or the intention to add more information during the rebuttal phase, but it is unfair to other works. Not being open-source limits the contribution to the community, and the inability to view the links makes it impossible to evaluate the workload.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)