def get_action(self, hist_graph, curr_graph, tag2unit_dict): assert isinstance( curr_graph, dgl.DGLGraph), "get action is designed to work on a single graph!" num_time_steps = hist_graph.batch_size hist_node_feature = hist_graph.ndata.pop('node_feature') curr_node_feature = curr_graph.ndata.pop('node_feature') maximum_num_enemy = get_largest_number_of_enemy_nodes([curr_graph]) nn_actions, info_dict = self.brain.get_action( num_time_steps, hist_graph, hist_node_feature, curr_graph, curr_node_feature, maximum_num_enemy) ally_tags = info_dict['ally_tags'] enemy_tags = info_dict['enemy_tags'] sc2_actions = nn_action_to_sc2_action(nn_actions=nn_actions, ally_tags=ally_tags, enemy_tags=enemy_tags, tag2unit_dict=tag2unit_dict) hist_graph.ndata['node_feature'] = hist_node_feature curr_graph.ndata['node_feature'] = curr_node_feature return nn_actions, sc2_actions, info_dict
ac_module = ActorCriticModule(node_input_dim=26) # ac_module = ActorCriticModule(node_input_dim=26) done_cnt = 0 i = 0 while True: # print("=========={} th iter ============".format(i)) cur_state = env.observe() # for tag, unit in cur_state['tag2unit_dict'].items(): # print("TAG: {} | NAME: {} | Health: {} | POS: {} ".format(tag, unit.name, unit.health, unit.position)) g = cur_state['g'] node_feature = g.ndata.pop('node_feature') num_enemy = get_largest_number_of_enemy_nodes([g]) nn_actions, info_dict = ac_module.get_action(g, node_feature, num_enemy) tag2unit_dict = cur_state['tag2unit_dict'] ally_tags = info_dict['ally_tags'] enemy_tags = info_dict['enemy_tags'] sc2_actions = nn_action_to_sc2_action(nn_actions=nn_actions, ally_tags=ally_tags, enemy_tags=enemy_tags, tag2unit_dict=tag2unit_dict) loss = ac_module(g, node_feature, num_enemy) next_state, reward, done = env.step(action=sc2_actions)
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) # 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) # casting actions to one torch tensor actions = torch.cat(actions).long() # prepare rewards rewards = torch.Tensor(rewards) # preparing dones dones = torch.Tensor(dones) 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
def fit(self): # 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) 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