Madrona is a research game engine designed specifically for creating learning environments that execute with extremely high throughput on a single GPU (up to millions of aggregate environment steps per second). It achieves this by running thousands of independent instances of an environment concurrently on the GPU. Madrona's goal is to enable orders of magnitude speedups for AI agent training (and other applications needing large amounts of simulated experience) by making it easier for researchers to create new high-performance environments for a wide range of tasks.

Demo learning environments that have been ported to Madrona:

• OpenAI's Hide and Seek (2 million steps per second)
• Overcooked-AI (40 million steps/second)
• Hanabi (20 million steps/sec)
• Madrona Escape Room

* Timings measured on a RTX 4090 GPU.

For more information, please see the following resources:

• Madrona GitHub repo
• Our short blog post introducing the engine
• Bidipta Sarkar's post on building an Overcooked-AI batch simulator using Madrona
• SIGGRAPH 2023 paper about the Madrona engine
• Madrona FAQ

Please refer to our technical paper for an in-depth description of core Madrona engine implementation details and for an analysis of the engine's performance.

Abstract

Training AI agents to perform complex tasks in simulated worlds requires millions to billions of steps of experience. To achieve high performance, today's fastest simulators for training AI agents adopt the idea of batch simulation: using a single simulation engine to simultaneously step many environments in parallel. We introduce a framework for productively authoring novel training environments (including custom logic for environment generation, environment time stepping, and generating agent observations and rewards) that execute as high-performance, GPU-accelerated batched simulators. Our key observation is that the Entity Component System (ECS) design pattern, popular for expressing CPU-side game logic today, is also well-suited for providing the structure needed for high-performance batched simulators. We contribute the first fully-GPU accelerated ECS implementation that natively supports batch environment simulation. We demonstrate how ECS abstractions impose structure on a training environment's logic and state that allows the system to efficiently manage state, amortize work, and identify GPU-friendly coherent parallel computations within and across different environments. We implement several learning environments in this framework, and demonstrate GPU speedups of two to three orders of magnitude over open source CPU baselines and 5x — 33x over strong baselines running on a 32-thread CPU. An implementation of the OpenAI hide and seek 3D environment written in our framework, which performs rigid body physics and ray tracing in each simulator step, achieves over 1.9 million environment steps per second on a single GPU.

Citation

@article{shacklett23madrona,
    title   = {An Extensible, Data-Oriented Architecture for
               High-Performance, Many-World Simulation},
    author  = {Brennan Shacklett and Luc Guy Rosenzweig and
               Zhiqiang Xie and Bidipta Sarkar and Andrew Szot and
               Erik Wijmans and Vladlen Koltun and Dhruv Batra and 
               Kayvon Fatahalian},
    journal = {ACM Trans. Graph.},
    volume  = {42},
    number  = {4},
    year    = {2023}
}

What's the best way to contact Madrona's developers?

Please contact us by opening a GitHub issue or directly reaching out via email to the first author of the SIGGRAPH 2023 Madrona publication.

What environments / tasks does Madrona support?

Madrona is not an RL environment simulator. It is a game engine / framework that makes it easier for developers (like RL researchers) to create their own custom learning environment simulators for their own custom tasks that achieve high throughput by running on GPUs. For example, instead of creating a new environment simulator using an existing game engine like Unity or Unreal, you might consider creating it using Madrona to achieve significantly higher experience generation throughput.

We try to maintain an updated list of environment simulators built with Madrona on the Madrona GitHub README. If you've made your own environment using Madrona and want to share, let us know and we'll add it to the list.

My learning task requires simulation of complex dynamical systems in contact-rich settings. (I currently use MuJoCo or IsaacGym.) Is Madrona a good fit for my work?

At the moment, probably not. Madrona currently provides a XPBD-based rigid body physics library for basic 3D collision and contact support. (This library is used by our 3D Hide and Seek and Escape Room example environments.) However, tasks that require advanced physics implementations (such as robotics manipulation tasks) are likely better off using other systems at this time. Hopefully in the future we can find ways to integrate libraries like MuJoCo or Isaac Gym into Madrona by allowing them to communicate with the rest of the game engine using ECS structures. This would enable the development of new environments that are unconstrained in their custom game logic and also benefit from world-class physics simulation.

Does Madrona have rendering support? Can I use Madrona for training "pixels to actions" agents?

Madrona currently has rendering support for visualization purposes only. While Madrona can simulate thousands of environments concurrently, the visualizer can only view a single environment at a time, and this output isn't exposed to PyTorch. Adding high-performance batch rendering support using ideas from our prior bps3D project is something we hope to do in the near future. Let us know if you have an immediate use case for it. The learning tasks we used as initial drivers of Madrona development involved agent policies that took internal game state to actions, so fast rendering has been lower priority for us.

I have an existing environment simulator written in Python. What is the easiest way to port it to Madrona?

At this time Madrona requires game logic to be written in C++, so using Madrona will require porting this logic to C++ and re-organizing your game state to leverage the entity component system architecture. Bidipta Sarkar has a great post on porting the Overcooked-AI environment to Madrona. That post describes the process of extracting the core simulation logic from a large Python code base and re-implementing it using Madrona's C++ APIs. A trivial port of Overcooked-AI was done in only a few days, and it increased simulator performance by over 1000X. The Madrona github repo also contains links to sample projects that can be used as starter code for creating new environments.

Can I use Madrona while keeping parts of my simulator (e.g., reward functions) in Python?

Yes. While Madrona currently requires the core custom environment simulation logic for a task to be written in C++, logic that executes at the end of the simulation step like reward functions can be implemented using PyTorch tensor operations. This is possible by exporting any of the internal simulation state (ECS components) needed as input to the reward functions as PyTorch tensors, and then computing rewards in PyTorch after Madrona finishes the current step. Madrona provides built-in support for exporting ECS components to PyTorch, making this relatively straightforward.

How is Madrona related to projects like NVIDIA Warp or Numba?

Projects like NVIDIA Warp and Numba (and others) aim to make GPU programming easier by bringing high-level languages (Python) to the GPU. These frameworks could provide a high-productivity way to author individual ECS systems for use within the Madrona engine. A useful future project would be to engineer the appropriate support for generating Madrona-compatible CUDA kernels from ECS systems written in these languages. (Contact us if you're interested in contributing!)

Doing so would allow developers to build complex batch simulators while realizing the productivity benefits of expressing environment logic in high-level languages and the performance benefits of using ECS APIs for efficiently managing state (custom components, creating and deleting entities) and gluing together different systems.

Can I use Madrona with a different learning framework than PyTorch?

Madrona exports simulation state to learning frameworks using nanobind's dlpack integration that supports most common frameworks. Currently, only PyTorch entry points for this functionality are exposed, because that is the framework we use and test. If you're interested in using Madrona with a different learning framework, please open a GitHub issue with a link to training code in your preferred framework for testing purposes and we can add the core engine support.