621 lines
24 KiB
Plaintext
621 lines
24 KiB
Plaintext
// 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
|
|
|
|
__constant__ float kEpsilon = 1.0e-10; // used to prevent division by 0
|
|
|
|
extern "C" {
|
|
// CUDA version of the scenario_step() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
|
|
// CUDA version of the sir_step() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
__device__ void cuda_sir_step(
|
|
float* susceptible,
|
|
float* infected,
|
|
float* recovered,
|
|
float* vaccinated,
|
|
float* deaths,
|
|
int* num_vaccines_available_t,
|
|
const int* kRealWorldStringencyPolicyHistory,
|
|
const float kStatePopulation,
|
|
const int kNumAgents,
|
|
const int kBetaDelay,
|
|
const float kBetaSlope,
|
|
const float kbetaIntercept,
|
|
int* stringency_level,
|
|
float* beta,
|
|
const float kGamma,
|
|
const float kDeathRate,
|
|
const int kEnvId,
|
|
const int kAgentId,
|
|
int timestep,
|
|
const int kEpisodeLength,
|
|
const int kArrayIdxCurrentTime,
|
|
const int kArrayIdxPrevTime,
|
|
const int kTimeIndependentArrayIdx
|
|
) {
|
|
float susceptible_fraction_vaccinated = min(
|
|
1.0,
|
|
num_vaccines_available_t[kTimeIndependentArrayIdx] /
|
|
(susceptible[kArrayIdxPrevTime] + kEpsilon));
|
|
float vaccinated_t = min(
|
|
static_cast<float>(num_vaccines_available_t[
|
|
kTimeIndependentArrayIdx]),
|
|
susceptible[kArrayIdxPrevTime]);
|
|
|
|
// (S/N) * I in place of (S*I) / N to prevent overflow
|
|
float neighborhood_SI_over_N = susceptible[kArrayIdxPrevTime] /
|
|
kStatePopulation * infected[kArrayIdxPrevTime];
|
|
int stringency_level_tmk;
|
|
if (timestep < kBetaDelay) {
|
|
stringency_level_tmk = kRealWorldStringencyPolicyHistory[
|
|
(timestep - 1) * (kNumAgents - 1) + kAgentId];
|
|
} else {
|
|
stringency_level_tmk = stringency_level[kEnvId * (
|
|
kEpisodeLength + 1) * (kNumAgents - 1) +
|
|
(timestep - kBetaDelay) * (kNumAgents - 1) + kAgentId];
|
|
}
|
|
beta[kTimeIndependentArrayIdx] = stringency_level_tmk *
|
|
kBetaSlope + kbetaIntercept;
|
|
|
|
float dS_t = -(neighborhood_SI_over_N * beta[
|
|
kTimeIndependentArrayIdx] *
|
|
(1 - susceptible_fraction_vaccinated) + vaccinated_t);
|
|
float dR_t = kGamma * infected[kArrayIdxPrevTime] + vaccinated_t;
|
|
float dI_t = - dS_t - dR_t;
|
|
|
|
susceptible[kArrayIdxCurrentTime] = max(
|
|
0.0,
|
|
susceptible[kArrayIdxPrevTime] + dS_t);
|
|
infected[kArrayIdxCurrentTime] = max(
|
|
0.0,
|
|
infected[kArrayIdxPrevTime] + dI_t);
|
|
recovered[kArrayIdxCurrentTime] = max(
|
|
0.0,
|
|
recovered[kArrayIdxPrevTime] + dR_t);
|
|
|
|
vaccinated[kArrayIdxCurrentTime] = vaccinated_t +
|
|
vaccinated[kArrayIdxPrevTime];
|
|
float recovered_but_not_vaccinated = recovered[kArrayIdxCurrentTime] -
|
|
vaccinated[kArrayIdxCurrentTime];
|
|
deaths[kArrayIdxCurrentTime] = recovered_but_not_vaccinated *
|
|
kDeathRate;
|
|
}
|
|
|
|
// CUDA version of the softplus() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
__device__ float softplus(float x) {
|
|
const float kBeta = 1.0;
|
|
const float kThreshold = 20.0;
|
|
if (kBeta * x < kThreshold) {
|
|
return 1.0 / kBeta * log(1.0 + exp(kBeta * x));
|
|
} else {
|
|
return x;
|
|
}
|
|
}
|
|
|
|
__device__ float signal2unemployment(
|
|
const int kEnvId,
|
|
const int kAgentId,
|
|
float* signal,
|
|
const float* kUnemploymentConvolutionalFilters,
|
|
const float kUnemploymentBias,
|
|
const int kNumAgents,
|
|
const int kFilterLen,
|
|
const int kNumFilters
|
|
) {
|
|
float unemployment = 0.0;
|
|
const int kArrayIndexOffset = kEnvId * (kNumAgents - 1) * kNumFilters *
|
|
kFilterLen + kAgentId * kNumFilters * kFilterLen;
|
|
for (int index = 0; index < (kFilterLen * kNumFilters); index ++) {
|
|
unemployment += signal[kArrayIndexOffset + index] *
|
|
kUnemploymentConvolutionalFilters[index];
|
|
}
|
|
return softplus(unemployment) + kUnemploymentBias;
|
|
}
|
|
|
|
// CUDA version of the unemployment_step() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
__device__ void cuda_unemployment_step(
|
|
float* unemployed,
|
|
int* stringency_level,
|
|
int* delta_stringency_level,
|
|
const float* kGroupedConvolutionalFilterWeights,
|
|
const float* kUnemploymentConvolutionalFilters,
|
|
const float* kUnemploymentBias,
|
|
float* convolved_signal,
|
|
const int kFilterLen,
|
|
const int kNumFilters,
|
|
const float kStatePopulation,
|
|
const int kNumAgents,
|
|
const int kEnvId,
|
|
const int kAgentId,
|
|
int timestep,
|
|
const int kArrayIdxCurrentTime,
|
|
const int kArrayIdxPrevTime
|
|
) {
|
|
// Shift array by kNumAgents - 1
|
|
for (int idx = 0; idx < kFilterLen - 1; idx ++) {
|
|
delta_stringency_level[
|
|
kEnvId * kFilterLen * (kNumAgents - 1) + idx *
|
|
(kNumAgents - 1) + kAgentId
|
|
] =
|
|
delta_stringency_level[
|
|
kEnvId * kFilterLen * (kNumAgents - 1) + (idx + 1) *
|
|
(kNumAgents - 1) + kAgentId
|
|
];
|
|
}
|
|
|
|
delta_stringency_level[
|
|
kEnvId * kFilterLen * (kNumAgents - 1) + (kFilterLen - 1) *
|
|
(kNumAgents - 1) + kAgentId
|
|
] = stringency_level[kArrayIdxCurrentTime] -
|
|
stringency_level[kArrayIdxPrevTime];
|
|
|
|
// convolved_signal refers to the convolution between the filter weights
|
|
// and the delta stringency levels
|
|
for (int filter_idx = 0; filter_idx < kNumFilters; filter_idx ++) {
|
|
for (int idx = 0; idx < kFilterLen; idx ++) {
|
|
convolved_signal[
|
|
kEnvId * (kNumAgents - 1) * kNumFilters * kFilterLen +
|
|
kAgentId * kNumFilters * kFilterLen +
|
|
filter_idx * kFilterLen +
|
|
idx
|
|
] =
|
|
delta_stringency_level[kEnvId * kFilterLen * (kNumAgents - 1) +
|
|
idx * (kNumAgents - 1) + kAgentId] *
|
|
kGroupedConvolutionalFilterWeights[kAgentId * kNumFilters +
|
|
filter_idx];
|
|
}
|
|
}
|
|
|
|
float unemployment_rate = signal2unemployment(
|
|
kEnvId,
|
|
kAgentId,
|
|
convolved_signal,
|
|
kUnemploymentConvolutionalFilters,
|
|
kUnemploymentBias[kAgentId],
|
|
kNumAgents,
|
|
kFilterLen,
|
|
kNumFilters);
|
|
|
|
unemployed[kArrayIdxCurrentTime] =
|
|
unemployment_rate * kStatePopulation / 100.0;
|
|
}
|
|
|
|
// CUDA version of the economy_step() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
__device__ void cuda_economy_step(
|
|
float* infected,
|
|
float* deaths,
|
|
float* unemployed,
|
|
float* incapacitated,
|
|
float* cant_work,
|
|
float* num_people_that_can_work,
|
|
const float kStatePopulation,
|
|
const float kInfectionTooSickToWorkRate,
|
|
const float kPopulationBetweenAge18And65,
|
|
const float kDailyProductionPerWorker,
|
|
float* productivity,
|
|
float* subsidy,
|
|
float* postsubsidy_productivity,
|
|
int timestep,
|
|
const int kArrayIdxCurrentTime,
|
|
int kTimeIndependentArrayIdx
|
|
) {
|
|
incapacitated[kTimeIndependentArrayIdx] =
|
|
kInfectionTooSickToWorkRate * infected[kArrayIdxCurrentTime] +
|
|
deaths[kArrayIdxCurrentTime];
|
|
cant_work[kTimeIndependentArrayIdx] =
|
|
incapacitated[kTimeIndependentArrayIdx] *
|
|
kPopulationBetweenAge18And65 + unemployed[kArrayIdxCurrentTime];
|
|
int num_workers = static_cast<int>(kStatePopulation) * kPopulationBetweenAge18And65;
|
|
num_people_that_can_work[kTimeIndependentArrayIdx] = max(
|
|
0.0,
|
|
num_workers - cant_work[kTimeIndependentArrayIdx]);
|
|
productivity[kArrayIdxCurrentTime] =
|
|
num_people_that_can_work[kTimeIndependentArrayIdx] *
|
|
kDailyProductionPerWorker;
|
|
|
|
postsubsidy_productivity[kArrayIdxCurrentTime] =
|
|
productivity[kArrayIdxCurrentTime] +
|
|
subsidy[kArrayIdxCurrentTime];
|
|
}
|
|
|
|
// CUDA version of crra_nonlinearity() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
__device__ float crra_nonlinearity(
|
|
float x,
|
|
const float kEta,
|
|
const int kNumDaysInAnYear
|
|
) {
|
|
float annual_x = kNumDaysInAnYear * x;
|
|
float annual_x_clipped = annual_x;
|
|
if (annual_x < 0.1) {
|
|
annual_x_clipped = 0.1;
|
|
} else if (annual_x > 3.0) {
|
|
annual_x_clipped = 3.0;
|
|
}
|
|
float annual_crra = 1 + (pow(annual_x_clipped, (1 - kEta)) - 1) /
|
|
(1 - kEta);
|
|
float daily_crra = annual_crra / kNumDaysInAnYear;
|
|
return daily_crra;
|
|
}
|
|
|
|
// CUDA version of min_max_normalization() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
__device__ float min_max_normalization(
|
|
float x,
|
|
const float kMinX,
|
|
const float kMaxX
|
|
) {
|
|
return (x - kMinX) / (kMaxX - kMinX + kEpsilon);
|
|
}
|
|
|
|
// CUDA version of get_rew() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
__device__ float get_rew(
|
|
const float kHealthIndexWeightage,
|
|
float health_index,
|
|
const float kEconomicIndexWeightage,
|
|
float economic_index
|
|
) {
|
|
return (
|
|
kHealthIndexWeightage * health_index
|
|
+ kEconomicIndexWeightage * economic_index) /
|
|
(kHealthIndexWeightage + kEconomicIndexWeightage);
|
|
}
|
|
|
|
// CUDA version of scenario_step() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
__global__ void CudaCovidAndEconomySimulationStep(
|
|
float* susceptible,
|
|
float* infected,
|
|
float* recovered,
|
|
float* deaths,
|
|
float* vaccinated,
|
|
float* unemployed,
|
|
float* subsidy,
|
|
float* productivity,
|
|
int* stringency_level,
|
|
const int kNumStringencyLevels,
|
|
float* postsubsidy_productivity,
|
|
int* num_vaccines_available_t,
|
|
const int* kRealWorldStringencyPolicyHistory,
|
|
const int kBetaDelay,
|
|
const float* kBetaSlopes,
|
|
const float* kbetaIntercepts,
|
|
float* beta,
|
|
const float kGamma,
|
|
const float kDeathRate,
|
|
float* incapacitated,
|
|
float* cant_work,
|
|
float* num_people_that_can_work,
|
|
const int* us_kStatePopulation,
|
|
const float kInfectionTooSickToWorkRate,
|
|
const float kPopulationBetweenAge18And65,
|
|
const int kFilterLen,
|
|
const int kNumFilters,
|
|
int* delta_stringency_level,
|
|
const float* kGroupedConvolutionalFilterWeights,
|
|
const float* kUnemploymentConvolutionalFilters,
|
|
const float* kUnemploymentBias,
|
|
float* signal,
|
|
const float kDailyProductionPerWorker,
|
|
const float* maximum_productivity,
|
|
float* obs_a_world_agent_state,
|
|
float* obs_a_world_agent_postsubsidy_productivity,
|
|
float* obs_a_world_lagged_stringency_level,
|
|
float* obs_a_time,
|
|
float* obs_p_world_agent_state,
|
|
float* obs_p_world_agent_postsubsidy_productivity,
|
|
float* obs_p_world_lagged_stringency_level,
|
|
float* obs_p_time,
|
|
int * env_timestep_arr,
|
|
const int kNumAgents,
|
|
const int kEpisodeLength
|
|
) {
|
|
const int kEnvId = blockIdx.x;
|
|
const int kAgentId = threadIdx.x;
|
|
|
|
assert(env_timestep_arr[kEnvId] > 0 &&
|
|
env_timestep_arr[kEnvId] <= kEpisodeLength);
|
|
assert (kAgentId <= kNumAgents - 1);
|
|
const int kNumFeatures = 6;
|
|
|
|
if (kAgentId < (kNumAgents - 1)) {
|
|
// Indices for time-dependent and time-independent arrays
|
|
// Time dependent arrays have shapes (num_envs,
|
|
// kEpisodeLength + 1, kNumAgents - 1)
|
|
// Time independent arrays have shapes (num_envs, kNumAgents - 1)
|
|
const int kArrayIndexOffset = kEnvId * (kEpisodeLength + 1) *
|
|
(kNumAgents - 1);
|
|
int kArrayIdxCurrentTime = kArrayIndexOffset +
|
|
env_timestep_arr[kEnvId] * (kNumAgents - 1) + kAgentId;
|
|
int kArrayIdxPrevTime = kArrayIndexOffset +
|
|
(env_timestep_arr[kEnvId] - 1) * (kNumAgents - 1) + kAgentId;
|
|
const int kTimeIndependentArrayIdx = kEnvId *
|
|
(kNumAgents - 1) + kAgentId;
|
|
|
|
const float kStatePopulation = static_cast<float>(us_kStatePopulation[kAgentId]);
|
|
|
|
cuda_sir_step(
|
|
susceptible,
|
|
infected,
|
|
recovered,
|
|
vaccinated,
|
|
deaths,
|
|
num_vaccines_available_t,
|
|
kRealWorldStringencyPolicyHistory,
|
|
kStatePopulation,
|
|
kNumAgents,
|
|
kBetaDelay,
|
|
kBetaSlopes[kAgentId],
|
|
kbetaIntercepts[kAgentId],
|
|
stringency_level,
|
|
beta,
|
|
kGamma,
|
|
kDeathRate,
|
|
kEnvId,
|
|
kAgentId,
|
|
env_timestep_arr[kEnvId],
|
|
kEpisodeLength,
|
|
kArrayIdxCurrentTime,
|
|
kArrayIdxPrevTime,
|
|
kTimeIndependentArrayIdx);
|
|
|
|
cuda_unemployment_step(
|
|
unemployed,
|
|
stringency_level,
|
|
delta_stringency_level,
|
|
kGroupedConvolutionalFilterWeights,
|
|
kUnemploymentConvolutionalFilters,
|
|
kUnemploymentBias,
|
|
signal,
|
|
kFilterLen,
|
|
kNumFilters,
|
|
kStatePopulation,
|
|
kNumAgents,
|
|
kEnvId,
|
|
kAgentId,
|
|
env_timestep_arr[kEnvId],
|
|
kArrayIdxCurrentTime,
|
|
kArrayIdxPrevTime);
|
|
|
|
cuda_economy_step(
|
|
infected,
|
|
deaths,
|
|
unemployed,
|
|
incapacitated,
|
|
cant_work,
|
|
num_people_that_can_work,
|
|
kStatePopulation,
|
|
kInfectionTooSickToWorkRate,
|
|
kPopulationBetweenAge18And65,
|
|
kDailyProductionPerWorker,
|
|
productivity,
|
|
subsidy,
|
|
postsubsidy_productivity,
|
|
env_timestep_arr[kEnvId],
|
|
kArrayIdxCurrentTime,
|
|
kTimeIndependentArrayIdx);
|
|
|
|
// CUDA version of generate observations
|
|
// Agents' observations
|
|
int kFeatureArrayIndexOffset = kEnvId * kNumFeatures *
|
|
(kNumAgents - 1) + kAgentId;
|
|
obs_a_world_agent_state[
|
|
kFeatureArrayIndexOffset + 0 * (kNumAgents - 1)
|
|
] = susceptible[kArrayIdxCurrentTime] / kStatePopulation;
|
|
obs_a_world_agent_state[
|
|
kFeatureArrayIndexOffset + 1 * (kNumAgents - 1)
|
|
] = infected[kArrayIdxCurrentTime] / kStatePopulation;
|
|
obs_a_world_agent_state[
|
|
kFeatureArrayIndexOffset + 2 * (kNumAgents - 1)
|
|
] = recovered[kArrayIdxCurrentTime] / kStatePopulation;
|
|
obs_a_world_agent_state[
|
|
kFeatureArrayIndexOffset + 3 * (kNumAgents - 1)
|
|
] = deaths[kArrayIdxCurrentTime] / kStatePopulation;
|
|
obs_a_world_agent_state[
|
|
kFeatureArrayIndexOffset + 4 * (kNumAgents - 1)
|
|
] = vaccinated[kArrayIdxCurrentTime] / kStatePopulation;
|
|
obs_a_world_agent_state[
|
|
kFeatureArrayIndexOffset + 5 * (kNumAgents - 1)
|
|
] = unemployed[kArrayIdxCurrentTime] / kStatePopulation;
|
|
|
|
for (int feature_id = 0; feature_id < kNumFeatures; feature_id ++) {
|
|
const int kIndex = feature_id * (kNumAgents - 1);
|
|
obs_p_world_agent_state[kFeatureArrayIndexOffset +
|
|
kIndex
|
|
] = obs_a_world_agent_state[kFeatureArrayIndexOffset +
|
|
kIndex];
|
|
}
|
|
|
|
obs_a_world_agent_postsubsidy_productivity[
|
|
kTimeIndependentArrayIdx
|
|
] = postsubsidy_productivity[kArrayIdxCurrentTime] /
|
|
maximum_productivity[kAgentId];
|
|
obs_p_world_agent_postsubsidy_productivity[
|
|
kTimeIndependentArrayIdx
|
|
] = obs_a_world_agent_postsubsidy_productivity[
|
|
kTimeIndependentArrayIdx
|
|
];
|
|
|
|
int t_beta = env_timestep_arr[kEnvId] - kBetaDelay + 1;
|
|
if (t_beta < 0) {
|
|
obs_a_world_lagged_stringency_level[
|
|
kTimeIndependentArrayIdx
|
|
] = kRealWorldStringencyPolicyHistory[
|
|
env_timestep_arr[kEnvId] * (kNumAgents - 1) + kAgentId
|
|
] / static_cast<float>(kNumStringencyLevels);
|
|
} else {
|
|
obs_a_world_lagged_stringency_level[
|
|
kTimeIndependentArrayIdx
|
|
] = stringency_level[
|
|
kArrayIndexOffset +
|
|
t_beta * (kNumAgents - 1) +
|
|
kAgentId
|
|
] / static_cast<float>(kNumStringencyLevels);
|
|
}
|
|
obs_p_world_lagged_stringency_level[
|
|
kTimeIndependentArrayIdx
|
|
] = obs_a_world_lagged_stringency_level[
|
|
kTimeIndependentArrayIdx];
|
|
// Below, we assume observation scaling = True
|
|
// (otherwise, 'obs_a_time[kTimeIndependentArrayIdx] =
|
|
// static_cast<float>(env_timestep_arr[kEnvId])
|
|
obs_a_time[kTimeIndependentArrayIdx] =
|
|
env_timestep_arr[kEnvId] / static_cast<float>(kEpisodeLength);
|
|
} else if (kAgentId == kNumAgents - 1) {
|
|
obs_p_time[kEnvId] = env_timestep_arr[kEnvId] /
|
|
static_cast<float>(kEpisodeLength);
|
|
}
|
|
}
|
|
|
|
// CUDA version of the compute_reward() in
|
|
// "ai_economist.foundation.scenarios.covid19_env.py"
|
|
__global__ void CudaComputeReward(
|
|
float* rewards_a,
|
|
float* rewards_p,
|
|
const int kNumDaysInAnYear,
|
|
const int kValueOfLife,
|
|
const float kRiskFreeInterestRate,
|
|
const float kEconomicRewardCrraEta,
|
|
const float* kMinMarginalAgentHealthIndex,
|
|
const float* kMaxMarginalAgentHealthIndex,
|
|
const float* kMinMarginalAgentEconomicIndex,
|
|
const float* kMaxMarginalAgentEconomicIndex,
|
|
const float kMinMarginalPlannerHealthIndex,
|
|
const float kMaxMarginalPlannerHealthIndex,
|
|
const float kMinMarginalPlannerEconomicIndex,
|
|
const float kMaxMarginalPlannerEconomicIndex,
|
|
const float* kWeightageOnMarginalAgentHealthIndex,
|
|
const float* kWeightageOnMarginalPlannerHealthIndex,
|
|
const float kWeightageOnMarginalAgentEconomicIndex,
|
|
const float kWeightageOnMarginalPlannerEconomicIndex,
|
|
const float* kAgentsHealthNorm,
|
|
const float* kAgentsEconomicNorm,
|
|
const float kPlannerHealthNorm,
|
|
const float kPlannerEconomicNorm,
|
|
float* deaths,
|
|
float* subsidy,
|
|
float* postsubsidy_productivity,
|
|
int* env_done_arr,
|
|
int* env_timestep_arr,
|
|
const int kNumAgents,
|
|
const int kEpisodeLength
|
|
) {
|
|
const int kEnvId = blockIdx.x;
|
|
const int kAgentId = threadIdx.x;
|
|
|
|
assert(env_timestep_arr[kEnvId] > 0 &&
|
|
env_timestep_arr[kEnvId] <= kEpisodeLength);
|
|
assert (kAgentId <= kNumAgents - 1);
|
|
|
|
const int kArrayIndexOffset = kEnvId * (kEpisodeLength + 1) *
|
|
(kNumAgents - 1);
|
|
if (kAgentId < (kNumAgents - 1)) {
|
|
// Agents' rewards
|
|
// Indices for time-dependent and time-independent arrays
|
|
// Time dependent arrays have shapes (num_envs,
|
|
// kEpisodeLength + 1, kNumAgents - 1)
|
|
// Time independent arrays have shapes (num_envs, kNumAgents - 1)
|
|
int kArrayIdxCurrentTime = kArrayIndexOffset +
|
|
env_timestep_arr[kEnvId] * (kNumAgents - 1) + kAgentId;
|
|
int kArrayIdxPrevTime = kArrayIndexOffset +
|
|
(env_timestep_arr[kEnvId] - 1) * (kNumAgents - 1) + kAgentId;
|
|
const int kTimeIndependentArrayIdx = kEnvId *
|
|
(kNumAgents - 1) + kAgentId;
|
|
|
|
float marginal_deaths = deaths[kArrayIdxCurrentTime] -
|
|
deaths[kArrayIdxPrevTime];
|
|
|
|
// Note: changing the order of operations to prevent overflow
|
|
float marginal_agent_health_index = - marginal_deaths /
|
|
(kAgentsHealthNorm[kAgentId] /
|
|
static_cast<float>(kValueOfLife));
|
|
|
|
float marginal_agent_economic_index = crra_nonlinearity(
|
|
postsubsidy_productivity[kArrayIdxCurrentTime] /
|
|
kAgentsEconomicNorm[kAgentId],
|
|
kEconomicRewardCrraEta,
|
|
kNumDaysInAnYear);
|
|
|
|
marginal_agent_health_index = min_max_normalization(
|
|
marginal_agent_health_index,
|
|
kMinMarginalAgentHealthIndex[kAgentId],
|
|
kMaxMarginalAgentHealthIndex[kAgentId]);
|
|
marginal_agent_economic_index = min_max_normalization(
|
|
marginal_agent_economic_index,
|
|
kMinMarginalAgentEconomicIndex[kAgentId],
|
|
kMaxMarginalAgentEconomicIndex[kAgentId]);
|
|
|
|
rewards_a[kTimeIndependentArrayIdx] = get_rew(
|
|
kWeightageOnMarginalAgentHealthIndex[kAgentId],
|
|
marginal_agent_health_index,
|
|
kWeightageOnMarginalPlannerHealthIndex[kAgentId],
|
|
marginal_agent_economic_index);
|
|
} else if (kAgentId == kNumAgents - 1) {
|
|
// Planner's rewards
|
|
float total_marginal_deaths = 0;
|
|
for (int ag_id = 0; ag_id < (kNumAgents - 1); ag_id ++) {
|
|
total_marginal_deaths += (
|
|
deaths[kArrayIndexOffset + env_timestep_arr[kEnvId] *
|
|
(kNumAgents - 1) + ag_id] -
|
|
deaths[kArrayIndexOffset + (env_timestep_arr[kEnvId] - 1) *
|
|
(kNumAgents - 1) + ag_id]);
|
|
}
|
|
// Note: changing the order of operations to prevent overflow
|
|
float marginal_planner_health_index = -total_marginal_deaths /
|
|
(kPlannerHealthNorm / static_cast<float>(kValueOfLife));
|
|
|
|
float total_subsidy = 0.0;
|
|
float total_postsubsidy_productivity = 0.0;
|
|
for (int ag_id = 0; ag_id < (kNumAgents - 1); ag_id ++) {
|
|
total_subsidy += subsidy[kArrayIndexOffset +
|
|
env_timestep_arr[kEnvId] * (kNumAgents - 1) + ag_id];
|
|
total_postsubsidy_productivity +=
|
|
postsubsidy_productivity[kArrayIndexOffset +
|
|
env_timestep_arr[kEnvId] * (kNumAgents - 1) + ag_id];
|
|
}
|
|
|
|
float cost_of_subsidy = (1 + kRiskFreeInterestRate) *
|
|
total_subsidy;
|
|
float marginal_planner_economic_index = crra_nonlinearity(
|
|
(total_postsubsidy_productivity - cost_of_subsidy) /
|
|
kPlannerEconomicNorm,
|
|
kEconomicRewardCrraEta,
|
|
kNumDaysInAnYear);
|
|
|
|
marginal_planner_health_index = min_max_normalization(
|
|
marginal_planner_health_index,
|
|
kMinMarginalPlannerHealthIndex,
|
|
kMaxMarginalPlannerHealthIndex);
|
|
marginal_planner_economic_index = min_max_normalization(
|
|
marginal_planner_economic_index,
|
|
kMinMarginalPlannerEconomicIndex,
|
|
kMaxMarginalPlannerEconomicIndex);
|
|
|
|
rewards_p[kEnvId] = get_rew(
|
|
kWeightageOnMarginalAgentEconomicIndex,
|
|
marginal_planner_health_index,
|
|
kWeightageOnMarginalPlannerEconomicIndex,
|
|
marginal_planner_economic_index);
|
|
}
|
|
|
|
// Wait here for all agents to finish computing rewards
|
|
__syncthreads();
|
|
|
|
// Use only agent 0's thread to set done_arr
|
|
if (kAgentId == 0) {
|
|
if (env_timestep_arr[kEnvId] == kEpisodeLength) {
|
|
env_timestep_arr[kEnvId] = 0;
|
|
env_done_arr[kEnvId] = 1;
|
|
}
|
|
}
|
|
}
|
|
}
|