1688 lines
70 KiB
Python
1688 lines
70 KiB
Python
# Copyright (c) 2021, salesforce.com, inc.
|
|
# All rights reserved.
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
# For full license text, see the LICENSE file in the repo root
|
|
# or https://opensource.org/licenses/BSD-3-Clause
|
|
|
|
import json
|
|
import os
|
|
from datetime import datetime, timedelta
|
|
|
|
import GPUtil
|
|
import numpy as np
|
|
|
|
from ai_economist.foundation.base.base_env import BaseEnvironment, scenario_registry
|
|
from ai_economist.foundation.utils import verify_activation_code
|
|
|
|
try:
|
|
num_gpus_available = len(GPUtil.getAvailable())
|
|
print(f"Inside covid19_env.py: {num_gpus_available} GPUs are available.")
|
|
if num_gpus_available == 0:
|
|
print("No GPUs found! Running the simulation on a CPU.")
|
|
else:
|
|
from warp_drive.utils.constants import Constants
|
|
from warp_drive.utils.data_feed import DataFeed
|
|
|
|
_OBSERVATIONS = Constants.OBSERVATIONS
|
|
_ACTIONS = Constants.ACTIONS
|
|
_REWARDS = Constants.REWARDS
|
|
except ModuleNotFoundError:
|
|
print(
|
|
"Warning: The 'WarpDrive' package is not found and cannot be used! "
|
|
"If you wish to use WarpDrive, please run "
|
|
"'pip install rl-warp-drive' first."
|
|
)
|
|
except ValueError:
|
|
print("No GPUs found! Running the simulation on a CPU.")
|
|
|
|
|
|
@scenario_registry.add
|
|
class CovidAndEconomyEnvironment(BaseEnvironment):
|
|
"""
|
|
A simulation to model health and economy dynamics amidst the COVID-19 pandemic.
|
|
The environment comprising 51 agents (each agent corresponding to a US state and
|
|
Washington D.C.) and the Federal Government (planner). The state agents decide the
|
|
stringency level of the policy response to the pandemic, while the federal
|
|
government provides subsidies to eligible individuals.
|
|
|
|
This simulation makes modeling assumptions. For details, see the technical paper:
|
|
https://arxiv.org/abs/2108.02904
|
|
|
|
Args:
|
|
use_real_world_data (bool): Replay what happened in the real world.
|
|
Real-world data comprises SIR (susceptible/infected/recovered),
|
|
unemployment, government policy, and vaccination numbers.
|
|
This setting also sets use_real_world_policies=True.
|
|
use_real_world_policies (bool): Run the environment with real-world policies
|
|
(stringency levels and subsidies). With this setting and
|
|
use_real_world_data=False, SIR and economy dynamics are still
|
|
driven by fitted models.
|
|
path_to_data_and_fitted_params (dirpath): Full path to the directory containing
|
|
the data, fitted parameters and model constants. This defaults to
|
|
"ai_economist/datasets/covid19_datasets/data_and_fitted_params".
|
|
For details on obtaining these parameters, please see the notebook
|
|
"ai-economist-foundation/ai_economist/datasets/covid19_datasets/
|
|
gather_real_world_data_and_fit_parameters.ipynb".
|
|
start_date (string): Date (YYYY-MM-DD) to start the simulation.
|
|
pop_between_age_18_65 (float): Fraction of the population between ages 18-65.
|
|
This is the subset of the population whose employment/unemployment affects
|
|
economic productivity.
|
|
Range: 0 <= pop_between_age_18_65 <= 1.
|
|
infection_too_sick_to_work_rate (float): Fraction of people infected with
|
|
COVID-19. Infected people don't work.
|
|
Range: 0 <= infection_too_sick_to_work_rate <= 1
|
|
risk_free_interest_rate (float): Percentage of interest paid by the federal
|
|
government to borrow money from the federal reserve for COVID-19 relief
|
|
(direct payments). Higher interest rates mean that direct payments
|
|
have a larger cost on the federal government's economic index.
|
|
Range: 0 <= risk_free_interest_rate
|
|
economic_reward_crra_eta (float): CRRA eta parameter for modeling the economic
|
|
reward non-linearity.
|
|
A useful reference: https://en.wikipedia.org/wiki/Isoelastic_utility
|
|
Range: 0 <= economic_reward_crra_eta
|
|
health_priority_scaling_agents (float): A factor indicating how much more the
|
|
states prioritize health (roughly speaking, loss of lives due to
|
|
opening up more) over the economy (roughly speaking, a loss in GDP
|
|
due to shutting down resulting in more unemployment) compared to the
|
|
real-world.
|
|
For example, a value of 1 corresponds to the real-world, while
|
|
a value of 2 means that states cared twice as much about public health
|
|
(preventing deaths), while a value of 0.5 means that states cared twice
|
|
as much about the economy (preventing GDP drops).
|
|
Range: 0 <= health_priority_scaling_agents
|
|
health_priority_scaling_planner (float): same as above,
|
|
but for the federal government.
|
|
Range: 0 <= health_priority_scaling_planner
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*base_env_args,
|
|
use_real_world_data=False,
|
|
use_real_world_policies=False,
|
|
path_to_data_and_fitted_params="",
|
|
start_date="2020-03-22",
|
|
pop_between_age_18_65=0.6,
|
|
infection_too_sick_to_work_rate=0.1,
|
|
risk_free_interest_rate=0.03,
|
|
economic_reward_crra_eta=2,
|
|
health_priority_scaling_agents=1,
|
|
health_priority_scaling_planner=1,
|
|
reward_normalization_factor=1,
|
|
**base_env_kwargs,
|
|
):
|
|
verify_activation_code()
|
|
|
|
# Used for datatype checks
|
|
self.np_float_dtype = np.float32
|
|
self.np_int_dtype = np.int32
|
|
|
|
# Flag to use real-world data or the fitted models instead
|
|
self.use_real_world_data = use_real_world_data
|
|
# Flag to use real-world policies (actions) or the supplied actions instead
|
|
self.use_real_world_policies = use_real_world_policies
|
|
|
|
# If we use real-world data, we also want to use the real-world policies
|
|
if self.use_real_world_data:
|
|
print(
|
|
"Using real-world data to initialize as well as to "
|
|
"step through the env."
|
|
)
|
|
# Note: under this setting, the real_world policies are also used.
|
|
assert self.use_real_world_policies, (
|
|
"Since the env. config. 'use_real_world_data' is True, please also "
|
|
"set 'use_real_world_policies' to True."
|
|
)
|
|
else:
|
|
print(
|
|
"Using the real-world data to only initialize the env, "
|
|
"and using the fitted models to step through the env."
|
|
)
|
|
|
|
# Load real-world date
|
|
if path_to_data_and_fitted_params == "":
|
|
current_dir = os.path.dirname(__file__)
|
|
self.path_to_data_and_fitted_params = os.path.join(
|
|
current_dir, "../../../datasets/covid19_datasets/data_and_fitted_params"
|
|
)
|
|
else:
|
|
self.path_to_data_and_fitted_params = path_to_data_and_fitted_params
|
|
|
|
print(
|
|
"Loading real-world data from {}".format(
|
|
self.path_to_data_and_fitted_params
|
|
)
|
|
)
|
|
real_world_data_npz = np.load(
|
|
os.path.join(self.path_to_data_and_fitted_params, "real_world_data.npz")
|
|
)
|
|
self._real_world_data = {}
|
|
for key in list(real_world_data_npz):
|
|
self._real_world_data[key] = real_world_data_npz[key]
|
|
|
|
# Load fitted parameters
|
|
print(
|
|
"Loading fit parameters from {}".format(self.path_to_data_and_fitted_params)
|
|
)
|
|
self.load_model_constants(self.path_to_data_and_fitted_params)
|
|
self.load_fitted_params(self.path_to_data_and_fitted_params)
|
|
|
|
try:
|
|
self.start_date = datetime.strptime(start_date, self.date_format)
|
|
except ValueError:
|
|
print(f"Incorrect data format, should be {self.date_format}")
|
|
|
|
# Start date should be beyond the date for which data is available
|
|
assert self.start_date >= self.policy_start_date
|
|
|
|
# Compute a start date index based on policy start date
|
|
self.start_date_index = (self.start_date - self.policy_start_date).days
|
|
assert 0 <= self.start_date_index < len(self._real_world_data["policy"])
|
|
|
|
# For date logging (This will be overwritten in additional_reset_steps;
|
|
# see below)
|
|
self.current_date = None
|
|
|
|
# When using real-world policy, limit the episode length
|
|
# to the length of the available policy.
|
|
if self.use_real_world_policies:
|
|
real_world_policy_length = (
|
|
len(self._real_world_data["policy"]) - self.start_date_index
|
|
)
|
|
print("Using real-world policies, ignoring external action inputs.")
|
|
assert base_env_kwargs["episode_length"] <= real_world_policy_length, (
|
|
f"The real-world policies are only available for "
|
|
f"{real_world_policy_length} timesteps; so the 'episode_length' "
|
|
f"in the environment configuration can only be at most "
|
|
f"{real_world_policy_length}"
|
|
)
|
|
else:
|
|
print("Using external action inputs.")
|
|
|
|
# US states and populations
|
|
self.num_us_states = len(self.us_state_population)
|
|
|
|
assert (
|
|
base_env_kwargs["n_agents"] == self.num_us_states
|
|
), "n_agents should be set to the number of US states, i.e., {}.".format(
|
|
self.num_us_states
|
|
)
|
|
# Note: For a faster environment step time, we collate all the individual agents
|
|
# into a single agent index "a" and we flatten the component action masks too.
|
|
assert base_env_kwargs[
|
|
"collate_agent_step_and_reset_data"
|
|
], "The env. config 'collate_agent_step_and_reset_data' should be set to True."
|
|
super().__init__(*base_env_args, **base_env_kwargs)
|
|
|
|
# Add attributes to self.world for use in components
|
|
self.world.us_state_population = self.us_state_population
|
|
self.world.us_population = self.us_population
|
|
self.world.start_date = self.start_date
|
|
self.world.n_stringency_levels = self.num_stringency_levels
|
|
self.world.use_real_world_policies = self.use_real_world_policies
|
|
if self.use_real_world_policies:
|
|
# Agent open/close stringency levels
|
|
self.world.real_world_stringency_policy = self._real_world_data["policy"][
|
|
self.start_date_index :
|
|
]
|
|
# Planner subsidy levels
|
|
self.world.real_world_subsidy = self._real_world_data["subsidy"][
|
|
self.start_date_index :
|
|
]
|
|
|
|
# Policy --> Unemployment
|
|
# For accurately modeling the state-wise unemployment, we convolve
|
|
# the current stringency policy with a family of exponential filters
|
|
# with separate means (lambdas).
|
|
# This code sets up things we will use in `unemployment_step()`,
|
|
# which includes a detailed breakdown of how the unemployment model is
|
|
# implemented.
|
|
self.stringency_level_history = None
|
|
# Each filter captures a temporally extended response to a stringency change.
|
|
self.num_filters = len(self.conv_lambdas)
|
|
self.f_ts = np.tile(
|
|
np.flip(np.arange(self.filter_len), (0,))[None, None],
|
|
(1, self.num_filters, 1),
|
|
).astype(self.np_float_dtype)
|
|
self.unemp_conv_filters = np.exp(-self.f_ts / self.conv_lambdas[None, :, None])
|
|
# Each state weights these filters differently.
|
|
self.repeated_conv_weights = np.repeat(
|
|
self.grouped_convolutional_filter_weights.reshape(
|
|
self.num_us_states, self.num_filters
|
|
)[:, :, np.newaxis],
|
|
self.filter_len,
|
|
axis=-1,
|
|
)
|
|
|
|
# For manually modulating SIR/Unemployment parameters
|
|
self._beta_intercepts_modulation = 1
|
|
self._beta_slopes_modulation = 1
|
|
self._unemployment_modulation = 1
|
|
|
|
# Economy-related
|
|
# Interest rate for borrowing money from the federal reserve
|
|
self.risk_free_interest_rate = self.np_float_dtype(risk_free_interest_rate)
|
|
|
|
# Compute each worker's daily productivity when at work (to match 2019 GDP)
|
|
# We assume the open/close stringency policy level was always at it's lowest
|
|
# value (i.e., 1) before the pandemic started.
|
|
num_unemployed_at_stringency_level_1 = self.unemployment_step(
|
|
np.ones(self.num_us_states)
|
|
)
|
|
workforce = (
|
|
self.us_population * pop_between_age_18_65
|
|
- np.sum(num_unemployed_at_stringency_level_1)
|
|
).astype(self.np_int_dtype)
|
|
workers_per_capita = (workforce / self.us_population).astype(
|
|
self.np_float_dtype
|
|
)
|
|
gdp_per_worker = (self.gdp_per_capita / workers_per_capita).astype(
|
|
self.np_float_dtype
|
|
)
|
|
self.num_days_in_an_year = 365
|
|
self.daily_production_per_worker = (
|
|
gdp_per_worker / self.num_days_in_an_year
|
|
).astype(self.np_float_dtype)
|
|
|
|
self.infection_too_sick_to_work_rate = self.np_float_dtype(
|
|
infection_too_sick_to_work_rate
|
|
)
|
|
assert 0 <= self.infection_too_sick_to_work_rate <= 1
|
|
|
|
self.pop_between_age_18_65 = self.np_float_dtype(pop_between_age_18_65)
|
|
assert 0 <= self.pop_between_age_18_65 <= 1
|
|
|
|
# Compute max possible productivity values (used for agent reward normalization)
|
|
max_productivity_t = self.economy_step(
|
|
self.us_state_population,
|
|
np.zeros((self.num_us_states), dtype=self.np_int_dtype),
|
|
np.zeros((self.num_us_states), dtype=self.np_int_dtype),
|
|
num_unemployed_at_stringency_level_1,
|
|
infection_too_sick_to_work_rate=self.infection_too_sick_to_work_rate,
|
|
population_between_age_18_65=self.pop_between_age_18_65,
|
|
)
|
|
self.maximum_productivity_t = max_productivity_t
|
|
|
|
# Economic reward non-linearity
|
|
self.economic_reward_crra_eta = self.np_float_dtype(economic_reward_crra_eta)
|
|
assert 0.0 <= self.economic_reward_crra_eta < 20.0
|
|
|
|
# Health indices are normalized by maximum annual GDP
|
|
self.agents_health_norm = self.maximum_productivity_t * self.num_days_in_an_year
|
|
self.planner_health_norm = np.sum(self.agents_health_norm)
|
|
|
|
# Economic indices are normalized by maximum annual GDP
|
|
self.agents_economic_norm = (
|
|
self.maximum_productivity_t * self.num_days_in_an_year
|
|
)
|
|
self.planner_economic_norm = np.sum(self.agents_economic_norm)
|
|
|
|
def scale_health_over_economic_index(health_priority_scaling, alphas):
|
|
"""
|
|
Given starting alpha(s), compute new alphas so that the
|
|
resulting alpha:1-alpha ratio is scaled by health_weightage
|
|
"""
|
|
z = alphas / (1 - alphas) # alphas = z / (1 + z)
|
|
scaled_z = health_priority_scaling * z
|
|
new_alphas = scaled_z / (1 + scaled_z)
|
|
return new_alphas
|
|
|
|
# Agents' health and economic index weightages
|
|
# fmt: off
|
|
self.weightage_on_marginal_agent_health_index = \
|
|
scale_health_over_economic_index(
|
|
health_priority_scaling_agents,
|
|
self.inferred_weightage_on_agent_health_index,
|
|
)
|
|
# fmt: on
|
|
assert (
|
|
(self.weightage_on_marginal_agent_health_index >= 0)
|
|
& (self.weightage_on_marginal_agent_health_index <= 1)
|
|
).all()
|
|
self.weightage_on_marginal_agent_economic_index = (
|
|
1 - self.weightage_on_marginal_agent_health_index
|
|
)
|
|
|
|
# Planner's health and economic index weightages
|
|
# fmt: off
|
|
self.weightage_on_marginal_planner_health_index = \
|
|
scale_health_over_economic_index(
|
|
health_priority_scaling_planner,
|
|
self.inferred_weightage_on_planner_health_index,
|
|
)
|
|
# fmt: on
|
|
assert 0 <= self.weightage_on_marginal_planner_health_index <= 1
|
|
self.weightage_on_marginal_planner_economic_index = (
|
|
1 - self.weightage_on_marginal_planner_health_index
|
|
)
|
|
|
|
# Normalization factor for the reward (often useful for RL training)
|
|
self.reward_normalization_factor = reward_normalization_factor
|
|
|
|
# CUDA-related attributes (for GPU simulations)
|
|
# Note: these will be set / overwritten via the env_wrapper
|
|
# use_cuda will be set to True (by the env_wrapper), if needed
|
|
# to be simulated on the GPU
|
|
self.use_cuda = False
|
|
self.cuda_data_manager = None
|
|
self.cuda_function_manager = None
|
|
self.cuda_step = lambda *args, **kwargs: None
|
|
self.cuda_compute_reward = lambda *args, **kwargs: None
|
|
|
|
# Adding use_cuda to self.world for use in components
|
|
self.world.use_cuda = self.use_cuda
|
|
self.world.cuda_data_manager = self.cuda_data_manager
|
|
self.world.cuda_function_manager = self.cuda_function_manager
|
|
|
|
name = "CovidAndEconomySimulation"
|
|
agent_subclasses = ["BasicMobileAgent", "BasicPlanner"]
|
|
|
|
required_entities = []
|
|
|
|
def reset_starting_layout(self):
|
|
pass
|
|
|
|
def reset_agent_states(self):
|
|
self.world.clear_agent_locs()
|
|
|
|
def get_data_dictionary(self):
|
|
"""
|
|
Create a dictionary of data to push to the GPU (device).
|
|
"""
|
|
data_dict = DataFeed()
|
|
# Global States
|
|
data_dict.add_data(
|
|
name="susceptible",
|
|
data=self.world.global_state["Susceptible"],
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="infected",
|
|
data=self.world.global_state["Infected"],
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="recovered",
|
|
data=self.world.global_state["Recovered"],
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="deaths",
|
|
data=self.world.global_state["Deaths"],
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="unemployed",
|
|
data=self.world.global_state["Unemployed"],
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="vaccinated",
|
|
data=self.world.global_state["Vaccinated"],
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
# Actions
|
|
data_dict.add_data(
|
|
name="stringency_level",
|
|
data=self.world.global_state["Stringency Level"].astype(self.np_int_dtype),
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="subsidy_level",
|
|
data=self.world.global_state["Subsidy Level"].astype(self.np_int_dtype),
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
# Economy-related
|
|
data_dict.add_data(
|
|
name="subsidy",
|
|
data=self.world.global_state["Subsidy"],
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="postsubsidy_productivity",
|
|
data=self.world.global_state["Postsubsidy Productivity"],
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="productivity",
|
|
data=np.zeros_like(
|
|
self.world.global_state["Susceptible"], dtype=self.np_float_dtype
|
|
),
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="incapacitated",
|
|
data=np.zeros((self.num_us_states), dtype=self.np_float_dtype),
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="cant_work",
|
|
data=np.zeros((self.num_us_states), dtype=self.np_float_dtype),
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="num_people_that_can_work",
|
|
data=np.zeros((self.num_us_states), dtype=self.np_float_dtype),
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="us_state_population",
|
|
data=self.us_state_population,
|
|
)
|
|
data_dict.add_data(
|
|
name="infection_too_sick_to_work_rate",
|
|
data=self.infection_too_sick_to_work_rate,
|
|
)
|
|
data_dict.add_data(
|
|
name="population_between_age_18_65",
|
|
data=self.pop_between_age_18_65,
|
|
)
|
|
data_dict.add_data(
|
|
name="daily_production_per_worker",
|
|
data=self.daily_production_per_worker,
|
|
)
|
|
data_dict.add_data(
|
|
name="maximum_productivity",
|
|
data=self.maximum_productivity_t,
|
|
)
|
|
# SIR-related
|
|
data_dict.add_data(
|
|
name="real_world_stringency_policy_history",
|
|
data=(
|
|
self._real_world_data["policy"][
|
|
self.start_date_index - self.beta_delay + 1 : self.start_date_index,
|
|
:,
|
|
]
|
|
).astype(self.np_int_dtype),
|
|
)
|
|
data_dict.add_data(
|
|
name="beta_delay",
|
|
data=self.beta_delay,
|
|
)
|
|
data_dict.add_data(
|
|
name="beta_slopes",
|
|
data=self.beta_slopes,
|
|
)
|
|
data_dict.add_data(
|
|
name="beta_intercepts",
|
|
data=self.beta_intercepts,
|
|
)
|
|
data_dict.add_data(
|
|
name="beta",
|
|
data=np.zeros((self.num_us_states), dtype=self.np_float_dtype),
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="gamma",
|
|
data=self.gamma,
|
|
)
|
|
data_dict.add_data(
|
|
name="death_rate",
|
|
data=self.death_rate,
|
|
)
|
|
# Unemployment fit parameters
|
|
data_dict.add_data(
|
|
name="filter_len",
|
|
data=self.filter_len,
|
|
)
|
|
data_dict.add_data(
|
|
name="num_filters",
|
|
data=self.num_filters,
|
|
)
|
|
data_dict.add_data(
|
|
name="delta_stringency_level",
|
|
data=(
|
|
self.stringency_level_history[1:] - self.stringency_level_history[:-1]
|
|
).astype(self.np_int_dtype),
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
data_dict.add_data(
|
|
name="grouped_convolutional_filter_weights",
|
|
data=self.grouped_convolutional_filter_weights,
|
|
)
|
|
data_dict.add_data(
|
|
name="unemp_conv_filters",
|
|
data=self.unemp_conv_filters,
|
|
)
|
|
data_dict.add_data(
|
|
name="unemployment_bias",
|
|
data=self.unemployment_bias,
|
|
)
|
|
data_dict.add_data(
|
|
name="signal",
|
|
data=np.zeros(
|
|
(self.n_agents, self.num_filters, self.filter_len),
|
|
dtype=self.np_float_dtype,
|
|
),
|
|
save_copy_and_apply_at_reset=True,
|
|
)
|
|
# Reward-related
|
|
data_dict.add_data(
|
|
name="min_marginal_agent_health_index",
|
|
data=self.min_marginal_agent_health_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="max_marginal_agent_health_index",
|
|
data=self.max_marginal_agent_health_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="min_marginal_agent_economic_index",
|
|
data=self.min_marginal_agent_economic_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="max_marginal_agent_economic_index",
|
|
data=self.max_marginal_agent_economic_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="min_marginal_planner_health_index",
|
|
data=self.min_marginal_planner_health_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="max_marginal_planner_health_index",
|
|
data=self.max_marginal_planner_health_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="min_marginal_planner_economic_index",
|
|
data=self.min_marginal_planner_economic_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="max_marginal_planner_economic_index",
|
|
data=self.max_marginal_planner_economic_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="weightage_on_marginal_agent_health_index",
|
|
data=self.weightage_on_marginal_agent_health_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="weightage_on_marginal_agent_economic_index",
|
|
data=self.weightage_on_marginal_agent_economic_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="weightage_on_marginal_planner_health_index",
|
|
data=self.weightage_on_marginal_planner_health_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="weightage_on_marginal_planner_economic_index",
|
|
data=self.weightage_on_marginal_planner_economic_index,
|
|
)
|
|
data_dict.add_data(
|
|
name="value_of_life",
|
|
data=self.value_of_life,
|
|
)
|
|
data_dict.add_data(
|
|
name="economic_reward_crra_eta",
|
|
data=self.economic_reward_crra_eta,
|
|
)
|
|
data_dict.add_data(
|
|
name="num_days_in_an_year",
|
|
data=self.num_days_in_an_year,
|
|
)
|
|
data_dict.add_data(
|
|
name="risk_free_interest_rate",
|
|
data=self.risk_free_interest_rate,
|
|
)
|
|
data_dict.add_data(
|
|
name="agents_health_norm",
|
|
data=self.agents_health_norm,
|
|
)
|
|
data_dict.add_data(
|
|
name="agents_economic_norm",
|
|
data=self.agents_economic_norm,
|
|
)
|
|
data_dict.add_data(
|
|
name="planner_health_norm",
|
|
data=self.planner_health_norm,
|
|
)
|
|
data_dict.add_data(
|
|
name="planner_economic_norm",
|
|
data=self.planner_economic_norm,
|
|
)
|
|
|
|
return data_dict
|
|
|
|
def get_tensor_dictionary(self):
|
|
"""
|
|
Create a dictionary of (Pytorch-accessible) data to push to the GPU (device).
|
|
"""
|
|
tensor_dict = DataFeed()
|
|
return tensor_dict
|
|
|
|
def scenario_step(self):
|
|
"""
|
|
Update the state of the USA based on the Covid-19 and Economy dynamics.
|
|
This internally implements three steps
|
|
- sir_step() - updates the susceptible, infected, recovered, deaths
|
|
and vaccination numbers based on the SIR equations
|
|
- unemployment_step() - uses the unemployment model to updates the unemployment
|
|
based on the stringency levels
|
|
- economy_step - computes the current producitivity numbers for the agents
|
|
"""
|
|
if self.use_cuda:
|
|
self.cuda_step(
|
|
self.cuda_data_manager.device_data("susceptible"),
|
|
self.cuda_data_manager.device_data("infected"),
|
|
self.cuda_data_manager.device_data("recovered"),
|
|
self.cuda_data_manager.device_data("deaths"),
|
|
self.cuda_data_manager.device_data("vaccinated"),
|
|
self.cuda_data_manager.device_data("unemployed"),
|
|
self.cuda_data_manager.device_data("subsidy"),
|
|
self.cuda_data_manager.device_data("productivity"),
|
|
self.cuda_data_manager.device_data("stringency_level"),
|
|
self.cuda_data_manager.device_data("num_stringency_levels"),
|
|
self.cuda_data_manager.device_data("postsubsidy_productivity"),
|
|
self.cuda_data_manager.device_data("num_vaccines_available_t"),
|
|
self.cuda_data_manager.device_data(
|
|
"real_world_stringency_policy_history"
|
|
),
|
|
self.cuda_data_manager.device_data("beta_delay"),
|
|
self.cuda_data_manager.device_data("beta_slopes"),
|
|
self.cuda_data_manager.device_data("beta_intercepts"),
|
|
self.cuda_data_manager.device_data("beta"),
|
|
self.cuda_data_manager.device_data("gamma"),
|
|
self.cuda_data_manager.device_data("death_rate"),
|
|
self.cuda_data_manager.device_data("incapacitated"),
|
|
self.cuda_data_manager.device_data("cant_work"),
|
|
self.cuda_data_manager.device_data("num_people_that_can_work"),
|
|
self.cuda_data_manager.device_data("us_state_population"),
|
|
self.cuda_data_manager.device_data("infection_too_sick_to_work_rate"),
|
|
self.cuda_data_manager.device_data("population_between_age_18_65"),
|
|
self.cuda_data_manager.device_data("filter_len"),
|
|
self.cuda_data_manager.device_data("num_filters"),
|
|
self.cuda_data_manager.device_data("delta_stringency_level"),
|
|
self.cuda_data_manager.device_data(
|
|
"grouped_convolutional_filter_weights"
|
|
),
|
|
self.cuda_data_manager.device_data("unemp_conv_filters"),
|
|
self.cuda_data_manager.device_data("unemployment_bias"),
|
|
self.cuda_data_manager.device_data("signal"),
|
|
self.cuda_data_manager.device_data("daily_production_per_worker"),
|
|
self.cuda_data_manager.device_data("maximum_productivity"),
|
|
self.cuda_data_manager.device_data(
|
|
f"{_OBSERVATIONS}_a_world-agent_state"
|
|
),
|
|
self.cuda_data_manager.device_data(
|
|
f"{_OBSERVATIONS}_a_world-agent_postsubsidy_productivity"
|
|
),
|
|
self.cuda_data_manager.device_data(
|
|
f"{_OBSERVATIONS}_a_world-lagged_stringency_level"
|
|
),
|
|
self.cuda_data_manager.device_data(f"{_OBSERVATIONS}_a_time"),
|
|
self.cuda_data_manager.device_data(
|
|
f"{_OBSERVATIONS}_p_world-agent_state"
|
|
),
|
|
self.cuda_data_manager.device_data(
|
|
f"{_OBSERVATIONS}_p_world-agent_postsubsidy_productivity"
|
|
),
|
|
self.cuda_data_manager.device_data(
|
|
f"{_OBSERVATIONS}_p_world-lagged_stringency_level"
|
|
),
|
|
self.cuda_data_manager.device_data(f"{_OBSERVATIONS}_p_time"),
|
|
self.cuda_data_manager.device_data("_timestep_"),
|
|
self.cuda_data_manager.meta_info("n_agents"),
|
|
self.cuda_data_manager.meta_info("episode_length"),
|
|
block=self.world.cuda_function_manager.block,
|
|
grid=self.world.cuda_function_manager.grid,
|
|
)
|
|
else:
|
|
prev_t = self.world.timestep - 1
|
|
curr_t = self.world.timestep
|
|
|
|
self.current_date += timedelta(days=1)
|
|
|
|
# SIR
|
|
# ---
|
|
if self.use_real_world_data:
|
|
_S_t = np.maximum(
|
|
self._real_world_data["susceptible"][
|
|
curr_t + self.start_date_index
|
|
],
|
|
0,
|
|
)
|
|
_I_t = np.maximum(
|
|
self._real_world_data["infected"][curr_t + self.start_date_index],
|
|
0,
|
|
)
|
|
_R_t = np.maximum(
|
|
self._real_world_data["recovered"][curr_t + self.start_date_index],
|
|
0,
|
|
)
|
|
_V_t = np.maximum(
|
|
self._real_world_data["vaccinated"][curr_t + self.start_date_index],
|
|
0,
|
|
)
|
|
_D_t = np.maximum(
|
|
self._real_world_data["deaths"][curr_t + self.start_date_index],
|
|
0,
|
|
)
|
|
|
|
else: # Use simulation logic
|
|
if curr_t - self.beta_delay < 0:
|
|
if self.start_date_index + curr_t - self.beta_delay < 0:
|
|
stringency_level_tmk = np.ones(self.num_us_states)
|
|
else:
|
|
stringency_level_tmk = self._real_world_data["policy"][
|
|
self.start_date_index + curr_t - self.beta_delay, :
|
|
]
|
|
else:
|
|
stringency_level_tmk = self.world.global_state["Stringency Level"][
|
|
curr_t - self.beta_delay
|
|
]
|
|
stringency_level_tmk = stringency_level_tmk.astype(self.np_int_dtype)
|
|
|
|
_S_tm1 = self.world.global_state["Susceptible"][prev_t]
|
|
_I_tm1 = self.world.global_state["Infected"][prev_t]
|
|
_R_tm1 = self.world.global_state["Recovered"][prev_t]
|
|
_V_tm1 = self.world.global_state["Vaccinated"][prev_t]
|
|
|
|
# Vaccination
|
|
# -----------
|
|
num_vaccines_available_t = np.zeros(
|
|
self.n_agents, dtype=self.np_int_dtype
|
|
)
|
|
for aidx, agent in enumerate(self.world.agents):
|
|
# "Load" the vaccines in the inventory into this vector.
|
|
num_vaccines_available_t[aidx] = agent.state["Vaccines Available"]
|
|
# Agents always use whatever vaccines they can, so this becomes 0:
|
|
agent.state["Total Vaccinated"] += agent.state["Vaccines Available"]
|
|
agent.state["Vaccines Available"] = 0
|
|
|
|
# SIR step
|
|
# --------
|
|
_dS, _dI, _dR, _dV = self.sir_step(
|
|
_S_tm1,
|
|
_I_tm1,
|
|
stringency_level_tmk,
|
|
num_vaccines_available_t,
|
|
)
|
|
_S_t = np.maximum(_S_tm1 + _dS, 0)
|
|
_I_t = np.maximum(_I_tm1 + _dI, 0)
|
|
_R_t = np.maximum(_R_tm1 + _dR, 0)
|
|
_V_t = np.maximum(_V_tm1 + _dV, 0)
|
|
|
|
num_recovered_but_not_vaccinated_t = _R_t - _V_t
|
|
_D_t = self.death_rate * num_recovered_but_not_vaccinated_t
|
|
|
|
# Update global state
|
|
# -------------------
|
|
self.world.global_state["Susceptible"][curr_t] = _S_t
|
|
self.world.global_state["Infected"][curr_t] = _I_t
|
|
self.world.global_state["Recovered"][curr_t] = _R_t
|
|
self.world.global_state["Deaths"][curr_t] = _D_t
|
|
self.world.global_state["Vaccinated"][curr_t] = _V_t
|
|
|
|
# Unemployment
|
|
# ------------
|
|
if self.use_real_world_data:
|
|
num_unemployed_t = self._real_world_data["unemployed"][
|
|
self.start_date_index + curr_t
|
|
]
|
|
else:
|
|
num_unemployed_t = self.unemployment_step(
|
|
current_stringency_level=self.world.global_state[
|
|
"Stringency Level"
|
|
][curr_t]
|
|
)
|
|
|
|
self.world.global_state["Unemployed"][curr_t] = num_unemployed_t
|
|
|
|
# Productivity
|
|
# ------------
|
|
productivity_t = self.economy_step(
|
|
self.us_state_population,
|
|
infected=_I_t,
|
|
deaths=_D_t,
|
|
unemployed=num_unemployed_t,
|
|
infection_too_sick_to_work_rate=self.infection_too_sick_to_work_rate,
|
|
population_between_age_18_65=self.pop_between_age_18_65,
|
|
)
|
|
|
|
# Subsidies
|
|
# ---------
|
|
# Add federal government subsidy to productivity
|
|
daily_statewise_subsidy_t = self.world.global_state["Subsidy"][curr_t]
|
|
postsubsidy_productivity_t = productivity_t + daily_statewise_subsidy_t
|
|
self.world.global_state["Postsubsidy Productivity"][
|
|
curr_t
|
|
] = postsubsidy_productivity_t
|
|
|
|
# Update agent state
|
|
# ------------------
|
|
current_date_string = datetime.strftime(
|
|
self.current_date, format=self.date_format
|
|
)
|
|
for agent in self.world.agents:
|
|
agent.state["Total Susceptible"] = _S_t[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
agent.state["New Infections"] = (
|
|
_I_t[agent.idx] - agent.state["Total Infected"]
|
|
).astype(self.np_int_dtype)
|
|
agent.state["Total Infected"] = _I_t[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
agent.state["Total Recovered"] = _R_t[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
agent.state["New Deaths"] = _D_t[agent.idx] - agent.state[
|
|
"Total Deaths"
|
|
].astype(self.np_int_dtype)
|
|
agent.state["Total Deaths"] = _D_t[agent.idx].astype(self.np_int_dtype)
|
|
agent.state["Total Vaccinated"] = _V_t[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
|
|
agent.state["Total Unemployed"] = num_unemployed_t[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
agent.state["New Subsidy Received"] = daily_statewise_subsidy_t[
|
|
agent.idx
|
|
]
|
|
agent.state["Postsubsidy Productivity"] = postsubsidy_productivity_t[
|
|
agent.idx
|
|
]
|
|
agent.state["Date"] = current_date_string
|
|
|
|
# Update planner state
|
|
# --------------------
|
|
self.world.planner.state["Total Susceptible"] = np.sum(_S_t).astype(
|
|
self.np_int_dtype
|
|
)
|
|
self.world.planner.state["New Infections"] = (
|
|
np.sum(_I_t) - self.world.planner.state["Total Infected"]
|
|
).astype(self.np_int_dtype)
|
|
self.world.planner.state["Total Infected"] = np.sum(_I_t).astype(
|
|
self.np_int_dtype
|
|
)
|
|
self.world.planner.state["Total Recovered"] = np.sum(_R_t).astype(
|
|
self.np_int_dtype
|
|
)
|
|
self.world.planner.state["New Deaths"] = (
|
|
np.sum(_D_t) - self.world.planner.state["Total Deaths"]
|
|
).astype(self.np_int_dtype)
|
|
self.world.planner.state["Total Deaths"] = np.sum(_D_t).astype(
|
|
self.np_int_dtype
|
|
)
|
|
self.world.planner.state["Total Vaccinated"] = np.sum(_V_t).astype(
|
|
self.np_int_dtype
|
|
)
|
|
self.world.planner.state["Total Unemployed"] = np.sum(
|
|
num_unemployed_t
|
|
).astype(self.np_int_dtype)
|
|
self.world.planner.state["New Subsidy Provided"] = np.sum(
|
|
daily_statewise_subsidy_t
|
|
)
|
|
self.world.planner.state["Postsubsidy Productivity"] = np.sum(
|
|
postsubsidy_productivity_t
|
|
)
|
|
self.world.planner.state["Date"] = current_date_string
|
|
|
|
def generate_observations(self):
|
|
"""
|
|
- Process agent-specific and planner-specific data into an observation.
|
|
- Observations contain only the relevant features for that actor.
|
|
:return: a dictionary of observations for each agent and planner
|
|
"""
|
|
redux_agent_global_state = None
|
|
for feature in [
|
|
"Susceptible",
|
|
"Infected",
|
|
"Recovered",
|
|
"Deaths",
|
|
"Vaccinated",
|
|
"Unemployed",
|
|
]:
|
|
if redux_agent_global_state is None:
|
|
redux_agent_global_state = self.world.global_state[feature][
|
|
self.world.timestep
|
|
]
|
|
else:
|
|
redux_agent_global_state = np.vstack(
|
|
(
|
|
redux_agent_global_state,
|
|
self.world.global_state[feature][self.world.timestep],
|
|
)
|
|
)
|
|
normalized_redux_agent_state = (
|
|
redux_agent_global_state / self.us_state_population[None]
|
|
)
|
|
|
|
# Productivity
|
|
postsubsidy_productivity_t = self.world.global_state[
|
|
"Postsubsidy Productivity"
|
|
][self.world.timestep]
|
|
normalized_postsubsidy_productivity_t = (
|
|
postsubsidy_productivity_t / self.maximum_productivity_t
|
|
)
|
|
|
|
# Let agents know about the policy about to affect SIR infection-rate beta
|
|
t_beta = self.world.timestep - self.beta_delay + 1
|
|
if t_beta < 0:
|
|
lagged_stringency_level = self._real_world_data["policy"][
|
|
self.start_date_index + t_beta
|
|
]
|
|
else:
|
|
lagged_stringency_level = self.world.global_state["Stringency Level"][
|
|
t_beta
|
|
]
|
|
|
|
normalized_lagged_stringency_level = (
|
|
lagged_stringency_level / self.num_stringency_levels
|
|
)
|
|
|
|
# To condition policy on agent id
|
|
agent_index = np.eye(self.n_agents, dtype=self.np_int_dtype)
|
|
|
|
# Observation dict - Agents
|
|
# -------------------------
|
|
obs_dict = dict()
|
|
obs_dict["a"] = {
|
|
"agent_index": agent_index,
|
|
"agent_state": normalized_redux_agent_state,
|
|
"agent_postsubsidy_productivity": normalized_postsubsidy_productivity_t,
|
|
"lagged_stringency_level": normalized_lagged_stringency_level,
|
|
}
|
|
|
|
# Observation dict - Planner
|
|
# --------------------------
|
|
obs_dict[self.world.planner.idx] = {
|
|
"agent_state": normalized_redux_agent_state,
|
|
"agent_postsubsidy_productivity": normalized_postsubsidy_productivity_t,
|
|
"lagged_stringency_level": normalized_lagged_stringency_level,
|
|
}
|
|
|
|
return obs_dict
|
|
|
|
def compute_reward(self):
|
|
"""
|
|
Compute the social welfare metrics for each agent and the planner.
|
|
:return: a dictionary of rewards for each agent in the simulation
|
|
"""
|
|
if self.use_cuda:
|
|
self.cuda_compute_reward(
|
|
self.cuda_data_manager.device_data(f"{_REWARDS}_a"),
|
|
self.cuda_data_manager.device_data(f"{_REWARDS}_p"),
|
|
self.cuda_data_manager.device_data("num_days_in_an_year"),
|
|
self.cuda_data_manager.device_data("value_of_life"),
|
|
self.cuda_data_manager.device_data("risk_free_interest_rate"),
|
|
self.cuda_data_manager.device_data("economic_reward_crra_eta"),
|
|
self.cuda_data_manager.device_data("min_marginal_agent_health_index"),
|
|
self.cuda_data_manager.device_data("max_marginal_agent_health_index"),
|
|
self.cuda_data_manager.device_data("min_marginal_agent_economic_index"),
|
|
self.cuda_data_manager.device_data("max_marginal_agent_economic_index"),
|
|
self.cuda_data_manager.device_data("min_marginal_planner_health_index"),
|
|
self.cuda_data_manager.device_data("max_marginal_planner_health_index"),
|
|
self.cuda_data_manager.device_data(
|
|
"min_marginal_planner_economic_index"
|
|
),
|
|
self.cuda_data_manager.device_data(
|
|
"max_marginal_planner_economic_index"
|
|
),
|
|
self.cuda_data_manager.device_data(
|
|
"weightage_on_marginal_agent_health_index"
|
|
),
|
|
self.cuda_data_manager.device_data(
|
|
"weightage_on_marginal_agent_economic_index"
|
|
),
|
|
self.cuda_data_manager.device_data(
|
|
"weightage_on_marginal_planner_health_index"
|
|
),
|
|
self.cuda_data_manager.device_data(
|
|
"weightage_on_marginal_planner_economic_index"
|
|
),
|
|
self.cuda_data_manager.device_data("agents_health_norm"),
|
|
self.cuda_data_manager.device_data("agents_economic_norm"),
|
|
self.cuda_data_manager.device_data("planner_health_norm"),
|
|
self.cuda_data_manager.device_data("planner_economic_norm"),
|
|
self.cuda_data_manager.device_data("deaths"),
|
|
self.cuda_data_manager.device_data("subsidy"),
|
|
self.cuda_data_manager.device_data("postsubsidy_productivity"),
|
|
self.cuda_data_manager.device_data("_done_"),
|
|
self.cuda_data_manager.device_data("_timestep_"),
|
|
self.cuda_data_manager.meta_info("n_agents"),
|
|
self.cuda_data_manager.meta_info("episode_length"),
|
|
block=self.world.cuda_function_manager.block,
|
|
grid=self.world.cuda_function_manager.grid,
|
|
)
|
|
return {} # Return empty dict. Reward arrays are updated in-place
|
|
rew = {"a": 0, "p": 0}
|
|
|
|
def crra_nonlinearity(x, eta):
|
|
# Reference: https://en.wikipedia.org/wiki/Isoelastic_utility
|
|
# To be applied to (marginal) economic indices
|
|
annual_x = self.num_days_in_an_year * x
|
|
annual_x_clipped = np.clip(annual_x, 0.1, 3)
|
|
annual_crra = 1 + (annual_x_clipped ** (1 - eta) - 1) / (1 - eta)
|
|
daily_crra = annual_crra / self.num_days_in_an_year
|
|
return daily_crra
|
|
|
|
def min_max_normalization(x, min_x, max_x):
|
|
eps = 1e-10
|
|
return (x - min_x) / (max_x - min_x + eps)
|
|
|
|
def get_weighted_average(
|
|
health_index_weightage,
|
|
health_index,
|
|
economic_index_weightage,
|
|
economic_index,
|
|
):
|
|
return (
|
|
health_index_weightage * health_index
|
|
+ economic_index_weightage * economic_index
|
|
) / (health_index_weightage + economic_index_weightage)
|
|
|
|
# Changes this last timestep:
|
|
marginal_deaths = (
|
|
self.world.global_state["Deaths"][self.world.timestep]
|
|
- self.world.global_state["Deaths"][self.world.timestep - 1]
|
|
)
|
|
|
|
subsidy_t = self.world.global_state["Subsidy"][self.world.timestep]
|
|
postsubsidy_productivity_t = self.world.global_state[
|
|
"Postsubsidy Productivity"
|
|
][self.world.timestep]
|
|
|
|
# Health index -- the cost equivalent (annual GDP) of covid deaths
|
|
# Note: casting deaths to float to prevent overflow issues
|
|
marginal_agent_health_index = (
|
|
-marginal_deaths.astype(self.np_float_dtype)
|
|
* self.value_of_life
|
|
/ self.agents_health_norm
|
|
).astype(self.np_float_dtype)
|
|
|
|
# Economic index -- fraction of annual GDP achieved
|
|
# Use a "crra" nonlinearity on the agent economic reward
|
|
marginal_agent_economic_index = crra_nonlinearity(
|
|
postsubsidy_productivity_t / self.agents_economic_norm,
|
|
self.economic_reward_crra_eta,
|
|
).astype(self.np_float_dtype)
|
|
|
|
# Min-max Normalization
|
|
marginal_agent_health_index = min_max_normalization(
|
|
marginal_agent_health_index,
|
|
self.min_marginal_agent_health_index,
|
|
self.max_marginal_agent_health_index,
|
|
).astype(self.np_float_dtype)
|
|
marginal_agent_economic_index = min_max_normalization(
|
|
marginal_agent_economic_index,
|
|
self.min_marginal_agent_economic_index,
|
|
self.max_marginal_agent_economic_index,
|
|
).astype(self.np_float_dtype)
|
|
|
|
# Agent Rewards
|
|
# -------------
|
|
agent_rewards = get_weighted_average(
|
|
self.weightage_on_marginal_agent_health_index,
|
|
marginal_agent_health_index,
|
|
self.weightage_on_marginal_agent_economic_index,
|
|
marginal_agent_economic_index,
|
|
)
|
|
rew["a"] = agent_rewards / self.reward_normalization_factor
|
|
|
|
# Update agent states
|
|
# -------------------
|
|
for agent in self.world.agents:
|
|
agent.state["Health Index"] += marginal_agent_health_index[agent.idx]
|
|
agent.state["Economic Index"] += marginal_agent_economic_index[agent.idx]
|
|
|
|
# National level
|
|
# --------------
|
|
# Health index -- the cost equivalent (annual GDP) of covid deaths
|
|
# Note: casting deaths to float to prevent overflow issues
|
|
marginal_planner_health_index = (
|
|
-np.sum(marginal_deaths).astype(self.np_float_dtype)
|
|
* self.value_of_life
|
|
/ self.planner_health_norm
|
|
)
|
|
|
|
# Economic index -- fraction of annual GDP achieved (minus subsidy cost)
|
|
cost_of_subsidy_t = (1 + self.risk_free_interest_rate) * np.sum(subsidy_t)
|
|
# Use a "crra" nonlinearity on the planner economic reward
|
|
marginal_planner_economic_index = crra_nonlinearity(
|
|
(np.sum(postsubsidy_productivity_t) - cost_of_subsidy_t)
|
|
/ self.planner_economic_norm,
|
|
self.economic_reward_crra_eta,
|
|
)
|
|
|
|
# Min-max Normalization
|
|
marginal_planner_health_index = min_max_normalization(
|
|
marginal_planner_health_index,
|
|
self.min_marginal_planner_health_index,
|
|
self.max_marginal_planner_health_index,
|
|
)
|
|
marginal_planner_economic_index = min_max_normalization(
|
|
marginal_planner_economic_index,
|
|
self.min_marginal_planner_economic_index,
|
|
self.max_marginal_planner_economic_index,
|
|
)
|
|
|
|
# Update planner states
|
|
# -------------------
|
|
self.world.planner.state["Health Index"] += marginal_planner_health_index
|
|
self.world.planner.state["Economic Index"] += marginal_planner_economic_index
|
|
|
|
# Planner Reward
|
|
# --------------
|
|
planner_rewards = get_weighted_average(
|
|
self.weightage_on_marginal_planner_health_index,
|
|
marginal_planner_health_index,
|
|
self.weightage_on_marginal_planner_economic_index,
|
|
marginal_planner_economic_index,
|
|
)
|
|
rew[self.world.planner.idx] = planner_rewards / self.reward_normalization_factor
|
|
|
|
return rew
|
|
|
|
def additional_reset_steps(self):
|
|
assert self.world.timestep == 0
|
|
|
|
# Reset current date
|
|
self.current_date = self.start_date
|
|
|
|
# SIR numbers at timestep 0
|
|
susceptible_0 = self._real_world_data["susceptible"][self.start_date_index]
|
|
infected_0 = self._real_world_data["infected"][self.start_date_index]
|
|
newly_infected_0 = (
|
|
infected_0
|
|
- self._real_world_data["infected"][max(0, self.start_date_index - 1)]
|
|
)
|
|
recovered_0 = self._real_world_data["recovered"][self.start_date_index]
|
|
deaths_0 = recovered_0 * self.death_rate
|
|
|
|
# Unemployment and vaccinated numbers at timestep 0
|
|
unemployed_0 = self._real_world_data["unemployed"][self.start_date_index]
|
|
vaccinated_0 = self._real_world_data["vaccinated"][self.start_date_index]
|
|
|
|
# Create a global state dictionary to save episode data
|
|
self.world.global_state = {}
|
|
self.set_global_state("Susceptible", susceptible_0, t=self.world.timestep)
|
|
self.set_global_state("Infected", infected_0, t=self.world.timestep)
|
|
self.set_global_state("Recovered", recovered_0, t=self.world.timestep)
|
|
self.set_global_state("Deaths", deaths_0, t=self.world.timestep)
|
|
|
|
self.set_global_state("Unemployed", unemployed_0, t=self.world.timestep)
|
|
self.set_global_state("Vaccinated", vaccinated_0, t=self.world.timestep)
|
|
|
|
new_deaths_0 = (
|
|
deaths_0
|
|
- self._real_world_data["recovered"][max(0, self.start_date_index - 1)]
|
|
* self.death_rate
|
|
)
|
|
|
|
# Reset stringency level history.
|
|
# Pad with stringency levels of 1 corresponding to states being fully open
|
|
# (as was the case before the pandemic).
|
|
self.stringency_level_history = np.pad(
|
|
self._real_world_data["policy"][: self.start_date_index + 1],
|
|
[(self.filter_len, 0), (0, 0)],
|
|
constant_values=1,
|
|
)[-(self.filter_len + 1) :]
|
|
|
|
# Set the stringency level based to the real-world policy
|
|
self.set_global_state(
|
|
"Stringency Level",
|
|
self._real_world_data["policy"][self.start_date_index],
|
|
t=self.world.timestep,
|
|
)
|
|
|
|
# All US states start with zero subsidy and zero Postsubsidy Productivity
|
|
self.set_global_state("Subsidy Level", dtype=self.np_float_dtype)
|
|
self.set_global_state("Subsidy", dtype=self.np_float_dtype)
|
|
self.set_global_state("Postsubsidy Productivity", dtype=self.np_float_dtype)
|
|
|
|
# Set initial agent states
|
|
# ------------------------
|
|
current_date_string = datetime.strftime(
|
|
self.current_date, format=self.date_format
|
|
)
|
|
|
|
for agent in self.world.agents:
|
|
agent.state["Total Susceptible"] = susceptible_0[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
agent.state["New Infections"] = newly_infected_0[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
agent.state["Total Infected"] = infected_0[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
agent.state["Total Recovered"] = recovered_0[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
agent.state["New Deaths"] = new_deaths_0[agent.idx].astype(
|
|
self.np_int_dtype
|
|
)
|
|
agent.state["Total Deaths"] = deaths_0[agent.idx].astype(self.np_int_dtype)
|
|
agent.state["Health Index"] = np.array([0]).astype(self.np_float_dtype)
|
|
agent.state["Economic Index"] = np.array([0]).astype(self.np_float_dtype)
|
|
agent.state["Date"] = current_date_string
|
|
|
|
# Planner state fields
|
|
self.world.planner.state["Total Susceptible"] = np.sum(
|
|
[agent.state["Total Susceptible"] for agent in self.world.agents]
|
|
).astype(self.np_int_dtype)
|
|
self.world.planner.state["New Infections"] = np.sum(
|
|
[agent.state["New Infections"] for agent in self.world.agents]
|
|
).astype(self.np_int_dtype)
|
|
self.world.planner.state["Total Infected"] = np.sum(
|
|
[agent.state["Total Infected"] for agent in self.world.agents]
|
|
).astype(self.np_int_dtype)
|
|
self.world.planner.state["Total Recovered"] = np.sum(
|
|
[agent.state["Total Recovered"] for agent in self.world.agents]
|
|
).astype(self.np_int_dtype)
|
|
self.world.planner.state["New Deaths"] = np.sum(
|
|
[agent.state["New Deaths"] for agent in self.world.agents]
|
|
).astype(self.np_int_dtype)
|
|
self.world.planner.state["Total Deaths"] = np.sum(
|
|
[agent.state["Total Deaths"] for agent in self.world.agents]
|
|
).astype(self.np_int_dtype)
|
|
self.world.planner.state["Total Vaccinated"] = np.sum(vaccinated_0).astype(
|
|
self.np_int_dtype
|
|
)
|
|
self.world.planner.state["Health Index"] = np.array([0]).astype(
|
|
self.np_float_dtype
|
|
)
|
|
self.world.planner.state["Economic Index"] = np.array([0]).astype(
|
|
self.np_float_dtype
|
|
)
|
|
|
|
self.world.planner.state["Date"] = current_date_string
|
|
|
|
# Reset any manually set parameter modulations
|
|
self._beta_intercepts_modulation = 1
|
|
self._beta_slopes_modulation = 1
|
|
self._unemployment_modulation = 1
|
|
|
|
def set_global_state(self, key=None, value=None, t=None, dtype=None):
|
|
# Use floats by default for the SIR dynamics
|
|
if dtype is None:
|
|
dtype = self.np_float_dtype
|
|
assert key in [
|
|
"Susceptible",
|
|
"Infected",
|
|
"Recovered",
|
|
"Deaths",
|
|
"Unemployed",
|
|
"Vaccinated",
|
|
"Stringency Level",
|
|
"Subsidy Level",
|
|
"Subsidy",
|
|
"Postsubsidy Productivity",
|
|
]
|
|
# If no values are passed, set everything to zeros.
|
|
if key not in self.world.global_state:
|
|
self.world.global_state[key] = np.zeros(
|
|
(self.episode_length + 1, self.num_us_states), dtype=dtype
|
|
)
|
|
|
|
if t is not None and value is not None:
|
|
assert isinstance(value, np.ndarray)
|
|
assert value.shape[0] == self.world.global_state[key].shape[1]
|
|
|
|
self.world.global_state[key][t] = value
|
|
else:
|
|
pass
|
|
|
|
def set_parameter_modulations(
|
|
self, beta_intercept=None, beta_slope=None, unemployment=None
|
|
):
|
|
"""
|
|
Apply parameter modulation, which will be in effect until the next env reset.
|
|
|
|
Each modulation term scales the associated set of model parameters by the
|
|
input value. This method is useful for performing a sensitivity analysis.
|
|
|
|
In effect, the transmission rate (beta) will be calculated as:
|
|
beta = (m_s * beta_slope)*lagged_stringency + (m_i * beta_intercept)
|
|
|
|
The unemployment rate (u) will be calculated as:
|
|
u = SOFTPLUS( m_u * SUM(u_filter_weight * u_filter_response) ) + u_0
|
|
|
|
Args:
|
|
beta_intercept: (float, >= 0) Modulation applied to the intercept term
|
|
of the beta model, m_i in above equations
|
|
beta_slope: (float, >= 0) Modulation applied to the slope term of the
|
|
beta model, m_s in above equations
|
|
unemployment: (float, >= 0) Modulation applied to the weighted sum of
|
|
unemployment filter responses, m_u in above equations.
|
|
|
|
Example:
|
|
# Reset the environment
|
|
env.reset()
|
|
|
|
# Increase the slope of the beta response by 15%
|
|
env.set_parameter_modulations(beta_slope=1.15)
|
|
|
|
# Run the environment (this example skips over action selection for brevity)
|
|
for t in range(env.episode_length):
|
|
env.step(actions[t])
|
|
"""
|
|
if beta_intercept is not None:
|
|
beta_intercept = float(beta_intercept)
|
|
assert beta_intercept >= 0
|
|
self._beta_intercepts_modulation = beta_intercept
|
|
|
|
if beta_slope is not None:
|
|
beta_slope = float(beta_slope)
|
|
assert beta_slope >= 0
|
|
self._beta_slopes_modulation = beta_slope
|
|
|
|
if unemployment is not None:
|
|
unemployment = float(unemployment)
|
|
assert unemployment >= 0
|
|
self._unemployment_modulation = unemployment
|
|
|
|
def unemployment_step(self, current_stringency_level):
|
|
"""
|
|
Computes unemployment given the current stringency level and past levels.
|
|
|
|
Unemployment is computed as follows:
|
|
1) For each of self.num_filters, an exponentially decaying filter is
|
|
convolved with the history of stringency changes. Responses move forward in
|
|
time, so a stringency change at time t-1 impacts the response at time t.
|
|
2) The filter responses at time t (the current timestep) are summed together
|
|
using state-specific weights.
|
|
3) The weighted sum is passed through a SOFTPLUS function to capture excess
|
|
unemployment due to stringency policy.
|
|
4) The excess unemployment is added to a state-specific baseline unemployment
|
|
level to get the total unemployment.
|
|
|
|
Note: Internally, unemployment is computed somewhat differently for speed.
|
|
In particular, no convolution is used. Instead the "filter response" at
|
|
time t is just a temporally discounted sum of past stringency changes,
|
|
with the discounting given by the filter decay rate.
|
|
"""
|
|
|
|
def softplus(x, beta=1, threshold=20):
|
|
"""
|
|
Numpy implementation of softplus. For reference, see
|
|
https://pytorch.org/docs/stable/generated/torch.nn.Softplus.html
|
|
"""
|
|
return 1 / beta * np.log(1 + np.exp(beta * x)) * (
|
|
beta * x <= threshold
|
|
) + x * (beta * x > threshold)
|
|
|
|
if (
|
|
self.world.timestep == 0
|
|
): # computing unemployment at closure policy "all ones"
|
|
delta_stringency_level = np.zeros((self.filter_len, self.num_us_states))
|
|
else:
|
|
self.stringency_level_history = np.concatenate(
|
|
(
|
|
self.stringency_level_history[1:],
|
|
current_stringency_level.reshape(1, -1),
|
|
)
|
|
)
|
|
delta_stringency_level = (
|
|
self.stringency_level_history[1:] - self.stringency_level_history[:-1]
|
|
)
|
|
|
|
# Rather than modulating the unemployment params,
|
|
# modulate the deltas (same effect)
|
|
delta_stringency_level = delta_stringency_level * self._unemployment_modulation
|
|
|
|
# Expand the [time, state] delta history to have a dimension for filter channel
|
|
x_data = delta_stringency_level[None].transpose(2, 0, 1)
|
|
|
|
# Apply the state-specific filter weights to each channel
|
|
weighted_x_data = x_data * self.repeated_conv_weights
|
|
|
|
# Compute the discounted sum of the weighted deltas, with each channel using
|
|
# a discounting rate reflecting the time constant of the filter channel. Also
|
|
# sum over channels and use a softplus to get excess unemployment.
|
|
excess_unemployment = softplus(
|
|
np.sum(weighted_x_data * self.unemp_conv_filters, axis=(1, 2)), beta=1
|
|
)
|
|
|
|
# Add excess unemployment to baseline unemployment
|
|
unemployment_rate = excess_unemployment + self.unemployment_bias
|
|
|
|
# Convert the rate (which is a percent) to raw numbers for output
|
|
num_unemployed_t = unemployment_rate * self.us_state_population / 100
|
|
return num_unemployed_t
|
|
|
|
# --- Scenario-specific ---
|
|
def economy_step(
|
|
self,
|
|
population,
|
|
infected,
|
|
deaths,
|
|
unemployed,
|
|
infection_too_sick_to_work_rate=0.05,
|
|
population_between_age_18_65=0.67,
|
|
):
|
|
"""
|
|
Computes how much production occurs.
|
|
|
|
Assumptions:
|
|
|
|
- People that cannot work: "infected + aware" and "unemployed" and "deaths".
|
|
- No life/death cycles.
|
|
|
|
See __init__() for pre-computation of each worker's daily productivity.
|
|
"""
|
|
|
|
incapacitated = (infection_too_sick_to_work_rate * infected) + deaths
|
|
cant_work = (incapacitated * population_between_age_18_65) + unemployed
|
|
|
|
num_workers = population * population_between_age_18_65
|
|
|
|
num_people_that_can_work = np.maximum(0, num_workers - cant_work)
|
|
|
|
productivity = (
|
|
num_people_that_can_work * self.daily_production_per_worker
|
|
).astype(self.np_float_dtype)
|
|
|
|
return productivity
|
|
|
|
def sir_step(self, S_tm1, I_tm1, stringency_level_tmk, num_vaccines_available_t):
|
|
"""
|
|
Simulates SIR infection model in the US.
|
|
"""
|
|
intercepts = self.beta_intercepts * self._beta_intercepts_modulation
|
|
slopes = self.beta_slopes * self._beta_slopes_modulation
|
|
beta_i = (intercepts + slopes * stringency_level_tmk).astype(
|
|
self.np_float_dtype
|
|
)
|
|
|
|
small_number = 1e-10 # used to prevent indeterminate cases
|
|
susceptible_fraction_vaccinated = np.minimum(
|
|
np.ones((self.num_us_states), dtype=self.np_int_dtype),
|
|
num_vaccines_available_t / (S_tm1 + small_number),
|
|
).astype(self.np_float_dtype)
|
|
vaccinated_t = np.minimum(num_vaccines_available_t, S_tm1)
|
|
|
|
# Record R0
|
|
R0 = beta_i / self.gamma
|
|
for agent in self.world.agents:
|
|
agent.state["R0"] = R0[agent.idx]
|
|
|
|
# S -> I; dS
|
|
neighborhood_SI_over_N = (S_tm1 / self.us_state_population) * I_tm1
|
|
dS_t = (
|
|
-beta_i * neighborhood_SI_over_N * (1 - susceptible_fraction_vaccinated)
|
|
- vaccinated_t
|
|
).astype(self.np_float_dtype)
|
|
|
|
# I -> R; dR
|
|
dR_t = (self.gamma * I_tm1 + vaccinated_t).astype(self.np_float_dtype)
|
|
|
|
# dI from d(S + I + R) = 0
|
|
# ------------------------
|
|
dI_t = -dS_t - dR_t
|
|
|
|
dV_t = vaccinated_t.astype(self.np_float_dtype)
|
|
|
|
return dS_t, dI_t, dR_t, dV_t
|
|
|
|
def load_model_constants(self, path_to_model_constants):
|
|
filename = "model_constants.json"
|
|
assert filename in os.listdir(path_to_model_constants), (
|
|
"Unable to locate '{}' in '{}'.\nPlease run the "
|
|
"'gather_real_world_data.ipynb' notebook first".format(
|
|
filename, path_to_model_constants
|
|
)
|
|
)
|
|
with open(os.path.join(path_to_model_constants, filename), "r") as fp:
|
|
model_constants_dict = json.load(fp)
|
|
fp.close()
|
|
|
|
self.date_format = model_constants_dict["DATE_FORMAT"]
|
|
self.us_state_idx_to_state_name = model_constants_dict[
|
|
"US_STATE_IDX_TO_STATE_NAME"
|
|
]
|
|
self.us_state_population = self.np_int_dtype(
|
|
model_constants_dict["US_STATE_POPULATION"]
|
|
)
|
|
self.us_population = self.np_int_dtype(model_constants_dict["US_POPULATION"])
|
|
self.num_stringency_levels = model_constants_dict["NUM_STRINGENCY_LEVELS"]
|
|
self.death_rate = self.np_float_dtype(model_constants_dict["SIR_MORTALITY"])
|
|
self.gamma = self.np_float_dtype(model_constants_dict["SIR_GAMMA"])
|
|
self.gdp_per_capita = self.np_float_dtype(
|
|
model_constants_dict["GDP_PER_CAPITA"]
|
|
)
|
|
|
|
def load_fitted_params(self, path_to_fitted_params):
|
|
filename = "fitted_params.json"
|
|
assert filename in os.listdir(path_to_fitted_params), (
|
|
"Unable to locate '{}' in '{}'.\nIf you ran the "
|
|
"'gather_real_world_data.ipynb' notebook to download the latest "
|
|
"real-world data, please also run the "
|
|
"'fit_parameters.ipynb' notebook.".format(filename, path_to_fitted_params)
|
|
)
|
|
with open(os.path.join(path_to_fitted_params, filename), "r") as fp:
|
|
fitted_params_dict = json.load(fp)
|
|
fp.close()
|
|
self.policy_start_date = datetime.strptime(
|
|
fitted_params_dict["POLICY_START_DATE"], self.date_format
|
|
)
|
|
self.value_of_life = self.np_int_dtype(fitted_params_dict["VALUE_OF_LIFE"])
|
|
self.beta_delay = self.np_int_dtype(fitted_params_dict["BETA_DELAY"])
|
|
self.beta_slopes = np.array(
|
|
fitted_params_dict["BETA_SLOPES"], dtype=self.np_float_dtype
|
|
)
|
|
self.beta_intercepts = np.array(
|
|
fitted_params_dict["BETA_INTERCEPTS"], dtype=self.np_float_dtype
|
|
)
|
|
self.min_marginal_agent_health_index = np.array(
|
|
fitted_params_dict["MIN_MARGINAL_AGENT_HEALTH_INDEX"],
|
|
dtype=self.np_float_dtype,
|
|
)
|
|
self.max_marginal_agent_health_index = np.array(
|
|
fitted_params_dict["MAX_MARGINAL_AGENT_HEALTH_INDEX"],
|
|
dtype=self.np_float_dtype,
|
|
)
|
|
self.min_marginal_agent_economic_index = np.array(
|
|
fitted_params_dict["MIN_MARGINAL_AGENT_ECONOMIC_INDEX"],
|
|
dtype=self.np_float_dtype,
|
|
)
|
|
self.max_marginal_agent_economic_index = np.array(
|
|
fitted_params_dict["MAX_MARGINAL_AGENT_ECONOMIC_INDEX"],
|
|
dtype=self.np_float_dtype,
|
|
)
|
|
self.min_marginal_planner_health_index = self.np_float_dtype(
|
|
fitted_params_dict["MIN_MARGINAL_PLANNER_HEALTH_INDEX"]
|
|
)
|
|
self.max_marginal_planner_health_index = self.np_float_dtype(
|
|
fitted_params_dict["MAX_MARGINAL_PLANNER_HEALTH_INDEX"]
|
|
)
|
|
self.min_marginal_planner_economic_index = self.np_float_dtype(
|
|
fitted_params_dict["MIN_MARGINAL_PLANNER_ECONOMIC_INDEX"]
|
|
)
|
|
self.max_marginal_planner_economic_index = self.np_float_dtype(
|
|
fitted_params_dict["MAX_MARGINAL_PLANNER_ECONOMIC_INDEX"]
|
|
)
|
|
self.inferred_weightage_on_agent_health_index = np.array(
|
|
fitted_params_dict["INFERRED_WEIGHTAGE_ON_AGENT_HEALTH_INDEX"],
|
|
dtype=self.np_float_dtype,
|
|
)
|
|
self.inferred_weightage_on_planner_health_index = self.np_float_dtype(
|
|
fitted_params_dict["INFERRED_WEIGHTAGE_ON_PLANNER_HEALTH_INDEX"]
|
|
)
|
|
self.filter_len = self.np_int_dtype(fitted_params_dict["FILTER_LEN"])
|
|
self.conv_lambdas = np.array(
|
|
fitted_params_dict["CONV_LAMBDAS"], dtype=self.np_float_dtype
|
|
)
|
|
self.unemployment_bias = np.array(
|
|
fitted_params_dict["UNEMPLOYMENT_BIAS"], dtype=self.np_float_dtype
|
|
)
|
|
self.grouped_convolutional_filter_weights = np.array(
|
|
fitted_params_dict["GROUPED_CONVOLUTIONAL_FILTER_WEIGHTS"],
|
|
dtype=self.np_float_dtype,
|
|
)
|
|
|
|
def scenario_metrics(self):
|
|
# End of episode metrics
|
|
# ----------------------
|
|
metrics_dict = {}
|
|
|
|
# State-level metrics
|
|
for agent in self.world.agents:
|
|
state_name = self.us_state_idx_to_state_name[str(agent.idx)]
|
|
|
|
for field in ["infected", "recovered", "deaths"]:
|
|
metric_key = "{}/{} (millions)".format(state_name, field)
|
|
metrics_dict[metric_key] = (
|
|
agent.state["Total " + field.capitalize()] / 1e6
|
|
)
|
|
|
|
metrics_dict["{}/mean_unemployment_rate (%)".format(state_name)] = (
|
|
np.mean(self.world.global_state["Unemployed"][1:, agent.idx], axis=0)
|
|
/ self.us_state_population[agent.idx]
|
|
* 100
|
|
)
|
|
|
|
metrics_dict[
|
|
"{}/mean_open_close_stringency_level".format(state_name)
|
|
] = np.mean(
|
|
self.world.global_state["Stringency Level"][1:, agent.idx], axis=0
|
|
)
|
|
|
|
metrics_dict["{}/total_productivity (billion $)".format(state_name)] = (
|
|
np.sum(
|
|
self.world.global_state["Postsubsidy Productivity"][1:, agent.idx]
|
|
)
|
|
/ 1e9
|
|
)
|
|
|
|
metrics_dict[
|
|
"{}/health_index_at_end_of_episode".format(state_name)
|
|
] = agent.state["Health Index"]
|
|
metrics_dict[
|
|
"{}/economic_index_at_end_of_episode".format(state_name)
|
|
] = agent.state["Economic Index"]
|
|
|
|
# USA-level metrics
|
|
metrics_dict["usa/vaccinated (% of population)"] = (
|
|
np.sum(self.world.global_state["Vaccinated"][self.world.timestep], axis=0)
|
|
/ self.us_population
|
|
* 100
|
|
)
|
|
metrics_dict["usa/deaths (thousands)"] = (
|
|
np.sum(self.world.global_state["Deaths"][self.world.timestep], axis=0) / 1e3
|
|
)
|
|
|
|
metrics_dict["usa/mean_unemployment_rate (%)"] = (
|
|
np.mean(
|
|
np.sum(self.world.global_state["Unemployed"][1:], axis=1)
|
|
/ self.us_population,
|
|
axis=0,
|
|
)
|
|
* 100
|
|
)
|
|
metrics_dict["usa/total_amount_subsidized (trillion $)"] = (
|
|
np.sum(self.world.global_state["Subsidy"][1:], axis=(0, 1)) / 1e12
|
|
)
|
|
metrics_dict["usa/total_productivity (trillion $)"] = (
|
|
np.sum(self.world.global_state["Postsubsidy Productivity"][1:], axis=(0, 1))
|
|
/ 1e12
|
|
)
|
|
|
|
metrics_dict["usa/health_index_at_end_of_episode"] = self.world.planner.state[
|
|
"Health Index"
|
|
]
|
|
metrics_dict["usa/economic_index_at_end_of_episode"] = self.world.planner.state[
|
|
"Economic Index"
|
|
]
|
|
|
|
return metrics_dict
|