Package: rl-warp-drive
WarpDrive: Extremely Fast End-to-End Deep Single or Multi-Agent Reinforcement Learning on a GPUWarpDrive is a flexible, lightweight, and easy-to-use open-source reinforcement learning (RL)
framework that implements end-to-end multi-agent RL on a single or multiple GPUs (Graphics Processing Unit).Using the extreme parallelization capability of GPUs, WarpDrive enables orders-of-magnitude
faster RL compared to CPU simulation + GPU model implementations. It is extremely efficient as it avoids back-and-forth data copying between the CPU and the GPU,
and runs simulations across multiple agents and multiple environment replicas in parallel.We have some main updates since its initial open source,version 1.3: provides the auto scaling tools to achieve the optimal throughput per device.version 1.4: supports the distributed asynchronous training among multiple GPU devices.version 1.6: supports the aggregation of multiple GPU blocks for one environment replica.version 2.0: supports the dual backends of both CUDA C and JIT compiled Numba.(Our Blog article)version 2.6: supports single agent environments, including Cartpole, MountainCar, AcrobotTogether, these allow the user to run thousands or even millions of concurrent simulations and train
on extremely large batches of experience, achieving at least 100x throughput over CPU-based counterparts.EnvironmentsWe include several default multi-agent environments
based on the game of "Tag" for benchmarking and testing. In the "Tag" games, taggers are trying to run after
and tag the runners. They are fairly complicated games where thread synchronization, shared memory, high-dimensional indexing for thousands of interacting agents are involved.Several more complex environments such as Covid-19 environment and climate change environment have been developed based on WarpDrive, you may see examples inReal-World Problems and Collaborations.We extend our efforts to some single agent environments includinggym.classic_control. Single-agent is a special case of multi-agent environment in WarpDrive. Since each environment only has one agent, the scalability is even higher.Below, we show multi-agent RL policies
trained for different tagger:runner speed ratios using WarpDrive.
These environments canrunatmillions of steps per second,
andtrainin just a fewhours, all on a single GPU!Throughput, Scalability and ConvergenceMulti AgentBelow, we compare the training speed on an N1 16-CPU
node versus a single A100 GPU (using WarpDrive), for the Tag environment with 100 runners and 5 taggers. With the same environment configuration and training parameters, WarpDrive on a GPU is about 10× faster. Both scenarios are with 60 environment replicas running in parallel. Using more environments on the CPU node is infeasible as data copying gets too expensive. With WarpDrive, it is possible to scale up the number of environment replicas at least 10-fold, for even faster training.Single AgentBelow, we compare the training speed on a single A100 GPU (using WarpDrive), for the Cartpole-v1 with 10, 100, 1K, and 10K environment replicas running in parallel for 3000 epochs (hyperperams are the same). You may not see elsewhere such an amazing convergence with the number of environments scaled to this large magnitude.Code StructureWarpDrive provides a CUDA (or Numba) + Python framework and quality-of-life tools, so you can quickly build fast, flexible and massively distributed multi-agent RL systems. The following figure illustrates a bottoms-up overview of the design and components of WarpDrive. The user only needs to write a CUDA or Numba step function at the CUDA environment layer, while the rest is a pure Python interface. We have step-by-step tutorials for you to master the workflow.Python InterfaceWarpDrive provides tools to build and train
multi-agent RL systems quickly with just a few lines of code.
Here is a short example to train tagger and runner agents:# Create a wrapped environment object via the EnvWrapper# Ensure that env_backend is set to 'pycuda' or 'numba' (in order to run on the GPU)env_wrapper=EnvWrapper(TagContinuous(**run_config["env"]),num_envs=run_config["trainer"]["num_envs"],env_backend="pycuda")# Agents can share policy models: this dictionary maps policy model names to agent ids.policy_tag_to_agent_id_map={"tagger":list(env_wrapper.env.taggers),"runner":list(env_wrapper.env.runners),}# Create the trainer objecttrainer=Trainer(env_wrapper=env_wrapper,config=run_config,policy_tag_to_agent_id_map=policy_tag_to_agent_id_map,)# Perform training!trainer.train()Papers and Citing WarpDriveOur paper published atJournal of Machine Learning Research(JMLR)https://jmlr.org/papers/v23/22-0185.html. You can also find more details in our white paper:https://arxiv.org/abs/2108.13976.If you're using WarpDrive in your research or applications, please cite using this BibTeX:@article{JMLR:v23:22-0185,
  author  = {Tian Lan and Sunil Srinivasa and Huan Wang and Stephan Zheng},
  title   = {WarpDrive: Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPU},
  journal = {Journal of Machine Learning Research},
  year    = {2022},
  volume  = {23},
  number  = {316},
  pages   = {1--6},
  url     = {http://jmlr.org/papers/v23/22-0185.html}
}

