NGC | Catalog
CatalogResourcesWarpDrive Sampler for Numba

WarpDrive Sampler for Numba

For downloads and more information, please view on a desktop device.
Logo for WarpDrive Sampler for Numba


Tutorial covering the basics of WarpDrive Sampler using Numba



Latest Version



April 4, 2023

Compressed Size

21.06 KB

This is the second tutorial on WarpDrive, a PyCUDA-based framework for extremely parallelized multi-agent reinforcement learning (RL) on a single graphics processing unit (GPU). At this stage, we assume you have read our first tutorial on WarpDrive basics.

In this tutorial, we describe CUDASampler, a lightweight and fast action sampler based on the policy distribution across several RL agents and environment replicas. CUDASampler utilizes the GPU to parallelize operations to efficiently sample a large number of actions in parallel.


  • It reads the distribution on the GPU through Pytorch and samples actions exclusively at the GPU. There is no data transfer.
  • It maximizes parallelism down to the individual thread level, i.e., each agent at each environment has its own random seed and independent random sampling process.