adding ai_economist for modding

ai_economist/foundation/scenarios/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2020, 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
+
+from ai_economist.foundation.base.base_env import scenario_registry
+
+from .covid19 import covid19_env
+from .one_step_economy import one_step_economy
+from .simple_wood_and_stone import dynamic_layout, layout_from_file
+
+# Import files that add Scenario class(es) to scenario_registry
+# -------------------------------------------------------------

ai_economist/foundation/scenarios/covid19/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2020, 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

ai_economist/foundation/scenarios/covid19/covid19_build.cu

@@ -0,0 +1,13 @@
+// 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
+
+#ifndef CUDA_INCLUDES_COVID19_CONST_H_
+#define CUDA_INCLUDES_COVID19_CONST_H_
+
+#include "../../components/covid19_components_step.cu"
+#include "covid19_env_step.cu"
+
+#endif  // CUDA_INCLUDES_COVID19_CONST_H_

ai_economist/foundation/scenarios/covid19/covid19_env_step.cu

@@ -0,0 +1,620 @@
+// 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;
+            }
+        }
+    }
+}

ai_economist/foundation/scenarios/covid19/key_to_check_activation_code_against

@@ -0,0 +1,27 @@
+-----BEGIN RSA PRIVATE KEY-----
+MIIEpgIBAAKCAQEAk1+Qz0/Qg4OOGrskBJnVI9KVGTEUvsldHUV4AzLeecYZSV5+
+FZUQpl8lq1mstUZZ0xMlGSHz2t+AAJxyEro8mAj9gAp1qeN58pAX2k29DOt4YRnp
+sTF1UG+nrV2aW+jfH16aeVsjWY+Nq+GxGyE3Q5bsxOhnOg0TUaB6RY8SBE/scTHn
+bfNsgTc5EuiAAGqYYYdu12n5zeyvfjGW7bBf4Q9t0F0bI+YdZQY9HD35KAoNcqFQ
+dvd2vKbojejkn+WyO1amnZxgAhVjpT61FV4u18jPN0Qrt0LHuF5kUVzYal+73ySY
+BbwEo4onEn9xvUlQGFJWmv4OPwbI3d4nLqP+mQIHK9xUXfK97QKCAQABeR2EO0uu
+ERyRXa5Mh7xsOEq/OJ9sQq+si8B5gDyyM1SW61wQMKF4Wiqw68bMCVvGRwScZD+T
+XwBEBJMm9lCVx/UfOWqYSNFCk/YBefv9AI0Kg5lfCMZQuTdjMcbJdjoR5xoiCbO1
+ya7oOU8mfWx/SV0o/698b/zMVBKBBQDNZaN9pmtTOgm3G1QnM9ZlmrdlKYpe9Ihs
+3sG4437QaPhumdZi8IoLBGMyYL2O38pG34LJjIkP8Efj1QVTndIIZX8CKghir++j
+nUAyofFt7/PBS2k7gQ/1gFISwHxKjmzl/Fc25o7ahlLbO+i2UnRiB9IXcmiGDXMv
+tY09oXhxCtTZAoGBAMEkMTzoiqKjXLwKLyFIF5QzXqQKcGqfC8NhQMsm43K0TgHg
+Sv1fLdnKw0FWSG30gppBorAY9p5FoI+AWwTSd+AJhz7T1y/shpJx1oBR8qKWO5kO
+gMru9kRRb0zb5hydakie3mujz7GUPiXrntKZjC4QYLar0USPulJnU+UTF6QjAoGB
+AMNWJqG1ybrk0sNkWJJDW+MnMT0T9o0E+CtbRHqMHh7K1LF9Sc/qh0gLfDo51+kr
+pscLaaJiF1Q8phzDhW9QDeNv+4lknNqMFBCFtzns1wVDlXL4U87oqhuBSs6IZAuO
+CGVefYKgefdwn64rcyRNala44BbiMJKwRoDvvgH1FvATAoGAV1YK9ZHB1RkXkZ5a
+uBePXvkScaujH4DxadMGf2tBuI1wIpVwhxOQ56yDwYoAuexXPUa8BAx2V69/LFo7
+H/yDYqzndA8WwZLy8oy7Ug+fFLtCp7VhkEwMPciBq6KjzUyShIBlgZOx5m5kTbfu
+Cs2JQU35YHeompcpLooRG1/cFZkCgYAyVlWABzmgSKJL9ohwlSBBZFCjQ1mjN6uc
+uRJxncqfCe3XQ5erFjuWMfPayWONBsWexNucJFc7Iz2LzCOXkUsftldEEET9f/2w
+PrbsEu8khNTLqUcow2Whz+A8C0dV6p2cqtTKR1XlSmNVqP30lmpHcmF+R3M/J1ON
+K7S9zJJ+zwKBgHIuCATGCGCOzAsUo80OQL46j74SxRV3H1CJASLKzatiTo54dbO6
+86w+N6BfYtYeRlnX1CTGl6bHqVUMBBlKws8Ig3gV3xFS8BiSav8zQ2m99JuhlVHF
+Ocfowmuad3WXYvYXQ5IeP2JM/3q7BoPLg1DKP4GGZlNbatMRI+H0HimV
+-----END RSA PRIVATE KEY-----

ai_economist/foundation/scenarios/one_step_economy/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2020, 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

ai_economist/foundation/scenarios/one_step_economy/one_step_economy.py

