My ultimate goal was to create a high-performance GPU implementation of Overcooked. Since Madrona is a C++ API, this requires porting game logic from Python to CUDA C++.

However, prior to attempting the port, my first step was to extract the smallest subset of code from the Python Overcooked repo needed to reproduce the MDP of the original game. This was the only code I planned to port to CUDA C++. I did not touch any code related to visualizing the game, as I could simply rely on the visualizations provided by the original implementation.

The most important part of this preparation step was replacing more complex Python data structures (like Python dictionaries) with multi-dimensional arrays, since multi-dimensional arrays have a straightforward mapping to the collection-oriented data structures required by Madrona (see step 2 below). I also replaced the Python function for computing the task’s reward with a lookup table. The result was 400 lines of clean Python code that executed the Overcooked MDP. Although my goal during this step was to distill the original Python implementation to a clean piece of code to make porting easier, this “simplified” version of the Python code actually ran 50% faster than the original. (3000 steps per second!)

An experienced GPU programmer knows that one way to port Overcooked to run N environments at once on the GPU would be to take the Python game logic from step 1, which defines how to step one environment, and directly port that logic to CUDA C++ while changing all game state variables to be arrays of length N. In this design, every GPU thread executes the logic for a single unique environment instance by reading and writing to these game state arrays at the index given by the current threadId. This approach has several drawbacks, which include:

Instead, Madrona requires simulators to adopt a software design pattern called the Entity Component System (ECS), which yields code that in most cases can be authored in “single world programming model” and is well-suited for high levels of optimization and fast parallel execution on the GPU. It also naturally extends to more complex scenarios when environments have different numbers of agents or objects, and must dynamically create and destroy objects as a result of play.

The main ideas of the ECS design pattern are simple:

The state of a game is organized as collections of game entities. For example, in Overcooked an entity might be a chef or a food item.

The state of each entity is determined by the value of the entity's components. For example, in Overcooked the chef has components representing the position, orientation, currently held object, and next chosen action.

Those familiar with databases (or dataframes) might think of a collection of entities as a logical table, where each row in the two corresponds to an entity, and each column stores the values for a certain component. Given a list of all a game’s entities and the components each of these entities requires, the Madrona runtime can store the component values for all entities with the same components, across all game instances, in a single contiguous table for efficient access and compact storage. Here’s an example of that table for Chefs:

Conveniently, the game developer doesn’t need to think about how Madrona lays out component data structures when expressing logic. They only need to write ECS systems that execute in a data-parallel manner to read and modify the values of components.

For example, here’s an example of a system that is run in a data-parallel manner once for each world. Each invocation of this function iterates over all pairs of agents, checking for collisions and updating the position and orientation of agents accordingly. Notice that this code uses accessors to get at component data without knowledge of the underlying data layout or even what world is currently executing.

 1: // Run once per agent
 2: inline void observationSystem(Engine &ctx, AgentID &id)
 3: {
 4:   WorldState &ws = ctx.singleton<WorldState>();
 5:   int32_t current_player = id.id;
 6: 
 7:   for (int loc = 0; loc < ws.size; loc++) {
 8:     LocationXObservation &obs = ctx.get<LocationXObservation>(ctx.data().locationXplayers[current_player * ws.size + loc]);
 9:     LocationData &dat = ctx.get<LocationData>(ctx.data().locations[loc]);
10:     fillLocInfo(obs, ws, dat);  // Constant-time operation
11:   }
12: 
13:   for (int i = 0; i < ws.num_players; i++) {
14:     PlayerState &ps = ctx.get<PlayerState>(ctx.data().agents[i]);
15:     int32_t pos = ps.position;
16:     LocationXObservation &obs = ctx.get<LocationXObservation>(ctx.data().locationXplayers[current_player * ws.size + pos]);
17:     fillPlayerInfo(obs, ws, current_player, ps, i);  // Constant-time operation
18:   }
19: }

Although the code I wrote above doesn’t look parallel, Madrona runs the code in parallel by executing thousands of environments concurrently on the GPU (parallelism across environments). The full logic of the Overcooked simulator is expressed as a set of system functions like the one above that need to be executed each simulation step.

By moving from python to C++, and executing many environments in parallel to utilize all the GPUs cores, this initial port yielded a substantial performance improvement: approximately two million steps per second with a thousand environments. This simple implementation required essentially no knowledge of parallel computing principles or CUDA; Madrona automatically handled GPU parallelism across multiple game environments.

Each agent is given an observation of each tile in the environment, represented as a 3D array. Specifically, at each (x, y) location, there are 26 features representing some information about the state of the world. For instance, the feature 0 indicates if the agent is present at this location, while feature 21 indicates whether there is a soup object at a location.

Constructing the observations for a single agent therefore requires iterating over each tile in the environment to construct each “feature” that the agent observes. The features for each tile could be processed independently for each agent, so the observations for the first tile could be generated independently from all other tiles. Therefore, I defined a new entity that represents an agent’s understanding of the location along with a system that generates its specific component of the observation. For instance, in the simple five by four grid with two agents, there are 20 of these “location” entities per agent, so 40 entities generate the observations in parallel.

