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

Abstract

Social dilemmas pose a significant challenge in the field of multi-agent reinforcement learning (MARL). Melting Pot is an extensive framework designed to evaluate social dilemma environments, providing an evaluation protocol that measures generalization to new social partners across various test scenarios. However, running reinforcement learning algorithms in the official Melting Pot environments demands substantial computational resources. In this paper, we introduce SocialJax, a suite of sequential social dilemma environments implemented in JAX. JAX is a high-performance numerical computing library for Python that enables significant improvements in the operational efficiency of SocialJax on GPUs and TPUs. Our experiments demonstrate that the training pipeline of SocialJax achieves a 50× speedup in real-time performance compared to Melting Pot’s RLlib baselines. Additionally, we validate the effectiveness of baseline algorithms within the SocialJax environments. Finally, we use Schelling diagrams to verify the social dilemma properties of these environments, ensuring they accurately capture the dynamics of social dilemmas.

Individual rewards and Collective rewards

Individual Rewards indicate that each agent is assigned its own reward, which inherently makes each agent selfish. Collective rewards refer to a scenario where all agents in a multi-agent system share a single, unified reward signal derived from the environment. This approach ensures that all agents are aligned towards achieving the same objective, promoting collaboration and coordination among them.

Clean Up

Clean Up is a multiplayer game where players earn rewards by collecting apples, each worth +1 reward. Apples grow in an orchard and their regrowth depends on the cleanliness of a nearby river. Pollution accumulates in the river at a constant rate and once pollution surpasses a certain threshold, the apple growth rate drops to zero. Players have the option to perform a cleaning action that removes small amounts of pollution from the river. For continuous apple growth, the group must keep river pollution levels consistently low over time.

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 of being assigned to the first player's color or the second player's color. An agent receives a reward of 1 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.

Individual Rewards

Collective Rewards

Individual Rewards

Collective Rewards

Clean Up: we can see that agents with Individual Rewards are unwilling to clean the river and instead focus on staying in certain areas of the map, waiting to harvest apples. In contrast, agents with Collective Rewards repeatedly clean the river to maximize the overall team reward.

Coins: Agents with Individual Rewards try to consume all colored coins whenever possible, whereas agents with Collective Rewards only collect coins that match their own color.

Commons Harvest: Open

There are several patches of apples in the room. The players can receive a reward of 1, 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.

Commons Harvest: Closed

Its mechanism is the same as Common Harvest: Open, but the difference is that there are two rooms in 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.

Individual Rewards

Collective rewards

Individual Rewards

Collective rewards

Commons Harvest: Open: We can observe that agents with Individual Rewards quickly consume apples without waiting for them to regrow. In contrast, agents with Collective Rewards regulate the pace of apple collection.

Commons Harvest: Closed: We can see that Individual Rewards agents quickly enter the room and eat all the apples. Although Collective Rewards agents have learned to control the pace of apple consumption to gain more long-term rewards, the agents outside haven't learned to enter rooms.

Coop Mining

In this environment, two types of ore spawn randomly in empty spaces. Players are equipped with a mining beam to extract the ore. Iron ore (gray) can be mined individually and provides a reward of +1 upon extraction. In contrast, gold ore (yellow) requires coordinated mining by exactly two players within a 3-step window, granting a reward of +8 to each participant. When a player begins to mine for gold ore, it flashes to indicate readiness for another player to assist. If no other player cooperates or too many players attempt to mine simultaneously, the ore reverts to its original state, and no reward is granted.

Territory

Territory is a competitive multi-agent game where players paint walls to capture territory and earn rewards. The players claim walls by painting them with their unique color, either by touching or flinging paint. The paint dries after 25 steps, brightens, and rewards are given. The more walls that are painted and dried, the greater the reward players can receive. Walls can be zapped twice to permanently destroy them, making them non-claimable and traversable. Players can also zap each other and a player is permanently removed after being hit twice.

Nine players compete to claim resources by painting them in their unique colors. Players spawn in an open area near each other, away from the resource walls. Resource distribution is uneven, with some areas richer in walls than others.

Individual Rewards

Collective rewards

Individual Rewards

Collective rewards

Coop Mining: Agents with Collective Rewards exhibit more cooperative behavior in acquiring gold mines compared to agents with Individual Rewards.

Territory: The situation is different in Territory. Since it is an open environment, agents with Individual Rewards do not destroy others' resources due to excessive competition.

Speed Benchmark Results

We compare the performance of our training pipeline from 1 environment to 256 environments with the training speed of the Stable-Baselines3 and RLlib, based on the publicly available code from Melting Pot 2.0. 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 and Territory, SocialJax is approximately 50-140 times faster than Melting Pot 2.0.

Table 1. Speed test of Training pipeline with IPPO

We also evaluate the speed of SocialJax across various environments and environment counts using random actions. We observe a significant speedup when running environments in parallel on a GPU compared to running a single environment. The original environments run on a single CPU, and achieving the parallel efficiency of SocialJax would require scaling to hundreds of CPUs.

Table2. Speed test using random actions

Schelling Diagrams

To validate the social dilemma properties of SocialJax, we depict a Schelling diagrams for each environment. The cooperative policies are 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 defector agents.

Training Curves

Training curves of Common Reward IPPO and Individual Reward IPPO. The former setting results in outcomes where collective interests are considered, while the latter mainly leads to selfish behaviour and lower returns.

Conclusion

The use of hardware acceleration has revolutionized MARL research, making it possible to overcome computational limitations, accelerate idea iteration. SocialJax is a library designed to leverage these advancements in the context of social dilemmas. It provides a collection of social dilemma environments and baseline algorithms implemented in JAX, offering the efficiency of hardware accelerators. By achieving remarkable speed-ups over traditional CPU-based implementations, SocialJax significantly reduces the time required for experiments. In addition, its unified codebase consolidates a diverse range of social dilemma environments, addressing long-standing issues in MARL evaluation standardization. By enabling faster, more reliable, and more thorough assessments, SocialJax aims to advance progress in MARL social dilemma research.

Page updated

Google Sites

Report abuse