def forward(self, graph, node_feature, update_node_type_indices,
update_edge_type_indices):
"""
:param graph:
:param node_feature:
:param update_node_type_indices:
:param update_edge_type_indices:
:return:
"""
graph.ndata['node_feature'] = node_feature
if self.neighbor_degree == 0: # Update features with only own features
for ntype_idx in update_node_type_indices:
node_index = get_filtered_node_index_by_type(graph, ntype_idx)
apply_func = partial(self.apply_node_function_no_neighbor,
ntype_idx=ntype_idx)
graph.apply_nodes(func=apply_func, v=node_index)
else: # Update features with own features and 1 hop neighbor features
for etype_idx in update_edge_type_indices:
edge_index = get_filtered_edge_index_by_type(graph, etype_idx)
graph.send_and_recv(edge_index,
message_func=self.message_function,
reduce_func=self.reduce_function)
for ntype_idx in update_node_type_indices:
node_index = get_filtered_node_index_by_type(graph, ntype_idx)
apply_func = partial(self.apply_node_function_yes_neighbor,
ntype_idx=ntype_idx)
graph.apply_nodes(func=apply_func, v=node_index)
updated_node_feature = graph.ndata.pop('node_feature')
if self.neighbor_degree >= 1:
graph.ndata.pop('aggregated_message')
return updated_node_feature
def forward(self, graph, node_feature, sub_q_tots):
graph.ndata['node_feature'] = node_feature
device = node_feature.device
len_groups = []
q_tot = torch.zeros(graph.batch_size, device=device)
for i, sub_q_tot in enumerate(sub_q_tots):
node_indices = get_filtered_node_index_by_assignment(graph, i)
len_groups.append(len(node_indices))
mask = torch.zeros(size=(node_feature.shape[0], 1), device=device)
mask[node_indices, :] = 1
graph.ndata[
'masked_node_feature'] = graph.ndata['node_feature'] * mask
w_input = dgl.sum_nodes(graph, 'masked_node_feature')
if self.rectifier == 'abs':
q_tot = q_tot + torch.abs(self.w(w_input)).view(-1) * sub_q_tot
elif self.rectifier == 'softplus':
q_tot = q_tot + F.softplus(
self.w(w_input)).view(-1) * sub_q_tot
else:
raise RuntimeError("Not implemented rectifier")
_ = graph.ndata.pop('masked_node_feature')
# testing
_ = graph.ndata.pop('node_feature')
ally_indices = get_filtered_node_index_by_type(graph, NODE_ALLY)
_v = torch.zeros(size=(graph.number_of_nodes(), node_feature.shape[1]),
device=device)
_v[ally_indices, :] = node_feature[ally_indices, :]
graph.ndata['node_feature'] = _v
# testing
v = self.v(dgl.sum_nodes(graph, 'node_feature')).view(-1)
q_tot = q_tot + v
_ = graph.ndata.pop('node_feature')
len_groups = np.array(len_groups)
ratio_groups = len_groups / np.sum(len_groups)
print("Num elements in groups {}".format(len_groups))
print("Num elements ratio {}".format(ratio_groups))
ally_indices = get_filtered_node_index_by_type(graph, NODE_ALLY)
target_assignment_weight = graph.ndata['normalized_score'][
ally_indices]
print("Average normalized scores {}".format(
target_assignment_weight.mean(0)))
return q_tot
def forward(self, graph, node_feature, qs, ally_node_type_index=NODE_ALLY):
assert isinstance(graph, dgl.BatchedDGLGraph)
w_emb = self.w_gn(graph, node_feature) # [# nodes x # node_dim]
w = torch.abs(self.w_ff(graph, w_emb)) # [# nodes x # 1]
ally_node_indices = get_filtered_node_index_by_type(
graph, ally_node_type_index)
device = w_emb.device
_qs = torch.zeros(size=(graph.number_of_nodes(), 1), device=device)
w = w[ally_node_indices, :] # [# allies x 1]
_qs[ally_node_indices, :] = w * qs.view(-1, 1)
graph.ndata['node_feature'] = _qs
q_tot = dgl.sum_nodes(graph, 'node_feature')
_ = graph.ndata.pop('node_feature')
v_emb = self.v_gn(graph, node_feature) # [# nodes x # node_dim]
v = self.v_ff(graph, v_emb) # [# nodes x # 1]
v = v[ally_node_indices, :] # [# allies x 1]
_v = torch.zeros(size=(graph.number_of_nodes(), 1), device=device)
_v[ally_node_indices, :] = v
graph.ndata['node_feature'] = _v
v = dgl.sum_nodes(graph, 'node_feature')
_ = graph.ndata.pop('node_feature')
q_tot = q_tot + v
return q_tot.view(-1)
def compute_probs(self,
graph,
node_feature,
maximum_num_enemy,
ally_node_type_index=NODE_ALLY,
attack_edge_type_index=EDGE_IN_ATTACK_RANGE):
# get logits of each action
move_arg, hold_arg, attack_arg, high_level_prob = self(
graph, node_feature, maximum_num_enemy, attack_edge_type_index)
device = move_arg.device
# Prepare un-normalized probability of attacks
unnormed_ps = torch.cat((move_arg, hold_arg, attack_arg),
dim=-1) # of all units including enemies
ally_node_indices = get_filtered_node_index_by_type(
graph, ally_node_type_index)
unnormed_ps = unnormed_ps[ally_node_indices, :] # of only ally units
ally_tags = graph.ndata['tag']
ally_tags = ally_tags[ally_node_indices]
if 'enemy_tag' in graph.ndata.keys():
enemy_tags = graph.ndata['enemy_tag']
else:
enemy_tags = torch.zeros_like(ally_tags,
device=device).view(-1, 1) # dummy
enemy_tags = enemy_tags[ally_node_indices, :]
ps = torch.nn.functional.softmax(unnormed_ps, dim=-1)
high_level_prob = high_level_prob[
ally_node_indices, :] # [#. ally units x 3]
attack_dim = ps.shape[-1] - self.move_dim - 1
num_reps = torch.tensor([self.move_dim, 1, attack_dim],
dtype=torch.long,
device=device)
high_level_prob = high_level_prob.repeat_interleave(num_reps, dim=1)
ps = high_level_prob * ps
log_ps = torch.log(ps + VERY_SMALL_NUMBER)
ally_entropy = -torch.sum(log_ps * ps, dim=-1) # per unit entropy
log_p_move, log_p_hold, log_p_attack = torch.split(
log_ps, [self.move_dim, 1, maximum_num_enemy], dim=1)
return_dict = dict()
# for RL training
return_dict['unnormed_ps'] = unnormed_ps
return_dict['probs'] = ps
return_dict['log_ps'] = log_ps
return_dict['log_p_move'] = log_p_move
return_dict['log_p_hold'] = log_p_hold
return_dict['log_p_attack'] = log_p_attack
return_dict['ally_entropy'] = ally_entropy
# for SC2 interfacing
return_dict['ally_tags'] = ally_tags
return_dict['enemy_tags'] = enemy_tags
return return_dict
def get_q(self, graph, node_feature, maximum_num_enemy, critic=None):
if critic is None:
critic = self.critic
ally_node_idx = get_filtered_node_index_by_type(graph, NODE_ALLY)
q_m, q_h, q_a = critic(graph, node_feature, maximum_num_enemy)
q = torch.cat([q_m, q_h, q_a], dim=-1)
q_ally = q[ally_node_idx, :]
return q_ally
def forward(self, graph, node_feature, update_node_type_indices):
graph.ndata['node_feature'] = node_feature
for ntype_idx in update_node_type_indices:
node_index = get_filtered_node_index_by_type(graph, ntype_idx)
apply_func = partial(self.apply_node_function, ntype_idx=ntype_idx)
graph.apply_nodes(func=apply_func, v=node_index)
_ = graph.ndata.pop('node_feature')
updated_node_feature = graph.ndata.pop('updated_node_feature')
return updated_node_feature
def eval_1_episode(self):
# expected return
# dictionary = {'name': (str), 'win': (bool), 'sum_reward': (float)}
running_wr = self.env.winning_ratio
env_name = self.env.name
traj = self.run_1_episode()
last_graph = traj[-1].state
num_allies = get_filtered_node_index_by_type(last_graph, NODE_ALLY).size()
num_enemies = get_filtered_node_index_by_type(last_graph, NODE_ENEMY).size()
sum_reward = np.sum([exp.reward for exp in traj._trajectory])
eval_dict = dict()
eval_dict['name'] = env_name
eval_dict['win'] = num_allies > num_enemies
eval_dict['sum_reward'] = sum_reward
self.env.winning_ratio = running_wr
return eval_dict
def forward(self, graph, node_feature, update_node_type_indices,
update_edge_type_indices):
if self.use_concat:
graph.ndata['node_feature'] = torch.cat(
[node_feature, graph.ndata['init_node_feature']], dim=1)
else:
graph.ndata['node_feature'] = node_feature
message_func = partial(
self.message_function,
update_edge_type_indices=update_edge_type_indices)
reduce_func = partial(
self.reduce_function,
update_edge_type_indices=update_edge_type_indices)
graph.send_and_recv(graph.edges(),
message_func=message_func,
reduce_func=reduce_func)
if not self.use_multi_node_types: # default behavior
for ntype_idx in update_node_type_indices:
node_indices = get_filtered_node_index_by_type(
graph, ntype_idx)
graph.apply_nodes(self.apply_node_function, v=node_indices)
else: # testing
for ntype_idx in update_node_type_indices:
node_indices = get_filtered_node_index_by_type(
graph, ntype_idx)
node_updater = self.node_updater['updater{}'.format(ntype_idx)]
apply_node_func = partial(self.apply_node_function_multi_type,
updater=node_updater)
graph.apply_nodes(apply_node_func, v=node_indices)
updated_node_feature = graph.ndata.pop('updated_node_feature')
_ = graph.ndata.pop('aggregated_node_feature')
_ = graph.ndata.pop('node_feature')
return updated_node_feature
def forward(self, graph, node_feature):
device = node_feature.device
graph.ndata['node_feature'] = node_feature
graph.apply_nodes(func=self.apply_node_function)
prob = graph.ndata.pop('prob')
ally_indices = get_filtered_node_index_by_type(graph, NODE_ALLY)
_assignment = graph.ndata.pop('assignment')
assignment = torch.ones_like(
_assignment,
device=device) * -1 # masking out enemy assignments as -1
assignment[ally_indices] = _assignment[ally_indices]
_normalized_score = graph.ndata.pop('normalized_score')
normalized_score = torch.ones_like(_normalized_score,
device=device) * -1
normalized_score[ally_indices] = _normalized_score[ally_indices]
return prob, assignment, normalized_score
def forward(self, graph, node_feature, update_node_type_indices, update_edge_type_indices):
"""
:param graph: structure only graph
:param node_feature:
:param update_node_type_indices:
:param update_edge_type_indices:
:return:
"""
graph.ndata['node_feature'] = node_feature
message_func = partial(self.message_function, update_edge_type_indices=update_edge_type_indices)
graph.send_and_recv(graph.edges(), message_func=message_func, reduce_func=self.reduce_function)
for ntype_idx in update_node_type_indices:
node_index = get_filtered_node_index_by_type(graph, ntype_idx)
graph.apply_nodes(self.apply_node_function, v=node_index)
updated_node_feature = graph.ndata.pop('node_feature')
graph.ndata.pop('z')
return updated_node_feature
def forward(self, graph, node_feature, qs, ally_node_type_index=NODE_ALLY):
assert isinstance(graph, dgl.BatchedDGLGraph)
w_emb = self.w_gn(graph, node_feature) # [# nodes x # node_dim]
if self.rectifier == 'abs':
w = torch.abs(self.w_ff(graph, w_emb)) # [# nodes x # 1]
elif self.rectifier == 'softplus':
w = F.softplus(self.w_ff(graph, w_emb)) # [# nodes x # 1]
else:
raise RuntimeError("Not supported rectifier")
# Curee's trick
ally_indices = get_filtered_node_index_by_type(graph, NODE_ALLY)
allies_assignment = graph.ndata['assignment'][ally_indices]
target_allies = allies_assignment == self.target_assignment
target_indices = torch.arange(target_allies.size(0))[target_allies]
device = w_emb.device
_qs = torch.zeros(size=(graph.number_of_nodes(), 1), device=device)
w = w[target_indices, :] # [# assignments x 1]
_qs[target_indices, :] = w * qs[target_indices].view(-1, 1)
graph.ndata['node_feature'] = _qs
q_tot = dgl.sum_nodes(graph, 'node_feature')
_ = graph.ndata.pop('node_feature')
v_emb = self.v_gn(graph, node_feature) # [# nodes x # node_dim]
v = self.v_ff(graph, v_emb) # [# nodes x # 1]
v = v[target_indices, :] # [# allies x 1]
_v = torch.zeros(size=(graph.number_of_nodes(), 1), device=device)
_v[target_indices, :] = v
graph.ndata['node_feature'] = _v
v = dgl.sum_nodes(graph, 'node_feature')
_ = graph.ndata.pop('node_feature')
q_tot = q_tot + v
return q_tot.view(-1)
def forward(self, graph, node_feature, qs, ally_node_type_index=NODE_ALLY):
assert isinstance(graph, dgl.BatchedDGLGraph)
# w_emb = self.w_gn(graph, node_feature) # [# nodes x # node_dim]
# if self.rectifier == 'abs':
# w = torch.abs(self.w_ff(graph, w_emb)) # [# nodes x # 1]
# elif self.rectifier == 'softplus':
# w = F.softplus(self.w_ff(graph, w_emb)) # [# nodes x # 1]
# Curee's trick
ally_indices = get_filtered_node_index_by_type(graph, NODE_ALLY)
target_assignment_weight = graph.ndata['normalized_score'][
ally_indices][:, self.target_assignment]
device = node_feature.device
_qs = torch.zeros(size=(graph.number_of_nodes(), 1), device=device)
# w = w[ally_indices, :] # [# assignments x 1]
_qs[ally_indices, :] = qs.view(-1, 1) * target_assignment_weight.view(
-1, 1)
graph.ndata['node_feature'] = _qs
q_tot = dgl.sum_nodes(graph, 'node_feature')
_ = graph.ndata.pop('node_feature')
v_emb = self.v_gn(graph, node_feature) # [# nodes x # node_dim]
v = self.v_ff(graph, v_emb) # [# nodes x # 1]
v = v[ally_indices, :] # [# allies x 1]
_v = torch.zeros(size=(graph.number_of_nodes(), 1), device=device)
_v[ally_indices, :] = v
graph.ndata['node_feature'] = _v
v = dgl.sum_nodes(graph, 'node_feature')
_ = graph.ndata.pop('node_feature')
q_tot = q_tot + v
return q_tot.view(-1)
def compute_qs(self,
graph,
node_feature,
maximum_num_enemy,
ally_node_type_index=NODE_ALLY,
attack_edge_type_index=None):
if attack_edge_type_index is None:
attack_edge_type_index = self.attack_edge_type_index
move_arg, hold_arg, attack_arg, assignment, normalized_score = self(
graph, node_feature, maximum_num_enemy, attack_edge_type_index)
qs = torch.cat((move_arg, hold_arg, attack_arg),
dim=-1) # of all units including enemies
ally_node_indices = get_filtered_node_index_by_type(
graph, ally_node_type_index)
qs = qs[ally_node_indices, :] # of only ally units
ally_tags = graph.ndata['tag']
ally_tags = ally_tags[ally_node_indices]
if 'enemy_tag' in graph.ndata.keys():
enemy_tags = graph.ndata['enemy_tag']
else:
enemy_tags = torch.zeros_like(ally_tags).view(-1, 1) # dummy
enemy_tags = enemy_tags[ally_node_indices, :]
return_dict = dict()
# for RL training
return_dict['qs'] = qs
return_dict['assignment'] = assignment
return_dict['normalized_score'] = normalized_score
# for SC2 interfacing
return_dict['ally_tags'] = ally_tags
return_dict['enemy_tags'] = enemy_tags
return return_dict
def fit(self, device='cpu'):
# the prefix 'c' indicates #current# time stamp inputs
# the prefix 'n' indicates #next# time stamp inputs
# expected specs:
# bs = batch_size, nt = hist_num_time_steps
# 'h_graph' = list of graph lists [[g_(0,0), g_(0,1), ... g_(0,nt)],
# [g_(1,0), g_(1,1), ..., g_(1,nt)],
# [g_(2,0), ..., g_(bs, 0), ... g_(bs, nt)]]
# 'graph' = list of graphs [g_(0), g_(1), ..., g_(bs)]
fit_conf = self.conf.fit_conf
batch_size = fit_conf['batch_size']
hist_num_time_steps = fit_conf['hist_num_time_steps']
c_h_graph, c_graph, actions, rewards, n_h_graph, n_graph, dones = self.buffer.sample(
batch_size)
c_maximum_num_enemy = get_largest_number_of_enemy_nodes(c_graph)
n_maximum_num_enemy = get_largest_number_of_enemy_nodes(n_graph)
# casting actions to one torch tensor
actions = torch.cat(actions).long()
# 'c_graph' is now list of graphs
c_ally_units = [
len(get_filtered_node_index_by_type(graph, NODE_ALLY))
for graph in c_graph
]
c_ally_units = torch.Tensor(c_ally_units).long()
# prepare rewards
rewards = torch.Tensor(rewards)
rewards = rewards.repeat_interleave(c_ally_units, dim=0)
# preparing dones
dones = torch.Tensor(dones)
dones = dones.repeat_interleave(c_ally_units, dim=0)
# batching graphs
list_c_h_graph = [g for L in c_h_graph for g in L]
list_n_h_graph = [g for L in n_h_graph for g in L]
c_hist_graph = dgl.batch(list_c_h_graph)
n_hist_graph = dgl.batch(list_n_h_graph)
c_curr_graph = dgl.batch(c_graph)
n_curr_graph = dgl.batch(n_graph)
if device != 'cpu':
c_hist_graph.to(torch.device('cuda'))
n_hist_graph.to(torch.device('cuda'))
c_curr_graph.to(torch.device('cuda'))
n_curr_graph.to(torch.device('cuda'))
actions = actions.to(torch.device('cuda'))
rewards = rewards.to(torch.device('cuda'))
dones = dones.to(torch.device('cuda'))
c_hist_feature = c_hist_graph.ndata.pop('node_feature')
c_curr_feature = c_curr_graph.ndata.pop('node_feature')
n_hist_feature = n_hist_graph.ndata.pop('node_feature')
n_curr_feature = n_curr_graph.ndata.pop('node_feature')
fit_return_dict = self.brain.fit(
num_time_steps=hist_num_time_steps,
c_hist_graph=c_hist_graph,
c_hist_feature=c_hist_feature,
c_curr_graph=c_curr_graph,
c_curr_feature=c_curr_feature,
c_maximum_num_enemy=c_maximum_num_enemy,
n_hist_graph=n_hist_graph,
n_hist_feature=n_hist_feature,
n_curr_graph=n_curr_graph,
n_curr_feature=n_curr_feature,
n_maximum_num_enemy=n_maximum_num_enemy,
actions=actions,
rewards=rewards,
dones=dones)
return fit_return_dict