For example, here’s a rewrite of the observations function that now defines what to do per-map location, not per world. (Notice there are no for loops over agents or locations in the code.)

 1: // Run once per location-agent pair
 2: inline void observationSystem(Engine &ctx, LocationXObservation &obs, LocationXID &id)
 3: {
 4:   WorldState &ws = ctx.singleton<WorldState>();
 5: 
 6:   int32_t loc = id.id % (ws.size);
 7:   int32_t current_player = id.id / (ws.size);
 8: 
 9:   int32_t shift = 5 * ws.num_players;
10:   LocationData &dat = ctx.get<LocationData>(ctx.data().locations[loc]);
11:   Object &obj = dat.object;
12: 
13:   fillLocInfo(obs, ws, dat);  // Constant-time operation
14: 
15:   if (dat.past_player != -1) {
16:     clearPastPlayerObs(obs, ws, dat, current_player);  // Constant-time operation
17:   }
18: 
19:   if (dat.current_player != -1) {
20:     PlayerState &ps = ctx.get<PlayerState>(ctx.data().agents[dat.current_player]);
21:     fillNewPlayerObs(obs, ws, current_player, ps);  // Constant-time operation
22:   }
23: }

Now, when the Overcooked simulator invokes the new observation system, the code is run in parallel across all location-agent pairs of all instances. For a 4x5 Overcooked map, when running a batch of 1000 environments, this is 40,000 parallel invocations of the observationSystem function.

The overcooked simulator must detect collisions between the agents in an environment. To detect if a collision has occurred, the original Overcooked-AI code loops over all pairs of agents, and checks if any pair of agents is trying to move to the same location or have “swapped” spots. If this ever happens, no agent is allowed to move to a new location for that turn, and only the agents’ orientations are updated. Otherwise, both the positions and orientations of all agents are updated. Checking collisions for all pairs can be expensive if an Overcooked environment involves many agents.

To reduce the cost of collision detection, I moved to an algorithm that takes a constant amount of time per agent by using some additional memory per tile. The algorithm is more advanced in that, unlike the previous examples, it requires the use of fine-grained synchronization (atomics) to prevent race conditions between data-parallel function invocations in the same system, but it significantly reduces the amount of computation performed as the number of agents increases.

I modified the implementation to save information about the agents’ positions and proposed positions as new components for the tiles. Specifically, each tile keeps track of the id of the current agent that is on it (if any), along with at most one of the agents that wants to move there. Collisions are resolved as follows:

First, each agent proposes a new tile to move to. Within this tile, there is an atomic component representing the id of at most one future agent that wants to move to it. We store the id of this agent using a relaxed atomic store. If multiple agents wish to move to the same tile, the relaxed store guarantees that one of them will be the final value of the tile at the end of the system. As python-like pseudocode:

 1: # run once per player
 2: def get_proposed_moves(world, player, action):
 3:     if action is INTERACT:
 4:           player.proposed_position = player.position
 5:           player.proposed_orientation = player.orientation
 6:     else:
 7:           player.proposed_position = position when moving according to action, handling world terrain
 8:           player.proposed_orientation = orientation corresponding to action
 9: 
10:     location_to_check = entity corresponding to player.proposed_position
11:     location_to_check.future_player.store_relaxed(player.id)

Next, in parallel each agent checks if a collision has occurred. If an agent sees that the “future agent” of the tile it wants to move to is not themself, then it determines a collision has occurred. We can also detect “swapping” behavior by checking the IDs of the agents at the current and future locations. If any collisions have occurred, we store that information into a singleton value for the environment using a relaxed write. As python-like pseudocode:

 1: # run once per player
 2: def check_collisions(world, player):
 3:     old_location = entity corresponding to player.position
 4:     new_location = entity corresponding to player.proposed_position
 5:     if (new_location.future_player != player.id):
 6:           world.should_update_pos.store_relaxed(False)
 7: 
 8:     # check swapping
 9:     other_player = new_location.current_player
10:     if other_player is not None and other_player is not player:
11:           if old_location.future_player == player.id:
12:               world.should_update_pos.store_relaxed(False)

Finally, each agent moves to the new tile if there is no collision, or stays on their original tile if a collision has occurred.

Agents in Overcooked have the ability to “interact” with their environment by facing a specific tile and choosing the INTERACT action. For instance, interacting with a pile of onions with an empty hand allows the agent to take an onion, and interacting with cooked soup while holding a plate fills the plate with soup.

A quirk of the original Overcooked implementation is that it processes the interactions of agents with tiles in the order of their IDs. For instance, if agent 1 has a plate, agent 2 has nothing, and both interact with an empty counter at the same time, it will seem like agent 1 passed the plate to agent 2, because agent 1 will first place down the plate and agent 2 will pick up that plate. However, if the scenario was flipped so that agent 2 has the plate to begin with, it will only seem like agent 2 placed down the plate while agent 1 did nothing. Since we want our simulation to be backwards-compatible with the original game, we need to be careful when parallelizing this process since the order of execution impacts simulation output.

The key observation is that at most four agents can interact with a single tile at once. This means that each tile can have an array of four agent IDs, and we can use atomic “fetch add” operations to have an unsorted list of agents that have interacted with that tile. Afterwards, each agent can determine how many other agents need to interact with the tile before it can perform its interaction, allowing for safe parallelization in an ECS system. This is implemented by calling four systems back-to-back, where the i’th system processes the interaction required of the i’th lowest ID of all players interacting with the tile.

When there are only two agents, the additional logic introduced by the techniques in steps 4 and 5 undermine the benefits of parallelizing the interactions over agents. When simulating 1000 environments of the Cramped Room layout, this additional parallelism only brings the steps-per-second from 2 million to 3.5 million.

The real benefit of intra-environment parallelism becomes clearer in layouts with large numbers of agents. When testing on a much larger grid with 30 agents, we find a 10 times increase, going from 7000 to 70,000 steps per second.