tutorials/multi_agent_gpu_training_with_warp_drive.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1349103",
   "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",
   "id": "9f928d76",
   "metadata": {},
   "source": [
    "### Colab\n",
    "\n",
    "Try this notebook on [Colab](http://colab.research.google.com/github/salesforce/ai-economist/blob/master/tutorials/multi_agent_gpu_training_with_warp_drive.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d85d358",
   "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": "markdown",
   "id": "073387a4",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "Welcome! In this tutorial, we detail how we train multi-agent economic simulations built using [Foundation](https://github.com/salesforce/ai-economist/tree/master/ai_economist/foundation) and train it using [WarpDrive](https://github.com/salesforce/warp-drive), an open-source library we built for extremely fast multi-agent reinforcement learning (MARL) on a single GPU. For the purposes of exposition, we specifically consider the [COVID-19 and economy simulation](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/scenarios/covid19/covid19_env.py). The COVID-19 and economy is a simulation to model health and economy dynamics amidst the COVID-19 pandemic and comprises 52 agents.\n",
    "\n",
    "We put together this tutorial with these goals in mind:\n",
    "- Describe how we train multi-agent simulations from scratch, starting with just a Python implementation of the environment on a CPU.\n",
    "- Provide reference starting code to help perform extremely fast MARL training so the AI Economist community can focus more towards contributing multi-agent simulations to Foundation.\n",
    "\n",
    "We will cover the following concepts:\n",
    "1. Building a GPU-compatible environment.\n",
    "2. CPU-GPU environment consistency checker.\n",
    "3. Adding an *environment wrapper*.\n",
    "4. Creating a *trainer* object, and perform training.\n",
    "5. Generate a rollout using the trainer object and visualize it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e761c0d",
   "metadata": {},
   "source": [
    "### Prerequisites\n",
    "It is helpful to be familiar with [Foundation](https://github.com/salesforce/ai-economist/tree/master/ai_economist/foundation), a multi-agent economic simulator, and also the COVID-19 and Economic simulation ([paper here](https://arxiv.org/abs/2108.02904)). We recommend taking a look at the following tutorials:\n",
    "\n",
    "- [Foundation: the Basics](https://github.com/salesforce/ai-economist/blob/master/tutorials/economic_simulation_basic.ipynb)\n",
    "- [Extending Foundation](https://github.com/salesforce/ai-economist/blob/master/tutorials/economic_simulation_advanced.ipynb)\n",
    "- [COVID-19 and Economic Simulation](https://github.com/salesforce/ai-economist/blob/master/tutorials/covid19_and_economic_simulation.ipynb)\n",
    "\n",
    "It is also important to get familiarized with [WarpDrive](https://github.com/salesforce/warp-drive), a framework we developed for extremely fast end-to-end reinforcement learning on a single GPU. We also have a detailed tutorial on on how to [create custom environments](https://github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-4-create_custom_environments.md) and integrate with WarpDrive."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b87df34",
   "metadata": {},
   "source": [
    "# Dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55bb9c3b",
   "metadata": {},
   "source": [
    "You will need to install the [AI Economist](https://github.com/salesforce/ai-economist) and [WarpDrive](https://github.com/salesforce/warp-drive) pip packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb4e9a57",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, signal, sys, time\n",
    "IN_COLAB = 'google.colab' in sys.modules\n",
    "\n",
    "if IN_COLAB:\n",
    "    !git clone https://github.com/salesforce/ai-economist.git\n",
    "\n",
    "    %cd ai-economist\n",
    "    !pip install -e .\n",
    "    \n",
    "    # Restart the Python runtime to automatically use the installed packages\n",
    "    print(\"\\n\\nRestarting the Python runtime! Please (re-)run the cells below.\")\n",
    "    time.sleep(1)\n",
    "    os.kill(os.getpid(), signal.SIGKILL)\n",
    "else:\n",
    "    !pip install ai-economist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56a22701",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ensure that a GPU is present.\n",
    "import GPUtil\n",
    "num_gpus_available = len(GPUtil.getAvailable())\n",
    "assert num_gpus_available > 0, \"This notebook needs a GPU machine to run!!\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9a36cdb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install rl-warp-drive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bff9430",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ai_economist\n",
    "import numpy as np\n",
    "import os\n",
    "import yaml\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import dates as mdates\n",
    "from datetime import timedelta\n",
    "from timeit import Timer\n",
    "\n",
    "from ai_economist.foundation.scenarios.covid19.covid19_env import (\n",
    "    CovidAndEconomyEnvironment,\n",
    ")\n",
    "from ai_economist.foundation.env_wrapper import FoundationEnvWrapper\n",
    "\n",
    "from warp_drive.env_cpu_gpu_consistency_checker import EnvironmentCPUvsGPU\n",
    "from warp_drive.training.trainer import Trainer\n",
    "from warp_drive.training.utils.data_loader import create_and_push_data_placeholders\n",
    "from warp_drive.utils.env_registrar import EnvironmentRegistrar\n",
    "\n",
    "_PATH_TO_AI_ECONOMIST_PACKAGE_DIR = ai_economist.__path__[0]\n",
    "\n",
    "# Set font size for the matplotlib figures\n",
    "plt.rcParams.update({'font.size': 22})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8278c6e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set logger level e.g., DEBUG, INFO, WARNING, ERROR\n",
    "import logging\n",
    "\n",
    "logging.getLogger().setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac34e846",
   "metadata": {},
   "source": [
    "# 1. Building a GPU-Compatible Environment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9cbb4a4",
   "metadata": {},
   "source": [
    "We start with a Python environment that has the [Gym](https://gym.openai.com/docs/)-style `__init__`, `reset` and `step` APIs. For example, consider the [COVID-19 economic simulation](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/scenarios/covid19/covid19_env.py). To build a GPU-compatible environment that can be trained with WarpDrive, you will need to first implement the simulation itself in CUDA C. While there are other alternatives for GPU-based simulations such as [Numba](https://numba.readthedocs.io/en/stable/cuda/index.html) and [JAX](https://jax.readthedocs.io/en/latest/), CUDA C provides the most flexibility for building complex multi-agent simulation logic, and also the fastest performance. However, implementing the simulation in\n",
    "CUDA C also requires the GPU memory and threads to be carefully managed. Some pointers on building and testing the simulation in CUDA C are provided in this WarpDrive [tutorial](https://github.com/salesforce/warp-drive/blob/master/tutorials/tutorial-4-create_custom_environments.md).\n",
    "\n",
    "Important: when writing the step function using CUDA C, the function names should follow the following convention so that they can be used with WarpDrive APIs.\n",
    "- The scenario class needs to have a 'name' attribute. The scenario step function requires to be named as \"Cuda{scenario_name}Step\".\n",
    "- Every component class needs to have a 'name' attribute. The step function for the component in the scenario requires to be named as \"Cuda{component_name}Step\".\n",
    "- The function used to compute the rewards requires to be named as \"CudaComputeReward\".\n",
    "\n",
    "The code for the COVID-19 economic simulation's step function is [here](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/scenarios/covid19/covid19_env_step.cu).\n",
    "\n",
    "To use an existing Python Environment with WarpDrive, one needs to add two augmentations (see below) to the Python code. First, a `get_data_dictionary()` method that pushes all the data arrays and environment parameters required to run the simulation to the GPU. Second, the step-function should invoke the `cuda_step` kernel with the data arrays that the CUDA C step function should have access to passed as arguments.\n",
    "\n",
    "```python\n",
    "class Env:\n",
    "    def __init__(self, **env_config):\n",
    "        ...\n",
    "\n",
    "    def reset(self):\n",
    "        ...\n",
    "        return obs\n",
    "\n",
    "    def get_data_dictionary(self):\n",
    "        # Specify the data that needs to be \n",
    "        # pushed to the GPU.\n",
    "        data_feed = DataFeed()\n",
    "        data_feed.add_data(\n",
    "            name=\"variable_name\",\n",
    "            data=self.variable,\n",
    "            save_copy_and_apply_at_reset\n",
    "            =True,\n",
    "        )\n",
    "        ...\n",
    "        return data_feed\n",
    "\n",
    "    def step(self, actions):\n",
    "        if self.use_cuda:\n",
    "            self.cuda_step(\n",
    "                # Pass the relevant data \n",
    "                # feed keys as arguments \n",
    "                # to cuda_step. \n",
    "                # Note: cuda_data_manager \n",
    "                # is created by the \n",
    "                # EnvWrapper.\n",
    "                self.cuda_data_manager.\n",
    "                device_data(...),\n",
    "                ...\n",
    "            )\n",
    "        else:\n",
    "            ...\n",
    "            return obs, rew, done, info\n",
    "```\n",
    "The complete Python code is [here](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/scenarios/covid19/covid19_env.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4645b43",
   "metadata": {},
   "source": [
    "# 2. CPU-GPU Environment Consistency Checker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01f9b8ae",
   "metadata": {},
   "source": [
    "Before we train the simulation on the GPU, we will need to ensure consistency between the Python and CUDA C versions of the simulation. For this purpose, Foundation provides an [EnvironmentCPUvsGPU class](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/env_cpu_gpu_consistency_checker.py). This module essentially instantiates environment objects corresponding to the two versions of the simulation. It then steps through the two environment objects for a specified number of environment replicas `num_envs` and a specified number of episodes `num_episodes`, and verifies that the observations, actions, rewards and the “done” flags are the same after each step. We have created a testing [script](https://github.com/salesforce/ai-economist/blob/master/tests/run_covid19_cpu_gpu_consistency_checks.py) that performs the consistency checks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a990b032",
   "metadata": {},
   "source": [
    "First, we will create an environment configuration to test with. For more details on what the configuration parameters mean, please refer to the simulation [code](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/scenarios/covid19/covid19_env.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1943025",
   "metadata": {},
   "outputs": [],
   "source": [
    "env_config = {\n",
    "    'collate_agent_step_and_reset_data': True,\n",
    "     'components': [\n",
    "         {'ControlUSStateOpenCloseStatus': {'action_cooldown_period': 28}},\n",
    "          {'FederalGovernmentSubsidy': {'num_subsidy_levels': 20,\n",
    "            'subsidy_interval': 90,\n",
    "            'max_annual_subsidy_per_person': 20000}},\n",
    "          {'VaccinationCampaign': {'daily_vaccines_per_million_people': 3000,\n",
    "            'delivery_interval': 1,\n",
    "            'vaccine_delivery_start_date': '2021-01-12'}}\n",
    "     ],\n",
    "     'economic_reward_crra_eta': 2,\n",
    "     'episode_length': 540,\n",
    "     'flatten_masks': True,\n",
    "     'flatten_observations': False,\n",
    "     'health_priority_scaling_agents': 0.3,\n",
    "     'health_priority_scaling_planner': 0.45,\n",
    "     'infection_too_sick_to_work_rate': 0.1,\n",
    "     'multi_action_mode_agents': False,\n",
    "     'multi_action_mode_planner': False,\n",
    "     'n_agents': 51,\n",
    "     'path_to_data_and_fitted_params': '',\n",
    "     'pop_between_age_18_65': 0.6,\n",
    "     'risk_free_interest_rate': 0.03,\n",
    "     'world_size': [1, 1],\n",
    "     'start_date': '2020-03-22',\n",
    "     'use_real_world_data': False,\n",
    "     'use_real_world_policies': False\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "364047c4",
   "metadata": {},
   "source": [
    "Next, we will need to register the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f5f0e1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "env_registrar = EnvironmentRegistrar()\n",
    "env_registrar.add_cuda_env_src_path(\n",
    "    CovidAndEconomyEnvironment.name,\n",
    "    os.path.join(\n",
    "        _PATH_TO_AI_ECONOMIST_PACKAGE_DIR, \n",
    "        \"foundation/scenarios/covid19/covid19_build.cu\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "213dbdf0",
   "metadata": {},
   "source": [
    "We will also need to set some variables: `policy_tag_to_agent_id_map`, `separate_placeholder_per_policy` (defaults to False) and `obs_dim_corresponding_to_num_agents` (defaults to \"first\"). The variable descriptions are below in comments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33d7556b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The policy_tag_to_agent_id_map dictionary maps\n",
    "# policy model names to agent ids.\n",
    "policy_tag_to_agent_id_map = {\n",
    "    \"a\": [str(agent_id) for agent_id in range(env_config[\"n_agents\"])],\n",
    "    \"p\": [\"p\"],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "186a933b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Flag indicating whether separate obs, actions and rewards placeholders have to be created for each policy.\n",
    "# Set \"create_separate_placeholders_for_each_policy\" to True here \n",
    "# since the agents and the planner have different observation and action spaces.\n",
    "separate_placeholder_per_policy = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cbd2468",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Flag indicating the observation dimension corresponding to 'num_agents'.\n",
    "# Note: WarpDrive assumes that all the observation are shaped\n",
    "# (num_agents, *feature_dim), i.e., the observation dimension\n",
    "# corresponding to 'num_agents' is the first one. Instead, if the\n",
    "# observation dimension corresponding to num_agents is the last one,\n",
    "# we will need to permute the axes to align with WarpDrive's assumption\n",
    "obs_dim_corresponding_to_num_agents = \"last\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8f128d8",
   "metadata": {},
   "source": [
    "The consistency tests may be performed using the `test_env_reset_and_step()` API as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1898dfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "EnvironmentCPUvsGPU(\n",
    "    dual_mode_env_class=CovidAndEconomyEnvironment,\n",
    "    env_configs={\"test\": env_config},\n",
    "    num_envs=3,\n",
    "    num_episodes=2,\n",
    "    env_wrapper=FoundationEnvWrapper,\n",
    "    env_registrar=env_registrar,\n",
    "    policy_tag_to_agent_id_map=policy_tag_to_agent_id_map,\n",
    "    create_separate_placeholders_for_each_policy=True,\n",
    "    obs_dim_corresponding_to_num_agents=\"last\"    \n",
    ").test_env_reset_and_step()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f867963d",
   "metadata": {},
   "source": [
    "If the two implementations are consistent, you should see  `The CPU and the GPU environment outputs are consistent within 1 percent` at the end of the previous cell run."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c83bf539",
   "metadata": {},
   "source": [
    "# 3. Adding an *Environment Wrapper*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5d7f438",
   "metadata": {},
   "source": [
    "Once the Python and CUDA C implementation are consistent with one another, we use an environment wrapper to wrap the environment object, and run the simulation on the GPU. Accordingly, we need to set the `use_cuda` argument to True. under this setting, only the first environment reset happens on the CPU. Following that, the data arrays created at reset and the simulation parameters are copied over (a one-time operation) to the GPU memory. All the subsequent steps (and resets) happen only on the GPU. In other words, there's no back-and-forth data copying between the CPU and the GPU, and all the data arrays on the GPU are modified in-place. The environment wrapper also uses the `num_envs` argument (defaults to $1$) to instantiate multiple replicas of the environment on the GPU.\n",
    "\n",
    "Note: for running the simulation on a CPU, simply set use_cuda=False, and it is no different than actually running the Python simulation on a CPU - the reset and step calls also happen on the CPU.\n",
    "\n",
    "The environment wrapper essentially performs the following tasks that are required to run the simulation on the GPU:\n",
    "- Registers the CUDA step kernel, so that the step function can be invoked from the CPU (host).\n",
    "- Pushes all the data listed in the data dictionary to the GPU when the environment is reset for the very frst time.\n",
    "- Automatically resets every environment when it reaches its done state.\n",
    "- Adds the observation and action spaces to the environment, which are required when training the environment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ede05b74",
   "metadata": {},
   "source": [
    "The CPU and GPU versions of the environment object may be created via setting the appropriate value of the `use_cuda` flag."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77db1010",
   "metadata": {},
   "outputs": [],
   "source": [
    "cpu_env = FoundationEnvWrapper(\n",
    "    CovidAndEconomyEnvironment(**env_config),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa259b4e",
   "metadata": {},
   "source": [
    "Instantiating the GPU environment also initializes the data function managers and loads the CUDA kernels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83e433e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu_env = FoundationEnvWrapper(\n",
    "    CovidAndEconomyEnvironment(**env_config),\n",
    "    use_cuda=True,\n",
    "    env_registrar=env_registrar,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a204d966",
   "metadata": {},
   "source": [
    "# 4. Creating a Trainer Object and Perform Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fbea778",
   "metadata": {},
   "source": [
    "Next, we will prepare the environment for training on a GPU.We will need to define a run configuration (which comprises the environment, training, policy and saving configurations), and create a *trainer* object.\n",
    "\n",
    "We will load the run configuration from a saved yaml file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c41ee947",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_path = os.path.join(\n",
    "    _PATH_TO_AI_ECONOMIST_PACKAGE_DIR,\n",
    "    \"training/run_configs/\",\n",
    "    f\"covid_and_economy_environment.yaml\",\n",
    ")\n",
    "with open(config_path, \"r\", encoding=\"utf8\") as f:\n",
    "    run_config = yaml.safe_load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed835bbd",
   "metadata": {},
   "source": [
    "### Instantiating the Trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c758763",
   "metadata": {},
   "source": [
    "Next, we will create and instantiate the trainer object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97f241ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    env_wrapper=gpu_env,\n",
    "    config=run_config,\n",
    "    policy_tag_to_agent_id_map=policy_tag_to_agent_id_map,\n",
    "    create_separate_placeholders_for_each_policy=separate_placeholder_per_policy,\n",
    "    obs_dim_corresponding_to_num_agents=obs_dim_corresponding_to_num_agents,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bb37f72",
   "metadata": {},
   "source": [
    "### CPU vs GPU Performance Comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a582b81",
   "metadata": {},
   "source": [
    "Before performing training, let us see how the simulation speed on the GPU compares with that of the CPU. We will generate a set of random actions, and step through both versions of the simulation a few times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50fdb713",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_random_actions(env):\n",
    "    actions = {\n",
    "            str(agent_id): np.random.randint(\n",
    "                low=0,\n",
    "                high=env.env.action_space[str(agent_id)].n,\n",
    "                dtype=np.int32,\n",
    "            )\n",
    "            for agent_id in range(env.n_agents-1)\n",
    "    }\n",
    "    actions[\"p\"] = np.random.randint(\n",
    "        low=0,\n",
    "        high=env.env.action_space[\"p\"].n,\n",
    "        dtype=np.int32,\n",
    "    )\n",
    "    return actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5cd89eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def env_reset_and_step(env):\n",
    "    env.reset()\n",
    "    actions = generate_random_actions(env)\n",
    "    for t in range(env.episode_length):\n",
    "        env.step(actions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4dd1cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "Timer(lambda: env_reset_and_step(gpu_env)).timeit(number=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ca9f41c",
   "metadata": {},
   "outputs": [],
   "source": [
    "Timer(lambda: env_reset_and_step(cpu_env)).timeit(number=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d626a15a",
   "metadata": {},
   "source": [
    "Notice that with just $1$ replica of the environment, the environment step on the GPU is over 5x faster (on an A100 machine). When running training, it is typical to use several environment replicas, and that provides an even higher performance boost for the GPU, since WarpDrive runs all the environment replicas in parallel on separate GPU blocks, and the CPU cannot achieve the same amount of parallelization."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28a6846a",
   "metadata": {},
   "source": [
    "### Perform Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f0aac10",
   "metadata": {},
   "source": [
    "We perform training by invoking `trainer.train()`. The speed performance stats and metrics for the trained policies are printed on screen.\n",
    "Note: In this notebook, we only run training for $200$ iterations. You may run it for longer by setting the `num_episodes` configuration parameter [here](https://github.com/salesforce/ai-economist/blob/master/ai_economist/training/run_configs/covid_and_economy_environment.yaml)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d64d7ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f71acf9",
   "metadata": {},
   "source": [
    "# 5. Visualize the Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac0973cb",
   "metadata": {},
   "source": [
    "Post training, it is useful to visualize some of the environment's actions and observations to gain more insight into the kinds of policies the RL agents learn, and the resulting environment dynamics. For the COVID-19 and economic simulation, the actions - \"stringency level\" and \"subsidy level\" control the observables such such as \"susceptible\", \"infected\", \"recovered\", \"deaths\", \"unemployed\", \"vaccinated\" and \"productivity\" for each of the US states.\n",
    "\n",
    "Incidentally, these actions and observables also correspond to the names of arrays that were pushed to the GPU after the very first environment reset. At any time, the arrays can be fetched back to the CPU via the WarpDrive trainer's API `fetch_episode_states`, and visualized for the duration of an entire episode. Below, we also provide a helper function to perform the visualizations. Note that in this notebook, we only performed a few iterations of training, so the policies will not be quite trained at this point, so the plots seen in the visualization seen are going to be arbitrary. You may run training with a different set of configurations or for longer by setting the `num_episodes` configuration parameter [here](https://github.com/salesforce/ai-economist/blob/master/ai_economist/training/run_configs/covid_and_economy_environment.yaml), and you can visualize the policies after, using the same code provided below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43f755a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function: visualizations\n",
    "\n",
    "def visualize_states(\n",
    "    entity=\"USA\",\n",
    "    episode_states=None,\n",
    "    trainer=None,\n",
    "    ax=None\n",
    "):\n",
    "    assert trainer is not None\n",
    "    assert episode_states is not None\n",
    "       \n",
    "    # US state names to index mapping\n",
    "    us_state_name_to_idx = {v: k for k, v in trainer.cuda_envs.env.us_state_idx_to_state_name.items()}\n",
    "    us_state_name_to_idx[\"USA\"] = \"p\"\n",
    "    agent_id = us_state_name_to_idx[entity]\n",
    "    \n",
    "    assert entity is not None\n",
    "    assert entity in us_state_name_to_idx.keys(), f\"entity should be in {list(us_state_name_to_idx.keys())}\"\n",
    "    \n",
    "    for key in episode_states:\n",
    "        # Use the collated data at the last valid time index\n",
    "        last_valid_time_index = np.isnan(np.sum(episode_states[key], axis = (1, 2))).argmax() - 1\n",
    "        episode_states[key] = episode_states[key][last_valid_time_index]\n",
    "    \n",
    "    if agent_id == \"p\":\n",
    "        for key in episode_states:\n",
    "                \n",
    "            if key in [\"subsidy_level\", \"stringency_level\"]:\n",
    "                episode_states[key] = np.mean(episode_states[key], axis=-1) # average across all the US states\n",
    "            else:\n",
    "                episode_states[key] = np.sum(episode_states[key], axis=-1) # sum across all the US states\n",
    "            episode_states[key] = episode_states[key].reshape(-1, 1) # putting back the agent_id dimension\n",
    "        agent_id = 0\n",
    "    else:\n",
    "        agent_id = int(agent_id)\n",
    "        \n",
    "    if ax is None:\n",
    "        if len(episode_states) < 3:\n",
    "            cols = len(episode_states)\n",
    "        else:\n",
    "            cols = 3\n",
    "        scale = 8\n",
    "        rows = int(np.ceil(len(episode_states) / cols))\n",
    "    \n",
    "        h, w = scale*max(rows, cols), scale*max(rows, rows)\n",
    "        fig, ax = plt.subplots(rows, cols, figsize=(h, w), sharex=True, squeeze=False)\n",
    "    else:\n",
    "        rows, cols = ax.shape\n",
    "        \n",
    "    start_date = trainer.cuda_envs.env.start_date\n",
    "    \n",
    "    for idx, key in enumerate(episode_states):\n",
    "        row = idx // cols\n",
    "        col = idx % cols\n",
    "\n",
    "        dates = [start_date + timedelta(day) for day in range(episode_states[key].shape[0])]\n",
    "        ax[row][col].plot(dates, episode_states[key][:, agent_id], linewidth=3)\n",
    "        ax[row][col].set_ylabel(key)\n",
    "        ax[row][col].grid(b=True)\n",
    "        ax[row][col].xaxis.set_major_locator(mdates.MonthLocator(interval=3))\n",
    "        ax[row][col].xaxis.set_major_formatter(mdates.DateFormatter(\"%b'%y\"))\n",
    "    fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aa83589",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch the key state indicators for an episode.\n",
    "episode_states = trainer.fetch_episode_states(\n",
    "    [\n",
    "        \"stringency_level\",\n",
    "        \"subsidy_level\",        \n",
    "        \"susceptible\",\n",
    "        \"infected\",\n",
    "        \"recovered\",\n",
    "        \"deaths\",\n",
    "        \"unemployed\",\n",
    "        \"vaccinated\",\n",
    "        \"productivity\",\n",
    "        \"postsubsidy_productivity\",\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Visualize the fetched states for the USA.\n",
    "# Feel free to modify the 'entity' argument to visualize the curves for the US states (e.g., California, Utah) too.\n",
    "visualize_states(\n",
    "    entity=\"USA\",\n",
    "    episode_states=episode_states,\n",
    "    trainer=trainer,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1649b9c5",
   "metadata": {},
   "source": [
    "And that's it for this tutorial. Happy training with Foundation and WarpDrive!"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}