1931 lines
73 KiB
Python
1931 lines
73 KiB
Python
# 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 os
|
|
import random
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pycuda
|
|
import pycuda.autoinit
|
|
import pycuda.driver as cuda_driver
|
|
import scipy
|
|
import scipy.stats
|
|
import torch
|
|
from pycuda.compiler import SourceModule
|
|
from torch.distributions import Categorical
|
|
from tqdm import tqdm
|
|
|
|
from .constants import (
|
|
consumer_state_scaling_factors,
|
|
firm_state_scaling_factors,
|
|
govt_state_scaling_factors,
|
|
)
|
|
from .networks import DeterministicPolicy, IndependentPolicyNet, PolicyNet
|
|
from .util import expand_to_digit_form, size_after_digit_expansion
|
|
|
|
_NP_DTYPE = np.float32
|
|
|
|
# the below line is 'strangely' necessary to make PyTorch work with PyCUDA
|
|
pytorch_cuda_init_success = torch.cuda.FloatTensor(8)
|
|
|
|
|
|
# for opening source files within module
|
|
module_path = Path(__file__).parent
|
|
|
|
|
|
def interval_list_contains(interval_list, step):
|
|
for (lower, upper_non_inclusive) in interval_list:
|
|
if lower <= step < upper_non_inclusive:
|
|
return True
|
|
return False
|
|
|
|
|
|
class NoOpOptimizer:
|
|
"""
|
|
Dummy Optimizer.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def step(self):
|
|
pass
|
|
|
|
|
|
def seed_everything(seed):
|
|
torch.manual_seed(seed)
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
|
|
def reverse_cumsum(x):
|
|
# assumes summing along episode iteration dim
|
|
return x + torch.sum(x, dim=-2, keepdims=True) - torch.cumsum(x, dim=-2)
|
|
|
|
|
|
def discounted_returns(rewards, gamma):
|
|
maxt = rewards.shape[-2]
|
|
cumulative_rewards = 0
|
|
returns = torch.zeros_like(rewards)
|
|
for t in reversed(range(maxt)):
|
|
returns[:, t, :] = rewards[:, t, :] + gamma * cumulative_rewards
|
|
cumulative_rewards = rewards[:, t, :] + cumulative_rewards
|
|
return returns
|
|
|
|
|
|
def compute_theta_coef(hparams_dict, episode):
|
|
anneal_dict = hparams_dict["agents"]["consumer_anneal_theta"]
|
|
if anneal_dict["anneal_on"]:
|
|
exp_decay_length_in_steps = anneal_dict["exp_decay_length_in_steps"]
|
|
theta_coef = np.float32(1.0 - (np.exp(-episode / exp_decay_length_in_steps)))
|
|
else:
|
|
return np.float32(1.0)
|
|
return theta_coef
|
|
|
|
|
|
def government_action_mask(hparams_dict, step):
|
|
government_actions_array = hparams_dict["agents"]["government_actions_array"]
|
|
tax_annealing_params = hparams_dict["agents"]["government_anneal_taxes"]
|
|
|
|
income_tax = torch.tensor(government_actions_array[:, 0]).cuda()
|
|
corporate_tax = torch.tensor(government_actions_array[:, 1]).cuda()
|
|
mask = torch.zeros(income_tax.shape[0]).cuda()
|
|
|
|
if not tax_annealing_params["anneal_on"]:
|
|
return None
|
|
a0 = tax_annealing_params["start"]
|
|
max_tax = tax_annealing_params["increase_const"] * step + a0
|
|
mask[(income_tax > max_tax) | (corporate_tax > max_tax)] -= 1000.0
|
|
|
|
return mask
|
|
|
|
|
|
def firm_action_mask(hparams_dict, step):
|
|
# pick out all firm actions where wage is the wrong height,
|
|
# and assign -1000.0 to those
|
|
firm_actions_array = hparams_dict["agents"]["firm_actions_array"]
|
|
wage_annealing_params = hparams_dict["agents"]["firm_anneal_wages"]
|
|
price_annealing_params = hparams_dict["agents"]["firm_anneal_prices"]
|
|
wages = torch.tensor(firm_actions_array[:, 1]).cuda()
|
|
prices = torch.tensor(firm_actions_array[:, 0]).cuda()
|
|
mask = torch.zeros(wages.shape[0]).cuda()
|
|
|
|
if not (wage_annealing_params["anneal_on"] or price_annealing_params["anneal_on"]):
|
|
return None
|
|
|
|
if wage_annealing_params["anneal_on"]:
|
|
a0 = wage_annealing_params["start"]
|
|
max_wage = wage_annealing_params["increase_const"] * step + a0
|
|
min_wage = -wage_annealing_params["decrease_const"] * step + a0
|
|
mask[(wages < min_wage) | (wages > max_wage)] -= 1000.0
|
|
if price_annealing_params["anneal_on"]:
|
|
a0 = price_annealing_params["start"]
|
|
max_price = price_annealing_params["increase_const"] * step + a0
|
|
min_price = -price_annealing_params["decrease_const"] * step + a0
|
|
mask[(prices < min_price) | (prices > max_price)] -= 1000.0
|
|
|
|
return mask
|
|
|
|
|
|
def get_cuda_code(rel_path_to_cu_file, **preprocessor_vars_to_replace):
|
|
with open(module_path / rel_path_to_cu_file) as cudasource:
|
|
code_string = cudasource.read()
|
|
|
|
# format for preprocessor macros in firm_rbc.cu is M_VARNAME.
|
|
# Specify all these as args to nvcc.
|
|
options_list = [
|
|
f"-D M_{k.upper()}={v}" for k, v in preprocessor_vars_to_replace.items()
|
|
]
|
|
|
|
return code_string, options_list
|
|
|
|
|
|
def add_penalty_for_no_ponzi(
|
|
states, rewards, budget_offset, penalty_coef=20.0, penalty_scale=100.0
|
|
):
|
|
budget_violations = -torch.clamp_max(states[..., budget_offset], 0.0)
|
|
rewards[:, -1, :] -= penalty_coef * budget_violations / penalty_scale
|
|
|
|
|
|
def update_government_rewards(
|
|
government_rewards, consumer_rewards, firm_rewards, cfg_dict
|
|
):
|
|
assert (
|
|
government_rewards == 0.0
|
|
).all() # govt should have been assigned exactly 0 in cuda step function
|
|
total_rewards = consumer_rewards.sum(dim=-1)
|
|
if cfg_dict["agents"]["government_counts_firm_reward"] == 1:
|
|
total_rewards = total_rewards + cfg_dict["agents"].get(
|
|
"firm_reward_for_government_factor", 1.0
|
|
) * firm_rewards.sum(dim=-1)
|
|
|
|
government_rewards[..., 0] = total_rewards[:] # one govt for now
|
|
|
|
|
|
def update_penalty_coef(
|
|
states,
|
|
budget_offset,
|
|
prev_penalty_coef,
|
|
penalty_step_size=0.01,
|
|
penalty_scale=100.0,
|
|
):
|
|
budget_violations = -torch.clamp_max(states[..., budget_offset], 0.0)
|
|
new_penalty_coef = (
|
|
prev_penalty_coef
|
|
+ penalty_step_size * (budget_violations / penalty_scale).mean().item()
|
|
)
|
|
return new_penalty_coef
|
|
|
|
|
|
def get_actions_from_inds(action_inds, agents_dict):
|
|
|
|
_action_inds = action_inds.cpu().to(torch.long)
|
|
|
|
consumption_action_tensor = torch.tensor(
|
|
agents_dict["consumer_consumption_actions_array"]
|
|
)
|
|
|
|
work_action_tensor = torch.tensor(agents_dict["consumer_work_actions_array"])
|
|
num_firms = agents_dict["num_firms"]
|
|
out_shape = _action_inds.shape[:-1] + (agents_dict["consumer_action_dim"],)
|
|
consumer_actions_out = torch.zeros(out_shape)
|
|
idx_hours_worked = num_firms
|
|
idx_which_firm = num_firms + 1
|
|
|
|
for i in range(num_firms):
|
|
consumer_actions_out[..., i] = consumption_action_tensor[
|
|
_action_inds[..., i], :
|
|
].squeeze(dim=-1)
|
|
|
|
consumer_actions_out[..., num_firms] = work_action_tensor[
|
|
_action_inds[..., idx_hours_worked], :
|
|
].squeeze(dim=-1)
|
|
|
|
consumer_actions_out[..., (num_firms + 1)] = _action_inds[..., idx_which_firm]
|
|
|
|
return consumer_actions_out
|
|
|
|
|
|
def anneal_entropy_coef(entropy_dict, step):
|
|
if entropy_dict is None:
|
|
return 1.0
|
|
|
|
if entropy_dict["anneal_on"]:
|
|
coef_floor = entropy_dict.get("coef_floor", 0.0)
|
|
return max(
|
|
np.exp(-step / entropy_dict["exp_decay_length_in_steps"]), coef_floor
|
|
)
|
|
return 1.0
|
|
|
|
|
|
def get_grad_norm(policy):
|
|
grad_norm = 0.0
|
|
for p in list(filter(lambda p: p.grad is not None, policy.parameters())):
|
|
grad_norm += (p.grad.data.norm(2).item()) ** 2
|
|
return grad_norm
|
|
|
|
|
|
def get_ev(adv, returns, cutoff=-1.0):
|
|
return max(cutoff, (1 - (adv.detach().var() / returns.detach().var())).item())
|
|
|
|
|
|
def consumer_ppo_step(
|
|
policy,
|
|
states,
|
|
actions,
|
|
rewards,
|
|
optimizer,
|
|
gamma_const,
|
|
entropy_val=0.0,
|
|
value_loss_weight=1.0,
|
|
ppo_num_updates=3,
|
|
reward_scale=1.0,
|
|
clip_grad_norm=None,
|
|
clip_param=0.1,
|
|
):
|
|
# Get initial policy predictions
|
|
multi_action_probs, old_value_preds = policy(states)
|
|
|
|
old_value_preds = old_value_preds.detach()
|
|
# Get returns
|
|
rescaled_rewards = rewards / reward_scale
|
|
G_discounted_returns = discounted_returns(rescaled_rewards, gamma_const)
|
|
|
|
# Value function loss
|
|
|
|
sum_old_log_probs = 0.0
|
|
for action_ind, probs in enumerate(multi_action_probs):
|
|
_CategoricalDist = Categorical(probs)
|
|
sum_old_log_probs += -1.0 * _CategoricalDist.log_prob(actions[..., action_ind])
|
|
sum_old_log_probs = sum_old_log_probs.detach()
|
|
|
|
assert not G_discounted_returns.requires_grad
|
|
assert not sum_old_log_probs.requires_grad
|
|
assert not old_value_preds.requires_grad
|
|
|
|
# Compute ppo loss
|
|
for _ in range(ppo_num_updates):
|
|
multi_action_probs, value_preds = policy(states)
|
|
get_huber_loss = torch.nn.SmoothL1Loss()
|
|
value_pred_clipped = old_value_preds + (value_preds - old_value_preds).clamp(
|
|
-clip_param, clip_param
|
|
)
|
|
value_loss_new = get_huber_loss(
|
|
value_preds.squeeze(dim=-1), G_discounted_returns
|
|
) # can use huber loss instead
|
|
value_loss_clipped = get_huber_loss(
|
|
value_pred_clipped.squeeze(dim=-1), G_discounted_returns
|
|
)
|
|
|
|
value_loss = torch.max(value_loss_new, value_loss_clipped).mean()
|
|
|
|
# Policy loss with value function baseline.
|
|
advantages = G_discounted_returns - value_preds.detach().squeeze(dim=-1)
|
|
# Don't propagate through to VF network.
|
|
assert not advantages.requires_grad
|
|
|
|
# Trick: standardize advantages
|
|
standardized_advantages = (advantages - advantages.mean()) / (
|
|
advantages.std() + 1e-6
|
|
)
|
|
sum_mean_entropy = 0.0 # mean over batch and agents
|
|
sum_neg_log_probs = 0.0
|
|
|
|
for action_ind, probs in enumerate(multi_action_probs):
|
|
_CategoricalDist = Categorical(probs)
|
|
sum_neg_log_probs += -1.0 * _CategoricalDist.log_prob(
|
|
actions[..., action_ind]
|
|
)
|
|
sum_mean_entropy += _CategoricalDist.entropy().mean()
|
|
|
|
assert sum_neg_log_probs.requires_grad
|
|
# note: log probs are negative, so negate again here
|
|
ratio = torch.exp(-sum_neg_log_probs + sum_old_log_probs)
|
|
surr1 = ratio * standardized_advantages
|
|
surr2 = (
|
|
torch.clamp(ratio, 1.0 - clip_param, 1.0 + clip_param)
|
|
* standardized_advantages
|
|
)
|
|
|
|
ppo_loss = -torch.min(surr1, surr2).mean()
|
|
|
|
loss = (
|
|
ppo_loss - entropy_val * sum_mean_entropy + value_loss_weight * value_loss
|
|
)
|
|
|
|
# Apply gradients
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
|
|
if clip_grad_norm is not None:
|
|
torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=clip_grad_norm)
|
|
|
|
optimizer.step()
|
|
|
|
|
|
def ppo_step(
|
|
policy,
|
|
states,
|
|
actions,
|
|
rewards,
|
|
optimizer,
|
|
gamma_const,
|
|
entropy_val=0.0,
|
|
value_loss_weight=1.0,
|
|
ppo_num_updates=3,
|
|
actions_mask=None,
|
|
reward_scale=1.0,
|
|
clip_grad_norm=None,
|
|
clip_param=0.1,
|
|
):
|
|
# Get initial policy predictions
|
|
probs, old_value_preds = policy(states, actions_mask=actions_mask)
|
|
old_value_preds = old_value_preds.detach()
|
|
|
|
# Get returns
|
|
rescaled_rewards = rewards / reward_scale
|
|
G_discounted_returns = discounted_returns(rescaled_rewards, gamma_const)
|
|
|
|
# Value function loss
|
|
|
|
_CategoricalDist = Categorical(probs)
|
|
old_log_probs = -1.0 * _CategoricalDist.log_prob(actions).detach()
|
|
|
|
assert not G_discounted_returns.requires_grad
|
|
assert not old_log_probs.requires_grad
|
|
assert not old_value_preds.requires_grad
|
|
|
|
# Compute ppo loss
|
|
for _ in range(ppo_num_updates):
|
|
probs, value_preds = policy(states, actions_mask=actions_mask)
|
|
get_huber_loss = torch.nn.SmoothL1Loss()
|
|
value_pred_clipped = old_value_preds + (value_preds - old_value_preds).clamp(
|
|
-clip_param, clip_param
|
|
)
|
|
value_loss_new = get_huber_loss(
|
|
value_preds.squeeze(dim=-1), G_discounted_returns
|
|
) # can use huber loss instead
|
|
value_loss_clipped = get_huber_loss(
|
|
value_pred_clipped.squeeze(dim=-1), G_discounted_returns
|
|
)
|
|
|
|
value_loss = torch.max(value_loss_new, value_loss_clipped).mean()
|
|
|
|
# Policy loss with value function baseline.
|
|
advantages = G_discounted_returns - value_preds.detach().squeeze(dim=-1)
|
|
# Don't propagate through to VF network.
|
|
assert not advantages.requires_grad
|
|
|
|
# Trick: standardize advantages
|
|
standardized_advantages = (advantages - advantages.mean()) / (
|
|
advantages.std() + 1e-6
|
|
)
|
|
|
|
_CategoricalDist = Categorical(probs)
|
|
neg_log_probs = -1.0 * _CategoricalDist.log_prob(actions)
|
|
mean_entropy = _CategoricalDist.entropy().mean()
|
|
|
|
assert neg_log_probs.requires_grad
|
|
# note: log probs are negative, so negate again here
|
|
ratio = torch.exp(-neg_log_probs + old_log_probs)
|
|
surr1 = ratio * standardized_advantages
|
|
surr2 = (
|
|
torch.clamp(ratio, 1.0 - clip_param, 1.0 + clip_param)
|
|
* standardized_advantages
|
|
)
|
|
|
|
ppo_loss = -torch.min(surr1, surr2).mean()
|
|
|
|
loss = ppo_loss - entropy_val * mean_entropy + value_loss_weight * value_loss
|
|
|
|
# Apply gradients
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
|
|
if clip_grad_norm is not None:
|
|
torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=clip_grad_norm)
|
|
|
|
optimizer.step()
|
|
|
|
|
|
def consumer_policy_gradient_step(
|
|
policy,
|
|
states,
|
|
actions,
|
|
rewards,
|
|
optimizer,
|
|
gamma_const,
|
|
entropy_val=0.0,
|
|
value_loss_weight=1.0,
|
|
reward_scale=1.0,
|
|
clip_grad_norm=None,
|
|
):
|
|
# Get policy and value predictions
|
|
multi_action_probs, value_preds = policy(states)
|
|
|
|
# Get returns
|
|
rescaled_rewards = rewards / reward_scale
|
|
G_discounted_returns = discounted_returns(rescaled_rewards, gamma_const)
|
|
|
|
# Value function loss
|
|
get_huber_loss = torch.nn.SmoothL1Loss()
|
|
value_loss = get_huber_loss(
|
|
value_preds.squeeze(dim=-1), G_discounted_returns
|
|
).mean() # can use huber loss instead
|
|
|
|
# Policy loss with value function baseline.
|
|
advantages = G_discounted_returns - value_preds.detach().squeeze(dim=-1)
|
|
# Don't propagate through to VF network.
|
|
assert not advantages.requires_grad
|
|
|
|
# Trick: standardize advantages
|
|
standardized_advantages = (advantages - advantages.mean()) / (
|
|
advantages.std() + 1e-6
|
|
)
|
|
|
|
# Compute policy loss
|
|
sum_mean_entropy = 0.0 # mean over batch and agents
|
|
sum_neg_log_probs = 0.0
|
|
|
|
for action_ind, probs in enumerate(multi_action_probs):
|
|
_CategoricalDist = Categorical(probs)
|
|
sum_neg_log_probs += -1.0 * _CategoricalDist.log_prob(actions[..., action_ind])
|
|
sum_mean_entropy += _CategoricalDist.entropy().mean()
|
|
|
|
pg_loss = (sum_neg_log_probs * standardized_advantages).mean()
|
|
assert sum_neg_log_probs.requires_grad
|
|
|
|
loss = pg_loss - entropy_val * sum_mean_entropy + value_loss_weight * value_loss
|
|
|
|
# Apply gradients
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
|
|
if clip_grad_norm is not None:
|
|
torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=clip_grad_norm)
|
|
|
|
optimizer.step()
|
|
|
|
|
|
def policy_gradient_step(
|
|
policy,
|
|
states,
|
|
actions,
|
|
rewards,
|
|
optimizer,
|
|
gamma_const,
|
|
entropy_val=0.0,
|
|
value_loss_weight=1.0,
|
|
actions_mask=None,
|
|
reward_scale=1.0,
|
|
clip_grad_norm=None,
|
|
):
|
|
|
|
# here, we must perform digit scaling
|
|
optimizer.zero_grad()
|
|
probs, value_preds = policy(states, actions_mask=actions_mask)
|
|
rewards = rewards / reward_scale
|
|
G_discounted_returns = discounted_returns(rewards, gamma_const)
|
|
get_huber_loss = torch.nn.SmoothL1Loss()
|
|
value_loss = get_huber_loss(
|
|
value_preds.squeeze(dim=-1), G_discounted_returns
|
|
).mean() # can use huber loss instead
|
|
advantages = G_discounted_returns - value_preds.detach().squeeze(
|
|
dim=-1
|
|
) # compute advantages (don't propagate through to VF network)
|
|
assert not advantages.requires_grad
|
|
# mean and standardize advantages
|
|
standardized_advantages = (advantages - advantages.mean()) / (
|
|
advantages.std() + 1e-6
|
|
)
|
|
assert not standardized_advantages.requires_grad
|
|
m = Categorical(probs)
|
|
pg_loss = (-m.log_prob(actions) * standardized_advantages).mean()
|
|
assert pg_loss.requires_grad
|
|
entropy_regularize = entropy_val * m.entropy().mean()
|
|
loss = pg_loss - entropy_regularize + value_loss_weight * value_loss
|
|
loss.backward()
|
|
|
|
if clip_grad_norm is not None:
|
|
torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=clip_grad_norm)
|
|
|
|
optimizer.step()
|
|
|
|
|
|
def save_dense_log(
|
|
save_dir,
|
|
epi,
|
|
agent_type_arrays,
|
|
agent_action_arrays,
|
|
agent_aux_arrays,
|
|
):
|
|
print(f"Saving dense log at episode {epi}")
|
|
for agent_type in ["consumer", "firm", "government"]:
|
|
states_batch, actions_batch, rewards_batch = agent_type_arrays[agent_type]
|
|
aux_array = agent_aux_arrays[agent_type]
|
|
if aux_array is not None:
|
|
aux_array = aux_array.cpu().numpy()
|
|
np.savez(
|
|
str(Path(save_dir) / Path(f"episode_{epi}_{agent_type}.npz")),
|
|
states=states_batch.cpu().numpy(),
|
|
actions=actions_batch.cpu().numpy(),
|
|
rewards=rewards_batch.cpu().numpy(),
|
|
action_array=agent_action_arrays[agent_type],
|
|
aux_array=aux_array,
|
|
)
|
|
|
|
|
|
def save_policy_parameters(
|
|
save_dir,
|
|
epi,
|
|
consumer_policy,
|
|
firm_policy,
|
|
government_policy,
|
|
freeze_firms,
|
|
freeze_govt,
|
|
):
|
|
print(f"saving model parameters at episode {epi}")
|
|
consumer_path = (
|
|
Path(save_dir) / Path("saved_models") / Path(f"consumer_policy_{epi}.pt")
|
|
)
|
|
|
|
# always save the latest, to be overwritten later
|
|
consumer_path_latest = (
|
|
Path(save_dir) / Path("saved_models") / Path("consumer_policy_latest.pt")
|
|
)
|
|
os.makedirs(consumer_path.parent, exist_ok=True)
|
|
torch.save(consumer_policy.state_dict(), consumer_path)
|
|
torch.save(consumer_policy.state_dict(), consumer_path_latest)
|
|
|
|
if freeze_firms is None:
|
|
firm_path = (
|
|
Path(save_dir) / Path("saved_models") / Path(f"firm_policy_{epi}.pt")
|
|
)
|
|
firm_path_latest = (
|
|
Path(save_dir) / Path("saved_models") / Path("firm_policy_latest.pt")
|
|
)
|
|
|
|
os.makedirs(firm_path.parent, exist_ok=True)
|
|
torch.save(firm_policy.state_dict(), firm_path)
|
|
torch.save(firm_policy.state_dict(), firm_path_latest)
|
|
if freeze_govt is None:
|
|
government_path = (
|
|
Path(save_dir) / Path("saved_models") / Path(f"government_policy_{epi}.pt")
|
|
)
|
|
government_path_latest = (
|
|
Path(save_dir) / Path("saved_models") / Path("government_policy_latest.pt")
|
|
)
|
|
|
|
os.makedirs(government_path.parent, exist_ok=True)
|
|
torch.save(government_policy.state_dict(), government_path)
|
|
torch.save(government_policy.state_dict(), government_path_latest)
|
|
|
|
|
|
class ConsumerFirmRunManagerBatchParallel:
|
|
"""
|
|
The Real Business Cycle Experiment Management Class.
|
|
"""
|
|
|
|
def __init__(self, cfg_dict, freeze_firms=None, freeze_govt=None):
|
|
self.cfg_dict = cfg_dict
|
|
self.train_dict = cfg_dict["train"]
|
|
self.agents_dict = cfg_dict["agents"]
|
|
self.world_dict = cfg_dict["world"]
|
|
self.save_dense_every = self.train_dict["save_dense_every"]
|
|
self.save_dir = self.train_dict["save_dir"]
|
|
|
|
self.freeze_firms = freeze_firms
|
|
self.freeze_govt = freeze_govt
|
|
|
|
self.__init_cuda_functions()
|
|
self.__init_cuda_data_structs()
|
|
self.__init_torch_data()
|
|
|
|
def __init_cuda_data_structs(self):
|
|
__td = self.train_dict
|
|
__ad = self.agents_dict
|
|
__wd = self.world_dict
|
|
batch_size = __td["batch_size"]
|
|
num_consumers = __ad["num_consumers"]
|
|
num_firms = __ad["num_firms"]
|
|
num_governments = __ad["num_governments"]
|
|
firm_action_dim = __ad["firm_action_dim"]
|
|
government_action_dim = __ad["government_action_dim"]
|
|
consumer_state_dim = __ad["consumer_state_dim"]
|
|
firm_state_dim = __ad["firm_state_dim"]
|
|
government_state_dim = __ad["government_state_dim"]
|
|
global_state_dim = __ad["global_state_dim"]
|
|
consumer_endowment = __wd["initial_consumer_endowment"]
|
|
firm_endowment = __wd["initial_firm_endowment"]
|
|
initial_stocks = __wd["initial_stocks"]
|
|
initial_wages = __wd["initial_wages"]
|
|
initial_prices = __wd["initial_prices"]
|
|
consumer_theta = __wd["consumer_theta"]
|
|
|
|
consumer_rewards = np.zeros((batch_size, num_consumers), dtype=_NP_DTYPE)
|
|
consumer_states = np.zeros(
|
|
(batch_size, num_consumers, consumer_state_dim), dtype=_NP_DTYPE
|
|
)
|
|
|
|
firm_action_indices = np.zeros((batch_size, num_firms), dtype=np.int32)
|
|
firm_actions = np.zeros(
|
|
(batch_size, num_firms, firm_action_dim), dtype=_NP_DTYPE
|
|
)
|
|
firm_rewards = np.zeros((batch_size, num_firms), dtype=_NP_DTYPE)
|
|
firm_states = np.zeros((batch_size, num_firms, firm_state_dim), dtype=_NP_DTYPE)
|
|
|
|
government_action_indices = np.zeros(
|
|
(batch_size, num_governments), dtype=np.int32
|
|
)
|
|
government_actions = np.zeros(
|
|
(batch_size, num_governments, government_action_dim), dtype=_NP_DTYPE
|
|
)
|
|
government_rewards = np.zeros((batch_size, num_governments), dtype=_NP_DTYPE)
|
|
government_states = np.zeros(
|
|
(batch_size, num_governments, government_state_dim), dtype=_NP_DTYPE
|
|
)
|
|
|
|
# initialize states to right values here
|
|
|
|
# global state init
|
|
# for consumers, firms, and governments
|
|
for state_arr in [consumer_states, firm_states, government_states]:
|
|
# set prices to 1.0
|
|
state_arr[:, :, 0:num_firms] = initial_prices
|
|
# set wages to 0.0
|
|
state_arr[:, :, num_firms : (2 * num_firms)] = initial_wages
|
|
# set stocks to 0.0
|
|
state_arr[:, :, (2 * num_firms) : (3 * num_firms)] = initial_stocks
|
|
# set goods overdemanded to 0.0
|
|
state_arr[:, :, (3 * num_firms) : (4 * num_firms)] = 0.0
|
|
# set taxes to 0.0
|
|
state_arr[:, :, (4 * num_firms)] = 0.0
|
|
state_arr[:, :, (4 * num_firms) + 1] = 0.0
|
|
|
|
# consumer states, set theta and initial budget
|
|
if "paretoscaletheta" in __wd:
|
|
pareto_vals = np.expand_dims(
|
|
scipy.stats.pareto.ppf(
|
|
(np.arange(num_consumers) / num_consumers), __wd["paretoscaletheta"]
|
|
),
|
|
axis=0,
|
|
)
|
|
consumer_states[:, :, consumer_state_dim - 1] = consumer_theta * (
|
|
1.0 / pareto_vals
|
|
)
|
|
else:
|
|
consumer_states[:, :, consumer_state_dim - 1] = consumer_theta
|
|
consumer_states[:, :, global_state_dim] = consumer_endowment
|
|
|
|
# firm states
|
|
# capital
|
|
if __wd.get("initial_capital", None) == "proportional":
|
|
for i in range(num_firms):
|
|
firm_states[:, i, global_state_dim + 1] = ((i + 1) / 10.0) * 2.0
|
|
elif __wd.get("initial_capital", None) == "twolevel":
|
|
for i in range(num_firms):
|
|
if i < (num_firms // 2):
|
|
firm_states[:, i, global_state_dim + 1] = 5000
|
|
else:
|
|
firm_states[:, i, global_state_dim + 1] = 10000
|
|
else:
|
|
firm_states[:, :, global_state_dim + 1] = 1.0
|
|
|
|
# production alpha
|
|
if __wd["production_alpha"] == "proportional":
|
|
half_firms = num_firms // 2
|
|
for i in range(num_firms):
|
|
firm_states[:, i, global_state_dim + 2] = ((i % half_firms) + 1) * 0.2
|
|
elif __wd["production_alpha"] == "fixed_array":
|
|
alpha_arr = [0.2, 0.3, 0.4, 0.6, 0.8, 0.2, 0.3, 0.4, 0.6, 0.8]
|
|
for i in range(num_firms):
|
|
firm_states[:, i, global_state_dim + 2] = alpha_arr[i]
|
|
else:
|
|
for i in range(num_firms):
|
|
firm_states[:, i, global_state_dim + 2] = __wd["production_alpha"]
|
|
|
|
# set one-hot fields correctly by index for each firm
|
|
onehot_rows = np.eye(num_firms)
|
|
firm_states[:, :, (global_state_dim + 3) :] = onehot_rows
|
|
firm_states[:, :, global_state_dim] = firm_endowment
|
|
|
|
# government states
|
|
# for now, nothing beyond global state
|
|
|
|
self.consumer_states_gpu_tensor = torch.from_numpy(consumer_states).cuda()
|
|
# these are now tensors bc sampling for consumers via pytorch
|
|
self.consumer_rewards_gpu_pycuda = cuda_driver.mem_alloc(
|
|
consumer_rewards.nbytes
|
|
)
|
|
self.consumer_states_checkpoint_gpu_pycuda = cuda_driver.mem_alloc(
|
|
consumer_states.nbytes
|
|
)
|
|
cuda_driver.memcpy_htod(self.consumer_rewards_gpu_pycuda, consumer_rewards)
|
|
cuda_driver.memcpy_htod(
|
|
self.consumer_states_checkpoint_gpu_pycuda, consumer_states
|
|
)
|
|
|
|
self.firm_states_gpu_tensor = torch.from_numpy(firm_states).cuda()
|
|
self.firm_action_indices_gpu_pycuda = cuda_driver.mem_alloc(
|
|
firm_action_indices.nbytes
|
|
)
|
|
self.firm_actions_gpu_pycuda = cuda_driver.mem_alloc(firm_actions.nbytes)
|
|
self.firm_rewards_gpu_pycuda = cuda_driver.mem_alloc(firm_rewards.nbytes)
|
|
self.firm_states_checkpoint_gpu_pycuda = cuda_driver.mem_alloc(
|
|
firm_states.nbytes
|
|
)
|
|
cuda_driver.memcpy_htod(
|
|
self.firm_action_indices_gpu_pycuda, firm_action_indices
|
|
)
|
|
cuda_driver.memcpy_htod(self.firm_actions_gpu_pycuda, firm_actions)
|
|
cuda_driver.memcpy_htod(self.firm_rewards_gpu_pycuda, firm_rewards)
|
|
cuda_driver.memcpy_htod(self.firm_states_checkpoint_gpu_pycuda, firm_states)
|
|
|
|
self.government_states_gpu_tensor = torch.from_numpy(government_states).cuda()
|
|
self.government_action_indices_gpu_pycuda = cuda_driver.mem_alloc(
|
|
government_action_indices.nbytes
|
|
)
|
|
self.government_actions_gpu_pycuda = cuda_driver.mem_alloc(
|
|
government_actions.nbytes
|
|
)
|
|
self.government_rewards_gpu_pycuda = cuda_driver.mem_alloc(
|
|
government_rewards.nbytes
|
|
)
|
|
self.government_states_checkpoint_gpu_pycuda = cuda_driver.mem_alloc(
|
|
government_states.nbytes
|
|
)
|
|
cuda_driver.memcpy_htod(
|
|
self.government_action_indices_gpu_pycuda, government_action_indices
|
|
)
|
|
cuda_driver.memcpy_htod(self.government_actions_gpu_pycuda, government_actions)
|
|
cuda_driver.memcpy_htod(self.government_rewards_gpu_pycuda, government_rewards)
|
|
cuda_driver.memcpy_htod(
|
|
self.government_states_checkpoint_gpu_pycuda, government_states
|
|
)
|
|
|
|
def __init_torch_data(self):
|
|
|
|
__td = self.train_dict
|
|
__ad = self.agents_dict
|
|
|
|
batch_size = __td["batch_size"]
|
|
num_consumers = __ad["num_consumers"]
|
|
num_firms = __ad["num_firms"]
|
|
num_governments = __ad["num_governments"]
|
|
consumer_action_dim = __ad["consumer_action_dim"]
|
|
consumer_state_dim = __ad["consumer_state_dim"]
|
|
firm_state_dim = __ad["firm_state_dim"]
|
|
government_state_dim = __ad["government_state_dim"]
|
|
num_iters = int(self.world_dict["maxtime"])
|
|
|
|
consumer_states_batch = torch.zeros(
|
|
batch_size,
|
|
num_iters,
|
|
num_consumers,
|
|
consumer_state_dim,
|
|
dtype=torch.float32,
|
|
device="cpu",
|
|
)
|
|
consumer_actions_single = torch.zeros(
|
|
batch_size,
|
|
num_consumers,
|
|
num_firms + 1 + 1,
|
|
dtype=torch.int32,
|
|
device="cpu",
|
|
)
|
|
consumer_actions_batch = torch.zeros(
|
|
batch_size,
|
|
num_iters,
|
|
num_consumers,
|
|
num_firms + 1 + 1,
|
|
dtype=torch.int32,
|
|
device="cpu",
|
|
)
|
|
|
|
# auxiliary state info that is not part of observables.
|
|
# currently just the realized consumption
|
|
consumer_aux_batch = torch.zeros(
|
|
batch_size,
|
|
num_iters,
|
|
num_consumers,
|
|
num_firms,
|
|
dtype=torch.float32,
|
|
device="cpu",
|
|
)
|
|
|
|
consumer_rewards_batch = torch.zeros(
|
|
batch_size, num_iters, num_consumers, dtype=torch.float32, device="cpu"
|
|
)
|
|
self.consumer_states_batch_gpu_tensor = consumer_states_batch.cuda()
|
|
self.consumer_actions_batch_gpu_tensor = consumer_actions_batch.cuda()
|
|
self.consumer_actions_index_single_gpu_tensor = consumer_actions_single.cuda()
|
|
self.consumer_actions_single_gpu_tensor = torch.zeros(
|
|
batch_size,
|
|
num_consumers,
|
|
consumer_action_dim,
|
|
dtype=torch.float32,
|
|
device="cpu",
|
|
).cuda()
|
|
self.consumer_rewards_batch_gpu_tensor = consumer_rewards_batch.cuda()
|
|
self.consumer_aux_batch_gpu_tensor = consumer_aux_batch.cuda()
|
|
|
|
firm_states_batch = torch.zeros(
|
|
batch_size,
|
|
num_iters,
|
|
num_firms,
|
|
firm_state_dim,
|
|
dtype=torch.float32,
|
|
device="cpu",
|
|
)
|
|
firm_actions_batch = torch.zeros(
|
|
batch_size, num_iters, num_firms, dtype=torch.int32, device="cpu"
|
|
)
|
|
firm_rewards_batch = torch.zeros(
|
|
batch_size, num_iters, num_firms, dtype=torch.float32, device="cpu"
|
|
)
|
|
firm_aux_batch = torch.zeros(
|
|
batch_size, num_iters, num_firms, dtype=torch.float32, device="cpu"
|
|
)
|
|
self.firm_states_batch = firm_states_batch.cuda()
|
|
self.firm_actions_batch = firm_actions_batch.cuda()
|
|
self.firm_rewards_batch = firm_rewards_batch.cuda()
|
|
self.firm_aux_batch = firm_aux_batch.cuda()
|
|
|
|
government_states_batch = torch.zeros(
|
|
batch_size,
|
|
num_iters,
|
|
num_governments,
|
|
government_state_dim,
|
|
dtype=torch.float32,
|
|
device="cpu",
|
|
)
|
|
government_actions_batch = torch.zeros(
|
|
batch_size, num_iters, num_governments, dtype=torch.int32, device="cpu"
|
|
)
|
|
government_rewards_batch = torch.zeros(
|
|
batch_size, num_iters, num_governments, dtype=torch.float32, device="cpu"
|
|
)
|
|
self.government_states_batch = government_states_batch.cuda()
|
|
self.government_actions_batch = government_actions_batch.cuda()
|
|
self.government_rewards_batch = government_rewards_batch.cuda()
|
|
|
|
def __init_cuda_functions(self):
|
|
|
|
__td = self.train_dict
|
|
__ad = self.agents_dict
|
|
__wd = self.world_dict
|
|
|
|
if self.freeze_firms is not None:
|
|
countfirmreward = 0
|
|
else:
|
|
countfirmreward = self.agents_dict["government_counts_firm_reward"]
|
|
|
|
code, compiler_options = get_cuda_code(
|
|
Path("cuda") / Path("firm_rbc.cu"),
|
|
batchsize=__td["batch_size"],
|
|
numconsumers=__ad["num_consumers"],
|
|
numfirms=__ad["num_firms"],
|
|
numgovernments=__ad["num_governments"],
|
|
maxtime=__wd["maxtime"],
|
|
# numactionsconsumer=__ad["consumer_num_actions"],
|
|
numactionsconsumer=__ad["consumer_num_work_actions"],
|
|
numactionsfirm=__ad["firm_num_actions"],
|
|
numactionsgovernment=__ad["government_num_actions"],
|
|
interestrate=__wd["interest_rate"],
|
|
crra_param=__wd["crra_param"],
|
|
shouldboostfirmreward=int(__td["should_boost_firm_reward"]),
|
|
boostfirmrewardfactor=__td["boost_firm_reward_factor"],
|
|
countfirmreward=countfirmreward,
|
|
importerprice=__wd["importer_price"],
|
|
importerquantity=__wd["importer_quantity"],
|
|
laborfloor=__wd.get("labor_floor", 0.0),
|
|
useimporter=__wd["use_importer"],
|
|
)
|
|
|
|
mod = SourceModule(code, options=compiler_options, no_extern_c=True)
|
|
self.mod = mod
|
|
|
|
# --------------------------------------------------------------------
|
|
# Define Consumer actions -- maanged in Pytorch
|
|
# --------------------------------------------------------------------
|
|
self.consumption_action_tensor = torch.tensor(
|
|
__ad["consumer_consumption_actions_array"].astype(_NP_DTYPE)
|
|
).cuda()
|
|
self.work_action_tensor = torch.tensor(
|
|
__ad["consumer_work_actions_array"].astype(_NP_DTYPE)
|
|
).cuda()
|
|
|
|
# --------------------------------------------------------------------
|
|
# Define Firm actions -- maanged in CUDA
|
|
# --------------------------------------------------------------------
|
|
firm_index_to_action_gpu, _ = mod.get_global("kFirmIndexToAction")
|
|
cuda_driver.memcpy_htod(
|
|
firm_index_to_action_gpu,
|
|
__ad["firm_actions_array"].astype(_NP_DTYPE),
|
|
)
|
|
|
|
# --------------------------------------------------------------------
|
|
# Define Govt actions -- maanged in CUDA
|
|
# --------------------------------------------------------------------
|
|
government_index_to_action_gpu, _ = mod.get_global("kGovernmentIndexToAction")
|
|
cuda_driver.memcpy_htod(
|
|
government_index_to_action_gpu,
|
|
__ad["government_actions_array"].astype(_NP_DTYPE),
|
|
)
|
|
|
|
# --------------------------------------------------------------------
|
|
# Get handles to CUDA methods
|
|
# --------------------------------------------------------------------
|
|
self.cuda_init_random = mod.get_function("CudaInitKernel")
|
|
self.cuda_reset_env = mod.get_function("CudaResetEnv")
|
|
self.cuda_sample_actions = mod.get_function(
|
|
"CudaSampleFirmAndGovernmentActions"
|
|
)
|
|
self.cuda_step = mod.get_function("CudaStep")
|
|
self.cuda_free_mem = mod.get_function("CudaFreeRand")
|
|
|
|
def _update_consumer_actions_inplace(self):
|
|
# call after consumer_actions_single is updated
|
|
__ad = self.agents_dict
|
|
|
|
# Add asserts when ``loading'' arrays
|
|
# assert consumption_action_array.shape == (1, 1, 1)
|
|
# assert len(consumption_action_array.shape) == 3
|
|
|
|
num_firms = __ad["num_firms"]
|
|
idx_hours = num_firms
|
|
idx_which_firm = num_firms + 1
|
|
for i in range(num_firms):
|
|
|
|
consumption_actions_at_firm_i = (
|
|
self.consumer_actions_index_single_gpu_tensor[..., i].to(torch.long)
|
|
)
|
|
|
|
self.consumer_actions_single_gpu_tensor[
|
|
..., i
|
|
] = self.consumption_action_tensor[
|
|
consumption_actions_at_firm_i, :
|
|
].squeeze(
|
|
dim=-1
|
|
)
|
|
|
|
consumer_hours_worked = self.consumer_actions_index_single_gpu_tensor[
|
|
..., idx_hours
|
|
].to(torch.long)
|
|
|
|
self.consumer_actions_single_gpu_tensor[
|
|
..., num_firms
|
|
] = self.work_action_tensor[consumer_hours_worked, :].squeeze(dim=-1)
|
|
|
|
self.consumer_actions_single_gpu_tensor[
|
|
..., num_firms + 1
|
|
] = self.consumer_actions_index_single_gpu_tensor[..., idx_which_firm]
|
|
|
|
def sample_consumer_actions_and_store(self, consumer_probs_list):
|
|
# Every consumer has A action heads, output as a list of tensors.
|
|
# Sample from each of these lists and store the results.
|
|
|
|
with torch.no_grad():
|
|
for i, probs in enumerate(consumer_probs_list):
|
|
dist = Categorical(probs)
|
|
samples = dist.sample()
|
|
self.consumer_actions_index_single_gpu_tensor[..., i] = samples
|
|
|
|
self._update_consumer_actions_inplace()
|
|
|
|
def consumers_will_train_this_episode(self, epi):
|
|
__ad = self.agents_dict
|
|
if "training_schedule_mod" in self.agents_dict:
|
|
mod_val = epi % __ad["training_schedule_mod"]
|
|
return mod_val <= __ad["consumer_mod_threshold"]
|
|
if "consumer_training_list" in self.agents_dict:
|
|
return interval_list_contains(__ad["consumer_training_list"], epi)
|
|
if "train_consumers_every" in self.agents_dict:
|
|
mod_val = epi % __ad["train_consumers_every"]
|
|
else:
|
|
mod_val = 0
|
|
return epi >= self.agents_dict.get("consumer_training_start", 0) and (
|
|
mod_val == 0
|
|
)
|
|
|
|
def firms_will_train_this_episode(self, epi):
|
|
__ad = self.agents_dict
|
|
if "training_schedule_mod" in self.agents_dict:
|
|
mod_val = epi % __ad["training_schedule_mod"]
|
|
return mod_val > __ad["consumer_mod_threshold"]
|
|
if "firm_training_list" in self.agents_dict:
|
|
return interval_list_contains(__ad["firm_training_list"], epi) and (
|
|
self.freeze_firms is None
|
|
)
|
|
if "train_firms_every" in self.agents_dict:
|
|
mod_val = epi % __ad["train_firms_every"]
|
|
else:
|
|
mod_val = 0
|
|
return (
|
|
(epi >= self.agents_dict.get("firm_training_start", 0))
|
|
and (self.freeze_firms is None)
|
|
and (mod_val == 0)
|
|
)
|
|
|
|
def governments_will_train_this_episode(self, epi):
|
|
__ad = self.agents_dict
|
|
if "government_training_list" in self.agents_dict:
|
|
return interval_list_contains(__ad["government_training_list"], epi) and (
|
|
self.freeze_govt is None
|
|
)
|
|
if "train_government_every" in self.agents_dict:
|
|
mod_val = epi % self.agents_dict["train_government_every"]
|
|
else:
|
|
mod_val = 0
|
|
return (
|
|
(epi >= self.agents_dict.get("government_training_start", 0))
|
|
and (self.freeze_govt is None)
|
|
and (mod_val == 0)
|
|
)
|
|
|
|
def bestresponse_train(
|
|
self, train_type, num_episodes, rollout_path, ep_str="latest", checkpoint=100
|
|
):
|
|
# train one single type only
|
|
# load all policies from state dict
|
|
# reset all the environment stuff
|
|
|
|
__td = self.train_dict
|
|
__ad = self.agents_dict
|
|
num_iters = int(self.world_dict["maxtime"])
|
|
num_consumers = __ad["num_consumers"]
|
|
num_firms = __ad["num_firms"]
|
|
num_governments = __ad["num_governments"]
|
|
num_agents = num_consumers + num_firms + num_governments
|
|
block = (num_agents, 1, 1)
|
|
grid = (__td["batch_size"], 1)
|
|
|
|
seed_everything(__td["seed"])
|
|
self.cuda_init_random(np.int32(__td["seed"]), block=block, grid=grid)
|
|
|
|
# --------------------------------------------
|
|
# Define Consumer policy + optimizers
|
|
# --------------------------------------------
|
|
lr = __td["lr"]
|
|
|
|
consumer_expanded_size = size_after_digit_expansion(
|
|
__ad["consumer_state_dim"],
|
|
__ad["consumer_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
)
|
|
|
|
consumer_policy = IndependentPolicyNet(
|
|
consumer_expanded_size,
|
|
[__ad["consumer_num_consume_actions"]] * num_firms
|
|
+ [
|
|
__ad["consumer_num_work_actions"],
|
|
__ad["consumer_num_whichfirm_actions"],
|
|
],
|
|
norm_consts=(
|
|
torch.zeros(consumer_expanded_size).cuda(), # don't center for now
|
|
consumer_state_scaling_factors(self.cfg_dict),
|
|
),
|
|
).to("cuda")
|
|
consumer_policy.load_state_dict(
|
|
torch.load(
|
|
rollout_path
|
|
/ Path("saved_models")
|
|
/ Path(f"consumer_policy_{ep_str}.pt")
|
|
)
|
|
)
|
|
|
|
consumer_optim = torch.optim.Adam(consumer_policy.parameters(), lr=lr)
|
|
firm_expanded_size = size_after_digit_expansion(
|
|
__ad["firm_state_dim"],
|
|
__ad["firm_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
)
|
|
firm_policy = PolicyNet(
|
|
firm_expanded_size,
|
|
__ad["firm_num_actions"],
|
|
norm_consts=(
|
|
torch.zeros(firm_expanded_size).cuda(),
|
|
firm_state_scaling_factors(self.cfg_dict),
|
|
),
|
|
).to("cuda")
|
|
|
|
firm_policy.load_state_dict(
|
|
torch.load(
|
|
rollout_path / Path("saved_models") / Path(f"firm_policy_{ep_str}.pt")
|
|
)
|
|
)
|
|
|
|
firm_optim = torch.optim.Adam(firm_policy.parameters(), lr=lr)
|
|
government_expanded_size = size_after_digit_expansion(
|
|
__ad["government_state_dim"],
|
|
__ad["government_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
)
|
|
|
|
government_policy = PolicyNet(
|
|
government_expanded_size,
|
|
__ad["government_num_actions"],
|
|
norm_consts=(
|
|
torch.zeros(government_expanded_size).cuda(),
|
|
govt_state_scaling_factors(self.cfg_dict),
|
|
),
|
|
).to("cuda")
|
|
|
|
government_policy.load_state_dict(
|
|
torch.load(
|
|
rollout_path
|
|
/ Path("saved_models")
|
|
/ Path(f"government_policy_{ep_str}.pt")
|
|
)
|
|
)
|
|
|
|
government_optim = torch.optim.Adam(government_policy.parameters(), lr=lr)
|
|
rewards = []
|
|
|
|
agent_type_arrays = {
|
|
"consumer": (
|
|
self.consumer_states_batch_gpu_tensor,
|
|
self.consumer_actions_batch_gpu_tensor,
|
|
self.consumer_rewards_batch_gpu_tensor,
|
|
),
|
|
"firm": (
|
|
self.firm_states_batch,
|
|
self.firm_actions_batch,
|
|
self.firm_rewards_batch,
|
|
),
|
|
"government": (
|
|
self.government_states_batch,
|
|
self.government_actions_batch,
|
|
self.government_rewards_batch,
|
|
),
|
|
}
|
|
|
|
agent_action_arrays = {
|
|
"consumer": __ad["consumer_work_actions_array"],
|
|
"firm": __ad["firm_actions_array"],
|
|
"government": __ad["government_actions_array"],
|
|
}
|
|
|
|
agent_aux_arrays = {
|
|
"consumer": (self.consumer_aux_batch_gpu_tensor),
|
|
"firm": (self.firm_aux_batch),
|
|
"government": None,
|
|
}
|
|
|
|
pbar = tqdm(range(num_episodes))
|
|
for epi in pbar:
|
|
annealed_entropy_coef = 0.1 # later, do some computation to anneal this
|
|
self.cuda_reset_env(
|
|
CudaTensorHolder(self.consumer_states_gpu_tensor),
|
|
CudaTensorHolder(self.firm_states_gpu_tensor),
|
|
CudaTensorHolder(self.government_states_gpu_tensor),
|
|
self.consumer_states_checkpoint_gpu_pycuda,
|
|
self.firm_states_checkpoint_gpu_pycuda,
|
|
self.government_states_checkpoint_gpu_pycuda,
|
|
np.float32(1.0),
|
|
block=block,
|
|
grid=grid,
|
|
)
|
|
|
|
for _iter in range(num_iters):
|
|
|
|
# ------------------------
|
|
# Run policy and get probs
|
|
# ------------------------
|
|
with torch.no_grad():
|
|
# here, we must perform digit scaling
|
|
consumer_probs_list, _ = consumer_policy(
|
|
expand_to_digit_form(
|
|
self.consumer_states_gpu_tensor,
|
|
__ad["consumer_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
)
|
|
)
|
|
firm_probs, _ = firm_policy(
|
|
expand_to_digit_form(
|
|
self.firm_states_gpu_tensor,
|
|
__ad["firm_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
actions_mask=None,
|
|
)
|
|
government_probs, _ = government_policy(
|
|
expand_to_digit_form(
|
|
self.government_states_gpu_tensor,
|
|
__ad["government_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
actions_mask=None,
|
|
)
|
|
|
|
# ------------------------
|
|
# Get action samples
|
|
# ------------------------
|
|
# Sample consumer actions using PyTorch here on GPU!
|
|
self.sample_consumer_actions_and_store(consumer_probs_list)
|
|
|
|
# Sample firms + govt actions using PyCUDA on GPU!
|
|
self.cuda_sample_actions(
|
|
CudaTensorHolder(firm_probs),
|
|
self.firm_action_indices_gpu_pycuda,
|
|
self.firm_actions_gpu_pycuda,
|
|
CudaTensorHolder(government_probs),
|
|
self.government_action_indices_gpu_pycuda,
|
|
self.government_actions_gpu_pycuda,
|
|
block=block,
|
|
grid=grid,
|
|
)
|
|
|
|
# ------------------------
|
|
# Step on GPU
|
|
# ------------------------
|
|
self.cuda_step(
|
|
CudaTensorHolder(
|
|
# size: batches x n_consumers x consumer_state float
|
|
self.consumer_states_gpu_tensor
|
|
),
|
|
CudaTensorHolder(
|
|
# size: batches x n_consumers x consumer_action_dim float
|
|
self.consumer_actions_single_gpu_tensor
|
|
),
|
|
# size: batches x n_consumers x 1 float
|
|
self.consumer_rewards_gpu_pycuda,
|
|
CudaTensorHolder(
|
|
self.consumer_states_batch_gpu_tensor
|
|
), # size: batches x episode x n_consumers x consumer_state float
|
|
CudaTensorHolder(self.consumer_rewards_batch_gpu_tensor),
|
|
CudaTensorHolder(self.firm_states_gpu_tensor),
|
|
self.firm_action_indices_gpu_pycuda,
|
|
self.firm_actions_gpu_pycuda,
|
|
self.firm_rewards_gpu_pycuda,
|
|
CudaTensorHolder(self.firm_states_batch),
|
|
CudaTensorHolder(self.firm_actions_batch),
|
|
CudaTensorHolder(self.firm_rewards_batch),
|
|
CudaTensorHolder(self.government_states_gpu_tensor),
|
|
self.government_action_indices_gpu_pycuda,
|
|
self.government_actions_gpu_pycuda,
|
|
self.government_rewards_gpu_pycuda,
|
|
CudaTensorHolder(self.government_states_batch),
|
|
CudaTensorHolder(self.government_actions_batch),
|
|
CudaTensorHolder(self.government_rewards_batch),
|
|
CudaTensorHolder(self.consumer_aux_batch_gpu_tensor),
|
|
CudaTensorHolder(self.firm_aux_batch),
|
|
np.int32(_iter),
|
|
block=block,
|
|
grid=grid,
|
|
)
|
|
self.consumer_actions_batch_gpu_tensor[
|
|
:, _iter, :, :
|
|
] = self.consumer_actions_index_single_gpu_tensor
|
|
update_government_rewards(
|
|
self.government_rewards_batch,
|
|
self.consumer_rewards_batch_gpu_tensor,
|
|
self.firm_rewards_batch,
|
|
self.cfg_dict,
|
|
)
|
|
if train_type == "consumer":
|
|
consumer_reward_scale = self.agents_dict.get(
|
|
"consumer_reward_scale", 1.0
|
|
)
|
|
consumer_policy_gradient_step(
|
|
consumer_policy,
|
|
expand_to_digit_form(
|
|
self.consumer_states_batch_gpu_tensor,
|
|
__ad["consumer_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
self.consumer_actions_batch_gpu_tensor,
|
|
self.consumer_rewards_batch_gpu_tensor,
|
|
consumer_optim,
|
|
__td["gamma"],
|
|
entropy_val=annealed_entropy_coef * __td["entropy"],
|
|
value_loss_weight=__td["value_loss_weight"],
|
|
reward_scale=consumer_reward_scale,
|
|
clip_grad_norm=self.train_dict.get("clip_grad_norm", None),
|
|
)
|
|
rewards.append(self.consumer_rewards_batch_gpu_tensor.mean().item())
|
|
elif train_type == "firm":
|
|
firm_reward_scale = self.agents_dict.get("firm_reward_scale", 1.0)
|
|
policy_gradient_step(
|
|
firm_policy,
|
|
expand_to_digit_form(
|
|
self.firm_states_batch,
|
|
__ad["firm_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
self.firm_actions_batch,
|
|
self.firm_rewards_batch,
|
|
firm_optim,
|
|
__td["gamma"],
|
|
entropy_val=annealed_entropy_coef * __td["entropy"],
|
|
value_loss_weight=__td["value_loss_weight"],
|
|
actions_mask=None,
|
|
reward_scale=firm_reward_scale,
|
|
clip_grad_norm=self.train_dict.get("clip_grad_norm", None),
|
|
)
|
|
rewards.append(self.firm_rewards_batch.mean().item())
|
|
elif train_type == "government":
|
|
government_reward_scale = self.agents_dict.get(
|
|
"government_reward_scale", 1.0
|
|
)
|
|
policy_gradient_step(
|
|
government_policy,
|
|
expand_to_digit_form(
|
|
self.government_states_batch,
|
|
__ad["government_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
self.government_actions_batch,
|
|
self.government_rewards_batch,
|
|
government_optim,
|
|
__td["gamma"],
|
|
entropy_val=annealed_entropy_coef * __td["entropy"],
|
|
value_loss_weight=__td["value_loss_weight"],
|
|
actions_mask=None,
|
|
reward_scale=government_reward_scale,
|
|
clip_grad_norm=self.train_dict.get("clip_grad_norm", None),
|
|
)
|
|
rewards.append(self.government_rewards_batch.mean().item())
|
|
pbar.set_postfix({"reward": rewards[-1]})
|
|
if (epi % checkpoint) == 0:
|
|
# save policy every checkpoint steps
|
|
save_policy_parameters(
|
|
str(Path(self.save_dir) / f"br{train_type}"),
|
|
epi,
|
|
consumer_policy,
|
|
firm_policy,
|
|
government_policy,
|
|
self.freeze_firms,
|
|
self.freeze_govt,
|
|
)
|
|
save_dense_log(
|
|
str(Path(self.save_dir) / f"br{train_type}"),
|
|
epi,
|
|
agent_type_arrays,
|
|
agent_action_arrays,
|
|
agent_aux_arrays,
|
|
)
|
|
|
|
print(
|
|
f"{train_type}: starting reward {rewards[0]}, "
|
|
f"ending reward {rewards[-1]}, "
|
|
f"improvement in reward after {num_episodes}: {rewards[-1] - rewards[0]}"
|
|
)
|
|
|
|
self.cuda_free_mem(block=block, grid=grid)
|
|
return rewards
|
|
|
|
def train(self):
|
|
|
|
__td = self.train_dict
|
|
__ad = self.agents_dict
|
|
|
|
# Create logdir
|
|
os.makedirs(__td["save_dir"], exist_ok=True)
|
|
|
|
# Constants
|
|
num_iters = int(self.world_dict["maxtime"])
|
|
num_consumers = __ad["num_consumers"]
|
|
num_firms = __ad["num_firms"]
|
|
num_governments = __ad["num_governments"]
|
|
num_agents = num_consumers + num_firms + num_governments
|
|
|
|
# CUDA params: defines data shape on the GPU
|
|
block = (num_agents, 1, 1)
|
|
grid = (__td["batch_size"], 1)
|
|
|
|
# Set seeds
|
|
seed_everything(__td["seed"])
|
|
self.cuda_init_random(np.int32(__td["seed"]), block=block, grid=grid)
|
|
|
|
# --------------------------------------------
|
|
# Define Consumer policy + optimizers
|
|
# --------------------------------------------
|
|
lr = __td["lr"]
|
|
|
|
consumer_expanded_size = size_after_digit_expansion(
|
|
__ad["consumer_state_dim"],
|
|
__ad["consumer_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
)
|
|
|
|
consumer_policy = IndependentPolicyNet(
|
|
consumer_expanded_size,
|
|
[__ad["consumer_num_consume_actions"]] * num_firms
|
|
+ [
|
|
__ad["consumer_num_work_actions"],
|
|
__ad["consumer_num_whichfirm_actions"],
|
|
],
|
|
norm_consts=(
|
|
torch.zeros(consumer_expanded_size).cuda(), # don't center for now
|
|
consumer_state_scaling_factors(self.cfg_dict),
|
|
),
|
|
).to("cuda")
|
|
|
|
consumer_optim = torch.optim.Adam(
|
|
consumer_policy.parameters(),
|
|
lr=lr * self.agents_dict.get("consumer_lr_multiple", 1.0),
|
|
)
|
|
|
|
# --------------------------------------------
|
|
# Define Firm policy + optimizers
|
|
# --------------------------------------------
|
|
firm_expanded_size = size_after_digit_expansion(
|
|
__ad["firm_state_dim"],
|
|
__ad["firm_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
)
|
|
|
|
if self.freeze_firms is not None:
|
|
firm_policy = DeterministicPolicy(
|
|
firm_expanded_size,
|
|
__ad["firm_num_actions"],
|
|
self.freeze_firms,
|
|
)
|
|
firm_optim = NoOpOptimizer()
|
|
else:
|
|
firm_policy = PolicyNet(
|
|
firm_expanded_size,
|
|
__ad["firm_num_actions"],
|
|
norm_consts=(
|
|
torch.zeros(firm_expanded_size).cuda(),
|
|
firm_state_scaling_factors(self.cfg_dict),
|
|
),
|
|
).to("cuda")
|
|
|
|
firm_optim = torch.optim.Adam(
|
|
firm_policy.parameters(),
|
|
lr=lr * self.agents_dict.get("firm_lr_multiple", 1.0),
|
|
)
|
|
|
|
# --------------------------------------------
|
|
# Define Government policy + optimizers
|
|
# --------------------------------------------
|
|
government_expanded_size = size_after_digit_expansion(
|
|
__ad["government_state_dim"],
|
|
__ad["government_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
)
|
|
|
|
if self.freeze_govt is not None:
|
|
government_policy = DeterministicPolicy(
|
|
government_expanded_size,
|
|
__ad["government_num_actions"],
|
|
self.freeze_govt,
|
|
)
|
|
government_optim = NoOpOptimizer()
|
|
else:
|
|
government_policy = PolicyNet(
|
|
government_expanded_size,
|
|
__ad["government_num_actions"],
|
|
norm_consts=(
|
|
torch.zeros(government_expanded_size).cuda(),
|
|
govt_state_scaling_factors(self.cfg_dict),
|
|
),
|
|
).to("cuda")
|
|
government_optim = torch.optim.Adam(
|
|
government_policy.parameters(),
|
|
lr=lr * self.agents_dict.get("government_lr_multiple", 1.0),
|
|
)
|
|
|
|
# --------------------------------------------
|
|
# Logging
|
|
# --------------------------------------------
|
|
# For looking up GPU tensors
|
|
# --------------------------------------------
|
|
agent_type_arrays = {
|
|
"consumer": (
|
|
self.consumer_states_batch_gpu_tensor,
|
|
self.consumer_actions_batch_gpu_tensor,
|
|
self.consumer_rewards_batch_gpu_tensor,
|
|
),
|
|
"firm": (
|
|
self.firm_states_batch,
|
|
self.firm_actions_batch,
|
|
self.firm_rewards_batch,
|
|
),
|
|
"government": (
|
|
self.government_states_batch,
|
|
self.government_actions_batch,
|
|
self.government_rewards_batch,
|
|
),
|
|
}
|
|
|
|
agent_action_arrays = {
|
|
"consumer": __ad["consumer_work_actions_array"],
|
|
"firm": __ad["firm_actions_array"],
|
|
"government": __ad["government_actions_array"],
|
|
}
|
|
|
|
agent_aux_arrays = {
|
|
"consumer": (self.consumer_aux_batch_gpu_tensor),
|
|
"firm": (self.firm_aux_batch),
|
|
"government": None,
|
|
}
|
|
|
|
# --------------------------------------------
|
|
# Training policy XYZ starts at which step?
|
|
# --------------------------------------------
|
|
firm_no_ponzi_coef = self.agents_dict.get("firm_noponzi_start", 0.0)
|
|
consumer_no_ponzi_coef = self.agents_dict.get("consumer_noponzi_start", 0.0)
|
|
lagr_num_steps = self.train_dict.get("lagr_num_steps", 1)
|
|
|
|
firm_training_start = self.agents_dict.get("firm_training_start", 0)
|
|
consumer_training_start = self.agents_dict.get("consumer_training_start", 0)
|
|
government_training_start = self.agents_dict.get("government_training_start", 0)
|
|
|
|
firm_action_start = self.agents_dict.get("firm_begin_anneal_action", 0)
|
|
government_action_start = self.agents_dict.get(
|
|
"government_begin_anneal_action", 0
|
|
)
|
|
|
|
# --------------------------------------------
|
|
# Training loop
|
|
# --------------------------------------------
|
|
if self.train_dict.get("infinite_episodes", False):
|
|
epi_iterator = itertools.count(0, 1)
|
|
else:
|
|
epi_iterator = range(__td["num_episodes"])
|
|
|
|
final_epi = None
|
|
for epi in tqdm(epi_iterator):
|
|
|
|
firm_actions_mask = firm_action_mask(
|
|
self.cfg_dict,
|
|
max(epi - firm_action_start, 0),
|
|
)
|
|
government_actions_mask = government_action_mask(
|
|
self.cfg_dict,
|
|
max(epi - government_action_start, 0),
|
|
)
|
|
theta_coef = compute_theta_coef(self.cfg_dict, epi)
|
|
|
|
# Reset environment for all agents
|
|
self.cuda_reset_env(
|
|
CudaTensorHolder(self.consumer_states_gpu_tensor),
|
|
CudaTensorHolder(self.firm_states_gpu_tensor),
|
|
CudaTensorHolder(self.government_states_gpu_tensor),
|
|
self.consumer_states_checkpoint_gpu_pycuda,
|
|
self.firm_states_checkpoint_gpu_pycuda,
|
|
self.government_states_checkpoint_gpu_pycuda,
|
|
theta_coef,
|
|
block=block,
|
|
grid=grid,
|
|
)
|
|
|
|
# Learning Loop
|
|
for _iter in range(num_iters):
|
|
|
|
# ------------------------
|
|
# Run policy and get probs
|
|
# ------------------------
|
|
with torch.no_grad():
|
|
# here, we must perform digit scaling
|
|
consumer_probs_list, _ = consumer_policy(
|
|
expand_to_digit_form(
|
|
self.consumer_states_gpu_tensor,
|
|
__ad["consumer_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
)
|
|
)
|
|
firm_probs, _ = firm_policy(
|
|
expand_to_digit_form(
|
|
self.firm_states_gpu_tensor,
|
|
__ad["firm_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
actions_mask=firm_actions_mask,
|
|
)
|
|
government_probs, _ = government_policy(
|
|
expand_to_digit_form(
|
|
self.government_states_gpu_tensor,
|
|
__ad["government_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
actions_mask=government_actions_mask,
|
|
)
|
|
|
|
# ------------------------
|
|
# Get action samples
|
|
# ------------------------
|
|
# Sample consumer actions using PyTorch here on GPU!
|
|
self.sample_consumer_actions_and_store(consumer_probs_list)
|
|
|
|
# Sample firms + govt actions using PyCUDA on GPU!
|
|
self.cuda_sample_actions(
|
|
CudaTensorHolder(firm_probs),
|
|
self.firm_action_indices_gpu_pycuda,
|
|
self.firm_actions_gpu_pycuda,
|
|
CudaTensorHolder(government_probs),
|
|
self.government_action_indices_gpu_pycuda,
|
|
self.government_actions_gpu_pycuda,
|
|
block=block,
|
|
grid=grid,
|
|
)
|
|
|
|
# ------------------------
|
|
# Step on GPU
|
|
# ------------------------
|
|
self.cuda_step(
|
|
CudaTensorHolder(
|
|
# size: batches x n_consumers x consumer_state float
|
|
self.consumer_states_gpu_tensor
|
|
),
|
|
CudaTensorHolder(
|
|
# size: batches x n_consumers x consumer_action_dim float
|
|
self.consumer_actions_single_gpu_tensor
|
|
),
|
|
# size: batches x n_consumers x 1 float
|
|
self.consumer_rewards_gpu_pycuda,
|
|
CudaTensorHolder(
|
|
# size: batches x episode x n_consumers x consumer_state float
|
|
self.consumer_states_batch_gpu_tensor
|
|
),
|
|
CudaTensorHolder(self.consumer_rewards_batch_gpu_tensor),
|
|
CudaTensorHolder(self.firm_states_gpu_tensor),
|
|
self.firm_action_indices_gpu_pycuda,
|
|
self.firm_actions_gpu_pycuda,
|
|
self.firm_rewards_gpu_pycuda,
|
|
CudaTensorHolder(self.firm_states_batch),
|
|
CudaTensorHolder(self.firm_actions_batch),
|
|
CudaTensorHolder(self.firm_rewards_batch),
|
|
CudaTensorHolder(self.government_states_gpu_tensor),
|
|
self.government_action_indices_gpu_pycuda,
|
|
self.government_actions_gpu_pycuda,
|
|
self.government_rewards_gpu_pycuda,
|
|
CudaTensorHolder(self.government_states_batch),
|
|
CudaTensorHolder(self.government_actions_batch),
|
|
CudaTensorHolder(self.government_rewards_batch),
|
|
CudaTensorHolder(self.consumer_aux_batch_gpu_tensor),
|
|
CudaTensorHolder(self.firm_aux_batch),
|
|
np.int32(_iter),
|
|
block=block,
|
|
grid=grid,
|
|
)
|
|
self.consumer_actions_batch_gpu_tensor[
|
|
:, _iter, :, :
|
|
] = self.consumer_actions_index_single_gpu_tensor
|
|
|
|
# ------------------------
|
|
# Add penalty for no-Ponzi
|
|
# ------------------------
|
|
add_penalty_for_no_ponzi(
|
|
self.firm_states_gpu_tensor,
|
|
self.firm_rewards_batch,
|
|
__ad["global_state_dim"],
|
|
penalty_coef=firm_no_ponzi_coef,
|
|
)
|
|
add_penalty_for_no_ponzi(
|
|
self.consumer_states_gpu_tensor,
|
|
self.consumer_rewards_batch_gpu_tensor,
|
|
__ad["global_state_dim"],
|
|
penalty_coef=consumer_no_ponzi_coef,
|
|
penalty_scale=__ad["consumer_penalty_scale"],
|
|
)
|
|
|
|
# add government rewards -- sum of consumer rewards
|
|
update_government_rewards(
|
|
self.government_rewards_batch,
|
|
self.consumer_rewards_batch_gpu_tensor,
|
|
self.firm_rewards_batch,
|
|
self.cfg_dict,
|
|
)
|
|
|
|
# Save dense logs
|
|
# ------------------------
|
|
if (epi % __td["save_model_every"]) == 0:
|
|
save_policy_parameters(
|
|
self.save_dir,
|
|
epi,
|
|
consumer_policy,
|
|
firm_policy,
|
|
government_policy,
|
|
self.freeze_firms,
|
|
self.freeze_govt,
|
|
)
|
|
if (epi % self.save_dense_every) == 0:
|
|
save_dense_log(
|
|
self.save_dir,
|
|
epi,
|
|
agent_type_arrays,
|
|
agent_action_arrays,
|
|
agent_aux_arrays,
|
|
)
|
|
|
|
# --------------------------------
|
|
# Curriculum: Train Consumers
|
|
# --------------------------------
|
|
if self.consumers_will_train_this_episode(epi):
|
|
consumer_entropy_coef = anneal_entropy_coef(
|
|
self.agents_dict.get("consumer_anneal_entropy", None),
|
|
epi - consumer_training_start,
|
|
)
|
|
consumer_reward_scale = self.agents_dict.get(
|
|
"consumer_reward_scale", 1.0
|
|
)
|
|
if __td["use_ppo"]:
|
|
consumer_ppo_step(
|
|
consumer_policy,
|
|
expand_to_digit_form(
|
|
self.consumer_states_batch_gpu_tensor,
|
|
__ad["consumer_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
self.consumer_actions_batch_gpu_tensor,
|
|
self.consumer_rewards_batch_gpu_tensor,
|
|
consumer_optim,
|
|
__td["gamma"],
|
|
entropy_val=consumer_entropy_coef * __td["entropy"],
|
|
value_loss_weight=__td["value_loss_weight"],
|
|
reward_scale=consumer_reward_scale,
|
|
ppo_num_updates=__td["ppo_num_updates"],
|
|
clip_param=__td["ppo_clip_param"],
|
|
clip_grad_norm=self.train_dict.get("clip_grad_norm", None),
|
|
)
|
|
else:
|
|
consumer_policy_gradient_step(
|
|
consumer_policy,
|
|
expand_to_digit_form(
|
|
self.consumer_states_batch_gpu_tensor,
|
|
__ad["consumer_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
self.consumer_actions_batch_gpu_tensor,
|
|
self.consumer_rewards_batch_gpu_tensor,
|
|
consumer_optim,
|
|
__td["gamma"],
|
|
entropy_val=consumer_entropy_coef * __td["entropy"],
|
|
value_loss_weight=__td["value_loss_weight"],
|
|
reward_scale=consumer_reward_scale,
|
|
clip_grad_norm=self.train_dict.get("clip_grad_norm", None),
|
|
)
|
|
if (epi % lagr_num_steps) == 0:
|
|
consumer_no_ponzi_coef = update_penalty_coef(
|
|
self.consumer_states_gpu_tensor,
|
|
__ad["global_state_dim"],
|
|
consumer_no_ponzi_coef,
|
|
penalty_step_size=__ad["consumer_noponzi_eta"],
|
|
penalty_scale=__ad["consumer_penalty_scale"],
|
|
)
|
|
else:
|
|
pass
|
|
|
|
# --------------------------------
|
|
# Curriculum: Train Firms
|
|
# --------------------------------
|
|
if self.firms_will_train_this_episode(epi):
|
|
firm_entropy_coef = anneal_entropy_coef(
|
|
self.agents_dict.get("firm_anneal_entropy", None),
|
|
epi - firm_training_start,
|
|
)
|
|
firm_reward_scale = self.agents_dict.get("firm_reward_scale", 1.0)
|
|
if __td["use_ppo"]:
|
|
ppo_step(
|
|
firm_policy,
|
|
expand_to_digit_form(
|
|
self.firm_states_batch,
|
|
__ad["firm_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
self.firm_actions_batch,
|
|
self.firm_rewards_batch,
|
|
firm_optim,
|
|
__td["gamma"],
|
|
entropy_val=firm_entropy_coef * __td["entropy"],
|
|
value_loss_weight=__td["value_loss_weight"],
|
|
actions_mask=firm_actions_mask,
|
|
reward_scale=firm_reward_scale,
|
|
ppo_num_updates=__td["ppo_num_updates"],
|
|
clip_param=__td["ppo_clip_param"],
|
|
clip_grad_norm=self.train_dict.get("clip_grad_norm", None),
|
|
)
|
|
else:
|
|
policy_gradient_step(
|
|
firm_policy,
|
|
expand_to_digit_form(
|
|
self.firm_states_batch,
|
|
__ad["firm_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
self.firm_actions_batch,
|
|
self.firm_rewards_batch,
|
|
firm_optim,
|
|
__td["gamma"],
|
|
entropy_val=firm_entropy_coef * __td["entropy"],
|
|
value_loss_weight=__td["value_loss_weight"],
|
|
actions_mask=firm_actions_mask,
|
|
reward_scale=firm_reward_scale,
|
|
clip_grad_norm=self.train_dict.get("clip_grad_norm", None),
|
|
)
|
|
|
|
if (epi % lagr_num_steps) == 0:
|
|
firm_no_ponzi_coef = update_penalty_coef(
|
|
self.firm_states_gpu_tensor,
|
|
__ad["global_state_dim"],
|
|
firm_no_ponzi_coef,
|
|
penalty_step_size=__ad["firm_noponzi_eta"],
|
|
)
|
|
else:
|
|
pass
|
|
|
|
# --------------------------------
|
|
# Curriculum: Train Governments
|
|
# --------------------------------
|
|
if self.governments_will_train_this_episode(epi):
|
|
government_entropy_coef = anneal_entropy_coef(
|
|
self.agents_dict.get("govt_anneal_entropy", None),
|
|
epi - government_training_start,
|
|
)
|
|
government_reward_scale = self.agents_dict.get(
|
|
"government_reward_scale", 1.0
|
|
)
|
|
if __td["use_ppo"]:
|
|
ppo_step(
|
|
government_policy,
|
|
expand_to_digit_form(
|
|
self.government_states_batch,
|
|
__ad["government_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
self.government_actions_batch,
|
|
self.government_rewards_batch,
|
|
government_optim,
|
|
__td["gamma"],
|
|
entropy_val=government_entropy_coef * __td["entropy"],
|
|
value_loss_weight=__td["value_loss_weight"],
|
|
actions_mask=government_actions_mask,
|
|
reward_scale=government_reward_scale,
|
|
ppo_num_updates=__td["ppo_num_updates"],
|
|
clip_param=__td["ppo_clip_param"],
|
|
clip_grad_norm=self.train_dict.get("clip_grad_norm", None),
|
|
)
|
|
else:
|
|
policy_gradient_step(
|
|
government_policy,
|
|
expand_to_digit_form(
|
|
self.government_states_batch,
|
|
__ad["government_digit_dims"],
|
|
__td["digit_representation_size"],
|
|
),
|
|
self.government_actions_batch,
|
|
self.government_rewards_batch,
|
|
government_optim,
|
|
__td["gamma"],
|
|
entropy_val=government_entropy_coef * __td["entropy"],
|
|
value_loss_weight=__td["value_loss_weight"],
|
|
actions_mask=government_actions_mask,
|
|
reward_scale=government_reward_scale,
|
|
clip_grad_norm=self.train_dict.get("clip_grad_norm", None),
|
|
)
|
|
else:
|
|
pass
|
|
|
|
# Store the value of the final episode
|
|
final_epi = epi
|
|
|
|
# ------------------------------------------------------------------
|
|
# Post-Training (may not reach this with an infinite training loop!)
|
|
# Save FINAL dense log.
|
|
# ------------------------------------------------------------------
|
|
save_dense_log(
|
|
self.save_dir,
|
|
"final",
|
|
agent_type_arrays,
|
|
agent_action_arrays,
|
|
agent_aux_arrays,
|
|
)
|
|
save_policy_parameters(
|
|
self.save_dir,
|
|
final_epi,
|
|
consumer_policy,
|
|
firm_policy,
|
|
government_policy,
|
|
self.freeze_firms,
|
|
self.freeze_govt,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Clean up
|
|
# ------------------------------------------------------------------
|
|
|
|
self.cuda_free_mem(block=block, grid=grid)
|
|
|
|
|
|
class CudaTensorHolder(pycuda.driver.PointerHolderBase):
|
|
"""
|
|
A class that facilitates casting tensors to pointers.
|
|
"""
|
|
|
|
def __init__(self, t):
|
|
super().__init__()
|
|
self.t = t
|
|
self.gpudata = t.data_ptr()
|
|
|
|
def get_pointer(self):
|
|
return self.t.data_ptr()
|