# WarpDrive: Extremely Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPU

arXiv: https://arxiv.org/abs/2108.13976 · HTML: https://mlsanity.com/paper/2108.13976/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 27 references that are in the map.
- First posted on arXiv 31 Aug 2021.

## Builds on (most cited first, 12 of 12)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Deep Reinforcement Learning for Robotic Manipulation with Asynchronous Off-Policy Updates](https://mlsanity.com/paper/1610.00633/) (2016)
- [Acme: A Research Framework for Distributed Reinforcement Learning](https://mlsanity.com/paper/2006.00979/) (2020)
- [Building a Foundation for Data-Driven, Interpretable, and Robust Policy Design using the AI Economist](https://mlsanity.com/paper/2108.02904/) (2021)
- [Megaverse: Simulating Embodied Agents at One Million Experiences per Second](https://mlsanity.com/paper/2107.08170/) (2021)
- [ELF: An Extensive, Lightweight and Flexible Research Platform for Real-time Strategy Games](https://mlsanity.com/paper/1707.01067/) (2017)
- [Mava: a research library for distributed multi-agent reinforcement learning in JAX](https://mlsanity.com/paper/2107.01460/) (2021)
- [The AI Economist: Optimal Economic Policy Design via Two-level Deep Reinforcement Learning](https://mlsanity.com/paper/2108.02755/) (2021)

## Built on by (most cited first, 10 of 10)

- [Heterogeneous-Agent Reinforcement Learning](https://mlsanity.com/paper/2304.09870/) (2023)
- [EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine](https://mlsanity.com/paper/2206.10558/) (2022)
- [Multi-agent Reinforcement Learning: A Comprehensive Survey](https://mlsanity.com/paper/2312.10256/) (2023)
- [Phantom -- A RL-driven multi-agent framework to model complex systems](https://mlsanity.com/paper/2210.06012/) (2022)
- [MSRL: Distributed Reinforcement Learning with Dataflow Fragments](https://mlsanity.com/paper/2210.00882/) (2022)
- [Analyzing Micro-Founded General Equilibrium Models with Many Agents using Deep Reinforcement Learning](https://mlsanity.com/paper/2201.01163/) (2022)
- [High Performance Simulation for Scalable Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2207.03945/) (2022)
- [Enabling High Data Throughput Reinforcement Learning on GPUs: A Domain Agnostic Framework for Data-Driven Scientific Research](https://mlsanity.com/paper/2408.00930/) (2024)
- [KAGE-Bench: Fast Known-Axis Visual Generalization Evaluation for Reinforcement Learning](https://mlsanity.com/paper/2601.14232/) (2026)
- [FALCON-S: Fixed-wing ground-effect Aerodynamics Simulator and Flight Control Learning Suite](https://mlsanity.com/paper/2609.06046/) (2026)

## Abstract (arXiv)

Deep reinforcement learning (RL) is a powerful framework to train decision-making models in complex environments. However, RL can be slow as it requires repeated interaction with a simulation of the environment. In particular, there are key system engineering bottlenecks when using RL in complex environments that feature multiple agents with high-dimensional state, observation, or action spaces. We present WarpDrive, a flexible, lightweight, and easy-to-use open-source RL framework that implements end-to-end deep multi-agent RL on a single GPU (Graphics Processing Unit), built on PyCUDA and PyTorch. Using the extreme parallelization capability of GPUs, WarpDrive enables orders-of-magnitude faster RL compared to common implementations that blend CPU simulations and GPU models. Our design runs simulations and the agents in each simulation in parallel. It eliminates data copying between CPU and GPU. It also uses a single simulation data store on the GPU that is safely updated in-place. WarpDrive provides a lightweight Python interface and flexible environment wrappers that are easy to use and extend. Together, this allows the user to easily run thousands of concurrent multi-agent simulations and train on extremely large batches of experience. Through extensive experiments, we verify that WarpDrive provides high-throughput and scales almost linearly to many agents and parallel environments. For example, WarpDrive yields 2.9 million environment steps/second with 2000 environments and 1000 agents (at least 100x higher throughput compared to a CPU implementation) in a benchmark Tag simulation. As such, WarpDrive is a fast and extensible multi-agent RL platform to significantly accelerate research and development.