def __init__(self,
args,
max_raw_chars,
max_instruction_span,
coach_mode,
num_resource_bin,
*,
num_unit_type=len(gc.UnitTypes),
num_cmd_type=len(gc.CmdTypes)):
super().__init__()
self.params = {
'args': args,
'max_raw_chars': max_raw_chars,
'max_instruction_span': max_instruction_span,
'coach_mode': coach_mode,
'num_resource_bin': num_resource_bin,
'num_unit_type': num_unit_type,
'num_cmd_type': num_cmd_type,
}
self.args = args
self.max_raw_chars = max_raw_chars
self.max_instruction_span = max_instruction_span
self.coach_mode = coach_mode
self.pos_candidate_inst = None
self.neg_candidate_inst = None
self.args.inst_dict_path = self.args.inst_dict_path.replace(
'scratch/rts_data', 'rts-replays')
self.inst_dict = self.load_inst_dict(self.args.inst_dict_path)
self.prev_inst_encoder = LSTMInstructionEncoder(
self.inst_dict.total_vocab_size,
self.args.word_emb_dim,
self.args.word_emb_dropout,
self.args.inst_hid_dim,
self.inst_dict.pad_word_idx,
)
self.glob_encoder = ConvGlobEncoder(args, num_unit_type, num_cmd_type,
num_resource_bin,
self.prev_inst_encoder)
# count encoders
self.count_encoder = MlpEncoder(self.args.num_count_channels * 2,
self.args.count_hid_dim,
self.args.count_hid_dim,
self.args.count_hid_layers - 1,
activate_out=True)
self.cons_count_encoder = MlpEncoder(num_unit_type,
self.args.count_hid_dim // 2,
self.args.count_hid_dim // 2,
self.args.count_hid_layers - 1,
activate_out=True)
self.moving_avg_encoder = MlpEncoder(num_unit_type,
self.args.count_hid_dim // 2,
self.args.count_hid_dim // 2,
self.args.count_hid_layers - 1,
activate_out=True)
self.frame_passed_encoder = nn.Embedding(
max_instruction_span + 2,
self.args.count_hid_dim // 2,
)
if self.args.glob_dropout > 0:
self.glob_dropout = nn.Dropout(self.args.glob_dropout)
self.glob_feat_dim = int(2.5 * self.args.count_hid_dim +
self.glob_encoder.glob_dim)
self.cont_cls = GlobClsHead(
self.glob_feat_dim,
self.args.inst_hid_dim, # for reducing hyper-parameter
2)
if self.coach_mode == 'rnn':
encoder = self.prev_inst_encoder
elif self.coach_mode == 'bow':
encoder = MeanBOWInstructionEncoder(
self.inst_dict.total_vocab_size, self.args.inst_hid_dim,
self.args.word_emb_dropout, self.inst_dict.pad_word_idx)
elif self.coach_mode == 'onehot' or self.coach_mode == 'rnn_gen':
pass
else:
assert False, 'unknown coach mode: %s' % self.coach_mode
if self.coach_mode == 'rnn' or self.coach_mode == 'bow':
self.inst_selector = RnnSelector(encoder, self.glob_feat_dim)
else:
self.inst_selector = None
self.value = nn.utils.weight_norm(nn.Linear(self.glob_feat_dim, 1),
dim=None)
self.sampler = ContSoftmaxSampler('cont', 'cont_pi', 'inst', 'inst_pi')
def __init__(
self,
args,
max_raw_chars,
max_instruction_span,
coach_mode,
num_resource_bin,
*,
num_unit_type=len(gc.UnitTypes),
num_cmd_type=len(gc.CmdTypes),
coach_rule_emb_size=0,
):
super().__init__()
self.params = {
"args": args,
"max_raw_chars": max_raw_chars,
"max_instruction_span": max_instruction_span,
"coach_mode": coach_mode,
"num_resource_bin": num_resource_bin,
"num_unit_type": num_unit_type,
"num_cmd_type": num_cmd_type,
}
self.num_unit_types = NUM_UNIT_TYPES
self.coach_rule_emb_size = coach_rule_emb_size
self.args = args
self.max_raw_chars = max_raw_chars
self.max_instruction_span = max_instruction_span
self.coach_mode = coach_mode
self.pos_candidate_inst = None
self.neg_candidate_inst = None
self.args.inst_dict_path = self.args.inst_dict_path.replace(
"scratch/rts_data", "rts-replays")
self.inst_dict = self.load_inst_dict(self.args.inst_dict_path)
self.prev_inst_encoder = LSTMInstructionEncoder(
self.inst_dict.total_vocab_size,
self.args.word_emb_dim,
self.args.word_emb_dropout,
self.args.inst_hid_dim,
self.inst_dict.pad_word_idx,
)
self.glob_encoder = ConvGlobEncoder(args, num_unit_type, num_cmd_type,
num_resource_bin,
self.prev_inst_encoder)
# count encoders
self.count_encoder = MlpEncoder(
self.args.num_count_channels * 2,
self.args.count_hid_dim,
self.args.count_hid_dim,
self.args.count_hid_layers - 1,
activate_out=True,
)
self.cons_count_encoder = MlpEncoder(
num_unit_type,
self.args.count_hid_dim // 2,
self.args.count_hid_dim // 2,
self.args.count_hid_layers - 1,
activate_out=True,
)
self.moving_avg_encoder = MlpEncoder(
num_unit_type,
self.args.count_hid_dim // 2,
self.args.count_hid_dim // 2,
self.args.count_hid_layers - 1,
activate_out=True,
)
self.frame_passed_encoder = nn.Embedding(
max_instruction_span + 2,
self.args.count_hid_dim // 2,
)
if self.args.glob_dropout > 0:
self.glob_dropout = nn.Dropout(self.args.glob_dropout)
self.glob_feat_dim = int(2.5 * self.args.count_hid_dim +
self.glob_encoder.glob_dim)
self.cont_cls = GlobClsHead(
self.glob_feat_dim,
self.args.inst_hid_dim, # for reducing hyper-parameter
2,
)
if self.coach_rule_emb_size > 0:
# TODO: Remove num units hardcoding
self.rule_emb = nn.Embedding(NUM_UNIT_TYPES,
self.coach_rule_emb_size)
self.rule_lstm = torch.nn.LSTM(self.coach_rule_emb_size,
self.glob_feat_dim,
batch_first=True)
if self.coach_mode == "rnn":
encoder = self.prev_inst_encoder
elif self.coach_mode == "bow":
encoder = MeanBOWInstructionEncoder(
self.inst_dict.total_vocab_size,
self.args.inst_hid_dim,
self.args.word_emb_dropout,
self.inst_dict.pad_word_idx,
)
elif self.coach_mode == "onehot" or self.coach_mode == "rnn_gen":
pass
else:
assert False, "unknown coach mode: %s" % self.coach_mode
if self.coach_mode == "rnn" or self.coach_mode == "bow":
self.inst_selector = RnnSelector(encoder, self.glob_feat_dim)
else:
self.inst_selector = None
self.value = nn.utils.weight_norm(nn.Linear(self.glob_feat_dim, 1),
dim=None)
self.sampler = ContSoftmaxSampler("cont", "cont_pi", "inst", "inst_pi")