adding ai_economist for modding

This commit is contained in:
2023-01-12 16:41:38 +01:00
parent 0479a4f6a4
commit f177f8f0ba
85 changed files with 19373 additions and 2 deletions
@@ -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
@@ -0,0 +1,638 @@
# 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 itertools
import numpy as np
import torch
_NP_DTYPE = np.float32
def all_agents_export_experiment_template(
NUMFIRMS, NUMCONSUMERS, NUMGOVERNMENTS, episodes_const=30000
):
consumption_choices = [
np.array([0.0 + 1.0 * c for c in range(11)], dtype=_NP_DTYPE)
]
work_choices = [
np.array([0.0 + 20 * 13 * h for h in range(5)], dtype=_NP_DTYPE)
] # specify dtype -- be consistent?
consumption_choices = np.array(
list(itertools.product(*consumption_choices)), dtype=_NP_DTYPE
)
work_choices = np.array(list(itertools.product(*work_choices)), dtype=_NP_DTYPE)
price_choices = np.array([0.0 + 500.0 * c for c in range(6)], dtype=_NP_DTYPE)
wage_choices = np.array([0.0, 11.0, 22.0, 33.0, 44.0], dtype=_NP_DTYPE)
capital_choices = np.array([0.1], dtype=_NP_DTYPE)
price_and_wage = np.array(
list(itertools.product(price_choices, wage_choices, capital_choices)),
dtype=_NP_DTYPE,
)
# government action discretization
income_taxation_choices = np.array(
[0.0 + 0.2 * c for c in range(6)], dtype=_NP_DTYPE
)
corporate_taxation_choices = np.array(
[0.0 + 0.2 * c for c in range(6)], dtype=_NP_DTYPE
)
tax_choices = np.array(
list(itertools.product(income_taxation_choices, corporate_taxation_choices)),
dtype=_NP_DTYPE,
)
global_state_dim = (
NUMFIRMS # prices
+ NUMFIRMS # wages
+ NUMFIRMS # stocks
+ NUMFIRMS # was good overdemanded
+ 2 * NUMGOVERNMENTS # tax rates
+ 1
) # time
global_state_digit_dims = list(
range(2 * NUMFIRMS, 3 * NUMFIRMS)
) # stocks are the only global state var that can get huge
consumer_state_dim = (
global_state_dim + 1 + 1
) # budget # theta, the disutility of work
firm_state_dim = (
global_state_dim
+ 1 # budget
+ 1 # capital
+ 1 # production alpha
+ NUMFIRMS # onehot specifying which firm
)
episodes_to_anneal_firm = 100000
episodes_to_anneal_government = 100000
government_phase1_start = 100000
government_state_dim = global_state_dim
DEFAULT_CFG_DICT = {
# actions_array key will be added below
"agents": {
"num_consumers": NUMCONSUMERS,
"num_firms": NUMFIRMS,
"num_governments": NUMGOVERNMENTS,
"global_state_dim": global_state_dim,
"consumer_state_dim": consumer_state_dim,
# action vectors are how much consume from each firm,
# how much to work, and which firm to choose
"consumer_action_dim": NUMFIRMS + 1 + 1,
"consumer_num_consume_actions": consumption_choices.shape[0],
"consumer_num_work_actions": work_choices.shape[0],
"consumer_num_whichfirm_actions": NUMFIRMS,
"firm_state_dim": firm_state_dim, # what are observations?
# actions are price and wage for own firm, and capital choices
"firm_action_dim": 3,
"firm_num_actions": price_and_wage.shape[0],
"government_state_dim": government_state_dim,
"government_action_dim": 2,
"government_num_actions": tax_choices.shape[0],
"max_possible_consumption": float(consumption_choices.max()),
"max_possible_hours_worked": float(work_choices.max()),
"max_possible_wage": float(wage_choices.max()),
"max_possible_price": float(price_choices.max()),
# these are dims which, due to being on a large scale,
# have to be expanded to a digit representation
"consumer_digit_dims": global_state_digit_dims
+ [global_state_dim], # global state + consumer budget
# global state + firm budget (do we need capital?)
"firm_digit_dims": global_state_digit_dims + [global_state_dim],
# govt only has global state
"government_digit_dims": global_state_digit_dims,
"firm_reward_scale": 10000,
"government_reward_scale": 100000,
"consumer_reward_scale": 50.0,
"firm_anneal_wages": {
"anneal_on": True,
"start": 22.0,
"increase_const": float(wage_choices.max() - 22.0)
/ (episodes_to_anneal_firm),
"decrease_const": (22.0) / episodes_to_anneal_firm,
},
"firm_anneal_prices": {
"anneal_on": True,
"start": 1000.0,
"increase_const": float(price_choices.max() - 1000.00)
/ episodes_to_anneal_firm,
"decrease_const": (1000.0) / episodes_to_anneal_firm,
},
"government_anneal_taxes": {
"anneal_on": True,
"start": 0.0,
"increase_const": 1.0 / episodes_to_anneal_government,
},
"firm_begin_anneal_action": 0,
"government_begin_anneal_action": government_phase1_start,
"consumer_anneal_theta": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
},
"consumer_anneal_entropy": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
"coef_floor": 0.1,
},
"firm_anneal_entropy": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
"coef_floor": 0.1,
},
"govt_anneal_entropy": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
"coef_floor": 0.1,
},
"consumer_noponzi_eta": 0.0,
"consumer_penalty_scale": 1.0,
"firm_noponzi_eta": 0.0,
"firm_training_start": episodes_to_anneal_firm,
"government_training_start": government_phase1_start
+ episodes_to_anneal_government,
"consumer_training_start": 0,
"government_counts_firm_reward": 0,
"should_boost_firm_reward": False,
"firm_reward_for_government_factor": 0.0025,
},
"world": {
"maxtime": 10,
"initial_firm_endowment": 22.0 * 1000 * NUMCONSUMERS,
"initial_consumer_endowment": 2000,
"initial_stocks": 0.0,
"initial_prices": 1000.0,
"initial_wages": 22.0,
"interest_rate": 0.1,
"consumer_theta": 0.01,
"crra_param": 0.1,
"production_alpha": "fixed_array", # only works for exactly 10 firms, kluge
"initial_capital": "twolevel",
"paretoscaletheta": 4.0,
"importer_price": 500.0,
"importer_quantity": 100.0,
"use_importer": 1,
},
"train": {
"batch_size": 8,
"base_seed": 1234,
"save_dense_every": 2000,
"save_model_every": 10000,
"num_episodes": 500000,
"infinite_episodes": False,
"lr": 0.01,
"gamma": 0.9999,
"entropy": 0.0,
"value_loss_weight": 1.0,
"digit_representation_size": 10,
"lagr_num_steps": 1,
"boost_firm_reward_factor": 1.0,
},
}
return (
DEFAULT_CFG_DICT,
consumption_choices,
work_choices,
price_and_wage,
tax_choices,
None,
None,
)
def all_agents_short_export_experiment_template(
NUMFIRMS, NUMCONSUMERS, NUMGOVERNMENTS, episodes_const=10000
):
consumption_choices = [
np.array([0.0 + 1.0 * c for c in range(11)], dtype=_NP_DTYPE)
]
work_choices = [
np.array([0.0 + 20 * 13 * h for h in range(5)], dtype=_NP_DTYPE)
] # specify dtype -- be consistent?
consumption_choices = np.array(
list(itertools.product(*consumption_choices)), dtype=_NP_DTYPE
)
work_choices = np.array(list(itertools.product(*work_choices)), dtype=_NP_DTYPE)
price_choices = np.array([0.0 + 500.0 * c for c in range(6)], dtype=_NP_DTYPE)
wage_choices = np.array([0.0, 11.0, 22.0, 33.0, 44.0], dtype=_NP_DTYPE)
capital_choices = np.array([0.1], dtype=_NP_DTYPE)
price_and_wage = np.array(
list(itertools.product(price_choices, wage_choices, capital_choices)),
dtype=_NP_DTYPE,
)
# government action discretization
income_taxation_choices = np.array(
[0.0 + 0.2 * c for c in range(6)], dtype=_NP_DTYPE
)
corporate_taxation_choices = np.array(
[0.0 + 0.2 * c for c in range(6)], dtype=_NP_DTYPE
)
tax_choices = np.array(
list(itertools.product(income_taxation_choices, corporate_taxation_choices)),
dtype=_NP_DTYPE,
)
global_state_dim = (
NUMFIRMS # prices
+ NUMFIRMS # wages
+ NUMFIRMS # stocks
+ NUMFIRMS # was good overdemanded
+ 2 * NUMGOVERNMENTS # tax rates
+ 1
) # time
global_state_digit_dims = list(
range(2 * NUMFIRMS, 3 * NUMFIRMS)
) # stocks are the only global state var that can get huge
consumer_state_dim = (
global_state_dim + 1 + 1
) # budget # theta, the disutility of work
firm_state_dim = (
global_state_dim
+ 1 # budget
+ 1 # capital
+ 1 # production alpha
+ NUMFIRMS # onehot specifying which firm
)
episodes_to_anneal_firm = 30000
episodes_to_anneal_government = 30000
government_phase1_start = 30000
government_state_dim = global_state_dim
DEFAULT_CFG_DICT = {
# actions_array key will be added below
"agents": {
"num_consumers": NUMCONSUMERS,
"num_firms": NUMFIRMS,
"num_governments": NUMGOVERNMENTS,
"global_state_dim": global_state_dim,
"consumer_state_dim": consumer_state_dim,
# action vectors are how much consume from each firm,
# how much to work, and which firm to choose
"consumer_action_dim": NUMFIRMS + 1 + 1,
"consumer_num_consume_actions": consumption_choices.shape[0],
"consumer_num_work_actions": work_choices.shape[0],
"consumer_num_whichfirm_actions": NUMFIRMS,
"firm_state_dim": firm_state_dim, # what are observations?
# actions are price and wage for own firm, and capital choices
"firm_action_dim": 3,
"firm_num_actions": price_and_wage.shape[0],
"government_state_dim": government_state_dim,
"government_action_dim": 2,
"government_num_actions": tax_choices.shape[0],
"max_possible_consumption": float(consumption_choices.max()),
"max_possible_hours_worked": float(work_choices.max()),
"max_possible_wage": float(wage_choices.max()),
"max_possible_price": float(price_choices.max()),
# these are dims which, due to being on a large scale,
# have to be expanded to a digit representation
"consumer_digit_dims": global_state_digit_dims
+ [global_state_dim], # global state + consumer budget
"firm_digit_dims": global_state_digit_dims
+ [global_state_dim], # global state + firm budget (do we need capital?)
# govt only has global state
"government_digit_dims": global_state_digit_dims,
"firm_reward_scale": 10000,
"government_reward_scale": 100000,
"consumer_reward_scale": 50.0,
"firm_anneal_wages": {
"anneal_on": True,
"start": 22.0,
"increase_const": float(wage_choices.max() - 22.0)
/ (episodes_to_anneal_firm),
"decrease_const": (22.0) / episodes_to_anneal_firm,
},
"firm_anneal_prices": {
"anneal_on": True,
"start": 1000.0,
"increase_const": float(price_choices.max() - 1000.00)
/ episodes_to_anneal_firm,
"decrease_const": (1000.0) / episodes_to_anneal_firm,
},
"government_anneal_taxes": {
"anneal_on": True,
"start": 0.0,
"increase_const": 1.0 / episodes_to_anneal_government,
},
"firm_begin_anneal_action": 0,
"government_begin_anneal_action": government_phase1_start,
"consumer_anneal_theta": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
},
"consumer_anneal_entropy": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
"coef_floor": 0.1,
},
"firm_anneal_entropy": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
"coef_floor": 0.1,
},
"govt_anneal_entropy": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
"coef_floor": 0.1,
},
"consumer_noponzi_eta": 0.0,
"consumer_penalty_scale": 1.0,
"firm_noponzi_eta": 0.0,
"firm_training_start": episodes_to_anneal_firm,
"government_training_start": government_phase1_start
+ episodes_to_anneal_government,
"consumer_training_start": 0,
"government_counts_firm_reward": 0,
"should_boost_firm_reward": False,
"firm_reward_for_government_factor": 0.0025,
},
"world": {
"maxtime": 10,
"initial_firm_endowment": 22.0 * 1000 * NUMCONSUMERS,
"initial_consumer_endowment": 2000,
"initial_stocks": 0.0,
"initial_prices": 1000.0,
"initial_wages": 22.0,
"interest_rate": 0.1,
"consumer_theta": 0.01,
"crra_param": 0.1,
"production_alpha": "fixed_array", # only works for exactly 10 firms, kluge
"initial_capital": "twolevel",
"paretoscaletheta": 4.0,
"importer_price": 500.0,
"importer_quantity": 100.0,
"use_importer": 1,
},
"train": {
"batch_size": 8,
"base_seed": 1234,
"save_dense_every": 2000,
"save_model_every": 10000,
"num_episodes": 200000,
"infinite_episodes": False,
"lr": 0.01,
"gamma": 0.9999,
"entropy": 0.0,
"value_loss_weight": 1.0,
"digit_representation_size": 10,
"lagr_num_steps": 1,
"boost_firm_reward_factor": 1.0,
},
}
return (
DEFAULT_CFG_DICT,
consumption_choices,
work_choices,
price_and_wage,
tax_choices,
None,
None,
)
def very_short_test_template(
NUMFIRMS, NUMCONSUMERS, NUMGOVERNMENTS, episodes_const=30000
):
consumption_choices = [
np.array([0.0 + 1.0 * c for c in range(11)], dtype=_NP_DTYPE)
]
work_choices = [
np.array([0.0 + 20 * 13 * h for h in range(5)], dtype=_NP_DTYPE)
] # specify dtype -- be consistent?
consumption_choices = np.array(
list(itertools.product(*consumption_choices)), dtype=_NP_DTYPE
)
work_choices = np.array(list(itertools.product(*work_choices)), dtype=_NP_DTYPE)
price_choices = np.array([0.0 + 500.0 * c for c in range(6)], dtype=_NP_DTYPE)
wage_choices = np.array([0.0, 11.0, 22.0, 33.0, 44.0], dtype=_NP_DTYPE)
capital_choices = np.array([0.1], dtype=_NP_DTYPE)
price_and_wage = np.array(
list(itertools.product(price_choices, wage_choices, capital_choices)),
dtype=_NP_DTYPE,
)
# government action discretization
income_taxation_choices = np.array(
[0.0 + 0.2 * c for c in range(6)], dtype=_NP_DTYPE
)
corporate_taxation_choices = np.array(
[0.0 + 0.2 * c for c in range(6)], dtype=_NP_DTYPE
)
tax_choices = np.array(
list(itertools.product(income_taxation_choices, corporate_taxation_choices)),
dtype=_NP_DTYPE,
)
global_state_dim = (
NUMFIRMS # prices
+ NUMFIRMS # wages
+ NUMFIRMS # stocks
+ NUMFIRMS # was good overdemanded
+ 2 * NUMGOVERNMENTS # tax rates
+ 1
) # time
global_state_digit_dims = list(
range(2 * NUMFIRMS, 3 * NUMFIRMS)
) # stocks are the only global state var that can get huge
consumer_state_dim = (
global_state_dim + 1 + 1
) # budget # theta, the disutility of work
firm_state_dim = (
global_state_dim
+ 1 # budget
+ 1 # capital
+ 1 # production alpha
+ NUMFIRMS # onehot specifying which firm
)
episodes_to_anneal_firm = 10
episodes_to_anneal_government = 10
government_phase1_start = 10
government_state_dim = global_state_dim
DEFAULT_CFG_DICT = {
# actions_array key will be added below
"agents": {
"num_consumers": NUMCONSUMERS,
"num_firms": NUMFIRMS,
"num_governments": NUMGOVERNMENTS,
"global_state_dim": global_state_dim,
"consumer_state_dim": consumer_state_dim,
# action vectors are how much consume from each firm,
# how much to work, and which firm to choose
"consumer_action_dim": NUMFIRMS + 1 + 1,
"consumer_num_consume_actions": consumption_choices.shape[0],
"consumer_num_work_actions": work_choices.shape[0],
"consumer_num_whichfirm_actions": NUMFIRMS,
"firm_state_dim": firm_state_dim, # what are observations?
# actions are price and wage for own firm, and capital choices
"firm_action_dim": 3,
"firm_num_actions": price_and_wage.shape[0],
"government_state_dim": government_state_dim,
"government_action_dim": 2,
"government_num_actions": tax_choices.shape[0],
"max_possible_consumption": float(consumption_choices.max()),
"max_possible_hours_worked": float(work_choices.max()),
"max_possible_wage": float(wage_choices.max()),
"max_possible_price": float(price_choices.max()),
# these are dims which, due to being on a large scale,
# have to be expanded to a digit representation
"consumer_digit_dims": global_state_digit_dims
+ [global_state_dim], # global state + consumer budget
"firm_digit_dims": global_state_digit_dims
+ [global_state_dim], # global state + firm budget (do we need capital?)
# govt only has global state
"government_digit_dims": global_state_digit_dims,
"firm_reward_scale": 10000,
"government_reward_scale": 100000,
"consumer_reward_scale": 50.0,
"firm_anneal_wages": {
"anneal_on": True,
"start": 22.0,
"increase_const": float(wage_choices.max() - 22.0)
/ (episodes_to_anneal_firm),
"decrease_const": (22.0) / episodes_to_anneal_firm,
},
"firm_anneal_prices": {
"anneal_on": True,
"start": 1000.0,
"increase_const": float(price_choices.max() - 1000.00)
/ episodes_to_anneal_firm,
"decrease_const": (1000.0) / episodes_to_anneal_firm,
},
"government_anneal_taxes": {
"anneal_on": True,
"start": 0.0,
"increase_const": 1.0 / episodes_to_anneal_government,
},
"firm_begin_anneal_action": 0,
"government_begin_anneal_action": government_phase1_start,
"consumer_anneal_theta": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
},
"consumer_anneal_entropy": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
"coef_floor": 0.1,
},
"firm_anneal_entropy": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
"coef_floor": 0.1,
},
"govt_anneal_entropy": {
"anneal_on": True,
"exp_decay_length_in_steps": episodes_const,
"coef_floor": 0.1,
},
"consumer_noponzi_eta": 0.0,
"consumer_penalty_scale": 1.0,
"firm_noponzi_eta": 0.0,
"firm_training_start": episodes_to_anneal_firm,
"government_training_start": government_phase1_start
+ episodes_to_anneal_government,
"consumer_training_start": 0,
"government_counts_firm_reward": 0,
"should_boost_firm_reward": False,
"firm_reward_for_government_factor": 0.0025,
"train_firms_every": 2,
"train_consumers_every": 1,
"train_government_every": 5,
},
"world": {
"maxtime": 10,
"initial_firm_endowment": 22.0 * 1000 * NUMCONSUMERS,
"initial_consumer_endowment": 2000,
"initial_stocks": 0.0,
"initial_prices": 1000.0,
"initial_wages": 22.0,
"interest_rate": 0.1,
"consumer_theta": 0.01,
"crra_param": 0.1,
"production_alpha": "fixed_array", # only works for exactly 10 firms, kluge
"initial_capital": "twolevel",
"paretoscaletheta": 4.0,
"importer_price": 500.0,
"importer_quantity": 100.0,
"use_importer": 1,
},
"train": {
"batch_size": 8,
"base_seed": 1234,
"save_dense_every": 2000,
"save_model_every": 10000,
"num_episodes": 100,
"infinite_episodes": False,
"lr": 0.01,
"gamma": 0.9999,
"entropy": 0.0,
"value_loss_weight": 1.0,
"digit_representation_size": 10,
"lagr_num_steps": 1,
"boost_firm_reward_factor": 1.0,
},
}
return (
DEFAULT_CFG_DICT,
consumption_choices,
work_choices,
price_and_wage,
tax_choices,
None,
None,
)
def global_state_scaling_factors(cfg_dict):
max_wage = cfg_dict["agents"]["max_possible_wage"]
max_price = cfg_dict["agents"]["max_possible_price"]
num_firms = cfg_dict["agents"]["num_firms"]
num_governments = cfg_dict["agents"]["num_governments"]
maxtime = cfg_dict["world"]["maxtime"]
digit_size = cfg_dict["train"]["digit_representation_size"]
return torch.tensor(
# prices, wages, stocks, overdemanded
([max_price] * num_firms)
+ ([max_wage] * num_firms)
+ ([1.0] * num_firms * digit_size) # stocks are expanded to digit form
+ ([1.0] * num_firms)
+ ([1.0] * (2 * num_governments))
+ [maxtime]
)
def consumer_state_scaling_factors(cfg_dict):
global_state_scales = global_state_scaling_factors(cfg_dict)
digit_size = cfg_dict["train"]["digit_representation_size"]
consumer_scales = torch.tensor(
([1.0] * digit_size) + [cfg_dict["world"]["consumer_theta"]]
)
return torch.cat((global_state_scales, consumer_scales)).cuda()
def firm_state_scaling_factors(cfg_dict):
num_firms = cfg_dict["agents"]["num_firms"]
global_state_scales = global_state_scaling_factors(cfg_dict)
digit_size = cfg_dict["train"]["digit_representation_size"]
# budget, capital, alpha, one-hot
firm_scales = torch.tensor(
([1.0] * digit_size) + [10000.0, 1.0] + ([1.0] * num_firms)
)
return torch.cat((global_state_scales, firm_scales)).cuda()
def govt_state_scaling_factors(cfg_dict):
return global_state_scaling_factors(cfg_dict).cuda()
@@ -0,0 +1,912 @@
// 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
// Real Business Chain implementation in CUDA C
#include <curand_kernel.h>
#include <math.h>
typedef enum {
kConsumerType,
kFirmType,
kGovernmentType,
} AgentType;
const size_t kBatchSize = M_BATCHSIZE;
const size_t kNumConsumers = M_NUMCONSUMERS;
const bool kCountFirmReward = M_COUNTFIRMREWARD;
const size_t kNumFirms = M_NUMFIRMS;
const size_t kNumGovts = M_NUMGOVERNMENTS;
const float kMaxTime = M_MAXTIME;
const float kCrraParam = M_CRRA_PARAM;
const float kInterestRate = M_INTERESTRATE;
const size_t kNumAgents = kNumConsumers + kNumFirms + kNumGovts;
const bool kIncentivizeFirmActivity = M_SHOULDBOOSTFIRMREWARD;
const float kFirmBoostRewardFactor = M_BOOSTFIRMREWARDFACTOR;
const bool kUseImporter = M_USEIMPORTER;
const float kImporterPrice = M_IMPORTERPRICE;
const float kImporterQuantity = M_IMPORTERQUANTITY;
const float kLaborFloor = M_LABORFLOOR;
// Global state =
const size_t kNumPrices = kNumFirms; // - prices,
const size_t kNumWages = kNumFirms; // - wages,
const size_t kNumInventories = kNumFirms; // - stocks,
const size_t kNumOverdemandFlags = kNumFirms; // - good overdemanded flag,
const size_t kNumCorporateTaxes = kNumGovts; // - corporate tax rate
const size_t kNumIncomeTaxes = kNumGovts; // - income tax rate
const size_t kNumTimeDimensions = 1; // - time step
const size_t kGlobalStateSize = kNumPrices + kNumWages + kNumInventories + kNumOverdemandFlags + kNumCorporateTaxes + kNumIncomeTaxes + kNumTimeDimensions;
const size_t kIdxPricesOffset = 0;
const size_t kIdxWagesOffset = kNumPrices;
const size_t kIdxStockOffset = kIdxWagesOffset + kNumInventories;
const size_t kIdxOverdemandOffset = kIdxStockOffset + kNumOverdemandFlags;
const size_t kIdxIncomeTaxOffset = kGlobalStateSize - 3;
const size_t kIdxCorporateTaxOffset = kGlobalStateSize - 2;
const size_t kIdxTimeOffset = kGlobalStateSize - 1;
// Consumer actions: consume, work, choose which firm to work for
const size_t kActionSizeConsumer = kNumFirms + 1 + 1;
// add budget and theta
const size_t kStateSizeConsumer = kGlobalStateSize + 1 + 1;
const size_t kIdxConsumerBudgetOffset = 0;
// offset from agent-specific state part of array
const size_t kIdxConsumerThetaOffset = 1;
// UNUSED for consumer. Actions are floats.
// __constant__ float cs_index_to_action[num_actions_consumer *
// kActionSizeConsumer]; const size_t kActionSizeConsumer = kNumFirms +
// kNumFirms; // consume + work
/*const size_t num_actions_consumer =
NUMACTIONSkConsumerType; // depends on discretization*/
// Firm actions: set wage, set price, invest in capital
const size_t kActionSizeFirm = 3;
// Number of actions depends on discretization of continuous action space.
const size_t kNumActionsFirm = M_NUMACTIONSFIRM;
// budget, capital, production alpha, and one-hot firm ID
const size_t kStateSizeFirm = kGlobalStateSize + 1 + 1 + 1 + kNumFirms;
// offset from agent-specific state part of array
const size_t kIdxFirmBudgetOffset = 0;
const size_t kIdxFirmCapitalOffset = 1;
const size_t kIdxFirmAlphaOffset = 2;
const size_t kIdxFirmOnehotOffset = 3;
// Constant memory available from ALL threads.
// See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#constant
__constant__ float kFirmIndexToAction[kNumActionsFirm * kActionSizeFirm];
// Corporate + income tax rates
const size_t kGovtActionSize = 2;
const size_t kNumActionsGovernment = M_NUMACTIONSGOVERNMENT;
const size_t kGovtStateSize = kGlobalStateSize;
__constant__ float
kGovernmentIndexToAction[kNumActionsGovernment * kGovtActionSize];
// One RNG state for each thread. Each thread is assigned to an agent in an env.
__device__ curandState_t
*rng_state_arr[kBatchSize * kNumAgents]; // not sure best way to do this
// Offsets into action vectors
const size_t kIdxConsumerDemandedOffset = 0;
const size_t kIdxConsumerWorkedOffset = kNumFirms;
const size_t kIdxConsumerWhichFirmOffset = kNumFirms + 1;
// currently 1 govt
const size_t kIdxThisThreadGovtId = 0;
extern "C" {
// ------------------
// CUDA C Utilities
// ------------------
__device__ void CopyFloatArraySlice(float *start_point, int num_elems,
float *destination) {
for (int i = 0; i < num_elems; i++) {
destination[i] = start_point[i];
}
}
__device__ void CopyIntArraySlice(int *start_point, int num_elems,
float *destination) {
for (int i = 0; i < num_elems; i++) {
destination[i] = start_point[i];
}
}
// unfortunately, you can't do templates with extern "C" linkage required for
// CUDA, so we have to define different functions for each case.
__device__ int *GetPointerFromMultiIndexFor3DIntTensor(int *array,
const dim3 &sizes,
const dim3 &index) {
unsigned int flat_index =
index.z + index.y * (sizes.z) + index.x * (sizes.z * sizes.y);
return &(array[flat_index]);
}
__device__ float *GetPointerFromMultiIndexFor3DFloatTensor(float *array,
const dim3 &sizes,
const dim3 &index) {
unsigned int flat_index =
index.z + index.y * (sizes.z) + index.x * (sizes.z * sizes.y);
return &(array[flat_index]);
}
__device__ float *
GetPointerFromMultiIndexFor4DTensor(float *array, const size_t *sizes,
const size_t *multi_index) {
// don't use this for arrays that arne't exactly size 4!!!
unsigned int flat_index = multi_index[3] + multi_index[2] * sizes[3] +
multi_index[1] * sizes[3] * sizes[2] +
multi_index[0] * sizes[3] * sizes[2] * sizes[1];
return &(array[flat_index]);
}
__global__ void CudaInitKernel(int seed) {
// we want to reset random seeds for all firms and consumers
int tidx = threadIdx.x;
const int kThisThreadGlobalArrayIdx = blockIdx.x * kNumAgents + threadIdx.x;
if (tidx < kNumAgents) {
curandState_t *s = new curandState_t;
if (s != 0) {
curand_init(seed, kThisThreadGlobalArrayIdx, 0, s);
}
rng_state_arr[kThisThreadGlobalArrayIdx] = s;
}
}
__global__ void CudaFreeRand() {
int tidx = threadIdx.x;
const int kThisThreadGlobalArrayIdx = blockIdx.x * kNumAgents + threadIdx.x;
if (tidx < kNumAgents) {
curandState_t *s = rng_state_arr[kThisThreadGlobalArrayIdx];
delete s;
}
}
__device__ int SearchIndex(float *distr, float p, int l, int r) {
int mid;
int left = l;
int right = r;
while (left <= right) {
mid = left + (right - left) / 2;
if (distr[mid] == p) {
return mid;
} else if (distr[mid] < p) {
left = mid + 1;
} else {
right = mid - 1;
}
}
return left > r ? r : left;
}
// --------------------
// Simulation Utilities
// --------------------
__device__ AgentType GetAgentType(const int agent_id) {
if (agent_id < kNumConsumers) {
return kConsumerType;
} else if (agent_id < (kNumConsumers + kNumFirms)) {
return kFirmType;
} else {
return kGovernmentType;
}
}
__device__ float GetCRRAUtil(float consumption, float crra_param) {
return (powf(consumption + 1, 1.0 - crra_param) - 1.0) / (1.0 - crra_param);
}
__global__ void CudaResetEnv(float *cs_state_arr, float *fm_state_arr,
float *govt_state_arr, float *cs_state_ckpt_arr,
float *fm_state_ckpt_arr,
float *govt_state_ckpt_arr,
float theta_anneal_factor) {
/*
Resets the environment by writing the initial state (checkpoint) into the
state array for this agent (thread).
*/
const int kBlockId = blockIdx.x;
const int kWithinBlockAgentId = threadIdx.x;
if (kWithinBlockAgentId >= kNumAgents) {
return;
}
AgentType ThisThreadAgentType = GetAgentType(kWithinBlockAgentId);
float *state_arr;
float *ckpt_arr;
dim3 my_state_shape, my_state_idx;
size_t my_state_size;
if (ThisThreadAgentType == kConsumerType) {
// This thread/agent is a consumer.
my_state_size = kStateSizeConsumer;
my_state_shape = {kBatchSize, kNumConsumers, kStateSizeConsumer};
my_state_idx = {kBlockId, kWithinBlockAgentId, 0};
state_arr = cs_state_arr;
ckpt_arr = cs_state_ckpt_arr;
} else if (ThisThreadAgentType == kFirmType) {
// This thread/agent is a firm.
my_state_size = kStateSizeFirm;
my_state_shape = {kBatchSize, kNumFirms, kStateSizeFirm};
my_state_idx = {kBlockId,
(unsigned int)(kWithinBlockAgentId - kNumConsumers), 0};
state_arr = fm_state_arr;
ckpt_arr = fm_state_ckpt_arr;
} else {
// This thread/agent is government.
my_state_size = kGovtStateSize;
my_state_shape = {kBatchSize, kNumGovts, kGovtStateSize};
my_state_idx = {
kBlockId,
(unsigned int)(kWithinBlockAgentId - kNumConsumers - kNumFirms), 0};
state_arr = govt_state_arr;
ckpt_arr = govt_state_ckpt_arr;
}
float *my_state_arr = GetPointerFromMultiIndexFor3DFloatTensor(
state_arr, my_state_shape, my_state_idx);
float *my_ckpt_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
ckpt_arr, my_state_shape, my_state_idx);
CopyFloatArraySlice(my_ckpt_ptr, my_state_size, my_state_arr);
// anneal theta
if (ThisThreadAgentType == kConsumerType) {
my_state_arr[kGlobalStateSize + kIdxConsumerThetaOffset] *=
theta_anneal_factor;
}
}
__device__ void GetAction(float *action_arr, float *index_to_action_arr,
int index, int agent_idx, int agent_action_size) {
// it needs to be possible to call this for either agents or firms
// Note: each thread is an agent.
for (int i = 0; i < agent_action_size; i++) {
action_arr[agent_idx * agent_action_size + i] =
index_to_action_arr[index * agent_action_size + i];
}
}
__global__ void CudaSampleFirmAndGovernmentActions(
float *fm_distr, int *fm_action_indices_arr, float *fm_actions_arr,
float *govt_distr, int *govt_action_indices_arr, float *govt_actions_arr) {
// Samples actions for firms and governments. Consumer actions are sampled in
// Pytorch...
const int kWithinBlockAgentId = threadIdx.x;
// Unused threads should not do anything.
if (threadIdx.x >= kNumAgents) {
return;
}
AgentType ThisThreadAgentType = GetAgentType(kWithinBlockAgentId);
// Index into rand states array
int kThisThreadGlobalArrayIdx = blockIdx.x * kNumAgents + threadIdx.x;
curandState_t rng_state = *rng_state_arr[kThisThreadGlobalArrayIdx];
*rng_state_arr[kThisThreadGlobalArrayIdx] = rng_state;
// float cs_cum_dist[num_actions_consumer];
float fm_cum_dist[kNumActionsFirm];
float govt_cum_dist[kNumActionsGovernment];
float *my_cumul_dist;
float *my_dist;
int *my_indices;
float *my_actions;
float *index_to_action;
int this_thread_global_array_idx;
size_t my_num_actions;
int my_action_size;
// Consumers have multiple action heads, hence sampling is more complicated.
if (ThisThreadAgentType == kConsumerType) {
return;
} else if (ThisThreadAgentType == kFirmType) {
// on firm thread
my_cumul_dist = fm_cum_dist;
my_dist = fm_distr;
my_indices = fm_action_indices_arr;
my_actions = fm_actions_arr;
my_num_actions = kNumActionsFirm;
index_to_action = kFirmIndexToAction;
my_action_size = kActionSizeFirm;
this_thread_global_array_idx =
(blockIdx.x * kNumFirms) + (threadIdx.x - kNumConsumers);
} else {
my_cumul_dist = govt_cum_dist;
my_dist = govt_distr;
my_indices = govt_action_indices_arr;
my_actions = govt_actions_arr;
my_num_actions = kNumActionsGovernment;
index_to_action = kGovernmentIndexToAction;
my_action_size = kGovtActionSize;
this_thread_global_array_idx =
(blockIdx.x * kNumGovts) + (threadIdx.x - kNumConsumers - kNumFirms);
}
// Compute CDF
my_cumul_dist[0] = my_dist[this_thread_global_array_idx * my_num_actions];
for (int i = 1; i < my_num_actions; i++) {
my_cumul_dist[i] =
my_dist[this_thread_global_array_idx * my_num_actions + i] +
my_cumul_dist[i - 1];
}
// Given sampled action which is a float in [0, 1], find the corresponding
// discrete action.
float sampled_float = curand_uniform(&rng_state);
const int index =
SearchIndex(my_cumul_dist, sampled_float, 0, (int)(my_num_actions - 1));
my_indices[this_thread_global_array_idx] = index;
GetAction(my_actions, index_to_action, index, this_thread_global_array_idx,
my_action_size);
}
__device__ float GetFirmProduction(float technology, float capital, float hours,
float alpha) {
if (hours < kLaborFloor) {
hours = 0.0;
}
return technology * powf(capital, 1.0 - alpha) * powf(hours, alpha);
}
// --------------------
// Simulation Logic
// --------------------
__global__ void
CudaStep(float *cs_state_arr, float *cs_actions_arr, float *cs_rewards_arr,
float *cs_state_arr_batch, float *cs_rewards_arr_batch,
float *fm_state_arr, int *fm_action_indices_arr, float *fm_actions_arr,
float *fm_rewards_arr, float *fm_state_arr_batch,
int *fm_actions_arr_batch, float *fm_rewards_arr_batch,
float *govt_state_arr, int *govt_action_indices_arr,
float *govt_actions_arr, float *govt_rewards_arr,
float *govt_state_arr_batch, int *govt_actions_arr_batch,
float *govt_rewards_arr_batch,
float *consumer_aux_batch,
float *firm_aux_batch,
int iter) {
// This function should be called with 1 block per copy of the environment.
// Within a block, each agent should have a thread.
const int kWithinBlockAgentId = threadIdx.x;
// return if we're on an extra thread not corresponding to an agent
if (kWithinBlockAgentId >= kNumAgents) {
return;
}
// -------------------------------------
// Start of variables and pointers defs.
// -------------------------------------
// __shared__ variables are block-local: can be seen by each thread ** in the
// block **
__shared__ float gross_demand_arr[kNumFirms];
__shared__ int num_consumer_demand_arr[kNumFirms];
__shared__ float hours_worked_arr[kNumFirms];
__shared__ float total_actually_consumer_arr[kNumFirms];
__shared__ float bought_by_importer_arr[kNumFirms];
__shared__ float next_global_state_arr[kGlobalStateSize];
__shared__ float tax_revenue_arr[kNumGovts];
__shared__ float total_utility_arr[kNumGovts];
__shared__ bool need_to_ration_this_good_arr[kNumFirms]; // whether or not to
// ration good i
float net_demand_arr[kNumFirms]; // amount demanded after budget constraints
// by a consumer (ignore for non-consumers)
int num_iter = (int)kMaxTime;
AgentType ThisThreadAgentType = GetAgentType(kWithinBlockAgentId);
float this_agent_reward = 0.0;
// pointer to start of state vector
// state vector consists of global state, then
// agent-specific state global part is of same size for
// all agents, but needs to be sliced out of different
// arrays depending on agent type
float *my_global_state_ptr;
float *my_action_arr;
// pointer to start of state vector in batch history
float *batch_state_ptr;
// sizes and indices for strided array access
dim3 my_state_shape, my_state_idx, action_shape;
// shape for batched array of scalars (action ind and reward)
dim3 batch_single_shape, single_idx;
// pointer to action index
int *batch_action_value_ptr;
// pointer to batch reward
float *batch_reward_value_ptr;
// pointers to action index for current arrays
int *my_action_value_ptr;
// pointers to reward index for current arrays
float *my_reward_value_ptr;
float *my_aux_batch_ptr;
// -----------------------------------
// End of variables and pointers defs.
// -----------------------------------
if (ThisThreadAgentType == kConsumerType) {
// get current state
my_state_shape = {kBatchSize, kNumConsumers, kStateSizeConsumer};
my_state_idx = {blockIdx.x, threadIdx.x, 0};
my_global_state_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
cs_state_arr, my_state_shape, my_state_idx); // index)
// get current action
action_shape = {kBatchSize, kNumConsumers, kActionSizeConsumer};
my_action_arr = GetPointerFromMultiIndexFor3DFloatTensor(
cs_actions_arr, action_shape, my_state_idx);
// index into the episode history and save prev state into it
size_t my_batch_state_shape[] = {kBatchSize, num_iter, kNumConsumers,
kStateSizeConsumer};
size_t my_batch_state_idx[] = {blockIdx.x, iter, threadIdx.x, 0};
batch_state_ptr = GetPointerFromMultiIndexFor4DTensor(
cs_state_arr_batch, my_batch_state_shape, my_batch_state_idx);
CopyFloatArraySlice(my_global_state_ptr, kStateSizeConsumer,
batch_state_ptr);
size_t my_aux_batch_shape[] = {kBatchSize, num_iter, kNumConsumers, kNumFirms};
my_aux_batch_ptr = GetPointerFromMultiIndexFor4DTensor(
consumer_aux_batch, my_aux_batch_shape, my_batch_state_idx
);
// Extract pointers to rewards, batch and current
batch_single_shape = {kBatchSize, (unsigned int)num_iter, kNumConsumers};
single_idx = {blockIdx.x, (unsigned int)iter, threadIdx.x};
batch_reward_value_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
cs_rewards_arr_batch, batch_single_shape, single_idx);
my_reward_value_ptr =
&(cs_rewards_arr[blockIdx.x * kNumConsumers + threadIdx.x]);
}
if (ThisThreadAgentType == kFirmType) {
// get current state
size_t this_thread_firm_id = (threadIdx.x - kNumConsumers);
my_state_shape = {kBatchSize, kNumFirms, kStateSizeFirm};
my_state_idx = {blockIdx.x, (unsigned int)this_thread_firm_id, 0};
my_global_state_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
fm_state_arr, my_state_shape, my_state_idx);
// get current action
action_shape = {kBatchSize, kNumFirms, kActionSizeFirm};
my_action_arr = GetPointerFromMultiIndexFor3DFloatTensor(
fm_actions_arr, action_shape, my_state_idx);
// index into the episode history and save prev state into it
size_t my_batch_state_shape[] = {kBatchSize, num_iter, kNumFirms,
kStateSizeFirm};
size_t my_batch_state_idx[] = {blockIdx.x, iter, this_thread_firm_id, 0};
batch_state_ptr = GetPointerFromMultiIndexFor4DTensor(
fm_state_arr_batch, my_batch_state_shape, my_batch_state_idx);
CopyFloatArraySlice(my_global_state_ptr, kStateSizeFirm, batch_state_ptr);
dim3 my_aux_batch_shape = {kBatchSize, num_iter, kNumFirms};
dim3 aux_batch_idx = {blockIdx.x, iter, this_thread_firm_id};
my_aux_batch_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
firm_aux_batch, my_aux_batch_shape, aux_batch_idx
);
// extract pointers to action indices and rewards, batch and current
batch_single_shape = {kBatchSize, (unsigned int)num_iter, kNumFirms};
single_idx = {blockIdx.x, (unsigned int)iter,
(unsigned int)this_thread_firm_id};
batch_action_value_ptr = GetPointerFromMultiIndexFor3DIntTensor(
fm_actions_arr_batch, batch_single_shape, single_idx);
batch_reward_value_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
fm_rewards_arr_batch, batch_single_shape, single_idx);
const int kThisThreadFirmIdx = blockIdx.x * kNumFirms + this_thread_firm_id;
my_action_value_ptr = &(fm_action_indices_arr[kThisThreadFirmIdx]);
my_reward_value_ptr = &(fm_rewards_arr[kThisThreadFirmIdx]);
}
if (ThisThreadAgentType == kGovernmentType) {
int this_thread_govt_id = (threadIdx.x - kNumConsumers - kNumFirms);
my_state_shape = {kBatchSize, kNumGovts, kGovtStateSize};
my_state_idx = {blockIdx.x, (unsigned int)this_thread_govt_id, 0};
my_global_state_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
govt_state_arr, my_state_shape, my_state_idx); // index)
// get current action
action_shape = {kBatchSize, kNumGovts, kGovtActionSize};
my_action_arr = GetPointerFromMultiIndexFor3DFloatTensor(
govt_actions_arr, action_shape, my_state_idx);
// index into the episode history and save prev state into it
size_t my_batch_state_shape[] = {kBatchSize, num_iter, kNumGovts,
kGovtStateSize};
size_t my_batch_state_idx[] = {blockIdx.x, iter, this_thread_govt_id, 0};
batch_state_ptr = GetPointerFromMultiIndexFor4DTensor(
govt_state_arr_batch, my_batch_state_shape, my_batch_state_idx);
CopyFloatArraySlice(my_global_state_ptr, kGovtStateSize, batch_state_ptr);
// extract pointers to action indices and rewards, batch and current
batch_single_shape = {kBatchSize, (unsigned int)num_iter, kNumGovts};
single_idx = {blockIdx.x, (unsigned int)iter,
(unsigned int)this_thread_govt_id};
batch_action_value_ptr = GetPointerFromMultiIndexFor3DIntTensor(
govt_actions_arr_batch, batch_single_shape, single_idx);
batch_reward_value_ptr = GetPointerFromMultiIndexFor3DFloatTensor(
govt_rewards_arr_batch, batch_single_shape, single_idx);
const int kThisThreadGovtIdx = blockIdx.x * kNumGovts + this_thread_govt_id;
my_action_value_ptr = &(govt_action_indices_arr[kThisThreadGovtIdx]);
my_reward_value_ptr = &(govt_rewards_arr[kThisThreadGovtIdx]);
}
// ----------------------------------------------
// State pointers and variables
// Create pointers to agent-specific state that will be updated by the
// simulation logic.
// ----------------------------------------------
float *my_state_arr = &(my_global_state_ptr[kGlobalStateSize]);
float *prices_arr = &(my_global_state_ptr[kIdxPricesOffset]);
float *wages_arr = &(my_global_state_ptr[kIdxWagesOffset]);
float *available_stock_arr = &(my_global_state_ptr[kIdxStockOffset]);
float time = my_global_state_ptr[kIdxTimeOffset];
float income_tax_rate = my_global_state_ptr[kIdxIncomeTaxOffset];
float corporate_tax_rate = my_global_state_ptr[kIdxCorporateTaxOffset];
// -------------------------------
// Safely initialize shared memory
// -------------------------------
if (ThisThreadAgentType == kFirmType) {
int this_thread_firm_id = threadIdx.x - kNumConsumers;
gross_demand_arr[this_thread_firm_id] = 0.0;
num_consumer_demand_arr[this_thread_firm_id] = 0;
hours_worked_arr[this_thread_firm_id] = 0.0;
total_actually_consumer_arr[this_thread_firm_id] = 0.0;
need_to_ration_this_good_arr[this_thread_firm_id] = false;
}
if (ThisThreadAgentType == kGovernmentType) {
tax_revenue_arr[0] = 0.0;
}
__syncthreads();
// -------------------------------------
// End - Safely initialize shared memory
// -------------------------------------
// -------------------------------------
// Process actions
// -------------------------------------
if (ThisThreadAgentType == kConsumerType) {
// amount demanded is just the first part of the action vector
float *this_agent_gross_demand_arr =
&(my_action_arr[kIdxConsumerDemandedOffset]);
const float this_agent_hours_worked =
my_action_arr[kIdxConsumerWorkedOffset];
const int worked_for_this_firm_id =
(int)my_action_arr[kIdxConsumerWhichFirmOffset];
// here, need to scale demands to meet the budget. put them in a local array
// *budgetDemanded logic should be: compute total expenditure given prices.
// if less than budget, copy existing demands
float __cost_of_demand = 0.0;
for (int i = 0; i < kNumFirms; i++) {
__cost_of_demand += this_agent_gross_demand_arr[i] * prices_arr[i];
}
// Scale demand to ensure that total demand at most equals total supply
// we want: my_state_arr being 0 always sends __scale_factor to 0
float __scale_factor = 1.0;
if ((__cost_of_demand > 0.0) && (__cost_of_demand > my_state_arr[0])) {
__scale_factor = my_state_arr[0] / __cost_of_demand;
}
// otherwise scale all demands down to meet budget
// copy them into a demanded array
for (int i = 0; i < kNumFirms; i++) {
net_demand_arr[i] = __scale_factor * this_agent_gross_demand_arr[i];
}
// adding up demand across threads **in the block**
// somehow store amount demanded per firm in an array demanded (copy from
// action) also store amount worked per firm in array worked
// Every thread executes atomicAdd_block in a memory-safe way.
for (int i = 0; i < kNumFirms; i++) {
// sum across threads in block
atomicAdd_block(&(gross_demand_arr[i]), net_demand_arr[i]);
// increment count of consumers who want good i
if (net_demand_arr[i] > 0) {
atomicAdd_block(&(num_consumer_demand_arr[i]), 1);
}
}
// increment total hours worked for firm i
atomicAdd_block(&(hours_worked_arr[worked_for_this_firm_id]),
this_agent_hours_worked);
}
// wait for everyone to finish tallying up their adding
__syncthreads();
if (ThisThreadAgentType == kFirmType) {
// check each firm if rationing needed
int this_thread_firm_id = threadIdx.x - kNumConsumers;
need_to_ration_this_good_arr[this_thread_firm_id] =
((gross_demand_arr[this_thread_firm_id] > 0.0) && (gross_demand_arr[this_thread_firm_id] >
available_stock_arr[this_thread_firm_id]));
}
// wait for single thread to finish checking demands
__syncthreads();
// ----------------------------------------
// Consumers: Rationing demand + Utility
// ----------------------------------------
// Logic:
// case 1: no overdemand
// case 2: overdemand, but some want less than 1/N -- fill everyone up to
// max(theirs, 1/N) case 3: overdemand, everyone wants more -- fill everyone
// up to max(theirs, 1/N)
float net_consumed_arr[kNumFirms]; // per consumer thread
// always add negligible positive money to avoid budgets becoming small
// negative numbers otherwise, when computing proportions, one may end up with
// negative stocks.
float cs_budget_delta = 0.01;
float fm_budget_delta = 0.01;
float capital_delta = 0.0;
if (ThisThreadAgentType == kConsumerType) {
// find out how much consumed
for (int i = 0; i < kNumFirms; i++) {
float __ration_factor = 1.0;
if (need_to_ration_this_good_arr[i]) {
// overdemanded
__ration_factor = available_stock_arr[i] / gross_demand_arr[i];
}
net_consumed_arr[i] = __ration_factor * net_demand_arr[i];
atomicAdd_block(&(total_actually_consumer_arr[i]), net_consumed_arr[i]);
}
// store amount actually consumed for this consumer
CopyFloatArraySlice(net_consumed_arr, kNumFirms, my_aux_batch_ptr);
// ----------------------------------------
// Compute consumer utility
// ----------------------------------------
float hours_worked = my_action_arr[kIdxConsumerWorkedOffset];
int worked_for_this_firm_id =
(int)my_action_arr[kIdxConsumerWhichFirmOffset];
// budget is first elem of consumer state, theta second
float __theta = my_state_arr[1];
float __this_consumer_util = 0.0;
float __total_hours_worked = 0.0;
float __gross_income = 0.0;
// Compute expenses
// Each consumer can consume from each firm, so loop over them.
for (int i = 0; i < kNumFirms; i++) {
__this_consumer_util += GetCRRAUtil(net_consumed_arr[i], kCrraParam);
cs_budget_delta -= prices_arr[i] * net_consumed_arr[i];
}
// Compute income
__total_hours_worked += hours_worked;
__gross_income += wages_arr[worked_for_this_firm_id] * hours_worked;
float __income_tax_paid = income_tax_rate * __gross_income;
cs_budget_delta += (__gross_income - __income_tax_paid);
// Update tax revenue (government)
atomicAdd_block(&(tax_revenue_arr[kIdxThisThreadGovtId]),
__income_tax_paid);
// Compute reward
this_agent_reward +=
__this_consumer_util - (__theta / 2.0) * (__total_hours_worked);
}
__syncthreads();
// ----------------------------------------
// Firms Exports: Add external consumption.
// ----------------------------------------
if (ThisThreadAgentType == kFirmType ) {
const int this_thread_firm_id = threadIdx.x - kNumConsumers;
if (kUseImporter) {
// sell remaining goods, if any, to importer, if price is high enough.
float __this_firm_price = prices_arr[this_thread_firm_id];
float __stock_after_consumers = available_stock_arr[this_thread_firm_id] - total_actually_consumer_arr[this_thread_firm_id];
if (__this_firm_price >= kImporterPrice) {
bought_by_importer_arr[this_thread_firm_id] = fmaxf(fminf(__stock_after_consumers, kImporterQuantity), 0.0); // floor to zero to avoid small negative floats
}
else {
bought_by_importer_arr[this_thread_firm_id] = 0.0;
}
}
else {
bought_by_importer_arr[this_thread_firm_id] = 0.0;
}
}
// ----------------------------------------
// Firms: Rationing demand + Utility
// ----------------------------------------
if (ThisThreadAgentType == kFirmType) {
const int this_thread_firm_id = threadIdx.x - kNumConsumers;
float __this_firm_revenue =
(total_actually_consumer_arr[this_thread_firm_id] + bought_by_importer_arr[this_thread_firm_id]) *
prices_arr[this_thread_firm_id];
float __wages_paid =
hours_worked_arr[this_thread_firm_id] * wages_arr[this_thread_firm_id];
// Firms can invest in new capital. This increases their production factor
// (see GetFirmProduction).
// here, after consumers consume, if price is >= than importer price, importer consumes up to their maximum of the goods, at the importer price
float __gross_income = __this_firm_revenue - __wages_paid;
capital_delta = fmaxf(my_action_arr[2] * __gross_income, 0.0);
float __gross_profit = __gross_income - capital_delta;
float __corp_tax_paid = corporate_tax_rate * fmaxf(__gross_profit, 0.0);
fm_budget_delta = (__gross_profit - __corp_tax_paid);
if (kIncentivizeFirmActivity) {
if ((fm_budget_delta + my_state_arr[0]) > 0.0) { // if positive budget
this_agent_reward += (kFirmBoostRewardFactor * __this_firm_revenue);
}
}
this_agent_reward += (__gross_profit - __corp_tax_paid);
atomicAdd_block(&(tax_revenue_arr[0]), __corp_tax_paid);
float __production = GetFirmProduction(0.01, my_state_arr[kIdxFirmCapitalOffset],
hours_worked_arr[this_thread_firm_id], my_state_arr[kIdxFirmAlphaOffset]);
// -------------------
// Update global state
// -------------------
// update prices in global state
next_global_state_arr[kIdxPricesOffset + this_thread_firm_id] =
my_action_arr[0];
// update wages in global state
next_global_state_arr[kIdxWagesOffset + this_thread_firm_id] =
my_action_arr[1];
// update stocks in global state
next_global_state_arr[kIdxStockOffset + this_thread_firm_id] =
available_stock_arr[this_thread_firm_id] -
total_actually_consumer_arr[this_thread_firm_id] - bought_by_importer_arr[this_thread_firm_id] + __production;
*my_aux_batch_ptr = bought_by_importer_arr[this_thread_firm_id];
// update overdemanded in global state
next_global_state_arr[kIdxOverdemandOffset + this_thread_firm_id] =
need_to_ration_this_good_arr[this_thread_firm_id] ? 1.0 : 0.0;
}
// -----------------
// Move time forward
// -----------------
// Let first firm tick time
if (ThisThreadAgentType == kFirmType) {
const int this_thread_firm_id = threadIdx.x - kNumConsumers;
if (this_thread_firm_id == 0) {
next_global_state_arr[kIdxTimeOffset] =
my_global_state_ptr[kIdxTimeOffset] + 1.0;
}
}
__syncthreads();
// ----------------------------------------
// Subsidies
// ----------------------------------------
// need to redistribute tax revenues
if (ThisThreadAgentType == kConsumerType) {
float __redistribution = tax_revenue_arr[0] / kNumConsumers;
cs_budget_delta += __redistribution;
}
// ----------------------------------------
// Social welfare
// ----------------------------------------
// After this point consumers and firms know their final reward, so can inform
// the government thread via shared memory
__syncthreads();
// ----------------------------------------
// Government sets taxes for the next round
// ----------------------------------------
if (ThisThreadAgentType == kGovernmentType) {
next_global_state_arr[kIdxIncomeTaxOffset] = my_action_arr[0];
next_global_state_arr[kIdxCorporateTaxOffset] = my_action_arr[1];
}
__syncthreads();
// -----------------------------------------------
// Copy next_global_state_arr into my global state
// -----------------------------------------------
// All agents need to see the updated global state
CopyFloatArraySlice(next_global_state_arr, kGlobalStateSize,
my_global_state_ptr);
// -----------------------------------------------
// Update budgets
// -----------------------------------------------
// Update budget (same for all agents)
if (ThisThreadAgentType == kConsumerType) {
my_state_arr[0] += cs_budget_delta;
}
if (ThisThreadAgentType == kFirmType) {
my_state_arr[0] += fm_budget_delta;
}
// Add interest rate on savings
if ((ThisThreadAgentType == kConsumerType) ||
(ThisThreadAgentType == kFirmType)) {
if (my_state_arr[0] > 0.0) {
my_state_arr[0] += my_state_arr[0] * kInterestRate;
}
}
// Add new capital
if (ThisThreadAgentType == kFirmType) {
my_state_arr[kIdxFirmCapitalOffset] += capital_delta;
}
// Add new capital
if ((ThisThreadAgentType == kFirmType) ||
(ThisThreadAgentType == kGovernmentType)) {
*batch_action_value_ptr = *my_action_value_ptr;
}
// Update rewards in global state
*my_reward_value_ptr = this_agent_reward;
*batch_reward_value_ptr = this_agent_reward;
}
// ************************
// End of extern "C" block.
// ************************
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,114 @@
# 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 torch
import torch.nn.functional as F
from torch import nn
class IndependentPolicyNet(nn.Module):
"""
Represents a policy network with separate heads for different types of actions.
Thus, the resulting policy will take the form
$pi(a | s) = pi_1(a_1 | s) pi_2(a_2 | s)...$
"""
def __init__(self, state_size, action_size_list, norm_consts=None):
super().__init__()
self.state_size = state_size
self.action_size_list = action_size_list
if norm_consts is not None:
self.norm_center, self.norm_scale = norm_consts
else:
self.norm_center = torch.zeros(self.state_size).cuda()
self.norm_scale = torch.ones(self.state_size).cuda()
self.fc1 = nn.Linear(state_size, 128)
self.fc2 = nn.Linear(128, 128)
# policy network head
self.action_heads = nn.ModuleList(
[nn.Linear(128, action_size) for action_size in action_size_list]
)
# value network head
self.fc4 = nn.Linear(128, 1)
def forward(self, x):
assert x.shape[-1] == self.state_size # Check if the last dimension matches
# Normalize the model input
new_shape = tuple(1 for _ in x.shape[:-1]) + (x.shape[-1],)
view_center = self.norm_center.view(new_shape)
view_scale = self.norm_scale.view(new_shape)
x = (x - view_center) / view_scale
# Feed forward
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
probs = [F.softmax(action_head(x), dim=-1) for action_head in self.action_heads]
vals = self.fc4(x)
return probs, vals
class PolicyNet(nn.Module):
"""
The policy network class to output acton probabilities and the value function.
"""
def __init__(self, state_size, action_size, norm_consts=None):
super().__init__()
self.state_size = state_size
self.action_size = action_size
if norm_consts is not None:
self.norm_center, self.norm_scale = norm_consts
else:
self.norm_center = torch.zeros(self.state_size).cuda()
self.norm_scale = torch.ones(self.state_size).cuda()
self.fc1 = nn.Linear(state_size, 128)
self.fc2 = nn.Linear(128, 128)
# policy network head
self.fc3 = nn.Linear(128, action_size)
# value network head
self.fc4 = nn.Linear(128, 1)
def forward(self, x, actions_mask=None):
# here, the action mask should be large negative constants for actions
# that shouldn't be allowed.
new_shape = tuple(1 for _ in x.shape[:-1]) + (x.shape[-1],)
view_center = self.norm_center.view(new_shape)
view_scale = self.norm_scale.view(new_shape)
x = (x - view_center) / view_scale
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
if actions_mask is not None:
probs = F.softmax(self.fc3(x) + actions_mask, dim=-1)
else:
probs = F.softmax(self.fc3(x), dim=-1)
vals = self.fc4(x)
return probs, vals
class DeterministicPolicy:
"""
A policy class that outputs deterministic actions.
"""
def __init__(self, state_size, action_size, action_choice):
self.state_size = state_size
self.action_size = action_size
self.action_choice = action_choice
self.actions_out = torch.zeros(action_size, device="cuda")
self.actions_out[self.action_choice] = 1.0
def __call__(self, x, actions_mask=None):
return self.forward(x)
def forward(self, x):
# output enough copies of the delta function
# distribution of the right size given x
x_batch_shapes = x.shape[:-1]
repeat_vals = x_batch_shapes + (1,)
return self.actions_out.repeat(*repeat_vals), None
@@ -0,0 +1,110 @@
# 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 torch
def dict_merge(dct, merge_dct):
"""Recursive dict merge. Inspired by :meth:``dict.update()``, instead of
updating only top-level keys, dict_merge recurses down into dicts nested
to an arbitrary depth, updating keys. The ``merge_dct`` is merged into
``dct``.
:param dct: dict onto which the merge is executed
:param merge_dct: dct merged into dct
:return: None
"""
for k, v in merge_dct.items():
# dct does not have k yet. Add it with value v.
if (k not in dct) and (not isinstance(dct, list)):
dct[k] = v
else:
# dct[k] and merge_dict[k] are both dictionaries. Recurse.
if isinstance(dct[k], (dict, list)) and isinstance(v, dict):
dict_merge(dct[k], merge_dct[k])
else:
# dct[k] and merge_dict[k] are both tuples or lists.
if isinstance(dct[k], (tuple, list)) and isinstance(v, (tuple, list)):
# They don't match. Overwrite with v.
if len(dct[k]) != len(v):
dct[k] = v
else:
for i, (d_val, v_val) in enumerate(zip(dct[k], v)):
if isinstance(d_val, dict) and isinstance(v_val, dict):
dict_merge(d_val, v_val)
else:
dct[k][i] = v_val
else:
dct[k] = v
def min_max_consumer_budget_delta(hparams_dict):
# largest single round changes
max_wage = hparams_dict["agents"]["max_possible_wage"]
max_hours = hparams_dict["agents"]["max_possible_hours_worked"]
max_price = hparams_dict["agents"]["max_possible_price"]
max_singlefirm_consumption = hparams_dict["agents"]["max_possible_consumption"]
num_firms = hparams_dict["agents"]["num_firms"]
min_budget = (
-max_singlefirm_consumption * max_price * num_firms
) # negative budget from consuming only
max_budget = max_hours * max_wage * num_firms # positive budget from only working
return min_budget, max_budget
def min_max_stock_delta(hparams_dict):
# for now, assuming 1.0 capital
max_hours = hparams_dict["agents"]["max_possible_hours_worked"]
max_singlefirm_consumption = hparams_dict["agents"]["max_possible_consumption"]
alpha = hparams_dict["world"]["production_alpha"]
if isinstance(alpha, str):
alpha = 0.8
num_consumers = hparams_dict["agents"]["num_consumers"]
max_delta = (max_hours * num_consumers) ** alpha
min_delta = -max_singlefirm_consumption * num_consumers
return min_delta, max_delta
def min_max_firm_budget(hparams_dict):
max_wage = hparams_dict["agents"]["max_possible_wage"]
max_hours = hparams_dict["agents"]["max_possible_hours_worked"]
max_singlefirm_consumption = hparams_dict["agents"]["max_possible_consumption"]
num_consumers = hparams_dict["agents"]["num_consumers"]
max_price = hparams_dict["agents"]["max_possible_price"]
max_delta = max_singlefirm_consumption * max_price * num_consumers
min_delta = -max_hours * max_wage * num_consumers
return min_delta, max_delta
def expand_to_digit_form(x, dims_to_expand, max_digits):
# first split x up
requires_grad = (
x.requires_grad
) # don't want to backprop through these ops, but do want
# gradients if x had them
with torch.no_grad():
tensor_pieces = []
expanded_digit_shape = x.shape[:-1] + (max_digits,)
for i in range(x.shape[-1]):
if i not in dims_to_expand:
tensor_pieces.append(x[..., i : i + 1])
else:
digit_entries = torch.zeros(expanded_digit_shape, device=x.device)
for j in range(max_digits):
digit_entries[..., j] = (x[..., i] % (10 ** (j + 1))) / (
10 ** (j + 1)
)
tensor_pieces.append(digit_entries)
output = torch.cat(tensor_pieces, dim=-1)
output.requires_grad_(requires_grad)
return output
def size_after_digit_expansion(existing_size, dims_to_expand, max_digits):
num_expanded = len(dims_to_expand)
# num non expanded digits, + all the expanded ones
return (existing_size - num_expanded) + (max_digits * num_expanded)