@@ -0,0 +1,336 @@
+# 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 numpy as np
+
+from ai_economist.foundation.base.base_env import BaseEnvironment, scenario_registry
+from ai_economist.foundation.scenarios.utils import rewards, social_metrics
+
+
+@scenario_registry.add
+class OneStepEconomy(BaseEnvironment):
+    """
+    A simple model featuring one "step" of setting taxes and earning income.
+
+    As described in https://arxiv.org/abs/2108.02755:
+        A simplified version of simple_wood_and_stone scenario where both the planner
+        and the agents each make a single decision: the planner setting taxes and the
+        agents choosing labor. Each agent chooses an amount of labor that optimizes
+        its post-tax utility, and this optimal labor depends on its skill and the tax
+        rates, and it does not depend on the labor choices of other agents. Before
+        the agents act, the planner sets the marginal tax rates in order to optimize
+        social welfare.
+
+    Note:
+        This scenario is intended to be used with the 'PeriodicBracketTax' and
+            'SimpleLabor' components.
+        It should use an episode length of 2. In the first step, taxes are set by
+            the planner via 'PeriodicBracketTax'. In the second, agents select how much
+            to work/earn via 'SimpleLabor'.
+
+    Args:
+        agent_reward_type (str): The type of utility function used to compute each
+            agent's reward. Defaults to "coin_minus_labor_cost".
+        isoelastic_eta (float): The shape parameter of the isoelastic function used
+            in the "isoelastic_coin_minus_labor" utility function.
+        labor_exponent (float): The labor exponent parameter used in the
+            "coin_minus_labor_cost" utility function.
+        labor_cost (float): The coefficient used to weight the cost of labor.
+        planner_reward_type (str): The type of social welfare function (SWF) used to
+            compute the planner's reward. Defaults to "inv_income_weighted_utility".
+        mixing_weight_gini_vs_coin (float): Must be between 0 and 1 (inclusive).
+            Controls the weighting of equality and productivity when using SWF
+            "coin_eq_times_productivity", where a value of 0 (default) yields equal
+            weighting, and 1 only considers productivity.
+    """
+
+    name = "one-step-economy"
+    agent_subclasses = ["BasicMobileAgent", "BasicPlanner"]
+    required_entities = ["Coin"]
+
+    def __init__(
+        self,
+        *base_env_args,
+        agent_reward_type="coin_minus_labor_cost",
+        isoelastic_eta=0.23,
+        labor_exponent=2.0,
+        labor_cost=1.0,
+        planner_reward_type="inv_income_weighted_utility",
+        mixing_weight_gini_vs_coin=0,
+        **base_env_kwargs
+    ):
+        super().__init__(*base_env_args, **base_env_kwargs)
+
+        self.num_agents = len(self.world.agents)
+
+        self.labor_cost = labor_cost
+        self.agent_reward_type = agent_reward_type
+        self.isoelastic_eta = isoelastic_eta
+        self.labor_exponent = labor_exponent
+        self.planner_reward_type = planner_reward_type
+        self.mixing_weight_gini_vs_coin = mixing_weight_gini_vs_coin
+        self.planner_starting_coin = 0
+
+        self.curr_optimization_metrics = {str(a.idx): 0 for a in self.all_agents}
+
+    # The following methods must be implemented for each scenario
+    # -----------------------------------------------------------
+    def reset_starting_layout(self):
+        """
+        Part 1/2 of scenario reset. This method handles resetting the state of the
+        environment managed by the scenario (i.e. resource & landmark layout).
+
+        Here, generate a resource source layout consistent with target parameters.
+        """
+
+    def reset_agent_states(self):
+        """
+        Part 2/2 of scenario reset. This method handles resetting the state of the
+        agents themselves (i.e. inventory, locations, etc.).
+
+        Here, empty inventories, give mobile agents any starting coin, and place them
+        in random accesible locations to start.
+        """
+        self.world.clear_agent_locs()
+
+        for agent in self.world.agents:
+            # Clear everything to start with
+            agent.state["inventory"] = {k: 0 for k in agent.state["inventory"].keys()}
+            agent.state["escrow"] = {k: 0 for k in agent.state["escrow"].keys()}
+            agent.state["endogenous"] = {k: 0 for k in agent.state["endogenous"].keys()}
+
+        self.world.planner.inventory["Coin"] = self.planner_starting_coin
+
+    def scenario_step(self):
+        """
+        Update the state of the world according to whatever rules this scenario
+        implements.
+
+        This gets called in the 'step' method (of base_env) after going through each
+        component step and before generating observations, rewards, etc.
+
+        NOTE: does not take agent actions into account.
+        """
+
+    def generate_observations(self):
+        """
+        Generate observations associated with this scenario.
+
+        A scenario does not need to produce observations and can provide observations
+        for only some agent types; however, for a given agent type, it should either
+        always or never yield an observation. If it does yield an observation,
+        that observation should always have the same structure/sizes!
+
+        Returns:
+            obs (dict): A dictionary of {agent.idx: agent_obs_dict}. In words,
+                return a dictionary with an entry for each agent (which can including
+                the planner) for which this scenario provides an observation. For each
+                entry, the key specifies the index of the agent and the value contains
+                its associated observation dictionary.
+
+        Here, non-planner agents receive spatial observations (depending on the env
+        config) as well as the contents of their inventory and endogenous quantities.
+        The planner also receives spatial observations (again, depending on the env
+        config) as well as the inventory of each of the mobile agents.
+        """
+        obs_dict = dict()
+        for agent in self.world.agents:
+            obs_dict[str(agent.idx)] = {}
+
+        coin_endowments = np.array(
+            [agent.total_endowment("Coin") for agent in self.world.agents]
+        )
+        equality = social_metrics.get_equality(coin_endowments)
+        productivity = social_metrics.get_productivity(coin_endowments)
+        normalized_per_capita_productivity = productivity / self.num_agents / 1000
+        obs_dict[self.world.planner.idx] = {
+            "normalized_per_capita_productivity": normalized_per_capita_productivity,
+            "equality": equality,
+        }
+
+        return obs_dict
+
+    def compute_reward(self):
+        """
+        Apply the reward function(s) associated with this scenario to get the rewards
+        from this step.
+
+        Returns:
+            rew (dict): A dictionary of {agent.idx: agent_obs_dict}. In words,
+                return a dictionary with an entry for each agent in the environment
+                (including the planner). For each entry, the key specifies the index of
+                the agent and the value contains the scalar reward earned this timestep.
+
+        Rewards are computed as the marginal utility (agents) or marginal social
+        welfare (planner) experienced on this timestep. Ignoring discounting,
+        this means that agents' (planner's) objective is to maximize the utility
+        (social welfare) associated with the terminal state of the episode.
+        """
+        curr_optimization_metrics = self.get_current_optimization_metrics(
+            self.world.agents,
+            isoelastic_eta=float(self.isoelastic_eta),
+            labor_exponent=float(self.labor_exponent),
+            labor_coefficient=float(self.labor_cost),
+        )
+        planner_agents_rew = {
+            k: v - self.curr_optimization_metrics[k]
+            for k, v in curr_optimization_metrics.items()
+        }
+        self.curr_optimization_metrics = curr_optimization_metrics
+        return planner_agents_rew
+
+    # Optional methods for customization
+    # ----------------------------------
+    def additional_reset_steps(self):
+        """
+        Extra scenario-specific steps that should be performed at the end of the reset
+        cycle.
+
+        For each reset cycle...
+            First, reset_starting_layout() and reset_agent_states() will be called.
+
+            Second, <component>.reset() will be called for each registered component.
+
+            Lastly, this method will be called to allow for any final customization of
+            the reset cycle.
+        """
+        self.curr_optimization_metrics = self.get_current_optimization_metrics(
+            self.world.agents,
+            isoelastic_eta=float(self.isoelastic_eta),
+            labor_exponent=float(self.labor_exponent),
+            labor_coefficient=float(self.labor_cost),
+        )
+
+    def scenario_metrics(self):
+        """
+        Allows the scenario to generate metrics (collected along with component metrics
+        in the 'metrics' property).
+
+        To have the scenario add metrics, this function needs to return a dictionary of
+        {metric_key: value} where 'value' is a scalar (no nesting or lists!)
+
+        Here, summarize social metrics, endowments, utilities, and labor cost annealing.
+        """
+        metrics = dict()
+
+        # Log social/economic indicators
+        coin_endowments = np.array(
+            [agent.total_endowment("Coin") for agent in self.world.agents]
+        )
+        pretax_incomes = np.array(
+            [agent.state["production"] for agent in self.world.agents]
+        )
+        metrics["social/productivity"] = social_metrics.get_productivity(
+            coin_endowments
+        )
+        metrics["social/equality"] = social_metrics.get_equality(coin_endowments)
+
+        utilities = np.array(
+            [self.curr_optimization_metrics[agent.idx] for agent in self.world.agents]
+        )
+        metrics[
+            "social_welfare/coin_eq_times_productivity"
+        ] = rewards.coin_eq_times_productivity(
+            coin_endowments=coin_endowments, equality_weight=1.0
+        )
+        metrics[
+            "social_welfare/inv_income_weighted_utility"
+        ] = rewards.inv_income_weighted_utility(
+            coin_endowments=pretax_incomes, utilities=utilities  # coin_endowments,
+        )
+
+        # Log average endowments, endogenous, and utility for agents
+        agent_endows = {}
+        agent_endogenous = {}
+        agent_utilities = []
+        for agent in self.world.agents:
+            for resource in agent.inventory.keys():
+                if resource not in agent_endows:
+                    agent_endows[resource] = []
+                agent_endows[resource].append(
+                    agent.inventory[resource] + agent.escrow[resource]
+                )
+
+            for endogenous, quantity in agent.endogenous.items():
+                if endogenous not in agent_endogenous:
+                    agent_endogenous[endogenous] = []
+                agent_endogenous[endogenous].append(quantity)
+
+            agent_utilities.append(self.curr_optimization_metrics[agent.idx])
+
+        for resource, quantities in agent_endows.items():
+            metrics["endow/avg_agent/{}".format(resource)] = np.mean(quantities)
+
+        for endogenous, quantities in agent_endogenous.items():
+            metrics["endogenous/avg_agent/{}".format(endogenous)] = np.mean(quantities)
+
+        metrics["util/avg_agent"] = np.mean(agent_utilities)
+
+        # Log endowments and utility for the planner
+        for resource, quantity in self.world.planner.inventory.items():
+            metrics["endow/p/{}".format(resource)] = quantity
+
+        metrics["util/p"] = self.curr_optimization_metrics[self.world.planner.idx]
+
+        return metrics
+
+    def get_current_optimization_metrics(
+        self, agents, isoelastic_eta=0.23, labor_exponent=2.0, labor_coefficient=0.1
+    ):
+        """
+        Compute optimization metrics based on the current state. Used to compute reward.
+
+        Returns:
+            curr_optimization_metric (dict): A dictionary of {agent.idx: metric}
+                with an entry for each agent (including the planner) in the env.
+        """
+        curr_optimization_metric = {}
+
+        coin_endowments = np.array([agent.total_endowment("Coin") for agent in agents])
+
+        pretax_incomes = np.array([agent.state["production"] for agent in agents])
+
+        # Optimization metric for agents:
+        for agent in agents:
+            if self.agent_reward_type == "isoelastic_coin_minus_labor":
+                assert 0.0 <= isoelastic_eta <= 1.0
+                curr_optimization_metric[
+                    agent.idx
+                ] = rewards.isoelastic_coin_minus_labor(
+                    coin_endowment=agent.total_endowment("Coin"),
+                    total_labor=agent.state["endogenous"]["Labor"],
+                    isoelastic_eta=isoelastic_eta,
+                    labor_coefficient=labor_coefficient,
+                )
+            elif self.agent_reward_type == "coin_minus_labor_cost":
+                assert labor_exponent > 1.0
+                curr_optimization_metric[agent.idx] = rewards.coin_minus_labor_cost(
+                    coin_endowment=agent.total_endowment("Coin"),
+                    total_labor=agent.state["endogenous"]["Labor"],
+                    labor_exponent=labor_exponent,
+                    labor_coefficient=labor_coefficient,
+                )
+        # Optimization metric for the planner:
+        if self.planner_reward_type == "coin_eq_times_productivity":
+            curr_optimization_metric[
+                self.world.planner.idx
+            ] = rewards.coin_eq_times_productivity(
+                coin_endowments=coin_endowments,
+                equality_weight=1 - self.mixing_weight_gini_vs_coin,
+            )
+        elif self.planner_reward_type == "inv_income_weighted_utility":
+            curr_optimization_metric[
+                self.world.planner.idx
+            ] = rewards.inv_income_weighted_utility(
+                coin_endowments=pretax_incomes,  # coin_endowments,
+                utilities=np.array(
+                    [curr_optimization_metric[agent.idx] for agent in agents]
+                ),
+            )
+        else:
+            print("No valid planner reward selected!")
+            raise NotImplementedError
+        return curr_optimization_metric

