def __init__(
self,
config: Config,
modules: typing.Dict[str, torch.nn.Module] = None,
):
self._config = config
self._device = torch.device(config.get('device'))
if modules is not None:
assert 'pE' in modules
assert 'pT' in modules
assert 'pH' in modules
self._modules = modules
else:
self._modules = {
'pE': E(self._config).to(self._device),
'pT': T(self._config).to(self._device),
'pH': PH(self._config).to(self._device),
}
def __init__(
self,
config: Config,
train_dataset: ProofTraceLMDataset,
):
self._config = config
self._action_coeff = config.get('prooftrace_lm_action_coeff')
self._value_coeff = config.get('prooftrace_lm_value_coeff')
self._device = torch.device(config.get('device'))
self._modules = {
'E': E(self._config).to(self._device),
'T': T(self._config).to(self._device),
'PH': PH(self._config).to(self._device),
'VH': VH(self._config).to(self._device),
}
self._ack = IOTAAck(
config.get('prooftrace_lm_iota_sync_dir'),
self._modules,
)
self._nll_loss = nn.NLLLoss()
self._mse_loss = nn.MSELoss()
self._train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=self._config.get('prooftrace_lm_batch_size'),
shuffle=True,
collate_fn=lm_collate,
)
Log.out('ACK initialization', {
"batch_size": self._config.get('prooftrace_lm_batch_size'),
})
self._train_batch = 0
def __init__(
self,
config: Config,
):
self._config = config
self._device = torch.device(config.get('device'))
self._modules = {
'E': E(self._config).to(self._device),
'T': T(self._config).to(self._device),
'PH': PH(self._config).to(self._device),
'VH': VH(self._config).to(self._device),
}
self._dataset_dir = os.path.join(
os.path.expanduser(config.get('prooftrace_dataset_dir')),
config.get('prooftrace_dataset_size'),
'test_traces',
)
with gzip.open(
os.path.join(
os.path.expanduser(config.get('prooftrace_dataset_dir')),
config.get('prooftrace_dataset_size'),
'traces.tokenizer',
), 'rb') as f:
self._tokenizer = pickle.load(f)
self._tst = IOTARun(
config.get('prooftrace_search_iota_sync_dir'),
'test',
self._modules,
)
self._test_gamma_size = config.get('prooftrace_search_test_gamma_size')
Log.out('TST initialization', {})
def __init__(
self,
config: Config,
):
self._config = config
self._learning_rate = config.get('prooftrace_lm_learning_rate')
self._min_update_count = \
config.get('prooftrace_lm_iota_min_update_count')
self._device = torch.device(config.get('device'))
self._save_dir = config.get('prooftrace_save_dir')
self._load_dir = config.get('prooftrace_load_dir')
self._epoch = 0
self._tb_writer = None
if self._config.get('tensorboard_log_dir'):
self._tb_writer = SummaryWriter(
self._config.get('tensorboard_log_dir'), )
self._modules = {
'E': E(self._config).to(self._device),
'T': T(self._config).to(self._device),
'PH': PH(self._config).to(self._device),
'VH': VH(self._config).to(self._device),
}
Log.out(
"SYN Initializing",
{
'parameter_count_E': self._modules['E'].parameters_count(),
'parameter_count_T': self._modules['T'].parameters_count(),
'parameter_count_PH': self._modules['PH'].parameters_count(),
'parameter_count_VH': self._modules['VH'].parameters_count(),
},
)
self._syn = IOTASyn(
config.get('prooftrace_lm_iota_sync_dir'),
self._modules,
)
self._optimizer = optim.Adam(
[
{
'params': self._modules['E'].parameters()
},
{
'params': self._modules['T'].parameters()
},
{
'params': self._modules['PH'].parameters()
},
{
'params': self._modules['VH'].parameters()
},
],
lr=self._learning_rate,
)
self._syn.broadcast({'config': self._config})
def __init__(
self,
config: Config,
):
self._config = config
self._rollout_size = config.get('prooftrace_ppo_rollout_size')
self._pool_size = config.get('prooftrace_env_pool_size')
self._epoch_count = config.get('prooftrace_ppo_epoch_count')
self._clip = config.get('prooftrace_ppo_clip')
self._grad_norm_max = config.get('prooftrace_ppo_grad_norm_max')
self._act_entropy_coeff = \
config.get('prooftrace_ppo_act_entropy_coeff')
self._ptr_entropy_coeff = \
config.get('prooftrace_ppo_ptr_entropy_coeff')
self._value_coeff = config.get('prooftrace_ppo_value_coeff')
self._learning_rate = config.get('prooftrace_ppo_learning_rate')
self._reset_gamma = \
config.get('prooftrace_ppo_reset_gamma')
self._fixed_gamma = \
config.get('prooftrace_ppo_fixed_gamma')
self._explore_alpha = \
config.get('prooftrace_ppo_explore_alpha')
self._explore_beta = \
config.get('prooftrace_ppo_explore_beta')
self._explore_beta_width = \
config.get('prooftrace_ppo_explore_beta_width')
self._step_reward_prob = \
config.get('prooftrace_ppo_step_reward_prob')
self._match_reward_prob = \
config.get('prooftrace_ppo_match_reward_prob')
self._device = torch.device(config.get('device'))
self._modules = {
'E': E(self._config).to(self._device),
'T': T(self._config).to(self._device),
'PH': PH(self._config).to(self._device),
'VH': VH(self._config).to(self._device),
}
self._ack = IOTAAck(
config.get('prooftrace_ppo_iota_sync_dir'),
self._modules,
)
self._rollouts = Rollouts(self._config)
self._pool = Pool(self._config, False)
self._rollouts.observations[0] = self._pool.reset(
self._reset_gamma,
self._fixed_gamma,
)
self._episode_stp_reward = [0.0] * self._pool_size
self._episode_mtc_reward = [0.0] * self._pool_size
self._episode_fnl_reward = [0.0] * self._pool_size
Log.out(
'ACK initialization', {
"pool_size": self._config.get('prooftrace_env_pool_size'),
"rollout_size":
self._config.get('prooftrace_ppo_rollout_size'),
"batch_size": self._config.get('prooftrace_ppo_batch_size'),
})