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

Abstract

Deep reinforcement learning (RL) is a powerful framework to train decision-making models in complex dynamical environments. However, RL can be slow as it learns through repeated interaction with a simulation of the environment. Accelerating RL requires both algorithmic and engineering innovations. In particular, there are key systems engineering bottlenecks when using RL in complex environments that feature multiple agents or highdimensional state, observation, or action spaces, for example. We present WarpDrive, a flexible, lightweight, and easy-to-use open-source RL framework that implements end-toend multi-agent RL on a single GPU (Graphics Processing Unit), building on PyCUDA and PyTorch. Using the extreme parallelization capability of GPUs, WarpDrive enables ordersof-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. Together, this allows the user to run thousands of concurrent multi-agent simulations and train on extremely large batches of experience. For example, WarpDrive yields 2.9 million environment steps/second with 2000 environments and 1000 agents (at least 100× higher throughput compared to a CPU implementation) in a benchmark Tag simulation. WarpDrive provides a lightweight Python interface and environment wrappers to simplify usage and promote flexibility and extensions. As such, WarpDrive provides a framework for building high-throughput RL systems. *: TL and SS contributed equally. TL and SS designed and developed WarpDrive. TL built the core CUDA library, including the DataManager and FunctionManagers. SS built the environment wrapper and the training pipeline. SS and TL wrote the simulation examples and unit tests. SS and TL ran RL experiments. SZ, SS, and TL drafted the paper. SZ conceived and directed the project. Code is available at https://www.github.com/ salesforce/warp-drive. We thank Alexander Trott for valuable comments on this paper.The name WarpDrive is inspired by the science fiction concept of a fictional superluminal spacecraft propulsion system. Moreover, at the time of writing, a warp is a group of 32 threads that are executing at the same time in (certain) GPUs.

EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine

“…A few recent works, e.g., Brax [7], Isaac Gym [19], and WarpDrive [16], use accelerators like GPUs and TPUs for the environment engine. Due to the highly parallel nature of the accelerators, numerous environments can be executed simultaneously.…”

Phantom -- A RL-driven multi-agent framework to model complex systems

“…MARL frameworks such as [10] and MAVA [15] are designed to enable easier and more efficient implementation of MARL algorithms. The former innovates by focusing on performance with the use of GPU and their parallelization power.…”