Madrona v0.1: A Game Engine for Training Agents

Posted on Sunday, August 6th 2023 by Brennan Shacklett

After about a year of work, I'm happy to announce the initial release of the Madrona Engine. Madrona is a new type of game engine designed from the ground up for the unique workload of training agents using algorithms like reinforcement learning. With this release, we hope to provide a useful platform for creating high-performance learning environments across a broad range of exciting applications such as learning complex game strategies, robotics / embodied AI, and game development.

What makes Madrona different from existing game engines?

Learning environments built on Madrona are able to simulate batches of thousands of environments simultaneously by default, enabling high-throughput experience generation for learning. This is in contrast to existing game engines that are designed to execute a single game instance at real-time rates. By treating parallel batch execution of many environments as a first class notion in the engine, Madrona is able to map this parallel workload entirely onto the GPU and achieve order of magnitude speedups over CPU execution.

Machine Learning Impact

For machine learning researchers, the massive speedups enabled by Madrona's architecture result in very tangible changes to learning workflows. For example, my colleague Bidipta Sarkar ported the Overcooked AI environment to Madrona and achieved a 60x speedup to end-to-end training. Concretely, his training runs dropped from over an hour to under a minute! In fact, you can use this Colab notebook to train Overcooked agents on a free GPU in just a couple minutes yourself. This level of speedup significantly improves research iteration time and open up new algorithmic possibilities. I highly recommend reading Bidipta's blog post for a detailed account of the impact on his research as well as the nitty gritty technical details of porting Overcooked AI to Madrona and optimizing it.

Engine Architecture Overview

The core technical idea behind Madrona is the realization that the Entity Component System (ECS) architecture from the game industry is a great fit for batch GPU execution of many learning environments. Specifically, the overall structure and separation of code and data provided by an ECS allow Madrona to efficiently map custom learning environment logic onto the GPU by maintaining coherent execution and access patterns while also managing the complexity of linking many diverse subsystems together efficiently. If you're not familiar with the ECS architecture, I highly recommend Sander Mertens' excellent ECS FAQ, as well as this GDC talk from one of the developers of Overwatch.

Expect the next post on this blog to be a deeper dive into my reasoning behind choosing the ECS architecture for Madrona and why I think the "game engine" abstraction is extremely useful for building learning environments.

Next Steps

Hopefully this short post has piqued your interest in learning more about the engine and perhaps playing around with it. If so, the following resources are good starting points:

If you're interested in using the Madrona engine to build a new learning environment, or just want to experiment on the ML side with the environments we've already built, feel free to reach out with ideas / feature requests, or for help with any issues you encounter. You can contact me at my (extremely inactive) twitter account: @shacklettbp, or via my stanford email: shown here. Remember that Madrona is a research project, so expect rough edges and missing features, especially in this initial release. If you find bugs or issues with the engine or any of the existing simulators, please create a GitHub issue on the appropriate repo!