tutorials/tutorial-7-training_with_warp_drive_and_pytorch_lightning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fast Multi-agent Reinforcement Learning on a GPU using WarpDrive and Pytorch Lightning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try this notebook on [Colab](http://colab.research.google.com/github/salesforce/warp-drive/blob/master/tutorials/tutorial-7-training_with_warp_drive_and_pytorch_lightning.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": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial shows how [WarpDrive](https://github.com/salesforce/warp-drive) can be used together with [PyTorch Lightning](https://github.com/PyTorchLightning/pytorch-lightning).\n",
    "\n",
    "WarpDrive is a flexible, lightweight, and easy-to-use RL framework that implements end-to-end deep multi-agent RL on a single GPU (Graphics Processing Unit). Using the extreme parallelization capability of GPUs, it enables [orders-of-magnitude faster RL](https://arxiv.org/abs/2108.13976) compared to common implementations that blend CPU simulations and GPU models. WarpDrive is extremely efficient as it runs simulations across multiple agents and multiple environment replicas in parallel and completely eliminates the back-and-forth data copying between the CPU and the GPU.\n",
    "\n",
    "Pytorch Lightning is a machine learning framework which [greatly reduces trainer boilerplate code](https://www.pytorchlightning.ai/), and improves training modularity and flexibility. It abstracts away most of the engineering pieces of code, so users can focus on research and building models, and iterate on experiments really fast. Pytorch Lightning also provides support for easily running the model on any hardware, performing distributed training, model checkpointing, performance profiling, logging and visualization.\n",
    "\n",
    "Below, we demonstrate how to use WarpDrive and PytorchLightning together to train a game of [Tag](https://github.com/salesforce/warp-drive/blob/master/example_envs/tag_continuous/tag_continuous.py) where multiple *tagger* agents are trying to run after and tag multiple other *runner* agents. As such, the Warpdrive framework comprises several utility functions that help easily implement any (OpenAI-)*gym-style* RL environment, and furthermore, provides quality-of-life tools to train it end-to-end using just a few lines of code. You may familiarize yourself with WarpDrive with the help of these [tutorials](https://github.com/salesforce/warp-drive/tree/master/tutorials).\n",
    "\n",
    "We invite everyone to **contribute to WarpDrive**, including adding new multi-agent environments, proposing new features and reporting issues on our open source [repository](https://github.com/salesforce/warp-drive)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook requires the `rl-warp-drive` as well as the `pytorch-lightning` packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -U rl_warp_drive\n",
    "\n",
    "! pip install 'pytorch_lightning>=1.8'\n",
    "\n",
    "# Also, install ffmpeg for visualizing animations\n",
    "! apt install ffmpeg --yes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "from example_envs.tag_continuous.tag_continuous import TagContinuous\n",
    "from example_envs.tag_continuous.generate_rollout_animation import (\n",
    "    generate_tag_env_rollout_animation,\n",
    ")\n",
    "from IPython.display import HTML\n",
    "from pytorch_lightning import Trainer\n",
    "from warp_drive.env_wrapper import EnvWrapper\n",
    "from warp_drive.training.pytorch_lightning import (\n",
    "    CUDACallback,\n",
    "    PerfStatsCallback,\n",
    "    WarpDriveModule,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytorch_lightning\n",
    "\n",
    "assert pytorch_lightning.__version__ >= \"1.8.0\", \"lightning versions < 1.8 have a few different callbacks\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set logger level e.g., DEBUG, INFO, WARNING, ERROR.\n",
    "logging.getLogger().setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Specify a set of run configurations for your experiments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The run configuration is a dictionary comprising the environment parameters, the trainer and the policy network settings, as well as configurations for saving.\n",
    "\n",
    "For our experiment, we consider an environment wherein $5$ taggers and $100$ runners play the game of [Tag](https://github.com/salesforce/warp-drive/blob/master/example_envs/tag_continuous/tag_continuous.py) on a $20 \\times 20$ plane. The game lasts $200$ timesteps. Each agent chooses it's own acceleration and turn actions at every timestep, and we use mechanics to determine how the agents move over the grid. When a tagger gets close to a runner, the runner is tagged, and is eliminated from the game. For the configuration below, the runners and taggers have the same unit skill levels, or top speeds.\n",
    "\n",
    "We train the agents using $50$ environments or simulations running in parallel. With WarpDrive, each simulation runs on sepate GPU blocks.\n",
    "\n",
    "There are two separate policy networks used for the tagger and runner agents. Each network is a fully-connected model with two layers each of $256$ dimensions. We use the Advantage Actor Critic (A2C) algorithm for training. WarpDrive also currently provides the option to use the Proximal Policy Optimization (PPO) algorithm instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_config = dict(\n",
    "    name=\"tag_continuous\",\n",
    "    # Environment settings.\n",
    "    env=dict(\n",
    "        num_taggers=5,  # number of taggers in the environment\n",
    "        num_runners=100,  # number of runners in the environment\n",
    "        grid_length=20.0,  # length of the (square) grid on which the game is played\n",
    "        episode_length=200,  # episode length in timesteps\n",
    "        max_acceleration=0.1,  # maximum acceleration\n",
    "        min_acceleration=-0.1,  # minimum acceleration\n",
    "        max_turn=2.35,  # 3*pi/4 radians\n",
    "        min_turn=-2.35,  # -3*pi/4 radians\n",
    "        num_acceleration_levels=10,  # number of discretized accelerate actions\n",
    "        num_turn_levels=10,  # number of discretized turn actions\n",
    "        skill_level_tagger=1.0,  # skill level for the tagger\n",
    "        skill_level_runner=1.0,  # skill level for the runner\n",
    "        use_full_observation=False,  # each agent only sees full or partial information\n",
    "        runner_exits_game_after_tagged=True,  # flag to indicate if a runner stays in the game after getting tagged\n",
    "        num_other_agents_observed=10,  # number of other agents each agent can see\n",
    "        tag_reward_for_tagger=10.0,  # positive reward for the tagger upon tagging a runner\n",
    "        tag_penalty_for_runner=-10.0,  # negative reward for the runner upon getting tagged\n",
    "        end_of_game_reward_for_runner=1.0,  # reward at the end of the game for a runner that isn't tagged\n",
    "        tagging_distance=0.02,  # margin between a tagger and runner to consider the runner as 'tagged'.\n",
    "    ),\n",
    "    # Trainer settings.\n",
    "    trainer=dict(\n",
    "        num_envs=50,  # number of environment replicas (number of GPU blocks used)\n",
    "        train_batch_size=10000,  # total batch size used for training per iteration (across all the environments)\n",
    "        num_episodes=500,  # total number of episodes to run the training for (can be arbitrarily high!)\n",
    "    ),\n",
    "    # Policy network settings.\n",
    "    policy=dict(\n",
    "        runner=dict(\n",
    "            to_train=True,  # flag indicating whether the model needs to be trained\n",
    "            algorithm=\"A2C\",  # algorithm used to train the policy\n",
    "            gamma=0.98,  # discount rate\n",
    "            lr=0.005,  # learning rate\n",
    "            model=dict(\n",
    "                type=\"fully_connected\", fc_dims=[256, 256], model_ckpt_filepath=\"\"\n",
    "            ),  # policy model settings\n",
    "        ),\n",
    "        tagger=dict(\n",
    "            to_train=True,\n",
    "            algorithm=\"A2C\",\n",
    "            gamma=0.98,\n",
    "            lr=0.002,\n",
    "            model=dict(\n",
    "                type=\"fully_connected\", fc_dims=[256, 256], model_ckpt_filepath=\"\"\n",
    "            ),\n",
    "        ),\n",
    "    ),\n",
    "    # Checkpoint saving setting.\n",
    "    saving=dict(\n",
    "        metrics_log_freq=10,  # how often (in iterations) to print the metrics\n",
    "        model_params_save_freq=5000,  # how often (in iterations) to save the model parameters\n",
    "        basedir=\"/tmp\",  # base folder used for saving\n",
    "        name=\"continuous_tag\",  # experiment name\n",
    "        tag=\"example\",  # experiment tag\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Instantiate the WarpDrive Module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to instantiate the WarpDrive module, we first use an environment wrapper to specify that the environment needs to be run on the GPU (via the `use_cuda` flag). Also, agents in the environment can share policy models; so we specify a dictionary to map each policy network model to the list of agent ids using that model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a wrapped environment object via the EnvWrapper\n",
    "# Ensure that env_backend is set to be \"pycuda\" (in order to run on the GPU)\n",
    "# WarpDrive v2 also supports JIT numba backend, \n",
    "# if you have installed Numba, you can set \"numba\" instead of \"pycuda\" too.\n",
    "env_wrapper = EnvWrapper(\n",
    "    TagContinuous(**run_config[\"env\"]),\n",
    "    num_envs=run_config[\"trainer\"][\"num_envs\"],\n",
    "    env_backend=\"pycuda\",\n",
    ")\n",
    "\n",
    "# Agents can share policy models: this dictionary maps policy model names to agent ids.\n",
    "policy_tag_to_agent_id_map = {\n",
    "    \"tagger\": list(env_wrapper.env.taggers),\n",
    "    \"runner\": list(env_wrapper.env.runners),\n",
    "}\n",
    "\n",
    "wd_module = WarpDriveModule(\n",
    "    env_wrapper=env_wrapper,\n",
    "    config=run_config,\n",
    "    policy_tag_to_agent_id_map=policy_tag_to_agent_id_map,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing an episode roll-out before training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use the `generate_tag_env_rollout_animation()` helper function in order to visualize an episode rollout. Internally, this function uses the WarpDrive module's `fetch_episode_states` to fetch the data arrays on the GPU for the duration of an entire episode. Specifically, we fetch the state arrays pertaining to agents' x and y locations on the plane and indicators on which agents are still active in the game, and will use these to visualize an episode roll-out. Note that this function may be invoked at any time during training, and it will use the state of the policy models at that time to sample actions and generate the visualization.\n",
    "\n",
    "The animation below shows a sample realization of the game episode before training, i.e., with randomly chosen agent actions. The $5$ taggers are marked in pink, while the $100$ blue agents are the runners. Both the taggers and runners move around randomly and about half the runners remain at the end of the episode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim = generate_tag_env_rollout_animation(wd_module, fps=25)\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create the Lightning Trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we create the trainer for training the WarpDrive model. We add the `performance stats` callbacks to the trainer to view the throughput performance of WarpDrive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_freq = run_config[\"saving\"][\"metrics_log_freq\"]\n",
    "\n",
    "# Define callbacks.\n",
    "cuda_callback = CUDACallback(module=wd_module)\n",
    "perf_stats_callback = PerfStatsCallback(\n",
    "    batch_size=wd_module.training_batch_size,\n",
    "    num_iters=wd_module.num_iters,\n",
    "    log_freq=log_freq,\n",
    ")\n",
    "\n",
    "# Instantiate the PytorchLightning trainer with the callbacks.\n",
    "# # Also, set the number of gpus to 1, since this notebook uses just a single GPU.\n",
    "num_gpus = 1\n",
    "num_episodes = run_config[\"trainer\"][\"num_episodes\"]\n",
    "episode_length = run_config[\"env\"][\"episode_length\"]\n",
    "training_batch_size = run_config[\"trainer\"][\"train_batch_size\"]\n",
    "num_epochs = int(num_episodes * episode_length / training_batch_size)\n",
    "\n",
    "# Set reload_dataloaders_every_n_epochs=1 to invoke\n",
    "# train_dataloader() each epoch.\n",
    "trainer = Trainer(\n",
    "    accelerator=\"gpu\",\n",
    "    devices=num_gpus,\n",
    "    callbacks=[cuda_callback, perf_stats_callback],\n",
    "    max_epochs=num_epochs,\n",
    "    reload_dataloaders_every_n_epochs=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start tensorboard.\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir lightning_logs/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the WarpDrive Module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we invoke training.\n",
    "\n",
    "Note: please scroll up to the tensorboard cell to visualize the curves during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "trainer.fit(wd_module)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize an episode-rollout after training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim = generate_tag_env_rollout_animation(wd_module, fps=25)\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: In the configuration above, we have set the trainer to only train on $50000$ rollout episodes, but you can increase the `num_episodes` configuration parameter to train further. As more training happens, the runners learn to escape the taggers, and the taggers learn to chase after the runner. Sometimes, the taggers also collaborate to team-tag runners. A good number of episodes to train on (for the configuration we have used) is $2$M or higher."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Finally, close the WarpDrive module to clear up the CUDA memory heap\n",
    "wd_module.graceful_close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learn More about WarpDrive and explore our tutorials!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "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
}