@misc{lan2021warpdrive,
  title={WarpDrive: Extremely Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPU}, 
  author={Tian Lan and Sunil Srinivasa and Huan Wang and Caiming Xiong and Silvio Savarese and Stephan Zheng},
  year={2021},
  eprint={2108.13976},
  archivePrefix={arXiv},
  primaryClass={cs.LG}
}Tutorials and Quick StartTutorialsFamiliarize yourself with WarpDrive by running these tutorials on Colab orNGC container!WarpDrive basics(Introdunction and PyCUDA)WarpDrive basics(Numba)WarpDrive sampler(PyCUDA)WarpDrive sampler(Numba)WarpDrive resetter and loggerCreate custom environments (PyCUDA)Create custom environments (Numba)Training with WarpDriveScaling Up training with WarpDriveTraining with WarpDrive + Pytorch LightningYou may also run thesetutorialslocally, but you will need a GPU machine with nvcc compiler installed
and a compatible Nvidia GPU driver. You will also needJupyter.
Seehttps://jupyter.readthedocs.io/en/latest/install.htmlfor installation instructionsExample Training ScriptWe provide some example scripts for you to quickly start the end-to-end training.
For example, if you want to train tag_continuous environment (10 taggers and 100 runners) with 2 GPUs and CUDA C backendpython example_training_script_pycuda.py -e tag_continuous -n 2or switch to JIT compiled Numba backend with 1 GPUpython example_training_script_numba.py -e tag_continuousYou can find full reference documentationhere.Real World Problems and CollaborationsAI Economist Covid Environment with WarpDrive: We train two-level multi-agent economic simulations usingAI-Economist Foundationand train it using WarpDrive. We specifically consider the COVID-19 and economy simulation in this example.Climate Change Cooperation Competitioncollaborated withMila. We provide the base version of the RICE (regional integrated climate environment)simulation environment.Pytorch Lightning Trainer with WarpDrive: We provide atutorial exampleand ablog articleof a multi-agent reinforcement learning training loop with WarpDrive andPytorch Lightning.NVIDIA NGC Catalog and Quick Deployment to VertexAI: WarpDrive image is hosted byNGC Catalog. The NGC catalog "hosts containers for the top AI and data science software, tuned, tested and optimized by NVIDIA". Our tutorials also enable the quick deployment to VertexAI supported by the NGC.Installation InstructionsTo get started, you'll need to havePython 3.7+and thenvcccompiler installed
with a compatible Nvidia GPU CUDA driver.CUDA (which includes nvcc) can be installed by following Nvidia's instructions here:https://developer.nvidia.com/cuda-downloads.Docker ImageV100 GPU: You can refer to theexample Dockerfileto configure your system.A100 GPU: Our latest image is published and maintained by NVIDIA NGC. We recommend you download the latest image fromNGC catalog.If you want to build your customized environment, we suggest you visitNvidia Docker Hubto download the CUDA and cuDNN images compatible with your system.
You should be able to use the command line utility to monitor the NVIDIA GPU devices in your system:nvidia-smiand see something like this+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   37C    P0    32W / 300W |      0MiB / 16160MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+In this snapshot, you can see we are using a Tesla V100 GPU and CUDA version 11.0.Installing using PipYou can install WarpDrive using the Python package manager:pip install rl_warp_driveInstalling from SourceClone this repository to your machine:git clone https://www.github.com/salesforce/warp-driveOptional, but recommended for first tries:Create a new conda environment (named "warp_drive" below) and activate it:conda create --name warp_drive python=3.7 --yes
conda activate warp_driveInstall as an editable Python package:cd warp_drive
pip install -e .Testing your InstallationYou can call directly from Python command to test all modules and the end-to-end training workflow.python warp_drive/utils/unittests/run_unittests_pycuda.py
python warp_drive/utils/unittests/run_unittests_numba.py
python warp_drive/utils/unittests/run_trainer_tests.pyLearn MoreFor more information, please check out ourblog,white paper, and codedocumentation.If you're interested in extending this framework, or have questions, join the
AI Economist Slack channel using thisinvite link.