tutorials/tutorial-2.b-warp_drive_sampler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) 2021, salesforce.com, inc. \\\n",
    "All rights reserved. \\\n",
    "SPDX-License-Identifier: BSD-3-Clause \\\n",
    "For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Try this notebook on [Colab](http://colab.research.google.com/github/salesforce/warp-drive/blob/master/tutorials/tutorial-2.b-warp_drive_sampler.ipynb)!**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ⚠️ PLEASE NOTE:\n",
    "This notebook runs on a GPU runtime.\\\n",
    "If running on Colab, choose Runtime > Change runtime type from the menu, then select `GPU` in the 'Hardware accelerator' dropdown menu."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "assert torch.cuda.device_count() > 0, \"This notebook needs a GPU to run!\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Welcome to WarpDrive!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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 for [introduction and pycuda](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-1.a-warp_drive_basics.ipynb), and [numba](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-1.b-warp_drive_basics.ipynb) on WarpDrive basics.\n",
    "\n",
    "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. \n",
    "\n",
    "Notably:\n",
    "\n",
    "1. It reads the distribution on the GPU through Pytorch and samples actions exclusively at the GPU. There is no data transfer. \n",
    "2. 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. \n",
    "3. It runs much faster than most GPU samplers. For example, it is significantly faster than Pytorch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can install the warp_drive package using\n",
    "\n",
    "- the pip package manager, OR\n",
    "- by cloning the warp_drive package and installing the requirements.\n",
    "\n",
    "We will install the latest version of WarpDrive using the pip package manager."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install -U rl_warp_drive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from warp_drive.managers.numba_managers.numba_function_manager import NumbaFunctionManager, NumbaSampler\n",
    "from warp_drive.managers.numba_managers.numba_data_manager import NumbaDataManager\n",
    "from warp_drive.utils.constants import Constants\n",
    "from warp_drive.utils.data_feed import DataFeed\n",
    "from warp_drive.utils.common import get_project_root\n",
    "\n",
    "_ACTIONS = Constants.ACTIONS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set logger level e.g., DEBUG, INFO, WARNING, ERROR\n",
    "import logging\n",
    "\n",
    "logging.getLogger().setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize NumbaSampler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first initialize the **NumbaDataManager** and **NumbaFunctionManager**. To illustrate the sampler, we first load a \"test_build_sampler.py\". Note that these low-level managers and modules will be hidden and called automatically by WarpDrive in any end-to-end training and simulation. In this and the next tutorials, we want to show how a few fundamental modules work and their performance, that is why some low-level APIs such as \"import_numba_from_source_code()\" are called. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda_data_manager = NumbaDataManager(num_agents=5, episode_length=10, num_envs=2)\n",
    "cuda_function_manager = NumbaFunctionManager(\n",
    "    num_agents=cuda_data_manager.meta_info(\"n_agents\"),\n",
    "    num_envs=cuda_data_manager.meta_info(\"n_envs\"),\n",
    ")\n",
    "\n",
    "_NUMBA_FILEPATH = f\"warp_drive.numba_includes\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This `_NUMBA_FILEPATH` includes several Numba core services provided by WarpDrive.  Since Numba uses JIT compiler, we do not need to compile an executable before loading it as we did in the PyCUDA mode, however, in most cases, the source code still needs some environment configuration file to populate a few global environment constants and settings to the source code, for example, number of environment replicas. WarpDrive provides such template configure file and automatically update placeholders according to the current environment. (import_numba_env_config() is a low-level API, user will not need call those internal APIs directly for any WarpDrive end-to-end simulation and training.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda_function_manager.import_numba_env_config(template_header_file=\"template_env_config.txt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use **NumbaFunctionManager** to load the source code. In this demo, we use `test_build.py` which collects those core services and includes the backend source code for `NumbaSampleController`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda_function_manager.import_numba_from_source_code(f\"{_NUMBA_FILEPATH}.test_build\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we initialize **NumbaSampler** and assign the random seed. `NumbaSampler` keeps independent randomness across all threads and blocks. Notice that `NumbaSampler` requires `NumbaFunctionManager` because `NumbaFunctionManager` manages all the CUDA function pointers including to the sampler. Also notice this test binary uses 2 environment replicas and 5 agents. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda_sampler = NumbaSampler(function_manager=cuda_function_manager)\n",
    "cuda_sampler.init_random(seed=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Actions Placeholder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we feed the **actions_a** placeholder into the GPU. It has the shape `(n_envs=2, n_agents=5)` as expected. Also we make it accessible by Pytorch, because during RL training, actions will be fed into the Pytorch trainer directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_feed = DataFeed()\n",
    "data_feed.add_data(name=f\"{_ACTIONS}_a\", data=[[[0], [0], [0], [0], [0]], [[0], [0], [0], [0], [0]]])\n",
    "cuda_data_manager.push_data_to_device(data_feed, torch_accessible=True)\n",
    "assert cuda_data_manager.is_data_on_device_via_torch(f\"{_ACTIONS}_a\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Action Sampled Distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define an action **distribution** here. During training, this distribution would be provided by the policy model implemented in Pytorch. The distribution has the shape `(n_envs, n_agents, **n_actions**)`. The last dimension `n_actions` defines the size of the action space for a particular *discrete* action. For example, if we have up, down, left, right and no-ops, `n_actions=5`.\n",
    "\n",
    "**n_actions** needs to be registered by the sampler so the sampler is able to pre-allocate a global memory space in GPU to speed up action sampling. This can be done by calling `sampler.register_actions()`.\n",
    "\n",
    "In this tutorial, we check if our sampled action distribution follows the given distribution. For example, the distribution [0.333, 0.333, 0.333] below suggests the 1st agent has 3 possible actions and each of them have equal probability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda_sampler.register_actions(\n",
    "    cuda_data_manager, action_name=f\"{_ACTIONS}_a\", num_actions=3\n",
    ")\n",
    "\n",
    "distribution = np.array(\n",
    "    [\n",
    "        [\n",
    "            [0.333, 0.333, 0.333],\n",
    "            [0.2, 0.5, 0.3],\n",
    "            [0.95, 0.02, 0.03],\n",
    "            [0.02, 0.95, 0.03],\n",
    "            [0.02, 0.03, 0.95],\n",
    "        ],\n",
    "        [\n",
    "            [0.1, 0.7, 0.2],\n",
    "            [0.7, 0.2, 0.1],\n",
    "            [0.5, 0.5, 0.0],\n",
    "            [0.0, 0.5, 0.5],\n",
    "            [0.5, 0.0, 0.5],\n",
    "        ],\n",
    "    ]\n",
    ")\n",
    "distribution = torch.from_numpy(distribution).float().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run 10000 times to collect statistics\n",
    "actions_batch = torch.from_numpy(np.empty((10000, 2, 5), dtype=np.int32)).cuda()\n",
    "\n",
    "for i in range(10000):\n",
    "    cuda_sampler.sample(cuda_data_manager, distribution, action_name=f\"{_ACTIONS}_a\")\n",
    "    actions_batch[i] = cuda_data_manager.data_on_device_via_torch(f\"{_ACTIONS}_a\")[:, :, 0]\n",
    "actions_batch_host = actions_batch.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actions_env_0 = actions_batch_host[:, 0]\n",
    "actions_env_1 = actions_batch_host[:, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    \"Sampled actions distribution versus the given distribution (in bracket) for env 0: \\n\"\n",
    ")\n",
    "for agent_id in range(5):\n",
    "    print(\n",
    "        f\"Sampled action distribution for agent_id: {agent_id}:\\n\"\n",
    "        f\"{(actions_env_0[:, agent_id] == 0).sum() / 10000.0}({distribution[0, agent_id, 0]}), \\n\"\n",
    "        f\"{(actions_env_0[:, agent_id] == 1).sum() / 10000.0}({distribution[0, agent_id, 1]}), \\n\"\n",
    "        f\"{(actions_env_0[:, agent_id] == 2).sum() / 10000.0}({distribution[0, agent_id, 2]})  \\n\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    \"Sampled actions distribution versus the given distribution (in bracket) for env 1: \"\n",
    ")\n",
    "\n",
    "for agent_id in range(5):\n",
    "    print(\n",
    "        f\"Sampled action distribution for agent_id: {agent_id}:\\n\"\n",
    "        f\"{(actions_env_1[:, agent_id] == 0).sum() / 10000.0}({distribution[1, agent_id, 0]}), \\n\"\n",
    "        f\"{(actions_env_1[:, agent_id] == 1).sum() / 10000.0}({distribution[1, agent_id, 1]}), \\n\"\n",
    "        f\"{(actions_env_1[:, agent_id] == 2).sum() / 10000.0}({distribution[1, agent_id, 2]})  \\n\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Action Randomness Across Threads"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another important validation is whether the sampler provides independent randomness across different agents and environment replicas. Given the same policy model for all the agents and environment replicas, we can check if the sampled actions are independently distributed. \n",
    "\n",
    "Here, we assign all agents across all envs the same distribution [0.25, 0.25, 0.25, 0.25]. It is equivalent to an uniform action distribution among all actions [0,1,2,3], across 5 agents and 2 envs. Then we check the standard deviation across the agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_feed = DataFeed()\n",
    "data_feed.add_data(name=f\"{_ACTIONS}_b\", data=[[[0], [0], [0], [0], [0]], [[0], [0], [0], [0], [0]]])\n",
    "cuda_data_manager.push_data_to_device(data_feed, torch_accessible=True)\n",
    "assert cuda_data_manager.is_data_on_device_via_torch(f\"{_ACTIONS}_b\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda_sampler.register_actions(\n",
    "    cuda_data_manager, action_name=f\"{_ACTIONS}_b\", num_actions=4\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distribution = np.array(\n",
    "    [\n",
    "        [\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "        ],\n",
    "        [\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "            [0.25, 0.25, 0.25, 0.25],\n",
    "        ],\n",
    "    ]\n",
    ")\n",
    "distribution = torch.from_numpy(distribution).float().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run 10000 times to collect statistics.\n",
    "actions_batch = torch.from_numpy(np.empty((10000, 2, 5), dtype=np.int32)).cuda()\n",
    "\n",
    "for i in range(10000):\n",
    "    cuda_sampler.sample(cuda_data_manager, distribution, action_name=f\"{_ACTIONS}_b\")\n",
    "    actions_batch[i] = cuda_data_manager.data_on_device_via_torch(f\"{_ACTIONS}_b\")[:, :, 0]\n",
    "actions_batch_host = actions_batch.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actions_batch_host"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actions_batch_host.std(axis=2).mean(axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check the independence of randomness among all threads, we can compare it with a Numpy implementation. Here we use `numpy.choice(4, 5)` to repeat the same process for an uniform action distribution among all actions [0,1,2,3], 5 agents and 2 envs. We should see that the variation of Numpy output is very close to our sampler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actions_batch_numpy = np.empty((10000, 2, 5), dtype=np.int32)\n",
    "for i in range(10000):\n",
    "    actions_batch_numpy[i, 0, :] = np.random.choice(4, 5)\n",
    "    actions_batch_numpy[i, 1, :] = np.random.choice(4, 5)\n",
    "actions_batch_numpy.std(axis=2).mean(axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running Speed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The total time for sampling includes receiving a new distribution and using this to sample. For Numba, it also includes the JIT time. In fact, JIT becomes the main time consumer in this simple example.\n",
    "Comparing our sampler with [torch.Categorical sampler](https://pytorch.org/docs/stable/distributions.html), \n",
    "we are almost the same. \n",
    "\n",
    "*Note: our sampler runs in parallel across threads, so this speed-up is almost constant when scaling up the number of agents or environment replicas, i.e., increasing the number of used threads.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.distributions import Categorical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distribution = np.array(\n",
    "    [\n",
    "        [\n",
    "            [0.333, 0.333, 0.333],\n",
    "            [0.2, 0.5, 0.3],\n",
    "            [0.95, 0.02, 0.03],\n",
    "            [0.02, 0.95, 0.03],\n",
    "            [0.02, 0.03, 0.95],\n",
    "        ],\n",
    "        [\n",
    "            [0.1, 0.7, 0.2],\n",
    "            [0.7, 0.2, 0.1],\n",
    "            [0.5, 0.5, 0.0],\n",
    "            [0.0, 0.5, 0.5],\n",
    "            [0.5, 0.0, 0.5],\n",
    "        ],\n",
    "    ]\n",
    ")\n",
    "distribution = torch.from_numpy(distribution).float().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_event = torch.cuda.Event(enable_timing=True)\n",
    "end_event = torch.cuda.Event(enable_timing=True)\n",
    "\n",
    "start_event.record()\n",
    "for _ in range(1000):\n",
    "    cuda_sampler.sample(cuda_data_manager, distribution, action_name=f\"{_ACTIONS}_a\")\n",
    "end_event.record()\n",
    "torch.cuda.synchronize()\n",
    "print(f\"time elapsed: {start_event.elapsed_time(end_event)} ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_event = torch.cuda.Event(enable_timing=True)\n",
    "end_event = torch.cuda.Event(enable_timing=True)\n",
    "\n",
    "start_event.record()\n",
    "for _ in range(1000):\n",
    "    Categorical(distribution).sample()\n",
    "end_event.record()\n",
    "torch.cuda.synchronize()\n",
    "print(f\"time elapsed: {start_event.elapsed_time(end_event)} ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learn More and Explore our Tutorials!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we suggest you check out our advanced [tutorial](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-3-warp_drive_reset_and_log.ipynb) on WarpDrive's reset and log controller!\n",
    "\n",
    "For your reference, all our tutorials are here:\n",
    "1. [WarpDrive basics(intro and pycuda)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-1.a-warp_drive_basics.ipynb)\n",
    "2. [WarpDrive basics(numba)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-1.b-warp_drive_basics.ipynb)\n",
    "3. [WarpDrive sampler(pycuda)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-2.a-warp_drive_sampler.ipynb)\n",
    "4. [WarpDrive sampler(numba)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-2.b-warp_drive_sampler.ipynb)\n",
    "5. [WarpDrive resetter and logger](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-3-warp_drive_reset_and_log.ipynb)\n",
    "6. [Create custom environments (pycuda)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-4.a-create_custom_environments_pycuda.md)\n",
    "7. [Create custom environments (numba)](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-4.b-create_custom_environments_numba.md)\n",
    "8. [Training with WarpDrive](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-5-training_with_warp_drive.ipynb)\n",
    "9. [Scaling Up training with WarpDrive](https://www.github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-6-scaling_up_training_with_warp_drive.md)\n",
    "10. [Training with WarpDrive + Pytorch Lightning](https://github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-7-training_with_warp_drive_and_pytorch_lightning.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}