ai_economist/foundation/scenarios/simple_wood_and_stone/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2020, 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

ai_economist/foundation/scenarios/simple_wood_and_stone/layout_from_file.py

@@ -0,0 +1,800 @@
+# Copyright (c) 2020, 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
+
+from copy import deepcopy
+from pathlib import Path
+
+import numpy as np
+from scipy import signal
+
+from ai_economist.foundation.base.base_env import BaseEnvironment, scenario_registry
+from ai_economist.foundation.scenarios.utils import rewards, social_metrics
+
+
+@scenario_registry.add
+class LayoutFromFile(BaseEnvironment):
+    """
+    World containing stone and wood with stochastic regeneration. Refers to a fixed
+    layout file (see ./map_txt/ for examples) to determine the spatial arrangement of
+    stone, wood, and water tiles.
+
+    Args:
+        planner_gets_spatial_obs (bool): Whether the planner agent receives spatial
+            observations from the world.
+        full_observability (bool): Whether the mobile agents' spatial observation
+            includes the full world view or is instead an egocentric view.
+        mobile_agent_observation_range (int): If not using full_observability,
+            the spatial range (on each side of the agent) that is visible in the
+            spatial observations.
+        env_layout_file (str): Name of the layout file in ./map_txt/ to use.
+            Note: The world dimensions of that layout must match the world dimensions
+            argument used to construct the environment.
+        resource_regen_prob (float): Probability that an empty source tile will
+            regenerate a new resource unit.
+        fixed_four_skill_and_loc (bool): Whether to use a fixed set of build skills and
+            starting locations, with agents grouped into starting locations based on
+            which skill quartile they are in. False, by default.
+            True, for experiments in https://arxiv.org/abs/2004.13332.
+            Note: Requires that the environment uses the "Build" component with
+            skill_dist="pareto".
+        starting_agent_coin (int, float): Amount of coin agents have at t=0. Defaults
+            to zero coin.
+        isoelastic_eta (float): Parameter controlling the shape of agent utility
+            wrt coin endowment.
+        energy_cost (float): Coefficient for converting labor to negative utility.
+        energy_warmup_constant (float): Decay constant that controls the rate at which
+            the effective energy cost is annealed from 0 to energy_cost. Set to 0
+            (default) to disable annealing, meaning that the effective energy cost is
+            always energy_cost. The units of the decay constant depend on the choice of
+            energy_warmup_method.
+        energy_warmup_method (str): How to schedule energy annealing (warmup). If
+            "decay" (default), use the number of completed episodes. If "auto",
+            use the number of timesteps where the average agent reward was positive.
+        planner_reward_type (str): The type of reward used for the planner. Options
+            are "coin_eq_times_productivity" (default),
+            "inv_income_weighted_coin_endowment", and "inv_income_weighted_utility".
+        mixing_weight_gini_vs_coin (float): Degree to which equality is ignored w/
+            "coin_eq_times_productivity". Default is 0, which weights equality and
+            productivity equally. If set to 1, only productivity is rewarded.
+    """
+
+    name = "layout_from_file/simple_wood_and_stone"
+    agent_subclasses = ["BasicMobileAgent", "BasicPlanner"]
+    required_entities = ["Wood", "Stone", "Water"]
+
+    def __init__(
+        self,
+        *base_env_args,
+        planner_gets_spatial_info=True,
+        full_observability=False,
+        mobile_agent_observation_range=5,
+        env_layout_file="quadrant_25x25_20each_30clump.txt",
+        resource_regen_prob=0.01,
+        fixed_four_skill_and_loc=False,
+        starting_agent_coin=0,
+        isoelastic_eta=0.23,
+        energy_cost=0.21,
+        energy_warmup_constant=0,
+        energy_warmup_method="decay",
+        planner_reward_type="coin_eq_times_productivity",
+        mixing_weight_gini_vs_coin=0.0,
+        **base_env_kwargs,
+    ):
+        super().__init__(*base_env_args, **base_env_kwargs)
+
+        # Whether agents receive spatial information in their observation tensor
+        self._planner_gets_spatial_info = bool(planner_gets_spatial_info)
+
+        # Whether the (non-planner) agents can see the whole world map
+        self._full_observability = bool(full_observability)
+
+        self._mobile_agent_observation_range = int(mobile_agent_observation_range)
+
+        # Load in the layout
+        path_to_layout_file = Path(f"{Path(__file__).parent}/map_txt/{env_layout_file}")
+
+        with open(path_to_layout_file, "r") as f:
+            self.env_layout_string = f.read()
+            self.env_layout = self.env_layout_string.split(";")
+
+        # Convert the layout to landmark maps
+        landmark_lookup = {"W": "Wood", "S": "Stone", "@": "Water"}
+        self._source_maps = {
+            r: np.zeros(self.world_size) for r in landmark_lookup.values()
+        }
+        for r, symbol_row in enumerate(self.env_layout):
+            for c, symbol in enumerate(symbol_row):
+                landmark = landmark_lookup.get(symbol, None)
+                if landmark:
+                    self._source_maps[landmark][r, c] = 1
+
+        # For controlling how resource regeneration behavior
+        self.layout_specs = dict(
+            Wood={
+                "regen_weight": float(resource_regen_prob),
+                "regen_halfwidth": 0,
+                "max_health": 1,
+            },
+            Stone={
+                "regen_weight": float(resource_regen_prob),
+                "regen_halfwidth": 0,
+                "max_health": 1,
+            },
+        )
+        assert 0 <= self.layout_specs["Wood"]["regen_weight"] <= 1
+        assert 0 <= self.layout_specs["Stone"]["regen_weight"] <= 1
+
+        # How much coin do agents begin with at upon reset
+        self.starting_agent_coin = float(starting_agent_coin)
+        assert self.starting_agent_coin >= 0.0
+
+        # Controls the diminishing marginal utility of coin.
+        # isoelastic_eta=0 means no diminishing utility.
+        self.isoelastic_eta = float(isoelastic_eta)
+        assert 0.0 <= self.isoelastic_eta <= 1.0
+
+        # The amount that labor is weighted in utility computation
+        # (once annealing is finished)
+        self.energy_cost = float(energy_cost)
+        assert self.energy_cost >= 0
+
+        # Which method to use for calculating the progress of energy annealing
+        # If method = 'decay': #completed episodes
+        # If method = 'auto' : #timesteps where avg. agent reward > 0
+        self.energy_warmup_method = energy_warmup_method.lower()
+        assert self.energy_warmup_method in ["decay", "auto"]
+        # Decay constant for annealing to full energy cost
+        # (if energy_warmup_constant == 0, there is no annealing)
+        self.energy_warmup_constant = float(energy_warmup_constant)
+        assert self.energy_warmup_constant >= 0
+        self._auto_warmup_integrator = 0
+
+        # Which social welfare function to use
+        self.planner_reward_type = str(planner_reward_type).lower()
+
+        # How much to weight equality if using SWF=eq*prod:
+        # 0 -> SWF=eq * prod
+        # 1 -> SWF=prod
+        self.mixing_weight_gini_vs_coin = float(mixing_weight_gini_vs_coin)
+        assert 0 <= self.mixing_weight_gini_vs_coin <= 1.0
+
+        # Use this to calculate marginal changes and deliver that as reward
+        self.init_optimization_metric = {agent.idx: 0 for agent in self.all_agents}
+        self.prev_optimization_metric = {agent.idx: 0 for agent in self.all_agents}
+        self.curr_optimization_metric = {agent.idx: 0 for agent in self.all_agents}
+
+        """
+        Fixed Four Skill and Loc
+        ------------------------
+        """
+        self.agent_starting_pos = {agent.idx: [] for agent in self.world.agents}
+
+        self.fixed_four_skill_and_loc = bool(fixed_four_skill_and_loc)
+        if self.fixed_four_skill_and_loc:
+            bm = self.get_component("Build")
+            assert bm.skill_dist == "pareto"
+            pmsm = bm.payment_max_skill_multiplier
+
+            # Temporarily switch to a fixed seed for controlling randomness
+            seed_state = np.random.get_state()
+            np.random.seed(seed=1)
+
+            # Generate a batch (100000) of num_agents (sorted/clipped) Pareto samples.
+            pareto_samples = np.random.pareto(4, size=(100000, self.n_agents))
+            clipped_skills = np.minimum(pmsm, (pmsm - 1) * pareto_samples + 1)
+            sorted_clipped_skills = np.sort(clipped_skills, axis=1)
+            # The skill level of the i-th skill-ranked agent is the average of the
+            # i-th ranked samples throughout the batch.
+            average_ranked_skills = sorted_clipped_skills.mean(axis=0)
+            self._avg_ranked_skill = average_ranked_skills * bm.payment
+
+            np.random.set_state(seed_state)
+
+            # Fill in the starting location associated with each skill rank
+            starting_ranked_locs = [
+                # Worst group of agents goes in top right
+                (0, self.world_size[1] - 1),
+                # Second-worst group of agents goes in bottom left
+                (self.world_size[0] - 1, 0),
+                # Second-best group of agents goes in top left
+                (0, 0),
+                # Best group of agents goes in bottom right
+                (self.world_size[1] - 1, self.world_size[1] - 1),
+            ]
+            self._ranked_locs = []
+
+            # Based on skill, assign each agent to one of the location groups
+            skill_groups = np.floor(
+                np.arange(self.n_agents) * (4 / self.n_agents),
+            ).astype(np.int)
+            n_in_group = np.zeros(4, dtype=np.int)
+            for g in skill_groups:
+                # The position within the group is given by the number of agents
+                # counted in the group thus far.
+                g_pos = n_in_group[g]
+
+                # Top right
+                if g == 0:
+                    r = starting_ranked_locs[g][0] + (g_pos // 4)
+                    c = starting_ranked_locs[g][1] - (g_pos % 4)
+                    self._ranked_locs.append((r, c))
+
+                # Bottom left
+                elif g == 1:
+                    r = starting_ranked_locs[g][0] - (g_pos // 4)
+                    c = starting_ranked_locs[g][1] + (g_pos % 4)
+                    self._ranked_locs.append((r, c))
+
+                # Top left
+                elif g == 2:
+                    r = starting_ranked_locs[g][0] + (g_pos // 4)
+                    c = starting_ranked_locs[g][1] + (g_pos % 4)
+                    self._ranked_locs.append((r, c))
+
+                # Bottom right
+                elif g == 3:
+                    r = starting_ranked_locs[g][0] - (g_pos // 4)
+                    c = starting_ranked_locs[g][1] - (g_pos % 4)
+                    self._ranked_locs.append((r, c))
+
+                else:
+                    raise ValueError
+
+                # Count the agent we just placed.
+                n_in_group[g] = n_in_group[g] + 1
+
+    @property
+    def energy_weight(self):
+        """
+        Energy annealing progress. Multiply with self.energy_cost to get the
+        effective energy coefficient.
+        """
+        if self.energy_warmup_constant <= 0.0:
+            return 1.0
+
+        if self.energy_warmup_method == "decay":
+            return float(1.0 - np.exp(-self._completions / self.energy_warmup_constant))
+
+        if self.energy_warmup_method == "auto":
+            return float(
+                1.0
+                - np.exp(-self._auto_warmup_integrator / self.energy_warmup_constant)
+            )
+
+        raise NotImplementedError
+
+    def get_current_optimization_metrics(self):
+        """
+        Compute optimization metrics based on the current state. Used to compute reward.
+
+        Returns:
+            curr_optimization_metric (dict): A dictionary of {agent.idx: metric}
+                with an entry for each agent (including the planner) in the env.
+        """
+        curr_optimization_metric = {}
+        # (for agents)
+        for agent in self.world.agents:
+            curr_optimization_metric[agent.idx] = rewards.isoelastic_coin_minus_labor(
+                coin_endowment=agent.total_endowment("Coin"),
+                total_labor=agent.state["endogenous"]["Labor"],
+                isoelastic_eta=self.isoelastic_eta,
+                labor_coefficient=self.energy_weight * self.energy_cost,
+            )
+        # (for the planner)
+        if self.planner_reward_type == "coin_eq_times_productivity":
+            curr_optimization_metric[
+                self.world.planner.idx
+            ] = rewards.coin_eq_times_productivity(
+                coin_endowments=np.array(
+                    [agent.total_endowment("Coin") for agent in self.world.agents]
+                ),
+                equality_weight=1 - self.mixing_weight_gini_vs_coin,
+            )
+        elif self.planner_reward_type == "inv_income_weighted_coin_endowments":
+            curr_optimization_metric[
+                self.world.planner.idx
+            ] = rewards.inv_income_weighted_coin_endowments(
+                coin_endowments=np.array(
+                    [agent.total_endowment("Coin") for agent in self.world.agents]
+                )
+            )
+        elif self.planner_reward_type == "inv_income_weighted_utility":
+            curr_optimization_metric[
+                self.world.planner.idx
+            ] = rewards.inv_income_weighted_utility(
+                coin_endowments=np.array(
+                    [agent.total_endowment("Coin") for agent in self.world.agents]
+                ),
+                utilities=np.array(
+                    [curr_optimization_metric[agent.idx] for agent in self.world.agents]
+                ),
+            )
+        else:
+            print("No valid planner reward selected!")
+            raise NotImplementedError
+        return curr_optimization_metric
+
+    # The following methods must be implemented for each scenario
+    # -----------------------------------------------------------
+
+    def reset_starting_layout(self):
+        """
+        Part 1/2 of scenario reset. This method handles resetting the state of the
+        environment managed by the scenario (i.e. resource & landmark layout).
+
+        Here, reset to the layout in the fixed layout file
+        """
+        self.world.maps.clear()
+        for landmark, landmark_map in self._source_maps.items():
+            self.world.maps.set(landmark, landmark_map)
+            if landmark in ["Stone", "Wood"]:
+                self.world.maps.set(landmark + "SourceBlock", landmark_map)
+
+    def reset_agent_states(self):
+        """
+        Part 2/2 of scenario reset. This method handles resetting the state of the
+        agents themselves (i.e. inventory, locations, etc.).
+
+        Here, empty inventories and place mobile agents in random, accessible
+        locations to start. Note: If using fixed_four_skill_and_loc, the starting
+        locations will be overridden in self.additional_reset_steps.
+        """
+        self.world.clear_agent_locs()
+        for agent in self.world.agents:
+            agent.state["inventory"] = {k: 0 for k in agent.inventory.keys()}
+            agent.state["escrow"] = {k: 0 for k in agent.inventory.keys()}
+            agent.state["endogenous"] = {k: 0 for k in agent.endogenous.keys()}
+            # Add starting coin
+            agent.state["inventory"]["Coin"] = float(self.starting_agent_coin)
+
+        self.world.planner.state["inventory"] = {
+            k: 0 for k in self.world.planner.inventory.keys()
+        }
+        self.world.planner.state["escrow"] = {
+            k: 0 for k in self.world.planner.escrow.keys()
+        }
+
+        for agent in self.world.agents:
+            r = np.random.randint(0, self.world_size[0])
+            c = np.random.randint(0, self.world_size[1])
+            n_tries = 0
+            while not self.world.can_agent_occupy(r, c, agent):
+                r = np.random.randint(0, self.world_size[0])
+                c = np.random.randint(0, self.world_size[1])
+                n_tries += 1
+                if n_tries > 200:
+                    raise TimeoutError
+            r, c = self.world.set_agent_loc(agent, r, c)
+
+    def scenario_step(self):
+        """
+        Update the state of the world according to whatever rules this scenario
+        implements.
+
+        This gets called in the 'step' method (of base_env) after going through each
+        component step and before generating observations, rewards, etc.
+
+        In this class of scenarios, the scenario step handles stochastic resource
+        regeneration.
+        """
+
+        resources = ["Wood", "Stone"]
+
+        for resource in resources:
+            d = 1 + (2 * self.layout_specs[resource]["regen_halfwidth"])
+            kernel = (
+                self.layout_specs[resource]["regen_weight"] * np.ones((d, d)) / (d ** 2)
+            )
+
+            resource_map = self.world.maps.get(resource)
+            resource_source_blocks = self.world.maps.get(resource + "SourceBlock")
+            spawnable = (
+                self.world.maps.empty + resource_map + resource_source_blocks
+            ) > 0
+            spawnable *= resource_source_blocks > 0
+
+            health = np.maximum(resource_map, resource_source_blocks)
+            respawn = np.random.rand(*health.shape) < signal.convolve2d(
+                health, kernel, "same"
+            )
+            respawn *= spawnable
+
+            self.world.maps.set(
+                resource,
+                np.minimum(
+                    resource_map + respawn, self.layout_specs[resource]["max_health"]
+                ),
+            )
+
+    def generate_observations(self):
+        """
+        Generate observations associated with this scenario.
+
+        A scenario does not need to produce observations and can provide observations
+        for only some agent types; however, for a given agent type, it should either
+        always or never yield an observation. If it does yield an observation,
+        that observation should always have the same structure/sizes!
+
+        Returns:
+            obs (dict): A dictionary of {agent.idx: agent_obs_dict}. In words,
+                return a dictionary with an entry for each agent (which can including
+                the planner) for which this scenario provides an observation. For each
+                entry, the key specifies the index of the agent and the value contains
+                its associated observation dictionary.
+
+        Here, non-planner agents receive spatial observations (depending on the env
+        config) as well as the contents of their inventory and endogenous quantities.
+        The planner also receives spatial observations (again, depending on the env
+        config) as well as the inventory of each of the mobile agents.
+        """
+        obs = {}
+        curr_map = self.world.maps.state
+
+        owner_map = self.world.maps.owner_state
+        loc_map = self.world.loc_map
+        agent_idx_maps = np.concatenate([owner_map, loc_map[None, :, :]], axis=0)
+        agent_idx_maps += 2
+        agent_idx_maps[agent_idx_maps == 1] = 0
+
+        agent_locs = {
+            str(agent.idx): {
+                "loc-row": agent.loc[0] / self.world_size[0],
+                "loc-col": agent.loc[1] / self.world_size[1],
+            }
+            for agent in self.world.agents
+        }
+        agent_invs = {
+            str(agent.idx): {
+                "inventory-" + k: v * self.inv_scale for k, v in agent.inventory.items()
+            }
+            for agent in self.world.agents
+        }
+
+        obs[self.world.planner.idx] = {
+            "inventory-" + k: v * self.inv_scale
+            for k, v in self.world.planner.inventory.items()
+        }
+        if self._planner_gets_spatial_info:
+            obs[self.world.planner.idx].update(
+                dict(map=curr_map, idx_map=agent_idx_maps)
+            )
+
+        # Mobile agents see the full map. Convey location info via one-hot map channels.
+        if self._full_observability:
+            for agent in self.world.agents:
+                my_map = np.array(agent_idx_maps)
+                my_map[my_map == int(agent.idx) + 2] = 1
+                sidx = str(agent.idx)
+                obs[sidx] = {"map": curr_map, "idx_map": my_map}
+                obs[sidx].update(agent_invs[sidx])
+
+        # Mobile agents only see within a window around their position
+        else:
+            w = (
+                self._mobile_agent_observation_range
+            )  # View halfwidth (only applicable without full observability)
+
+            padded_map = np.pad(
+                curr_map,
+                [(0, 1), (w, w), (w, w)],
+                mode="constant",
+                constant_values=[(0, 1), (0, 0), (0, 0)],
+            )
+
+            padded_idx = np.pad(
+                agent_idx_maps,
+                [(0, 0), (w, w), (w, w)],
+                mode="constant",
+                constant_values=[(0, 0), (0, 0), (0, 0)],
+            )
+
+            for agent in self.world.agents:
+                r, c = [c + w for c in agent.loc]
+                visible_map = padded_map[
+                    :, (r - w) : (r + w + 1), (c - w) : (c + w + 1)
+                ]
+                visible_idx = np.array(
+                    padded_idx[:, (r - w) : (r + w + 1), (c - w) : (c + w + 1)]
+                )
+
+                visible_idx[visible_idx == int(agent.idx) + 2] = 1
+
+                sidx = str(agent.idx)
+
+                obs[sidx] = {"map": visible_map, "idx_map": visible_idx}
+                obs[sidx].update(agent_locs[sidx])
+                obs[sidx].update(agent_invs[sidx])
+
+                # Agent-wise planner info (gets crunched into the planner obs in the
+                # base scenario code)
+                obs["p" + sidx] = agent_invs[sidx]
+                if self._planner_gets_spatial_info:
+                    obs["p" + sidx].update(agent_locs[sidx])
+
+        return obs
+
+    def compute_reward(self):
+        """
+        Apply the reward function(s) associated with this scenario to get the rewards
+        from this step.
+
+        Returns:
+            rew (dict): A dictionary of {agent.idx: agent_obs_dict}. In words,
+                return a dictionary with an entry for each agent in the environment
+                (including the planner). For each entry, the key specifies the index of
+                the agent and the value contains the scalar reward earned this timestep.
+
+        Rewards are computed as the marginal utility (agents) or marginal social
+        welfare (planner) experienced on this timestep. Ignoring discounting,
+        this means that agents' (planner's) objective is to maximize the utility
+        (social welfare) associated with the terminal state of the episode.
+        """
+
+        # "curr_optimization_metric" hasn't been updated yet, so it gives us the
+        # utility from the last step.
+        utility_at_end_of_last_time_step = deepcopy(self.curr_optimization_metric)
+
+        # compute current objectives and store the values
+        self.curr_optimization_metric = self.get_current_optimization_metrics()
+
+        # reward = curr - prev objectives
+        rew = {
+            k: float(v - utility_at_end_of_last_time_step[k])
+            for k, v in self.curr_optimization_metric.items()
+        }
+
+        # store the previous objective values
+        self.prev_optimization_metric.update(utility_at_end_of_last_time_step)
+
+        # Automatic Energy Cost Annealing
+        # -------------------------------
+        avg_agent_rew = np.mean([rew[a.idx] for a in self.world.agents])
+        # Count the number of timesteps where the avg agent reward was > 0
+        if avg_agent_rew > 0:
+            self._auto_warmup_integrator += 1
+
+        return rew
+
+    # Optional methods for customization
+    # ----------------------------------
+
+    def additional_reset_steps(self):
+        """
+        Extra scenario-specific steps that should be performed at the end of the reset
+        cycle.
+
+        For each reset cycle...
+            First, reset_starting_layout() and reset_agent_states() will be called.
+
+            Second, <component>.reset() will be called for each registered component.
+
+            Lastly, this method will be called to allow for any final customization of
+            the reset cycle.
+
+        For this scenario, this method resets optimization metric trackers. If using
+        fixed_four_skill_and_loc, this is where each agent gets assigned to one of
+        the four fixed skill/loc combinations. The agent-->skill/loc assignment is
+        permuted so that all four skill/loc combinations are used.
+        """
+        if self.fixed_four_skill_and_loc:
+            self.world.clear_agent_locs()
+            for i, agent in enumerate(self.world.get_random_order_agents()):
+                self.world.set_agent_loc(agent, *self._ranked_locs[i])
+                agent.state["build_payment"] = self._avg_ranked_skill[i]
+
+        # compute current objectives
+        curr_optimization_metric = self.get_current_optimization_metrics()
+
+        self.curr_optimization_metric = deepcopy(curr_optimization_metric)
+        self.init_optimization_metric = deepcopy(curr_optimization_metric)
+        self.prev_optimization_metric = deepcopy(curr_optimization_metric)
+
+    def scenario_metrics(self):
+        """
+        Allows the scenario to generate metrics (collected along with component metrics
+        in the 'metrics' property).
+
+        To have the scenario add metrics, this function needs to return a dictionary of
+        {metric_key: value} where 'value' is a scalar (no nesting or lists!)
+
+        Here, summarize social metrics, endowments, utilities, and labor cost annealing.
+        """
+        metrics = dict()
+
+        coin_endowments = np.array(
+            [agent.total_endowment("Coin") for agent in self.world.agents]
+        )
+        metrics["social/productivity"] = social_metrics.get_productivity(
+            coin_endowments
+        )
+        metrics["social/equality"] = social_metrics.get_equality(coin_endowments)
+
+        utilities = np.array(
+            [self.curr_optimization_metric[agent.idx] for agent in self.world.agents]
+        )
+        metrics[
+            "social_welfare/coin_eq_times_productivity"
+        ] = rewards.coin_eq_times_productivity(
+            coin_endowments=coin_endowments, equality_weight=1.0
+        )
+        metrics[
+            "social_welfare/inv_income_weighted_coin_endow"
+        ] = rewards.inv_income_weighted_coin_endowments(coin_endowments=coin_endowments)
+        metrics[
+            "social_welfare/inv_income_weighted_utility"
+        ] = rewards.inv_income_weighted_utility(
+            coin_endowments=coin_endowments, utilities=utilities
+        )
+
+        for agent in self.all_agents:
+            for resource, quantity in agent.inventory.items():
+                metrics[
+                    "endow/{}/{}".format(agent.idx, resource)
+                ] = agent.total_endowment(resource)
+
+            if agent.endogenous is not None:
+                for resource, quantity in agent.endogenous.items():
+                    metrics["endogenous/{}/{}".format(agent.idx, resource)] = quantity
+
+            metrics["util/{}".format(agent.idx)] = self.curr_optimization_metric[
+                agent.idx
+            ]
+
+        # Labor weight
+        metrics["labor/weighted_cost"] = self.energy_cost * self.energy_weight
+        metrics["labor/warmup_integrator"] = int(self._auto_warmup_integrator)
+
+        return metrics
+
+
+@scenario_registry.add
+class SplitLayout(LayoutFromFile):
+    """
+    Extends layout_from_file/simple_wood_and_stone to impose a row of water midway
+    through the map, uses a fixed set of pareto-distributed building skills (requires a
+    Build component), and places agents in the top/bottom depending on skill rank.
+
+    Args:
+        water_row (int): Row of the map where the water barrier is placed. Defaults
+        to half the world height.
+        skill_rank_of_top_agents (int, float, tuple, list): Index/indices specifying
+            which agent(s) to place in the top of the map. Indices refer to the skill
+            ranking, with 0 referring to the highest-skilled agent. Defaults to only
+            the highest-skilled agent in the top.
+        planner_gets_spatial_obs (bool): Whether the planner agent receives spatial
+            observations from the world.
+        full_observability (bool): Whether the mobile agents' spatial observation
+            includes the full world view or is instead an egocentric view.
+        mobile_agent_observation_range (int): If not using full_observability,
+            the spatial range (on each side of the agent) that is visible in the
+            spatial observations.
+        env_layout_file (str): Name of the layout file in ./map_txt/ to use.
+            Note: The world dimensions of that layout must match the world dimensions
+            argument used to construct the environment.
+        resource_regen_prob (float): Probability that an empty source tile will
+            regenerate a new resource unit.
+        starting_agent_coin (int, float): Amount of coin agents have at t=0. Defaults
+            to zero coin.
+        isoelastic_eta (float): Parameter controlling the shape of agent utility
+            wrt coin endowment.
+        energy_cost (float): Coefficient for converting labor to negative utility.
+        energy_warmup_constant (float): Decay constant that controls the rate at which
+            the effective energy cost is annealed from 0 to energy_cost. Set to 0
+            (default) to disable annealing, meaning that the effective energy cost is
+            always energy_cost. The units of the decay constant depend on the choice of
+            energy_warmup_method.
+        energy_warmup_method (str): How to schedule energy annealing (warmup). If
+            "decay" (default), use the number of completed episodes. If "auto",
+            use the number of timesteps where the average agent reward was positive.
+        planner_reward_type (str): The type of reward used for the planner. Options
+            are "coin_eq_times_productivity" (default),
+            "inv_income_weighted_coin_endowment", and "inv_income_weighted_utility".
+        mixing_weight_gini_vs_coin (float): Degree to which equality is ignored w/
+            "coin_eq_times_productivity". Default is 0, which weights equality and
+            productivity equally. If set to 1, only productivity is rewarded.
+    """
+
+    name = "split_layout/simple_wood_and_stone"
+
+    def __init__(
+        self,
+        *args,
+        water_row=None,
+        skill_rank_of_top_agents=None,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        if self.fixed_four_skill_and_loc:
+            raise ValueError(
+                "The split layout scenario does not support "
+                "fixed_four_skill_and_loc. Set this to False."
+            )
+
+        # Augment the fixed layout to include a row of water through the middle
+        if water_row is None:
+            self._water_line = self.world_size[0] // 2
+        else:
+            self._water_line = int(water_row)
+            assert 0 < self._water_line < self.world_size[0] - 1
+        for landmark, landmark_map in self._source_maps.items():
+            landmark_map[self._water_line, :] = 1 if landmark == "Water" else 0
+            self._source_maps[landmark] = landmark_map
+
+        # Controls logic for which agents (by skill rank) get placed on the top
+        if skill_rank_of_top_agents is None:
+            skill_rank_of_top_agents = [0]
+
+        if isinstance(skill_rank_of_top_agents, (int, float)):
+            self.skill_rank_of_top_agents = [int(skill_rank_of_top_agents)]
+        elif isinstance(skill_rank_of_top_agents, (tuple, list)):
+            self.skill_rank_of_top_agents = list(set(skill_rank_of_top_agents))
+        else:
+            raise TypeError(
+                "skill_rank_of_top_agents must be a scalar "
+                "index, or a list of scalar indices."
+            )
+        for rank in self.skill_rank_of_top_agents:
+            assert 0 <= rank < self.n_agents
+        assert 0 < len(self.skill_rank_of_top_agents) < self.n_agents
+
+        # Set the skill associated with each skill rank
+        bm = self.get_component("Build")
+        assert bm.skill_dist == "pareto"
+        pmsm = bm.payment_max_skill_multiplier
+        # Generate a batch (100000) of num_agents (sorted/clipped) Pareto samples.
+        pareto_samples = np.random.pareto(4, size=(100000, self.n_agents))
+        clipped_skills = np.minimum(pmsm, (pmsm - 1) * pareto_samples + 1)
+        sorted_clipped_skills = np.sort(clipped_skills, axis=1)
+        # The skill level of the i-th skill-ranked agent is the average of the
+        # i-th ranked samples throughout the batch.
+        average_ranked_skills = sorted_clipped_skills.mean(axis=0)
+        self._avg_ranked_skill = average_ranked_skills * bm.payment
+        # Reverse the order so index 0 is the highest-skilled
+        self._avg_ranked_skill = self._avg_ranked_skill[::-1]
+
+    def additional_reset_steps(self):
+        """
+        Extra scenario-specific steps that should be performed at the end of the reset
+        cycle.
+
+        For each reset cycle...
+            First, reset_starting_layout() and reset_agent_states() will be called.
+
+            Second, <component>.reset() will be called for each registered component.
+
+            Lastly, this method will be called to allow for any final customization of
+            the reset cycle.
+
+        For this scenario, this method resets optimization metric trackers. This is
+        where each agent gets assigned to one of the skills and the starting
+        locations are reset according to self.skill_rank_of_top_agents.
+        """
+        self.world.clear_agent_locs()
+        for i, agent in enumerate(self.world.get_random_order_agents()):
+            agent.state["build_payment"] = self._avg_ranked_skill[i]
+            if i in self.skill_rank_of_top_agents:
+                r_min, r_max = 0, self._water_line
+            else:
+                r_min, r_max = self._water_line + 1, self.world_size[0]
+
+            r = np.random.randint(r_min, r_max)
+            c = np.random.randint(0, self.world_size[1])
+            n_tries = 0
+            while not self.world.can_agent_occupy(r, c, agent):
+                r = np.random.randint(r_min, r_max)
+                c = np.random.randint(0, self.world_size[1])
+                n_tries += 1
+                if n_tries > 200:
+                    raise TimeoutError
+            self.world.set_agent_loc(agent, r, c)
+
+        # compute current objectives
+        curr_optimization_metric = self.get_current_optimization_metrics()
+
+        self.curr_optimization_metric = deepcopy(curr_optimization_metric)
+        self.init_optimization_metric = deepcopy(curr_optimization_metric)
+        self.prev_optimization_metric = deepcopy(curr_optimization_metric)

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/closed_quadrant_25x25_20each_30clump.txt

@@ -0,0 +1 @@
+WWWWW  W    @ W          ;WW W        @ W      W WW; W          @         W W;SW S S      @            ;            @   W      W ; SS         @            ; S          @            ;            @            ;            @            ;S      S    @            ;            @            ;            @            ;@@@@@@@@@@@@@@@@@@@@@@@@@; S          @            ;S S         @            ;   S        @            ;SS          @            ; SS         @            ;SS          @            ;  S         @            ;            @            ;            @            ;            @            ;            @            ;S           @            ;

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/env-pure_and_mixed-15x15.txt

@@ -0,0 +1 @@
+WWW    @       ;WSS    @       ;WWW    @       ;WWW    @       ;WSS            ;SWS            ;       @       ;@@@@  @@@@  @@@;       @       ;WWW    @    S  ; WW         SS ;WWW          SS;W W    @       ;W      @       ;       @     S ;

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/env-pure_and_mixed-25x25.txt

@@ -0,0 +1 @@
+ WW W       @            ; SWWW       @            ;SSWW        @            ;WSSSW       @            ;WSSWW                    ;WS WS                    ;  WWS                    ;SWW S       @            ;  S W       @            ; WS W       @            ;            @            ;            @            ;@@@@   @@@@@@@@@@@@   @@@;            @            ;            @            ;  W         @       SSSSS; WW         @       SSS  ;            @        S SS; WW W       @       SSS S;W WW                SSSS ;WWW                 S  SS; WWWW                S S ;WW  W       @       S    ;    W       @          S ;   W        @       S SSS;

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/env-pure_and_mixed-40x40.txt

@@ -0,0 +1 @@
+   W  WW                                ;         W                              ; W  W     W                             ;      W                                 ;   W        W                           ;  W    W W                              ;       W W                              ;      WW                                ;                                        ;                                        ;                                        ;                                        ;                                        ;                                        ;                                        ;                                        ;                          SS       W    ;                                WW     W;                                        ;                             S      S  S;                               W  WW S  ;                              S  S  W   ;                           W   WS   S S ;                                        ;                                        ;                                        ;                                        ;                                        ;                                        ;                                        ;                                        ;   S                                    ;S   S                                   ;                                        ;   SS                                   ;                                        ;     S  SS  S                           ;   SS S  S                              ;S SS    S   S                           ;SSSS S  SS                              ;

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/quadrant_25x25_20each_30clump.txt

@@ -0,0 +1 @@
+WWWWW  W    @ W          ;WW W        @ W      W WW; W          @         W W;SW S S      @            ;            @   W      W ; SS                      ; S                       ;                         ;                         ;S      S    @            ;            @            ;            @            ;@@@@@    @@@@@@@    @@@@@; S          @            ;S S         @            ;   S        @            ;SS                       ; SS                      ;SS                       ;  S                      ;            @            ;            @            ;            @            ;            @            ;S           @            ;

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/quadrant_25x25_20each_30clump_no_water.txt

@@ -0,0 +1 @@
+WWWWW  W      W          ;WW W          W      W WW; W                    W W;SW S S                   ;                W      W ; SS                      ; S                       ;                         ;                         ;S      S                 ;                         ;                         ;                         ; S                       ;S S                      ;   S                     ;SS                       ; SS                      ;SS                       ;  S                      ;                         ;                         ;                         ;                         ;S                        ;

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/quadrant_40x40_50each.txt

@@ -0,0 +1 @@
+WWWWWWWW   WW      @@ WW                ;WWWWWWWW   WW      @@ WW                ;WWW  W             @@ WW          W  WWW;  W                @@              WW WW;  W                @@              WW WW;SSW  S  SS         @@                   ;                   @@     W          W  ;                   @@     W          W  ;  SSS                                   ;  SSS                                   ;  S                                     ;                                        ;                                        ;                                        ;SS         SS      @@                   ;SS         SS      @@                   ;                   @@                   ;                   @@                   ;                   @@                   ;@@@@@@@@      @@@@@@@@@@@@      @@@@@@@@;@@@@@@@@      @@@@@@@@@@@@      @@@@@@@@;  S                @@                   ;SS SS              @@                   ;SS SS              @@                   ;     S             @@                   ;     S             @@                   ;SSS                                     ;  SSS                                   ;  SSS                                   ;SSS                                     ;   SS                                   ;   SS                                   ;                   @@                   ;                   @@                   ;                   @@                   ;                   @@                   ;                   @@                   ;                   @@                   ;SS                 @@                   ;SS                 @@                   ;

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/quadrant_40x40_50each_no_water.txt

@@ -0,0 +1 @@
+WWWWWWWW   WW         WW                ;WWWWWWWW   WW         WW                ;WWW  W                WW          W  WWW;  W                                WW WW;  W                                WW WW;SSW  S  SS                              ;                          W          W  ;                          W          W  ;  SSS                                   ;  SSS                                   ;  S                                     ;                                        ;                                        ;                                        ;SS         SS                           ;SS         SS                           ;                                        ;                                        ;                                        ;                                        ;                                        ;  S                                     ;SS SS                                   ;SS SS                                   ;     S                                  ;     S                                  ;SSS                                     ;  SSS                                   ;  SSS                                   ;SSS                                     ;   SS                                   ;   SS                                   ;                                        ;                                        ;                                        ;                                        ;                                        ;                                        ;SS                                      ;SS                                      ;

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/quadrant_8x8_4each_8clump.txt

@@ -0,0 +1 @@
+WWWW@WW W; WW @ WWW;SW S@ S  ;S   @    ;@@ @@@ @@; S  @    ;S S @    ;S   @    ;   S@    ;

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/top_wood_bottom_stone_14x14.txt

@@ -0,0 +1 @@
+  WW       W  ;   A  WW    A ;WW WA  W      ;       A  A   ;              ;  AA    A     ;         A    ;              ;              ;              ;  A           ;  A    SS    S;SA  S       AS;S     SA  SA  

ai_economist/foundation/scenarios/simple_wood_and_stone/map_txt/uniform_25x25_25each_65clump.txt

@@ -0,0 +1 @@
+                    SSSSS;                SS   SSSS;                  SS SSSS;                   S SSSS;                      SS ;                         ;                         ;                         ;                         ;                         ;                         ;                         ;                         ;                         ;                         ;                         ;                         ;                         ;                         ;                         ;                        W;                 W     WW;                      WW ;W     WW W          W WWW;WWWWW WW       W    WWWWW;

ai_economist/foundation/scenarios/utils/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2020, 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

ai_economist/foundation/scenarios/utils/rewards.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2020, 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 numpy as np
+
+from ai_economist.foundation.scenarios.utils import social_metrics
+
+
+def isoelastic_coin_minus_labor(
+    coin_endowment, total_labor, isoelastic_eta, labor_coefficient
+):
+    """Agent utility, concave increasing in coin and linearly decreasing in labor.
+
+    Args:
+        coin_endowment (float, ndarray): The amount of coin owned by the agent(s).
+        total_labor (float, ndarray): The amount of labor performed by the agent(s).
+        isoelastic_eta (float): Constant describing the shape of the utility profile
+            with respect to coin endowment. Must be between 0 and 1. 0 yields utility
+            that increases linearly with coin. 1 yields utility that increases with
+            log(coin). Utility from coin uses:
+                https://en.wikipedia.org/wiki/Isoelastic_utility
+        labor_coefficient (float): Constant describing the disutility experienced per
+            unit of labor performed. Disutility from labor equals:
+                labor_coefficient * total_labor
+
+    Returns:
+        Agent utility (float) or utilities (ndarray).
+    """
+    # https://en.wikipedia.org/wiki/Isoelastic_utility
+    assert np.all(coin_endowment >= 0)
+    assert 0 <= isoelastic_eta <= 1.0
+
+    # Utility from coin endowment
+    if isoelastic_eta == 1.0:  # dangerous
+        util_c = np.log(np.max(1, coin_endowment))
+    else:  # isoelastic_eta >= 0
+        util_c = (coin_endowment ** (1 - isoelastic_eta) - 1) / (1 - isoelastic_eta)
+
+    # disutility from labor
+    util_l = total_labor * labor_coefficient
+
+    # Net utility
+    util = util_c - util_l
+
+    return util
+
+
+def coin_minus_labor_cost(
+    coin_endowment, total_labor, labor_exponent, labor_coefficient
+):
+    """Agent utility, linearly increasing in coin and decreasing as a power of labor.
+
+    Args:
+        coin_endowment (float, ndarray): The amount of coin owned by the agent(s).
+        total_labor (float, ndarray): The amount of labor performed by the agent(s).
+        labor_exponent (float): Constant describing the shape of the utility profile
+            with respect to total labor. Must be between >1.
+        labor_coefficient (float): Constant describing the disutility experienced per
+            unit of labor performed. Disutility from labor equals:
+                labor_coefficient * total_labor.
+
+    Returns:
+        Agent utility (float) or utilities (ndarray).
+    """
+    # https://en.wikipedia.org/wiki/Isoelastic_utility
+    assert np.all(coin_endowment >= 0)
+    assert labor_exponent > 1
+
+    # Utility from coin endowment
+    util_c = coin_endowment
+
+    # Disutility from labor
+    util_l = (total_labor ** labor_exponent) * labor_coefficient
+
+    # Net utility
+    util = util_c - util_l
+
+    return util
+
+
+def coin_eq_times_productivity(coin_endowments, equality_weight):
+    """Social welfare, measured as productivity scaled by the degree of coin equality.
+
+    Args:
+        coin_endowments (ndarray): The array of coin endowments for each of the
+            agents in the simulated economy.
+        equality_weight (float): Constant that determines how productivity is scaled
+            by coin equality. Must be between 0 (SW = prod) and 1 (SW = prod * eq).
+
+    Returns:
+        Product of coin equality and productivity (float).
+    """
+    n_agents = len(coin_endowments)
+    prod = social_metrics.get_productivity(coin_endowments) / n_agents
+    equality = equality_weight * social_metrics.get_equality(coin_endowments) + (
+        1 - equality_weight
+    )
+    return equality * prod
+
+
+def inv_income_weighted_coin_endowments(coin_endowments):
+    """Social welfare, as weighted average endowment (weighted by inverse endowment).
+
+    Args:
+        coin_endowments (ndarray): The array of coin endowments for each of the
+            agents in the simulated economy.
+
+    Returns:
+        Weighted average coin endowment (float).
+    """
+    pareto_weights = 1 / np.maximum(coin_endowments, 1)
+    pareto_weights = pareto_weights / np.sum(pareto_weights)
+    return np.sum(coin_endowments * pareto_weights)
+
+
+def inv_income_weighted_utility(coin_endowments, utilities):
+    """Social welfare, as weighted average utility (weighted by inverse endowment).
+
+    Args:
+        coin_endowments (ndarray): The array of coin endowments for each of the
+            agents in the simulated economy.
+        utilities (ndarray): The array of utilities for each of the agents in the
+            simulated economy.
+
+    Returns:
+        Weighted average utility (float).
+    """
+    pareto_weights = 1 / np.maximum(coin_endowments, 1)
+    pareto_weights = pareto_weights / np.sum(pareto_weights)
+    return np.sum(utilities * pareto_weights)

ai_economist/foundation/scenarios/utils/social_metrics.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2020, 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 numpy as np
+
+
+def get_gini(endowments):
+    """Returns the normalized Gini index describing the distribution of endowments.
+
+    https://en.wikipedia.org/wiki/Gini_coefficient
+
+    Args:
+        endowments (ndarray): The array of endowments for each of the agents in the
+            simulated economy.
+
+    Returns:
+        Normalized Gini index for the distribution of endowments (float). A value of 1
+            indicates everything belongs to 1 agent (perfect inequality), whereas a
+            value of 0 indicates all agents have equal endowments (perfect equality).
+
+    Note:
+        Uses a slightly different method depending on the number of agents. For fewer
+        agents (<30), uses an exact but slow method. Switches to using a much faster
+        method for more agents, where both methods produce approximately equivalent
+        results.
+    """
+    n_agents = len(endowments)
+
+    if n_agents < 30:  # Slower. Accurate for all n.
+        diff_ij = np.abs(
+            endowments.reshape((n_agents, 1)) - endowments.reshape((1, n_agents))
+        )
+        diff = np.sum(diff_ij)
+        norm = 2 * n_agents * endowments.sum(axis=0)
+        unscaled_gini = diff / (norm + 1e-10)
+        gini = unscaled_gini / ((n_agents - 1) / n_agents)
+        return gini
+
+    # Much faster. Slightly overestimated for low n.
+    s_endows = np.sort(endowments)
+    return 1 - (2 / (n_agents + 1)) * np.sum(
+        np.cumsum(s_endows) / (np.sum(s_endows) + 1e-10)
+    )
+
+
+def get_equality(endowments):
+    """Returns the complement of the normalized Gini index (equality = 1 - Gini).
+
+    Args:
+        endowments (ndarray): The array of endowments for each of the agents in the
+            simulated economy.
+
+    Returns:
+        Normalized equality index for the distribution of endowments (float). A value
+            of 0 indicates everything belongs to 1 agent (perfect inequality),
+            whereas a value of 1 indicates all agents have equal endowments (perfect
+            equality).
+    """
+    return 1 - get_gini(endowments)
+
+
+def get_productivity(coin_endowments):
+    """Returns the total coin inside the simulated economy.
+
+    Args:
+        coin_endowments (ndarray): The array of coin endowments for each of the
+            agents in the simulated economy.
+
+    Returns:
+        Total coin endowment (float).
+    """
+    return np.sum(coin_endowments)