SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lianhao Yin
lianhao@mit.edu
Tim Seyde
tseyde@mit.edu
Tsun-Hsuan Wang
tsunw@mit.edu
Wei Xiao
weixy@mit.edu
Ramin Hasani
rhasani@mit.edu
Joshua Rountree
rountree@mit.edu
Daniela Rus
rus@csail.mit.edu
Multi-agent reinforcement learning (MARL) research is faced with a trade-off: it either uses complex environments requiring large compute resources, which makes it inaccessible to researchers with limited resources, or relies on simpler dynamics for faster execution, which makes the transferability of the results to more realistic tasks challenging. Motivated by these challenges, we present Gigastep, a fully vectorizable, MARL environment implemented in JAX, capable of executing up to one billion environment steps per second on consumer-grade hardware. Its design allows for comprehensive MARL experimentation, including a complex, high-dimensional space defined by 3D dynamics, stochasticity, and partial observations. Gigastep supports both collaborative and adversarial tasks, continuous and discrete action spaces, and provides RGB image and feature vector observations, allowing the evaluation of a wide range of MARL algorithms. We validate Gigastep's usability through an extensive set of experiments, underscoring its role in widening participation and promoting inclusivity in the MARL research community. MIT licensed code is available at https://github.com/mlech261/gigastep.
Decentralized partially observable Markov decision processes (Dec-POMDPs) [4, 42] represent a practically relevant yet largely unsolved challenge within the domain of deep reinforcement learning (Deep RL). These decision-making frameworks involve multiple agents operating in an environment where each agent has access to only partial observations, making it necessary to reason about both the global state and the private information of individual agents. Solving Dec-POMDPs efficiently and effectively has profound implications for real-world applications [9], such as human-machine interaction [30], multi-robot systems [13], autonomous driving [28, 29, 10], and decentralized control of complex networks [71, 11].
To address this challenge, numerous environments have been proposed, each with unique strengths and weaknesses. For instance, the Starcraft Multi-Agent Challenge (SMAC) [50] emerged as a popular platform for benchmarking Dec-POMDP algorithms. However, the SMAC v2 [15] revealed that this environment is fundamentally deterministic, allowing agents to memorize action sequences rather than developing genuinely generalizable behavior, thus providing a skewed perspective on algorithmic performance. Consequently, it is essential to question the robustness and effectiveness of the solutions that have been evaluated only in single environments.
Gigastep Overview
Features
Heterogenous Agents
Objectives
Asymmetric Teams
Waypoint following
Figure 1: An overview of Gigastep's features and characteristics
Hide & seek
Observability
Fully observable
Observation space
Vector
XN X
Partially observable + stochasticity
Action Space
Furthermore, the current landscape of Dec-POMDP research exhibits a stark dichotomy: on the one hand, we have environments that, although computationally accessible, lack the complexity to adequately represent the intricacies of real-world Dec-POMDP scenarios. These are often overly simplistic, two-dimensional, and void of the partial observability that characterizes Dec-POMDPs [37]. On the other hand, we have environments that are sufficiently complex but require such a high degree of resources that exclude many researchers, particularly those from underrepresented or underfunded institutions, thus hindering the democratization and inclusive progress of the field [55].
Motivated by these concerns, our paper aims to bridge this gap by introducing a novel benchmark environment called Giagstep for Dec-POMDPs that balances difficulty, realism, and computational accessibility. By addressing the identified limitations of existing methods, we hope to foster a more inclusive, comprehensive, and realistic exploration of Dec-POMDPs in deep reinforcement learning research. In Gigastep, a set of airborne agents separated into two teams navigate in the 3-dimensional space and need to achieve predefined objectives, such as following waypoints and tagging, chasing, or avoiding agents of the other team. Each agent perceives their neighborhood and shares information with allied teammates stochastically, making it an ideal instance of benchmark Dec-POMDPs. We summarize our contributions in three points:
- Gigastep is designed with efficiency in mind enabling MARL on consumer-grade hardware
- Gigastep is fully open source allowing customization, specialization, and further advancements by the research community
- Gigastep captures the open challenges in multi-agent reinforcement learning in decentralized partially observable Markov decision processes (Dec-POMDP)
In this section, we list related RL and MARL benchmark environments and describe their differences from Gigastep. Moreover, in Table 1, we summarize key features and lack thereof of the environments.
A wide range of simulators are designed to help develop reinforcement learning algorithms to solve problems that can be treated as Dec-POMDPs or general MDPs. These simulators focus on different aspects such as multi-player computer games [49, 15], sports [23], social interactions [34], life evaluations [60], particle simulation [64, 73], differentiable rigid body dynamics [17, 69], drones [54, 58, 43, 18, 65], etc. We elaborate on their advantages and limitations in the following paragraphs.
Drones simulators - Multiple drone simulators are developed for different research questions [54, 18, 58, 43]. Some of them focus on photo-realistic simulation such as Airsim [54], flight goggles [18], flightmare [58], etc. Some of them aim to simulate multi-drones for reinforcement learning such as Pybullet-drones [43] and Airsim [54]. However, they are not designed to be vectorized and run on GPU. The speed slows down significantly when the number of drones increases.
StarCraft Multi-Agent Challenge — The StarCraft Multi-Agent Challenge (SMAC) [49] and SMACv2 [15], are of most prominent MARL benchmarks. SMAC offers a complex, real-time, and partially observable platform for MARL research based on the popular game StarCraft II.
However, there are certain limitations to the SMAC environment that Gigastep aims to address. Firstly, while SMAC does offer multiple units and scenarios, the game-based nature inherently limits the variability and control over the environment's parameters. Our environment is fully open source and designed to be more flexible, enabling researchers to customize and design scenarios that match their specific research goals. Secondly, SMAC's computational efficiency is tied to the StarCraft II game engine, which can be a bottleneck for large-scale experiments and severely limits its access to underrepresented researchers with little compute infrastructure. Our JAX-based environment is designed from the ground up to exploit hardware accelerators like GPUs and TPUs, ensuring highly efficient simulations. Lastly, SMAC's interface to the environment, through PySC2, can be cumbersome to navigate for newcomers to the field. Our environment is engineered to be user-friendly, with an intuitive API that allows easy integration with various reinforcement learning algorithms and custom scenarios.
PettingZoo — PettingZoo is a shared Python interface for a range of multi-agent reinforcement learning (MARL) challenges. It includes an extensive collection of benchmark environments such as the multi-agent Arcade learning environment (ALE) [3], the Multi Particle Environments (MPE) [36], SISL (Stanford Intelligent Systems Laboratory) environments [19] as well as a few of other environments. While PettingZoo [64] provides a valuable set of MARL environments and tools, its environments include a few limitations that Gigastap aims to resolve. For instance, the multi-agent ALE environments comprise of only 1vs1 player games, whereas Gigastep allows teams of more than a thousand agents to play simultaneously. Moreover, compared to Gigastep, the MPE and the SISL environments possess only basic 2D dynamics, lack support for high-dimensional RGB image observations, and contain at most 8 agents.
brax and gymnax — Brax [17] and Gymnax [25] are single-agent RL environments that are entirely implemented in JAX. Consequently, these environments can fully leverage the computational power of GPU and other accelerator devices to gather experience data for RL training in parallel, enabling fast RL experimentation. However, they lack native support for multi-agent settings. Gigastep instead aims to bridge this gap and replicate the advantages of both Brax and Gymnax in multi-agent settings.
Magent – Agent [73] is a research platform for multi-agent reinforcement learning. It has a specific aim to observe and understand the social behaviour of agents and individual behaviour through interaction and communications. However, it only provides adversarial scenarios such as pursuit, gathering, and battle. It has no built-in vectorization and can only be parallel executed by creating workers, which is not computationally efficient for state-of-the-art reinforcement learning algorithms.
Google football - Google football [23] is design to simulate football game in particular. It is limited to the number of team members in the cooperative and competitive teams which represents a sports game. Therefore, the scalability in both the number of agents and scenarios is limited.
Melting Pot - Melting Pot [34] has the capability to evaluate the generalization of algorithms to novel social situations such as cooperation, competition, deception, reciprocation, trust, stubbornness, and so on. However, it is limited in the number of agents and can not be accelerated in GPU.
Multi-agent Mujoco - Multi-agent Mujoco [44] is a research platform that treats each joint of a robot as an individual agent, upon which, one can develop multi-agent algorithms to control a full robot. However, there is no higher level and task-aware cooperation and competition between multi-agents.
Multi-agent RL algorithms - Effective MARL algorithms need to scale across high-dimensional action spaces and enable coordination among distributed agents. Independent learners [61] provide scalability at the cost of exacerbating coordination challenges, particularly under partial observability [12, 39]. Fully centralized learners can become computationally infeasible in high-dimensional action spaces, and distributed methods may improve learning [51, 47]. A common approach within Q-learning is the decomposition of the joint value function into local utilities [45, 57, 63, 53], which similarly transfers to the actor-critic setting [66, 59, 44] and hierarchical control [67]. More generally, the concept of centralized training with decentralized execution enables effective training of distributed teams under partial observability. Interactions can further be limited to only local interplay among agents [70, 5]. Alternatively, evolutionary approaches side-step the need for dense gradient information and may improve coordination under sparse team-centric reward feedback [38, 62].
Table 1: Comparison of different RL and multi-agent RL environments
Environment	Adversarial multi-agent	Stochastic dynamics	Heterogeneous agents	>1000 agents	Open source	Vector observations	RGB image observations	3D dynamics	Built-in vectorization	GPU acceleration	Customize environment
gym [8]	✓	✓	✓	-	-	✓	✓	✓	✓	-	✓
gymnax [25]	-	-	✓	-	-	✓	✓	✓	-	✓	-
brax [17]	-	-	-	-	-	✓	✓	-	✓	✓	✓
SMAC [49]	✓	✓	✓	-	✓	-	✓	-	-	-	-
SMACv2 [15]	✓	✓	✓	-	✓	-	✓	-	-	-	-
ALE [3]	-	✓	-	-	-	✓	✓	✓	-	-	-
Minecraft MALMO [21]	-	-	✓	-	✓	-	-	✓	✓	-	-
Deepmind Lab [2]	-	-	✓	-	-	✓	✓	✓	-	-	✓
FortAttack [14]	✓	✓	-	-	✓	-	✓	-	-	-	-
Derk's Gym [41]	✓	✓	-	-	-	-	✓	-	-	✓	-
Hanabi [1]	✓	-	✓	✓	-	-	✓	✓	N/A	-	-
Neural MMO [60]	✓	✓	✓	-	✓	✓	-	✓	-	-	✓
MPE in PettingZoo [64]	✓	✓	✓	-	-	-	✓	✓	-	-	-
SISL in PettingZoo[19]	✓	✓	-	-	✓	-	✓	✓	-	-	-
jumanji [6]	✓	✓	✓	✓	-	-	✓	✓	-	✓	-
Google football [23]	✓	✓	-	✓	-	-	✓	✓	✓	-	✓
Multi-agent Mujoco [44]	✓	-	-	-	-	-	✓	✓	-	✓	-
Pomerman [46]	✓	✓	✓	-	✓	-	✓	✓	-	-	-
Melting Pot [34]	✓	✓	✓	✓	-	✓	-	✓	-	-	-
Magent [73]	✓	✓	-	-	✓	✓	✓	-	-	✓	✓
Gym-pybullet-drones [43]	✓	-	-	-	-	-	✓	✓	✓	✓	✓
Air-sim [54]	✓	-	-	-	-	-	✓	✓	✓	-	-
Flightmare [58]	✓	-	-	-	-	-	✓	✓	✓	✓	-
gigastep (ours)	✓	✓	✓	✓	✓	✓	✓	✓	✓	✓	✓
3 Designing Gigastep
In this section, we first provide a formal definition of Dec-POMDPs, followed by a detailed description of the design choices made with Gigastep.
A Decentralized Partially Observable Markov Decision Process (Dec-POMDP) is an extension of a Markov Decision Process (MDP) that facilitates the modeling of multi-agent systems within partially observable environments. Specifically, a Dec-POMDP introduces decentralization, in which each agent makes decisions based on its individual actions and observations without direct communication or knowledge of other agents' observations or actions.
Formally, a Dec-POMDP is defined by a tuple  $G = (n,S,A,T,\Omega ,O,R,\gamma)$  where:  $n$  is the number of agents.  $S$  is a finite set of states of the environment.  $A = \{A_{1},A_{2},\ldots A_{n}\}$  is a set of action sets, where  $A_{i}$  is the set of actions available to agent  $i$ . The combined action space of all agents is  $A = A_{1}\times A_{2}\times \dots \times A_{n}$ .  $T:S\times A\to \Delta (S)$  is the state transition function, where  $\Delta (X)$  denotes a probability distribution of the set  $X$ . For a given state  $s\in S$  and joint action  $a\in A$ ,  $T(s,a)$  is the probability distribution representing probability  $\mathbb{P}(s^{\prime}|a,s)$  of transitioning to state  $s^{\prime}$  from state  $s$  when joint action  $a$  is taken.  $\Omega = \{\Omega_1,\Omega_2,\dots ,\Omega_n\}$  is a set of observations, where  $\Omega_{i}$  is the set of observations for agent  $i$ . The combined observation space of all agents is  $\Omega = \Omega_{1}\times \Omega_{2}\times \dots \times \Omega_{n}$ .
Maps + Waypoints
Figure 2: Examples of the RGB image observation mode.
Stochastic observation cone
1000 vs 1000 agents
Heterogenous agent types
$O: S \times A \to \Delta(\Omega)$  is the observation function. For a given joint action  $a \in A$  and state  $s$ ,  $O(s, a)$  is the probability  $\mathbb{P}(o|a, s)$  that agents observe  $o$  after taking joint action  $a$  and in state  $s$ .  $R: S \times A \to R$  is the reward function with  $R(s, a)$  representing the immediate reward for taking joint action  $a$  in state  $s$ .  $\gamma \in [0,1]$  is the discount factor, which determines the present value of future rewards.
Each agent  $i \in \{1, 2, \ldots, n\}$  selects actions based on its policy  $\pi_i : o_i \to \Delta(A_i)$ , which is a mapping from its observation to a distribution over its action set  $A_i$ . The joint policy of all agents is then  $\pi = (\pi_1, \pi_2, \dots, \pi_n)$ . The goal in a Dec-POMDP is to find an optimal joint policy  $\pi^*$  that maximizes the expected discounted sum of rewards  $\mathbb{E}(\sum_{t=0}^{\infty} \gamma^t R(s_t, a_t) | \pi, s_0 = s)$ , where  $s_t$  and  $a_t$  are the states and joint action at time  $t$ , respectively.
3.2 Design considerations
Here, we discuss the design choices and constraints made in Gigastep.
Vectorizable environment stepping RL algorithms learn from their experience of interacting with the environment. Consequently, their training is limited by the amount of experience data that is available. However, generating this data can be a slow process, especially in complex environments where each interaction could take a considerable amount of computing time. This bottleneck significantly impacts the efficiency of experimentation with RL algorithms and hinders research progress. The limitation is further amplified for underrepresented researchers who may lack the funding for vast computer resources.
Therefore, there is a pressing need for a fast and vectorizable environment that can handle environment rollouts in parallel. Such an environment would enable multi-agent reinforcement learning algorithms to generate and learn from a larger amount of experience data in a shorter span of time, thereby significantly accelerating research progress. A vectorizable environment would also allow these algorithms to fully exploit the capabilities of modern GPUs, further enhancing their computational efficiency and speed.
We resolved this issue by implementing Gigastep fully in JAX [7], an accelerated linear algebra (XLA) package that allows compiled programs to run in parallel on CPU or GPU devices. Moreover, compared to existing approaches Gigastep support rending the agents observations as RGB images in a fully vectorized and jax . jit compilable fashion. Contrarily, existing benchmarks often provide only vectorized feature vector observations and rely on third-party game libraries for rending the agents' state and observation.
Collaborative and adversarial components An ideal MARL benchmark integrates both collaborative and adversarial components within the environment. The simultaneous existence of these aspects within an environment enables a more comprehensive emulation of complex real-world scenarios, thereby fostering enhanced learning capabilities in agents. The collaborative component encourages cooperative strategies, enhancing the ability of agents to work together towards common goals, while the adversarial aspect fosters competitiveness, which hones an agent's ability to adapt to an ever-changing environment and devise individualistic strategies. Together, these components not only enrich the learning experience but also improve the robustness of the agent's decision-making process, fostering a more nuanced understanding of their environment, consequently leading to improved performance and adaptability in diverse and dynamic contexts.
Stochastic dynamics In MARL, it is crucial to design environments that possess a level of stochasticity in their dynamics. This element of randomness is important to prevent agents from learning and exploiting deterministic action sequences, thereby encouraging more robust and adaptable learning. Moreover, stochasticity is a fundamental aspect in many real-world applications [32, 33], thus critical if MARL algorithms need to transfer from simulation to real-world deployment. The Starcraft Multi-Agent Challenge (SMAC) v1 is a prime example of a deterministic environment where agents were able to learn simple, reproducible action patterns that led to an overestimation of their abilities [49, 15]. The introduction of stochasticity in SMAC v2 [15], has considerably improved the testing of learned policies by introducing uncertainty and forcing agents to devise strategies that are resilient to changing conditions. Hence, incorporating stochasticity into MARL environments is a necessary step toward achieving more comprehensive and realistic evaluations of agent competencies.
3.3 Environment description
In this section, we discuss various features and characteristics of Gigastep. An overview of these features is illustrated in Fig. 1. In particular, Gigastep was inspired by multiplayer online battle arena (MOBA) style games, where two teams of agents compete with each other. The agents have heterogeneous capabilities, which enables the emergence of complex team dynamics. Complexity is often further added by partial observability of the game state and stochastic or low-bandwidth communication with a team.
Agents' features. The environment features agents that move in 3D rectangular landscapes. The agents are heterogeneous in their skill sets. For instance, agents could be faster, more resilient, and possess a longer field of view. The environment allows for symmetric and asymmetric team configurations where one team has more agents than the opponent.
Maps. Gigastep comes with customize maps in both 2D and 3D.
Objectives. The agents are split into two teams with configurable objectives such as tagging, chasing waypoints, and hide & seek. For the tagging objective, the goal revolves around tagging opponents to eliminate them, achieving victory when all enemy agents are tagged. For the chasing waypoint objective, random waypoints appear on the map, and points are awarded to the team whose agent reaches the waypoint first. Notably, Gigastep implements a physicality to the environment, where agents can collide with each other or the map boundaries, leading to their elimination for the rest of the episode. For the hide-and-seek objective, agents of team one attempt to catch agents of team two. Team one agents get rewards when they catch opponent agents, while team two agents are rewarded for staying alive the longest.
Observability. Moreover, Gigastep offers both full and partial observability (PO) schemes. In the PO setting, agents have a visibility cone within which they stochastically perceive opponent agents, with increasing detection probability the closer the foes are. Moreover, Gigastep incorporates a stochastic communication system whereby allied agents share information about detected enemies based on distance-related probability. This complex environment fosters intricate multi-agent strategies, providing a rich testbed for RL methodologies.
Action spaces. The action space of each agent consists of three dimensions: heading, acceleration, and climbing. With the heading dimension, the agent can change its movement direction on the xy-plane. With acceleration, the agent controls its velocity. Finally, with the climbing dimension, the agent can move along the z-axis in the physical space.
Gigastep provides both continuous and discrete action spaces. In the continuous setting, each action is clamped to the interval  $[-1, 1]$ . For the discrete action space setting we discretized the actions to  $\{-1, 0, +1\}^3$ , i.e., 27 possible actions.
Observation spaces. Gigastep supports two formats for presenting observations: feature vectors and RGB images. The feature vector format constructs a fixed-length vector of numerical features that represent the agent's own 3D state as well as the state of the  $k$  nearest other agents that are seen by the ego agent. The state data contains variables such as the location, heading angle, agent type, and health status. The feature vector observation format allows the straightforward use of multi-layer perceptron (MLP) or recurrent neural network (RNN) policies [31, 20, 27].
The RGB image observation format renders the agent's 3D view of the environment as a 2-dimensional array of pixels, each containing three color channels of 8-bit resolution. To render the  $x, y$ , and  $z$
Table 2: Selection of the scenarios included in Gigastep. The objectives are waypoint following (WF), hide and seek (HF), and tagging (TA). Build-in agent types are D (default), H (increased health), F (faster speed), R (increased beam range), and S (all stats increased).
Scenario name and Objective	Description
waypoint_3_vs_3	WF
waypoint_5_vs_3	WF	Asymmetric teams (5×D vs 3×F)
hide_and-seeking_5_vs_5	HS
hide_and-seeking_5_vs_10	HS	Asymmetric teams (5×D vs 10×D)
identical_1_vs_1	TA
identical_5_vs_5	TA
identical_5_vs_3	TA	Asymmetric teams (5×D vs 3×S)
special_5_vs_5	TA	Heterogeneous agents 1 × {D, H, F, R, S} per team
special_20_vs_20	TA	Heterogeneous agents 5 × {D, H, F, R} per team
...	...
location, the heading, altitude, and agent type in the RGB image, we took inspiration from the representation of the traffic collision avoidance system (TCAS) [22] which is used in commercial avionics to visualize multiple aircraft on a small screen, e.g., see Figure 2 for example renderings of Gigastep. In essence, TCAS uses color and shape to visualize three or more dimensional features on a two-dimensional image. We chose a default resolution of 84-by-84 pixels to be compatible with the resolution typically used with the arcade learning environment's (ALE)[3, 40] and many RL frameworks such as rlib [35]. Unlike existing MARL benchmarks, the rendering process is fully implemented in JAX and can be compiled using jax.jit and vectorized using jax.vmap.