tutorials/covid19_and_economic_simulation.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": [
    "### Colab\n",
    "\n",
    "Try this notebook on [Colab](http://colab.research.google.com/github/salesforce/ai-economist/blob/master/tutorials/covid19_and_economic_simulation.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial demonstrates how the covid and economic simulation can be used to simulate different health and economic policies during the COVID-19 pandemic in the US and to study their effect on social objectives that combine public health and economic productivity.\n",
    "\n",
    "We begin with a brief introduction to the problem and then show how we implement it in simulation.  \n",
    "\n",
    "For further reading, check out our [paper](https://arxiv.org/abs/2108.02904), [web demo](https://einstein.ai/the-ai-economist/ai-policy-foundation-and-covid-case-study), and [blog post](https://blog.einstein.ai/ai-economist-covid-case-study-ethics )."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The COVID-19 and economic simulation environment studies the effect of health and economic governmental policies on the spread of the COVID-19 pandemic and the economy, and captures the tradeoffs therein. It is built on top of our economic simulation framework [Foundation](https://github.com/salesforce/ai-economist). The environment comprises 52 entities overall, including 51 state governors (agents) - one each for the 50 US states and the District of Columbia, and the federal government (planner).\n",
    "\n",
    "In our simulation, the US state governors determine the stringency levels (which represents a combination of several containment and closure policies), while the federal government determines how much to subsidize each of the US states (via relief payments). Notably, this set of actions captures several interesting health-economy tradeoffs. For instance,\n",
    "- If the states shut down more, there will be fewer cases and deaths, but the unemployment rate also increases, hurting the economy.\n",
    "- If the federal government subsidizes more, it helps alleviate the economy of the US states, and, in turn, incentivizes them to shut down more in order to bring down the cases and deaths. However the federal government needs to borrow the subsidy money at a certain interest rate, and the cost of borrowing lowers its economic index, as does the lost economic output resulting from the additional shut downs.\n",
    "\n",
    "Each of the states and the federal government try to balance public health and the economy in order to improve **social welfare**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Social Welfare"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Social welfare is a (weighted) combination of the health index $H$ and the economic index $E$:\n",
    "$$SW = \\alpha H + (1-\\alpha)E.$$\n",
    "$\\alpha$ is a weighting term ($0 \\leq \\alpha \\leq 1$). By varying $\\alpha$, we can define a family of social welfare functions that prioritize between the health and the economy, ranging from focusing fully on economy ($\\alpha=0$) to fully on health ($\\alpha=1$).\n",
    "\n",
    "The health index $H$ is proportional to the negative of COVID-19-related deaths: fewer the deaths, higher is the health index. The economic index $E$ is a function of the annual GDP, unemployment and federal subsidies; so higher GDP implies a higher economic index and vice versa. For more mathematical details on $H$, $E$ and $SW$, please see our [paper](https://arxiv.org/abs/2108.02904)\n",
    "\n",
    "Our simulation can be used to study the effects of various US state and federal government policies on COVID-19 deaths, unemployment and the economy, and can be used in conjunction with reinforcement learning to optimize the policies for varying definitions of social welfare."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can install the ai-economist package using the pip package manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, signal, sys, time\n",
    "IN_COLAB = 'google.colab' in sys.modules\n",
    "\n",
    "if IN_COLAB:\n",
    "    !pip install ai-economist\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,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ai_economist import foundation\n",
    "from datetime import datetime\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import dates as mdates\n",
    "\n",
    "# Set font size for the matplotlib figures\n",
    "plt.rcParams.update({'font.size': 26})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### In order to download and use the latest COVID-19 data in our economic simulation, please run the following notebooks in \"ai_economist/datasets/covid19_datasets\".\n",
    "\n",
    "1. gather_real_world_data.ipynb\n",
    "2. fit_model_parameters.ipynb\n",
    "\n",
    "Upon running these notebooks, you will fetch the latest real-world data, and use the data to fit models that will be used in the COVID-19 and economy simulation. From the notebooks, you will need to record the \"path_to_data_and_fitted_params\" to set into the environment config below\n",
    "\n",
    "Note: If you do not wish to download the real-world data, you can still run this notebook as is, and it will use the data saved in the ai_economist/datasets/covid19_datasets/path_to_data_and_fitted_params directory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a simulation environment\n",
    "\n",
    "To create the covid-19 and economic simulation, we first define the configuration of the environment that will be built."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_config = {\n",
    "    # Scenario name - determines which scenario class to use\n",
    "    \"scenario_name\": \"CovidAndEconomySimulation\",\n",
    "    \n",
    "    # The list of components in this simulation\n",
    "    \"components\": [\n",
    "        {\"ControlUSStateOpenCloseStatus\": {\n",
    "            # action cooldown period in days.\n",
    "            # Once a stringency level is set, the state(s) cannot switch to another level\n",
    "            # for a certain number of days (referred to as the \"action_cooldown_period\")\n",
    "            \"action_cooldown_period\": 28\n",
    "        }},\n",
    "        {\"FederalGovernmentSubsidy\": {\n",
    "            # The number of subsidy levels.\n",
    "            \"num_subsidy_levels\": 20,\n",
    "            # The number of days over which the total subsidy amount is evenly rolled out.\n",
    "            \"subsidy_interval\": 90,\n",
    "            # The maximum annual subsidy that may be allocated per person.\n",
    "            \"max_annual_subsidy_per_person\": 20000,\n",
    "        }},\n",
    "        {\"VaccinationCampaign\": {\n",
    "            # The number of vaccines available per million people everyday.\n",
    "            \"daily_vaccines_per_million_people\": 3000,\n",
    "            # The number of days between vaccine deliveries.\n",
    "            \"delivery_interval\": 1,\n",
    "            # The date (YYYY-MM-DD) when vaccination begins\n",
    "            \"vaccine_delivery_start_date\": \"2021-01-12\",\n",
    "        }},\n",
    "    ],\n",
    "\n",
    "    # Date (YYYY-MM-DD) to start the simulation.\n",
    "    \"start_date\": \"2020-03-22\",\n",
    "    # How long to run the simulation for (in days)\n",
    "    \"episode_length\": 405,\n",
    "    \n",
    "    # use_real_world_data (bool): Replay what happened in the real world.\n",
    "    # Real-world data comprises SIR (susceptible/infected/recovered),\n",
    "    # unemployment, government policy, and vaccination numbers.\n",
    "    # This setting also sets use_real_world_policies=True.\n",
    "    \"use_real_world_data\": False,\n",
    "    # use_real_world_policies (bool): Run the environment with real-world policies\n",
    "    # (stringency levels and subsidies). With this setting and\n",
    "    # use_real_world_data=False, SIR and economy dynamics are still\n",
    "    # driven by fitted models.\n",
    "    \"use_real_world_policies\": False,\n",
    "    \n",
    "    # A factor indicating how much more the\n",
    "    # states prioritize health (roughly speaking, loss of lives due to\n",
    "    # opening up more) over the economy (roughly speaking, a loss in GDP\n",
    "    # due to shutting down resulting in more unemployment) compared to the\n",
    "    # real-world.\n",
    "    # For example, a value of 1 corresponds to the health weight that \n",
    "    # maximizes social welfare under the real-world policy, while\n",
    "    # a value of 2 means that states care twice as much about public health\n",
    "    # (preventing deaths), while a value of 0.5 means that states care twice\n",
    "    # as much about the economy (preventing GDP drops).\n",
    "    \"health_priority_scaling_agents\": 1,\n",
    "    # Same as above for the planner\n",
    "    \"health_priority_scaling_planner\": 1,\n",
    "    \n",
    "    # Full path to the directory containing\n",
    "    # the data, fitted parameters and model constants. This defaults to\n",
    "    # \"ai_economist/datasets/covid19_datasets/data_and_fitted_params\".\n",
    "    # For details on obtaining these parameters, please see the notebook\n",
    "    # \"ai-economist-foundation/ai_economist/datasets/covid19_datasets/\n",
    "    # gather_real_world_data_and_fit_parameters.ipynb\".\n",
    "    \"path_to_data_and_fitted_params\": \"\",\n",
    "    \n",
    "    # Economy-related parameters\n",
    "    # Fraction of people infected with COVID-19. Infected people don't work.\n",
    "    \"infection_too_sick_to_work_rate\": 0.1,\n",
    "    # Fraction of the population between ages 18-65.\n",
    "    # This is the subset of the population whose employment/unemployment affects\n",
    "    # economic productivity.\n",
    "    \"pop_between_age_18_65\": 0.6,\n",
    "    # Percentage of interest paid by the federal\n",
    "    # government to borrow money from the federal reserve for COVID-19 relief\n",
    "    # (direct payments). Higher interest rates mean that direct payments\n",
    "    # have a larger cost on the federal government's economic index.\n",
    "    \"risk_free_interest_rate\": 0.03,\n",
    "    # CRRA eta parameter for modeling the economic reward non-linearity.\n",
    "    \"economic_reward_crra_eta\": 2,\n",
    "       \n",
    "    # Number of agents in the simulation (50 US states + Washington DC)\n",
    "    \"n_agents\": 51,    \n",
    "    # World size: Not relevant to this simulation, but needs to be set for Foundation\n",
    "    \"world_size\": [1, 1],\n",
    "    # Flag to collate all the agents' observations, rewards and done flags into a single matrix\n",
    "    \"collate_agent_step_and_reset_data\": True,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to create an environment instance with a configuration, we can use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "env = foundation.make_env_instance(**env_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All of the US states can control their stringency levels, while the federal government can subsidize the states. We can view the action spaces for the agents and the planner that shows how many discretized levels are possible for each agent's action.\\\n",
    "For the states' stringency actions, an action value of $1$ corresponds to fully open, and higher values mean the states are more closed (maxes out at \"fully closed\"). For the federal government subsidies, $0$ corresponds to no subsidies, and higher values mean more subsidies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(env.get_agent(agent_idx=\"0\").action_dim, env.get_agent(agent_idx=\"p\").action_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Studying the effect of different government policies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's interact with the simulation by setting different state and federal government actions and see how the year 2020 (and beyond) would have panned out.\n",
    "\n",
    "First, let's compare the actual real-world data to our fitted simulation with real-world policies. These two scenarios can be studied by simply setting a couple of the environment configuration parameters -\n",
    "1. Set `\"use_real_world_policies\": True` to use the real-world policies (stringency levels and subsidies) to step through the environment.\n",
    "2. Set `\"use_real_world_data\": True` to replay *all* the real-world data through the environment.\n",
    "\n",
    "Also note that for these scenarios, we do not need to explicitly provide external action inputs for the step()."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a dense log dictionary to save the dense logs for different scenarios\n",
    "# Note: \"dense logs\" are basically logs of the environment's states, actions and rewards for a full episode. \n",
    "# They can be used to visualize an episode.\n",
    "\n",
    "dense_logs = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function: useful for generating arbitrary actions\n",
    "\n",
    "def generate_actions(env, type=\"random\", episode_length=None, seed=None):\n",
    "    if episode_length is None:\n",
    "        episode_length = env.episode_length\n",
    "    if seed is not None:\n",
    "        np.random.rand(seed)\n",
    "\n",
    "    action_seq = [None for _ in range(episode_length)]\n",
    "    num_agents = env.n_agents\n",
    "    agent_action_spaces = env.all_agents[0].action_spaces\n",
    "    planner_action_spaces = env.all_agents[-1].action_spaces\n",
    "\n",
    "    for timestep in range(episode_length):\n",
    "\n",
    "        actions = {}\n",
    "        if type == \"real_world\":\n",
    "            # For these cases, we do not need to explicitly provide external actions.\n",
    "            pass\n",
    "        \n",
    "        elif type == \"random\":\n",
    "            actions = {str(agent_id): np.random.randint(agent_action_spaces) \n",
    "                       for agent_id in range(num_agents)}\n",
    "            actions['p'] = np.random.randint(planner_action_spaces)\n",
    "            \n",
    "        elif type == \"states_open_no_subsidies\":\n",
    "            actions = {str(agent_id): np.array([1]) for agent_id in range(num_agents)}\n",
    "            actions['p'] = np.zeros_like(planner_action_spaces)\n",
    "            \n",
    "        elif type == \"states_closed_full_subsidies\":\n",
    "            actions = {str(agent_id): np.array([agent_action_spaces - 1]) \n",
    "                             for agent_id in range(num_agents)}\n",
    "            actions['p'] = np.array(planner_action_spaces) - 1\n",
    "            \n",
    "        elif type == \"states_closed_6_months_no_subsidies\":\n",
    "            if timestep < 6 * 30:\n",
    "                actions = {str(agent_id): np.array([agent_action_spaces - 1]) \n",
    "                                 for agent_id in range(num_agents)}\n",
    "            else:\n",
    "                actions = {str(agent_id): np.array([1]) for agent_id in range(num_agents)}\n",
    "            actions['p'] = np.zeros_like(planner_action_spaces)\n",
    "\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "\n",
    "        action_seq[timestep] = actions\n",
    "\n",
    "    return action_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function to fetch environment dense logs\n",
    "\n",
    "def fetch_env_dense_log(\n",
    "    env_config,\n",
    "    action_type=\"real_world\"\n",
    "):\n",
    "    env = foundation.make_env_instance(**env_config)\n",
    "    env.reset(force_dense_logging=True)\n",
    "    \n",
    "    action_seq = generate_actions(env, action_type)\n",
    "\n",
    "    for t in range(env.episode_length):\n",
    "        env.step(action_seq[t]);\n",
    "    return env._dense_log"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 1: real-world"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# minimize print clutter\n",
    "\n",
    "real_world_env_config = env_config.copy()\n",
    "real_world_env_config.update(\n",
    "    {\n",
    "        \"use_real_world_data\": True,\n",
    "        \"use_real_world_policies\": True   \n",
    "    }\n",
    ")\n",
    "dense_logs[\"real_world\"] = fetch_env_dense_log(\n",
    "    real_world_env_config,\n",
    "    action_type=\"real_world\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 2: simulation with real-world policies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# minimize print clutter\n",
    "\n",
    "sim_with_real_world_policies_env_config = env_config.copy()\n",
    "sim_with_real_world_policies_env_config.update(\n",
    "    {\n",
    "        \"use_real_world_policies\": True\n",
    "    }\n",
    ")\n",
    "dense_logs[\"sim_with_real_world_policies\"] = fetch_env_dense_log(\n",
    "    sim_with_real_world_policies_env_config,\n",
    "    action_type=\"real_world\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing some of the environment states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function: visualizations\n",
    "\n",
    "def visualize_states(\n",
    "    entity=\"USA\",\n",
    "    dense_logs = {},\n",
    "    ax=None,\n",
    "    states=None\n",
    "):\n",
    "    if states is None:\n",
    "        # Relevant states for visualizations\n",
    "        states = [\n",
    "            # actions\n",
    "            'Current Open Close Stringency Level',\n",
    "            'Current Subsidy Level',\n",
    "            # global states\n",
    "            'Total Susceptible',\n",
    "            'Total Infected',\n",
    "            'Total Recovered',\n",
    "            'Total Deaths',\n",
    "            'New Infections',\n",
    "            'New Deaths',    \n",
    "            'Total Vaccinated',\n",
    "            'Total Unemployed',\n",
    "            'Subsidy',\n",
    "            'Postsubsidy Productivity',\n",
    "            # rewards\n",
    "            'Health Index',\n",
    "            'Economic Index',\n",
    "            'Social Welfare',\n",
    "        ]\n",
    "\n",
    "    # US state names to index mapping\n",
    "    us_state_name_to_idx = {v: k for k, v in env.us_state_idx_to_state_name.items()}\n",
    "    us_state_name_to_idx[\"USA\"] = \"p\"\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",
    "    agent_id = us_state_name_to_idx[entity]\n",
    "    assert len(dense_logs) > 0  # at least one valid dense logs\n",
    "    \n",
    "    if ax is None:\n",
    "        if len(states) < 3:\n",
    "            cols = len(states)\n",
    "        else:\n",
    "            cols = 3\n",
    "        scale = 8\n",
    "        rows = int(np.ceil(len(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",
    "    for scenario in dense_logs:\n",
    "        dense_log = dense_logs[scenario]\n",
    "        \n",
    "        dates = [\n",
    "            datetime.strptime(dense_logs[scenario][\"states\"][t][agent_id][\"Date\"], \"%Y-%m-%d\") \\\n",
    "            for t in range(len(dense_logs[scenario][\"states\"]) - 1)\n",
    "        ]\n",
    "        \n",
    "        # Compute reward\n",
    "        if entity == \"USA\":\n",
    "            reward = [dense_log[\"rewards\"][t][agent_id]\n",
    "                      for t in range(len(dense_log[\"states\"]) - 1)]                \n",
    "        else:\n",
    "            # Fetch agent-specific reward from the collated rewards\n",
    "            reward = [dense_log[\"rewards\"][t][\"a\"][int(agent_id)] \n",
    "                      for t in range(len(dense_log[\"states\"]) - 1)]        \n",
    "\n",
    "        for idx, plot_key in enumerate(states):\n",
    "            row = idx // cols\n",
    "            col = idx % cols\n",
    "            \n",
    "            if plot_key == \"Current Open Close Stringency Level\":\n",
    "                if entity == \"USA\":\n",
    "                    # Average across all the US states\n",
    "                    values = np.mean(np.array([\n",
    "                        [dense_log[\"states\"][t][str(agent_id)].get(plot_key, np.nan) \n",
    "                         for t in range(len(dense_log[\"states\"]) - 1)] \n",
    "                        for agent_id in range(51)]), axis=0)\n",
    "                    plot_key = \"(Average) Current Open Close Stringency Level\"\n",
    "                else:\n",
    "                    values = [dense_log[\"states\"][t][agent_id].get(plot_key, np.nan) \n",
    "                              for t in range(len(dense_log[\"states\"]) - 1)]\n",
    "            elif plot_key == \"Current Subsidy Level\":\n",
    "                # Use the subsidy level set by the planner\n",
    "                values = [dense_log[\"states\"][t][\"p\"].get(plot_key, np.nan) \n",
    "                      for t in range(len(dense_log[\"states\"]) - 1)]\n",
    "            elif plot_key == \"Subsidy\":\n",
    "                if entity == \"USA\":\n",
    "                    plot_key = \"New Subsidy Provided\"\n",
    "                else:\n",
    "                    plot_key = \"New Subsidy Received\"\n",
    "                values = [dense_log[\"states\"][t][agent_id].get(plot_key, np.nan) \n",
    "                          for t in range(len(dense_log[\"states\"]) - 1)]\n",
    "            elif plot_key == \"Social Welfare\":\n",
    "                values = reward\n",
    "            else:\n",
    "                values = [dense_log[\"states\"][t][str(agent_id)].get(plot_key, np.nan) \n",
    "                          for t in range(len(dense_log[\"states\"]) - 1)]                \n",
    "\n",
    "            ax[row][col].plot(dates, values, linewidth=5, label=scenario+f\" (SW={np.nansum(reward):3.1f})\")\n",
    "            ax[row][col].set_ylabel(plot_key)\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",
    "            ax[row][col].grid(b=True)\n",
    "            \n",
    "        ax[0][0].legend(bbox_to_anchor=(0, 1.1, 1.125 * cols, 0.2), loc=\"best\",\n",
    "                        mode=\"expand\", borderaxespad=0, ncol=cols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the environment states for the USA.\n",
    "# You can set US state names to the entity keyword and visualize the state specific values as well ! \n",
    "\n",
    "visualize_states(\n",
    "    entity=\"USA\",\n",
    "    dense_logs = dense_logs,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Comparing the curves for the real-world and the simulation with real-world policies, we see that we are able to fit the real-world data reasonably well.\n",
    "\n",
    "We can also use this simulation to study the effect of any other (government) policy, for example,\n",
    "- states open + no subsidies\n",
    "- states closed + full subsidies\n",
    "- states closed for two months (and open after) + no subsidies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 3: states open + no subsidies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# minimize print clutter\n",
    "\n",
    "dense_logs[\"sim_with_states_open_no_subsidies\"] = fetch_env_dense_log(\n",
    "    env_config,\n",
    "    action_type=\"states_open_no_subsidies\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 4: states closed + full subsidies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# minimize print clutter\n",
    "\n",
    "dense_logs[\"sim_with_states_closed_full_subsidies\"] = fetch_env_dense_log(\n",
    "    env_config,\n",
    "    action_type=\"states_closed_full_subsidies\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 5: states closed for 6 months (and open after) + no subsidies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# minimize print clutter\n",
    "\n",
    "dense_logs[\"sim_with_states_closed_6_months_no_subsidies\"] = fetch_env_dense_log(\n",
    "    env_config,\n",
    "    action_type=\"states_closed_6_months_no_subsidies\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_states(\n",
    "    entity=\"USA\",\n",
    "    dense_logs = dense_logs,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is evident that different government policies result in very different behaviors and result in varying values of social welfare (SW) (the SW numbers are shown in parentheses in the legend):\\\n",
    "Real-World: 347.6\\\n",
    "Simulation with real-world policies: 349.1\\\n",
    "Simulation with states open, and no subsidies: 183.8\\\n",
    "Simulation with states closed, and full subsidies: 317.8\\\n",
    "Simulation with states closed for 6 months and open after, and no subsidies: 334.2\n",
    "\n",
    "Some of the key takeaways here are:\n",
    "- Higher stringency results in fewer infections and deaths, but also increased unemployment. This leads to the highest health index, but the lowest economic index.\n",
    "- Lower stringency results in exactly the opposite trend - the least unemployment, since everything is open, but also, leads to the highest deaths.\n",
    "- Interestingly, note that the real-world policy (```'sim_with_real_world_policies'```) leads to the highest social welfare, so (according to our simulation) the real-world actions were better than these 3 simple alternatives!\n",
    "\n",
    "We note that the welfare numbers provided here are for the USA (federal government). Feel free to play around with the ```entity``` argument in the cell above to visualize the curves for different states too."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Varying social welfare objectives"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our simulation is also flexible to allow for any set of health and economic priority scalings / multipliers for the agents and the planner, resulting in varying social welfare functions. Note that these scalings multiply on top of the health and economic priorities inferred from the real data (for details, see this [notebook](https://github.com/salesforce/ai-economist/blob/master/ai_economist/datasets/covid19_datasets/fit_model_parameters.ipynb)).\n",
    "\n",
    "We refer to the tuple $(m_a = 1, m_p = 1)$ as our base scenario, wherein the priority multipliers for both the agents (states) and the planner (the USA) is 1. Accordingly, $(m_a = 2, m_p = 1)$ means (all) the states prioritize health about twice as as the data would suggest, $(m_a = 0.5, m_p = 1)$ means the states prioritize *the economy* twice as much, and $(m_a = 1, m_p = 2)$ means the *federal government* prioritizes health twice as much.\n",
    "\n",
    "Let us generate some dense logs with different health and economy priority scalings and compare them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's consider different agent priorities and visualize the social welfare for the simulation with the set of actions where the states are always open and the federal government does not subsidize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# minimize print clutter\n",
    "\n",
    "# Social welfare comparisons for the state of California\n",
    "\n",
    "dense_logs = {}\n",
    "\n",
    "for m_a, m_p in [\n",
    "    (0.25, 1),\n",
    "    (0.5, 1),\n",
    "    (1, 1),    \n",
    "    (2, 1),\n",
    "    (4, 1)\n",
    "]:\n",
    "    print(f\"\\nRunning the simulation for priority {(m_a, m_p)}... \\n\")\n",
    "    prioritized_env_config = env_config.copy()\n",
    "    prioritized_env_config.update(\n",
    "        {\n",
    "            \"health_priority_scaling_agents\": m_a,\n",
    "            \"health_priority_scaling_planner\": m_p\n",
    "        }\n",
    "    )\n",
    "    dense_logs[f\"$(m_a={m_a}, m_p={m_p})$\"] = fetch_env_dense_log(\n",
    "        prioritized_env_config,\n",
    "        action_type=\"states_open_no_subsidies\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_states(\n",
    "    entity=\"California\",\n",
    "    dense_logs = dense_logs,\n",
    "    states = [\"Social Welfare\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looking at the social welfare numbers, we see that when the states prioritize the health the least, the social welfare is the highest, and the social welfare decreases as the states choose to prioritize health more and more. This ties back to the \"always open\" actions taken by the states here, as in - when the states do not prioritize health over the economy, keeping them always open yields the best social welfare"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's consider different planner priorities and visualize the social welfares for the simulation with the set of actions where the states are always closed and the federal government subsidizes the full (maximum) amount."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# minimize print clutter\n",
    "\n",
    "dense_logs = {}\n",
    "\n",
    "for m_a, m_p in [\n",
    "    (1, 0.25),\n",
    "    (1, 0.5),\n",
    "    (1, 1),\n",
    "    (1, 2),\n",
    "    (1, 4),\n",
    "]:\n",
    "    print(f\"\\nRunning the simulation for priority {(m_a, m_p)}... \\n\")\n",
    "    prioritized_env_config = env_config.copy()\n",
    "    prioritized_env_config.update(\n",
    "        {\n",
    "            \"health_priority_scaling_agents\": m_a,\n",
    "            \"health_priority_scaling_planner\": m_p\n",
    "        }\n",
    "    )\n",
    "    dense_logs[f\"$(m_a={m_a}, m_p={m_p})$\"] = fetch_env_dense_log(\n",
    "        prioritized_env_config,\n",
    "        action_type=\"states_closed_full_subsidies\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_states(\n",
    "    entity=\"USA\",\n",
    "    dense_logs = dense_logs,\n",
    "    states = [\"Social Welfare\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The social welfare for the planner (when it subsidizes the full amount) is highest also when it prioritizes the economy the most, and least when it prioritizes the economy the least, understandably so.\n",
    "\n",
    "In general, when states prioritize health, they like to shut down more , thus the health index improves at the cost of the economic index (due to increased unemployment from shutting down). When the federal government prioritizes health, it tends to spend more in subsidy to help states economically during shut downs (so they can afford to shut down more), which results in a higher health index. However, subsidy-driven shutdown has a high economic cost for the federal government. On the other hand, subsidies increase the state-level economic index. In our [paper](https://arxiv.org/abs/2108.02904), we use multi-agent reinforcement learning agent to optimize the state and federal government policies for several different social welfare objectives."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that's it for this tutorial. Feel free to play around with the environment and visualize the effects for different entities, action types and health/economy priorities."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "948099c6ab02a15a055545cbca87e716aee9b3e6a51d0f68b03005577a92a5b6"
  },
  "kernelspec": {
   "display_name": "Python 3.7.11 64-bit ('ai-economist': conda)",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}