worker.py
CFM 上游 rollout 与 success_rate 定义
原文件 · 224 行 · 11392 字节
third_party/heteromrta/worker.py
SHA-256: b1a285190d25a191a88634a5e177245f252cf5679095523173139daef2499dc3
1import pickle2import time3import torch4import numpy as np5import random6from env.task_env import TaskEnv7from attention import AttentionNet8import scipy.signal as signal9from parameters import *10import copy11from torch.nn import functional as F12from torch.distributions import Categorical131415def discount(x, gamma):16 """从后向前计算折扣累计值。"""17 return signal.lfilter([1], [1, -gamma], x[::-1], axis=0)[::-1]181920def zero_padding(x, padding_size, length):21 pad = torch.nn.ZeroPad2d((0, 0, 0, padding_size - length))22 x = pad(x)23 return x242526class Worker:27 """运行一个 CFM 环境副本并收集任务选择记录。"""28 def __init__(self, mete_agent_id, local_network, local_baseline, global_step,29 device='cuda', save_image=False, seed=None, env_params=None):3031 self.device = device32 self.metaAgentID = mete_agent_id33 self.global_step = global_step34 self.save_image = save_image35 if env_params is None:36 env_params = [EnvParams.SPECIES_AGENTS_RANGE, EnvParams.SPECIES_RANGE, EnvParams.TASKS_RANGE]37 self.env = TaskEnv(*env_params, EnvParams.TRAIT_DIM, EnvParams.DECISION_DIM, seed=seed, plot_figure=save_image)38 self.baseline_env = copy.deepcopy(self.env)39 self.local_baseline = local_baseline40 self.local_net = local_network41 self.experience = {idx:[] for idx in range(7)}42 self.episode_number = None43 self.perf_metrics = {}44 self.p_rnn_state = {}45 self.max_time = EnvParams.MAX_TIME46 # 默认最多记录 300 次有效选择；短测试可以通过实例字段调低上限。47 self.max_decisions = 30048 self.max_decision_steps = None4950 def run_episode(self, training=True, sample=False, max_waiting=False):51 """按环境给出的决策时间依次处理可行动的智能体。"""52 buffer_dict = {idx:[] for idx in range(7)}53 perf_metrics = {}54 current_action_index = 055 decision_step = 056 while (57 not self.env.finished58 and self.env.current_time < EnvParams.MAX_TIME59 and current_action_index < self.max_decisions60 and (61 self.max_decision_steps is None62 or decision_step < self.max_decision_steps63 )64 ):65 # next_decision 跳到下一个有智能体需要重新分配的时间点。66 with torch.no_grad():67 release_agents, current_time = self.env.next_decision()68 self.env.current_time = current_time69 random.shuffle(release_agents[0])70 finished_task = []71 while release_agents[0] or release_agents[1]:72 # 同一时刻可能有多个智能体需要选择任务，逐个调用网络并更新环境。73 agent_id = release_agents[0].pop(0) if release_agents[0] else release_agents[1].pop(0)74 agent = self.env.agent_dic[agent_id]75 task_info, total_agents, mask = self.convert_torch(self.env.agent_observe(agent_id, max_waiting))76 block_flag = mask[0, 1:].all().item()77 # 没有可行任务时跳过本次选择，环境继续推进时间。78 if block_flag and not np.all(self.env.get_matrix(self.env.task_dic, 'feasible_assignment')):79 agent['no_choice'] = block_flag80 continue81 elif block_flag and np.all(self.env.get_matrix(self.env.task_dic, 'feasible_assignment')) and agent['current_task'] < 0:82 continue83 if training:84 task_info, total_agents, mask = self.obs_padding(task_info, total_agents, mask)85 index = torch.LongTensor([agent_id]).reshape(1, 1, 1).to(self.device)86 probs, _ = self.local_net(task_info, total_agents, mask, index)87 if training:88 action = Categorical(probs).sample()89 while action.item() > self.env.tasks_num:90 action = Categorical(probs).sample()91 else:92 if sample:93 action = Categorical(probs).sample()94 else:95 action = torch.argmax(probs, dim=1)96 r, doable, f_t = self.env.agent_step(agent_id, action.item(), decision_step)97 agent['current_action_index'] = current_action_index98 finished_task.append(f_t)99 if training and doable:100 # 只记录环境接受的动作，稍后统一填入整条 rollout 的回报差值。101 buffer_dict[0] += total_agents102 buffer_dict[1] += task_info103 buffer_dict[2] += action.unsqueeze(0)104 buffer_dict[3] += mask105 buffer_dict[4] += torch.FloatTensor([[0]]).to(self.device) # reward106 buffer_dict[5] += index107 buffer_dict[6] += torch.FloatTensor([[0]]).to(self.device)108 current_action_index += 1109 self.env.finished = self.env.check_finished()110 decision_step += 1111112 terminal_reward, finished_tasks = self.env.get_episode_reward(self.max_time)113114 # 指标在 episode 结束后统一从环境状态中读取。115 perf_metrics['success_rate'] = [np.sum(finished_tasks)/len(finished_tasks)]116 perf_metrics['makespan'] = [self.env.current_time]117 perf_metrics['time_cost'] = [np.nanmean(self.env.get_matrix(self.env.task_dic, 'time_start'))]118 perf_metrics['waiting_time'] = [np.mean(self.env.get_matrix(self.env.agent_dic, 'sum_waiting_time'))]119 perf_metrics['travel_dist'] = [np.sum(self.env.get_matrix(self.env.agent_dic, 'travel_dist'))]120 perf_metrics['efficiency'] = [self.env.get_efficiency()]121 return terminal_reward, buffer_dict, perf_metrics122123 def baseline_test(self):124 """在相同初始场景上用 baseline 网络贪心选择任务。"""125 self.baseline_env.plot_figure = False126 perf_metrics = {}127 current_action_index = 0128 start = time.time()129 while not self.baseline_env.finished and self.baseline_env.current_time < self.max_time and current_action_index < 300:130 with torch.no_grad():131 release_agents, current_time = self.baseline_env.next_decision()132 random.shuffle(release_agents[0])133 self.baseline_env.current_time = current_time134 if time.time() - start > 30:135 break136 while release_agents[0] or release_agents[1]:137 agent_id = release_agents[0].pop(0) if release_agents[0] else release_agents[1].pop(0)138 agent = self.baseline_env.agent_dic[agent_id]139 task_info, total_agents, mask = self.convert_torch(self.baseline_env.agent_observe(agent_id, False))140 return_flag = mask[0, 1:].all().item()141 if return_flag and not np.all(self.baseline_env.get_matrix(self.baseline_env.task_dic, 'feasible_assignment')): ## add condition on returning to depot142 self.baseline_env.agent_dic[agent_id]['no_choice'] = return_flag143 continue144 elif return_flag and np.all(self.baseline_env.get_matrix(self.baseline_env.task_dic, 'feasible_assignment')) and agent['current_task'] < 0:145 continue146 task_info, total_agents, mask = self.obs_padding(task_info, total_agents, mask)147 index = torch.LongTensor([agent_id]).reshape(1, 1, 1).to(self.device)148 probs, _ = self.local_baseline(task_info, total_agents, mask, index)149 action = torch.argmax(probs, 1)150 self.baseline_env.agent_step(agent_id, action.item(), None)151 current_action_index += 1152 self.baseline_env.finished = self.baseline_env.check_finished()153154 reward, finished_tasks = self.baseline_env.get_episode_reward(self.max_time)155 return reward156157 def work(self, episode_number):158 """在同一设置下运行多份样本，并用组内平均回报作为 baseline。"""159 baseline_rewards = []160 buffers = []161 metrics = []162 max_waiting = TrainParams.FORCE_MAX_OPEN_TASK163 for _ in range(TrainParams.POMO_SIZE):164 # 多次重置到同一分布并采样动作，得到一组可相互比较的回报。165 self.env.init_state()166 terminal_reward, buffer, perf_metrics = self.run_episode(episode_number,True, max_waiting)167 if terminal_reward is np.nan:168 max_waiting = True169 continue170 baseline_rewards.append(terminal_reward)171 buffers.append(buffer)172 metrics.append(perf_metrics)173 baseline_reward = np.nanmean(baseline_rewards)174175 for idx, buffer in enumerate(buffers):176 # 第 6 项保存当前样本回报减去组内平均回报，训练时作为权重使用。177 for key in buffer.keys():178 if key == 6:179 for i in range(len(buffer[key])):180 buffer[key][i] += baseline_rewards[idx] - baseline_reward181 if key not in self.experience.keys():182 self.experience[key] = buffer[key]183 else:184 self.experience[key] += buffer[key]185186 for metric in metrics:187 for key in metric.keys():188 if key not in self.perf_metrics.keys():189 self.perf_metrics[key] = metric[key]190 else:191 self.perf_metrics[key] += metric[key]192193 if self.save_image:194 try:195 self.env.plot_animation(SaverParams.GIFS_PATH, episode_number)196 except:197 pass198 self.episode_number = episode_number199200 def convert_torch(self, args):201 data = []202 for arg in args:203 data.append(torch.tensor(arg, dtype=torch.float).to(self.device))204 return data205206 @staticmethod207 def obs_padding(task_info, agents, mask):208 """将任务和智能体数量补到网络初始化时使用的固定上限。"""209 task_info = F.pad(task_info, (0, 0, 0, EnvParams.TASKS_RANGE[1] + 1 - task_info.shape[1]), 'constant', 0)210 agents = F.pad(agents, (0, 0, 0, EnvParams.SPECIES_AGENTS_RANGE[1] * EnvParams.SPECIES_RANGE[1] - agents.shape[1]), 'constant', 0)211 mask = F.pad(mask, (0, EnvParams.TASKS_RANGE[1] + 1 - mask.shape[1]), 'constant', 1)212 return task_info, agents, mask213214215if __name__ == '__main__':216 device = torch.device('cuda')217 # torch.manual_seed(9)218 # checkpoint = torch.load(SaverParams.MODEL_PATH + '/checkpoint.pth')219 localNetwork = AttentionNet(TrainParams.AGENT_INPUT_DIM, TrainParams.TASK_INPUT_DIM, TrainParams.EMBEDDING_DIM).to(device)220 # localNetwork.load_state_dict(checkpoint['best_model'])221 for i in range(100):222 worker = Worker(1, localNetwork, localNetwork, 0, device=device, seed=i, save_image=False)223 worker.work(i)224 print(i)