# 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