adding ai_economist for modding
This commit is contained in:
@@ -0,0 +1,242 @@
|
||||
# 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 hashlib
|
||||
import itertools
|
||||
import json
|
||||
import os
|
||||
import pickle
|
||||
import struct
|
||||
import time
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
# defaults
|
||||
_NUM_FIRMS = 10
|
||||
|
||||
|
||||
def _bigint_from_bytes(num_bytes):
|
||||
"""
|
||||
See https://github.com/openai/gym/blob/master/gym/utils/seeding.py.
|
||||
"""
|
||||
sizeof_int = 4
|
||||
padding = sizeof_int - len(num_bytes) % sizeof_int
|
||||
num_bytes += b"\0" * padding
|
||||
int_count = int(len(num_bytes) / sizeof_int)
|
||||
unpacked = struct.unpack("{}I".format(int_count), num_bytes)
|
||||
accum = 0
|
||||
for i, val in enumerate(unpacked):
|
||||
accum += 2 ** (sizeof_int * 8 * i) * val
|
||||
return accum
|
||||
|
||||
|
||||
def seed_from_base_seed(base_seed):
|
||||
"""
|
||||
Hash base seed to reduce correlation.
|
||||
"""
|
||||
max_bytes = 4
|
||||
hash_func = hashlib.sha512(str(base_seed).encode("utf8")).digest()
|
||||
|
||||
return _bigint_from_bytes(hash_func[:max_bytes])
|
||||
|
||||
|
||||
def hash_from_dict(d):
|
||||
d_copy = deepcopy(d)
|
||||
del (d_copy["train"])["base_seed"]
|
||||
d_string = json.dumps(d_copy, sort_keys=True)
|
||||
return int(hashlib.sha256(d_string.encode("utf8")).hexdigest()[:8], 16)
|
||||
|
||||
|
||||
def cfg_dict_from_yaml(
|
||||
hparams_path,
|
||||
consumption_choices,
|
||||
work_choices,
|
||||
price_and_wage,
|
||||
tax_choices,
|
||||
group_name=None,
|
||||
):
|
||||
with open(hparams_path) as f:
|
||||
d = yaml.safe_load(f)
|
||||
|
||||
if group_name is not None:
|
||||
d["metadata"]["group_name"] = group_name
|
||||
d["metadata"]["hparamhash"] = hash_from_dict(d)
|
||||
d["agents"][
|
||||
"consumer_consumption_actions_array"
|
||||
] = consumption_choices # Note: hardcoded
|
||||
d["agents"]["consumer_work_actions_array"] = work_choices # Note: hardcoded
|
||||
d["agents"]["firm_actions_array"] = price_and_wage # Note: hardcoded
|
||||
d["agents"]["government_actions_array"] = tax_choices
|
||||
d["train"]["save_dir"] = str(hparams_path.absolute().parent)
|
||||
d["train"]["seed"] = seed_from_base_seed(d["train"]["base_seed"])
|
||||
return d
|
||||
|
||||
|
||||
def run_experiment_batch_parallel(
|
||||
experiment_dir,
|
||||
consumption_choices,
|
||||
work_choices,
|
||||
price_and_wage,
|
||||
tax_choices,
|
||||
group_name=None,
|
||||
consumers_only=False,
|
||||
no_firms=False,
|
||||
default_firm_action=None,
|
||||
default_government_action=None,
|
||||
):
|
||||
hparams_path = Path(experiment_dir) / Path("hparams.yaml")
|
||||
hparams_dict = cfg_dict_from_yaml(
|
||||
hparams_path,
|
||||
consumption_choices,
|
||||
work_choices,
|
||||
price_and_wage,
|
||||
tax_choices,
|
||||
group_name=group_name,
|
||||
)
|
||||
print(f"hparams_dict {hparams_dict}")
|
||||
# import this here so rest of file still imports without cuda installed
|
||||
from rbc.cuda_manager import ConsumerFirmRunManagerBatchParallel
|
||||
|
||||
if consumers_only:
|
||||
m = ConsumerFirmRunManagerBatchParallel(
|
||||
hparams_dict,
|
||||
freeze_firms=default_firm_action,
|
||||
freeze_govt=default_government_action,
|
||||
)
|
||||
elif no_firms:
|
||||
m = ConsumerFirmRunManagerBatchParallel(
|
||||
hparams_dict,
|
||||
freeze_firms=default_firm_action,
|
||||
)
|
||||
else:
|
||||
m = ConsumerFirmRunManagerBatchParallel(hparams_dict)
|
||||
m.train()
|
||||
|
||||
|
||||
def compare_global_states_within_type(states, global_state_size):
|
||||
# every agent within a batch should have the same global state
|
||||
first_agent_global = states[:, :, :1, :global_state_size]
|
||||
all_agents_global = states[:, :, :, :global_state_size]
|
||||
return np.isclose(all_agents_global, first_agent_global).all()
|
||||
|
||||
|
||||
def compare_global_states_across_types(
|
||||
consumer_states, firm_states, government_states, global_state_size
|
||||
):
|
||||
first_agent_global = consumer_states[:, :, :1, :global_state_size]
|
||||
return (
|
||||
np.isclose(firm_states[:, :, :, :global_state_size], first_agent_global).all(),
|
||||
np.isclose(
|
||||
government_states[:, :, :, :global_state_size], first_agent_global
|
||||
).all(),
|
||||
np.isclose(
|
||||
consumer_states[:, :, :, :global_state_size], first_agent_global
|
||||
).all(),
|
||||
)
|
||||
|
||||
|
||||
def check_no_negative_stocks(state, stock_offset, stock_size):
|
||||
stocks = state[:, :, :, stock_offset : (stock_offset + stock_size)]
|
||||
return (stocks >= -1.0e-3).all()
|
||||
|
||||
|
||||
train_param_sweeps = {
|
||||
"lr": [0.005, 0.001],
|
||||
"entropy": [0.01],
|
||||
"base_seed": [2596],
|
||||
"batch_size": [64],
|
||||
"clip_grad_norm": [1.0, 2.0, 5.0],
|
||||
}
|
||||
|
||||
# Other param sweeps
|
||||
agent_param_sweeps = {
|
||||
# "consumer_noponzi_eta": [0.1,0.05]
|
||||
}
|
||||
|
||||
world_param_sweeps = {
|
||||
# "interest_rate": [0.02, 0.0]
|
||||
}
|
||||
|
||||
|
||||
def add_all(d, keys_list, target_val):
|
||||
for k in keys_list:
|
||||
d[k] = target_val
|
||||
|
||||
|
||||
def sweep_cfg_generator(
|
||||
base_cfg,
|
||||
tr_param_sweeps=None,
|
||||
ag_param_sweeps=None,
|
||||
wld_param_sweeps=None,
|
||||
seed_from_timestamp=False,
|
||||
group_name=None,
|
||||
):
|
||||
# train_param_sweeps
|
||||
if tr_param_sweeps is None:
|
||||
tr_param_sweeps = {}
|
||||
# agent_param_sweeps
|
||||
if ag_param_sweeps is None:
|
||||
ag_param_sweeps = {}
|
||||
# world_param_sweeps
|
||||
if wld_param_sweeps is None:
|
||||
wld_param_sweeps = {}
|
||||
|
||||
assert isinstance(tr_param_sweeps, dict)
|
||||
assert isinstance(ag_param_sweeps, dict)
|
||||
assert isinstance(wld_param_sweeps, dict)
|
||||
|
||||
key_dict = {} # tells which key goes to which dict, e.g. "lr" -> "train", etc.
|
||||
if len(tr_param_sweeps) > 0:
|
||||
train_k, train_v = zip(*tr_param_sweeps.items())
|
||||
else:
|
||||
train_k, train_v = (), ()
|
||||
add_all(key_dict, train_k, "train")
|
||||
if len(ag_param_sweeps) > 0:
|
||||
agent_k, agent_v = zip(*ag_param_sweeps.items())
|
||||
else:
|
||||
agent_k, agent_v = (), ()
|
||||
add_all(key_dict, agent_k, "agents")
|
||||
if len(wld_param_sweeps) > 0:
|
||||
world_k, world_v = zip(*wld_param_sweeps.items())
|
||||
else:
|
||||
world_k, world_v = (), ()
|
||||
add_all(key_dict, world_k, "world")
|
||||
|
||||
k = train_k + agent_k + world_k
|
||||
v = train_v + agent_v + world_v
|
||||
|
||||
# have a "reverse lookup" dictionary for each key name
|
||||
for combination in itertools.product(*v):
|
||||
values_to_substitute = dict(zip(k, combination))
|
||||
out = deepcopy(base_cfg)
|
||||
for key, value in values_to_substitute.items():
|
||||
out[key_dict[key]][key] = value
|
||||
if seed_from_timestamp:
|
||||
int_timestamp = int(
|
||||
time.time() * 1000
|
||||
) # time.time() returns float, multiply 1000 for higher resolution
|
||||
out["train"]["base_seed"] += int_timestamp
|
||||
if group_name is not None:
|
||||
out["metadata"]["group"] = group_name
|
||||
yield out
|
||||
|
||||
|
||||
def create_job_dir(experiment_dir, job_name_base, cfg=None, action_arrays=None):
|
||||
unique_id = time.time()
|
||||
dirname = f"{job_name_base}-{unique_id}".replace(".", "-")
|
||||
dir_path = Path(experiment_dir) / Path(dirname)
|
||||
os.makedirs(str(dir_path), exist_ok=True)
|
||||
cfg["metadata"]["dirname"] = dirname
|
||||
cfg["metadata"]["group"] = str(Path(experiment_dir).name)
|
||||
with open(dir_path / Path("hparams.yaml"), "w") as f:
|
||||
f.write(yaml.dump(cfg))
|
||||
|
||||
if action_arrays is not None:
|
||||
with open(dir_path / Path("action_arrays.pickle"), "wb") as f:
|
||||
pickle.dump(action_arrays, f)
|
||||
Reference in New Issue
Block a user