class QLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.params = list(mac.parameters()) self.last_target_update_episode = 0 self.mixer = None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC self.target_mac = copy.deepcopy(mac) self.log_stats_t = -self.args.learner_log_interval - 1 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # Calculate estimated Q-Values mac_out = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.mac.forward(batch, t=t) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1) # Concat over time # Pick the Q-Values for the actions taken by each agent chosen_action_qvals = th.gather(mac_out[:, :-1], dim=3, index=actions).squeeze(3) # Remove the last dim # Calculate the Q-Values necessary for the target target_mac_out = [] self.target_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_outs = self.target_mac.forward(batch, t=t) target_mac_out.append(target_agent_outs) # We don't need the first timesteps Q-Value estimate for calculating targets target_mac_out = th.stack(target_mac_out[1:], dim=1) # Concat across time # Mask out unavailable actions target_mac_out[avail_actions[:, 1:] == 0] = -9999999 # Max over target Q-Values if self.args.double_q: # Get actions that maximise live Q (for double q-learning) mac_out_detach = mac_out.clone().detach() mac_out_detach[avail_actions == 0] = -9999999 cur_max_actions = mac_out_detach[:, 1:].max(dim=3, keepdim=True)[1] target_max_qvals = th.gather(target_mac_out, 3, cur_max_actions).squeeze(3) else: target_max_qvals = target_mac_out.max(dim=3)[0] # Mix if self.mixer is not None: chosen_action_qvals = self.mixer(chosen_action_qvals, batch["state"][:, :-1]) target_max_qvals = self.target_mixer(target_max_qvals, batch["state"][:, 1:]) # Calculate 1-step Q-Learning targets targets = rewards + self.args.gamma * (1 - terminated) * target_max_qvals # Td-error td_error = (chosen_action_qvals - targets.detach()) mask = mask.expand_as(td_error) # 0-out the targets that came from padded data masked_td_error = td_error * mask # Normal L2 loss, take mean over actual data loss = (masked_td_error ** 2).sum() / mask.sum() # Optimise self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.optimiser.step() if (episode_num - self.last_target_update_episode) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat("td_error_abs", (masked_td_error.abs().sum().item()/mask_elems), t_env) self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item()/(mask_elems * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item()/(mask_elems * self.args.n_agents), t_env) self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict(th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict(th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class SLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.n_actions_levin = args.n_actions self.params = list(mac.parameters()) self.last_target_update_episode = 0 self.mixer = None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.params += list(self.mixer.parameters()) if not args.SubAVG_Mixer_flag: self.target_mixer = copy.deepcopy(self.mixer) elif args.mixer == "qmix": self.target_mixer_list = [] for i in range(self.args.SubAVG_Mixer_K): self.target_mixer_list.append(copy.deepcopy(self.mixer)) self.levin_iter_target_mixer_update = 0 self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC if not self.args.SubAVG_Agent_flag: self.target_mac = copy.deepcopy(mac) else: self.target_mac_list = [] for i in range(self.args.SubAVG_Agent_K): self.target_mac_list.append(copy.deepcopy(mac)) self.levin_iter_target_update = 0 self.log_stats_t = -self.args.learner_log_interval - 1 # ====== levin ===== self.number = 0 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int, epsilon_levin=None): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # Calculate estimated Q-Values mac_out = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.mac.forward(batch, t=t) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1) # Concat over time # Pick the Q-Values for the actions taken by each agent chosen_action_qvals = th.gather(mac_out[:, :-1], dim=3, index=actions).squeeze(3) # Calculate the Q-Values necessary for the target target_mac_out = [] if not self.args.SubAVG_Agent_flag: self.target_mac.init_hidden(batch.batch_size) else: for i in range(self.args.SubAVG_Agent_K): self.target_mac_list[i].init_hidden(batch.batch_size) for t in range(batch.max_seq_length): if not self.args.SubAVG_Agent_flag: target_agent_outs = self.target_mac.forward(batch, t=t) # exp:使用 average DQN的target_mac else: target_agent_outs = 0 self.target_agent_out_list = [] for i in range(self.args.SubAVG_Agent_K): target_agent_out = self.target_mac_list[i].forward(batch, t=t) target_agent_outs = target_agent_outs + target_agent_out if self.args.SubAVG_Agent_flag_select: self.target_agent_out_list.append(target_agent_out) target_agent_outs = target_agent_outs / self.args.SubAVG_Agent_K if self.args.SubAVG_Agent_flag_select: if self.args.SubAVG_Agent_name_select_replacement == 'mean': target_out_select_sum = 0 for i in range(self.args.SubAVG_Agent_K): if self.args.SubAVG_Agent_flag_select > 0: target_out_select = th.where( self.target_agent_out_list[i] < target_agent_outs, target_agent_outs, self.target_agent_out_list[i]) else: target_out_select = th.where( self.target_agent_out_list[i] > target_agent_outs, target_agent_outs, self.target_agent_out_list[i]) target_out_select_sum = target_out_select_sum + target_out_select target_agent_outs = target_out_select_sum / self.args.SubAVG_Agent_K elif self.args.SubAVG_Agent_name_select_replacement == 'zero': target_out_select_sum = 0 target_select_bool_sum = 0 for i in range(self.args.SubAVG_Agent_K): if self.args.SubAVG_Agent_flag_select > 0: target_select_bool = ( self.target_agent_out_list[i] > target_agent_outs).float() target_out_select = th.where( self.target_agent_out_list[i] > target_agent_outs, self.target_agent_out_list[i], th.full_like(target_agent_outs, 0)) else: target_select_bool = ( self.target_agent_out_list[i] < target_agent_outs).float() target_out_select = th.where( self.target_agent_out_list[i] < target_agent_outs, self.target_agent_out_list[i], th.full_like(target_agent_outs, 0)) target_select_bool_sum = target_select_bool_sum + target_select_bool target_out_select_sum = target_out_select_sum + target_out_select if self.levin_iter_target_update < 2: pass # print("using average directly") else: target_agent_outs = target_out_select_sum / target_select_bool_sum target_mac_out.append(target_agent_outs) # We don't need the first timesteps Q-Value estimate for calculating targets target_mac_out = th.stack(target_mac_out, dim=1) # Concat across time # Mask out unavailable actions target_chosen_action_qvals = th.gather(target_mac_out, 3, batch['actions']).squeeze(-1) # Mix if self.mixer is None: target_qvals = target_chosen_action_qvals else: chosen_action_qvals = self.mixer(chosen_action_qvals, batch["state"][:, :-1]) if not self.args.SubAVG_Mixer_flag: target_qvals = self.target_mixer(target_chosen_action_qvals, batch['state']) elif self.args.mixer == "qmix": target_max_qvals_sum = 0 self.target_mixer_out_list = [] for i in range(self.args.SubAVG_Mixer_K): targe_mixer_out = self.target_mixer_list[i]( target_chosen_action_qvals, batch['state']) target_max_qvals_sum = target_max_qvals_sum + targe_mixer_out if self.args.SubAVG_Mixer_flag_select: self.target_mixer_out_list.append(targe_mixer_out) target_max_qvals = target_max_qvals_sum / self.args.SubAVG_Mixer_K # levin: mixer select if self.args.SubAVG_Mixer_flag_select: if self.args.SubAVG_Mixer_name_select_replacement == 'mean': target_mixer_select_sum = 0 for i in range(self.args.SubAVG_Mixer_K): if self.args.SubAVG_Mixer_flag_select > 0: target_mixer_select = th.where( self.target_mixer_out_list[i] < target_max_qvals, target_max_qvals, self.target_mixer_out_list[i]) else: target_mixer_select = th.where( self.target_mixer_out_list[i] > target_max_qvals, target_max_qvals, self.target_mixer_out_list[i]) target_mixer_select_sum = target_mixer_select_sum + target_mixer_select target_max_qvals = target_mixer_select_sum / self.args.SubAVG_Mixer_K elif self.args.SubAVG_Mixer_name_select_replacement == 'zero': target_mixer_select_sum = 0 target_mixer_select_bool_sum = 0 for i in range(self.args.SubAVG_Mixer_K): if self.args.SubAVG_Mixer_flag_select > 0: target_mixer_select_bool = ( self.target_mixer_out_list[i] > target_max_qvals).float() target_mixer_select = th.where( self.target_mixer_out_list[i] > target_max_qvals, self.target_mixer_out_list[i], th.full_like(target_max_qvals, 0)) else: target_mixer_select_bool = ( self.target_mixer_out_list[i] < target_max_qvals).float() target_mixer_select = th.where( self.target_mixer_out_list[i] < target_max_qvals, self.target_mixer_out_list[i], th.full_like(target_max_qvals, 0)) target_mixer_select_bool_sum = target_mixer_select_bool_sum + target_mixer_select_bool target_mixer_select_sum = target_mixer_select_sum + target_mixer_select if self.levin_iter_target_mixer_update < 2: pass # print("using average-mix directly") else: target_max_qvals = target_mixer_select_sum / target_mixer_select_bool_sum target_qvals = target_max_qvals if self.args.td_lambda <= 1 and self.args.td_lambda > 0: targets = build_td_lambda_targets(rewards, terminated, mask, target_qvals, self.args.n_agents, self.args.gamma, self.args.td_lambda) else: if self.args.td_lambda == 0: n = 1 # 1-step TD else: n = self.args.td_lambda targets = th.zeros_like(batch['reward']) targets += batch['reward'] for i in range(1, n): targets[:, :-i] += (self.args.gamma**i) * ( 1 - terminated[:, i - 1:]) * batch['reward'][:, i:] targets[:, :-n] += (self.args.gamma**n) * ( 1 - terminated[:, n - 1:]) * target_qvals[:, n:] targets = targets[:, :-1] # Td-error td_error = (chosen_action_qvals - targets.detach()) mask = mask.expand_as(td_error) # 0-out the targets that came from padded data masked_td_error = td_error * mask # Normal L2 loss, take mean over actual data loss = (masked_td_error**2).sum() / mask.sum() * 2 # Optimise self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.optimiser.step() if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", loss.item(), t_env) # self.logger.log_stat("loss_levin", loss_levin.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat( "td_error_abs", (masked_td_error.abs().sum().item() / mask_elems), t_env) self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.log_stats_t = t_env def _update_targets(self): if not self.args.SubAVG_Agent_flag: self.target_mac.load_state(self.mac) else: self.number = self.levin_iter_target_update % self.args.SubAVG_Agent_K self.target_mac_list[self.number].load_state(self.mac) self.levin_iter_target_update = self.levin_iter_target_update + 1 if self.mixer is not None: if not self.args.SubAVG_Mixer_flag: self.target_mixer.load_state_dict(self.mixer.state_dict()) elif self.args.mixer == "qmix": mixer_number = self.levin_iter_target_mixer_update % self.args.SubAVG_Mixer_K self.target_mixer_list[mixer_number].load_state_dict( self.mixer.state_dict()) self.levin_iter_target_mixer_update = self.levin_iter_target_mixer_update + 1 self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() if not self.args.SubAVG_Agent_flag: self.target_mac.cuda() else: for i in range(self.args.SubAVG_Agent_K): self.target_mac_list[i].cuda() if self.mixer is not None: self.mixer.cuda() if not self.args.SubAVG_Mixer_flag: self.target_mixer.cuda() elif self.args.mixer == "qmix": for i in range(self.args.SubAVG_Mixer_K): self.target_mixer_list[i].cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks if not self.args.SubAVG_Agent_flag: self.target_mac.load_models(path) else: for i in range(self.args.SubAVG_Agent_K): self.target_mac_list[i].load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class OffPGLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.n_agents = args.n_agents self.n_actions = args.n_actions self.mac = mac self.logger = logger self.last_target_update_step = 0 self.critic_training_steps = 0 self.log_stats_t = -self.args.learner_log_interval - 1 self.critic = OffPGCritic(scheme, args) self.mixer = QMixer(args) self.target_critic = copy.deepcopy(self.critic) self.target_mixer = copy.deepcopy(self.mixer) self.agent_params = list(mac.parameters()) self.critic_params = list(self.critic.parameters()) self.mixer_params = list(self.mixer.parameters()) self.params = self.agent_params + self.critic_params self.c_params = self.critic_params + self.mixer_params self.agent_optimiser = RMSprop(params=self.agent_params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) self.critic_optimiser = RMSprop(params=self.critic_params, lr=args.critic_lr, alpha=args.optim_alpha, eps=args.optim_eps) self.mixer_optimiser = RMSprop(params=self.mixer_params, lr=args.critic_lr, alpha=args.optim_alpha, eps=args.optim_eps) def train(self, batch: EpisodeBatch, t_env: int, log): # Get the relevant quantities bs = batch.batch_size max_t = batch.max_seq_length actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() avail_actions = batch["avail_actions"][:, :-1] mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) mask = mask.repeat(1, 1, self.n_agents).view(-1) states = batch["state"][:, :-1] #build q inputs = self.critic._build_inputs(batch, bs, max_t) q_vals = self.critic.forward(inputs).detach()[:, :-1] mac_out = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length - 1): agent_outs = self.mac.forward(batch, t=t) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1) # Concat over time # Mask out unavailable actions, renormalise (as in action selection) mac_out[avail_actions == 0] = 0 mac_out = mac_out / mac_out.sum(dim=-1, keepdim=True) mac_out[avail_actions == 0] = 0 # Calculated baseline q_taken = th.gather(q_vals, dim=3, index=actions).squeeze(3) pi = mac_out.view(-1, self.n_actions) baseline = th.sum(mac_out * q_vals, dim=-1).view(-1).detach() # Calculate policy grad with mask pi_taken = th.gather(pi, dim=1, index=actions.reshape(-1, 1)).squeeze(1) pi_taken[mask == 0] = 1.0 log_pi_taken = th.log(pi_taken) coe = self.mixer.k(states).view(-1) advantages = (q_taken.view(-1) - baseline).detach() coma_loss = -( (coe * advantages * log_pi_taken) * mask).sum() / mask.sum() # Optimise agents self.agent_optimiser.zero_grad() coma_loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.agent_params, self.args.grad_norm_clip) self.agent_optimiser.step() #compute parameters sum for debugging p_sum = 0. for p in self.agent_params: p_sum += p.data.abs().sum().item() / 100.0 if t_env - self.log_stats_t >= self.args.learner_log_interval: ts_logged = len(log["critic_loss"]) for key in [ "critic_loss", "critic_grad_norm", "td_error_abs", "q_taken_mean", "target_mean", "q_max_mean", "q_min_mean", "q_max_var", "q_min_var" ]: self.logger.log_stat(key, sum(log[key]) / ts_logged, t_env) self.logger.log_stat("q_max_first", log["q_max_first"], t_env) self.logger.log_stat("q_min_first", log["q_min_first"], t_env) #self.logger.log_stat("advantage_mean", (advantages * mask).sum().item() / mask.sum().item(), t_env) self.logger.log_stat("coma_loss", coma_loss.item(), t_env) self.logger.log_stat("agent_grad_norm", grad_norm, t_env) self.logger.log_stat("pi_max", (pi.max(dim=1)[0] * mask).sum().item() / mask.sum().item(), t_env) self.log_stats_t = t_env def train_critic(self, on_batch, best_batch=None, log=None): bs = on_batch.batch_size max_t = on_batch.max_seq_length rewards = on_batch["reward"][:, :-1] actions = on_batch["actions"][:, :] terminated = on_batch["terminated"][:, :-1].float() mask = on_batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = on_batch["avail_actions"][:] states = on_batch["state"] #build_target_q target_inputs = self.target_critic._build_inputs(on_batch, bs, max_t) target_q_vals = self.target_critic.forward(target_inputs).detach() targets_taken = self.target_mixer( th.gather(target_q_vals, dim=3, index=actions).squeeze(3), states) target_q = build_td_lambda_targets(rewards, terminated, mask, targets_taken, self.n_agents, self.args.gamma, self.args.td_lambda).detach() inputs = self.critic._build_inputs(on_batch, bs, max_t) mac_out = [] self.mac.init_hidden(bs) for i in range(max_t): agent_outs = self.mac.forward(on_batch, t=i) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1).detach() # Mask out unavailable actions, renormalise (as in action selection) mac_out[avail_actions == 0] = 0 mac_out = mac_out / mac_out.sum(dim=-1, keepdim=True) mac_out[avail_actions == 0] = 0 if best_batch is not None: best_target_q, best_inputs, best_mask, best_actions, best_mac_out = self.train_critic_best( best_batch) log["best_reward"] = th.mean( best_batch["reward"][:, :-1].squeeze(2).sum(-1), dim=0) target_q = th.cat((target_q, best_target_q), dim=0) inputs = th.cat((inputs, best_inputs), dim=0) mask = th.cat((mask, best_mask), dim=0) actions = th.cat((actions, best_actions), dim=0) states = th.cat((states, best_batch["state"]), dim=0) mac_out = th.cat((mac_out, best_mac_out), dim=0) #train critic mac_out = mac_out.detach() for t in range(max_t - 1): mask_t = mask[:, t:t + 1] if mask_t.sum() < 0.5: continue k = self.mixer.k(states[:, t:t + 1]).unsqueeze(3) #b = self.mixer.b(states[:, t:t+1]) q_vals = self.critic.forward(inputs[:, t:t + 1]) q_ori = q_vals q_vals = th.gather(q_vals, 3, index=actions[:, t:t + 1]).squeeze(3) q_vals = self.mixer.forward(q_vals, states[:, t:t + 1]) target_q_t = target_q[:, t:t + 1].detach() q_err = (q_vals - target_q_t) * mask_t critic_loss = (q_err**2).sum() / mask_t.sum() #Here introduce the loss for Qi v_vals = th.sum(q_ori * mac_out[:, t:t + 1], dim=3, keepdim=True) ad_vals = q_ori - v_vals goal = th.sum(k * v_vals, dim=2, keepdim=True) + k * ad_vals goal_err = (goal - q_ori) * mask_t goal_loss = 0.1 * (goal_err** 2).sum() / mask_t.sum() / self.args.n_actions #critic_loss += goal_loss self.critic_optimiser.zero_grad() self.mixer_optimiser.zero_grad() critic_loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.c_params, self.args.grad_norm_clip) self.critic_optimiser.step() self.mixer_optimiser.step() self.critic_training_steps += 1 log["critic_loss"].append(critic_loss.item()) log["critic_grad_norm"].append(grad_norm) mask_elems = mask_t.sum().item() log["td_error_abs"].append((q_err.abs().sum().item() / mask_elems)) log["target_mean"].append( (target_q_t * mask_t).sum().item() / mask_elems) log["q_taken_mean"].append( (q_vals * mask_t).sum().item() / mask_elems) log["q_max_mean"].append( (th.mean(q_ori.max(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems) log["q_min_mean"].append( (th.mean(q_ori.min(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems) log["q_max_var"].append( (th.var(q_ori.max(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems) log["q_min_var"].append( (th.var(q_ori.min(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems) if (t == 0): log["q_max_first"] = ( th.mean(q_ori.max(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems log["q_min_first"] = ( th.mean(q_ori.min(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems #update target network if (self.critic_training_steps - self.last_target_update_step ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_step = self.critic_training_steps def train_critic_best(self, batch): bs = batch.batch_size max_t = batch.max_seq_length rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"][:] states = batch["state"] # pr for all actions of the episode mac_out = [] self.mac.init_hidden(bs) for i in range(max_t): agent_outs = self.mac.forward(batch, t=i) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1).detach() # Mask out unavailable actions, renormalise (as in action selection) mac_out[avail_actions == 0] = 0 mac_out = mac_out / mac_out.sum(dim=-1, keepdim=True) mac_out[avail_actions == 0] = 0 critic_mac = th.gather(mac_out, 3, actions).squeeze(3).prod(dim=2, keepdim=True) #target_q take target_inputs = self.target_critic._build_inputs(batch, bs, max_t) target_q_vals = self.target_critic.forward(target_inputs).detach() targets_taken = self.target_mixer( th.gather(target_q_vals, dim=3, index=actions).squeeze(3), states) #expected q exp_q = self.build_exp_q(target_q_vals, mac_out, states).detach() # td-error targets_taken[:, -1] = targets_taken[:, -1] * (1 - th.sum(terminated, dim=1)) exp_q[:, -1] = exp_q[:, -1] * (1 - th.sum(terminated, dim=1)) targets_taken[:, :-1] = targets_taken[:, :-1] * mask exp_q[:, :-1] = exp_q[:, :-1] * mask td_q = (rewards + self.args.gamma * exp_q[:, 1:] - targets_taken[:, :-1]) * mask #compute target target_q = build_target_q(td_q, targets_taken[:, :-1], critic_mac, mask, self.args.gamma, self.args.tb_lambda, self.args.step).detach() inputs = self.critic._build_inputs(batch, bs, max_t) return target_q, inputs, mask, actions, mac_out def build_exp_q(self, target_q_vals, mac_out, states): target_exp_q_vals = th.sum(target_q_vals * mac_out, dim=3) target_exp_q_vals = self.target_mixer.forward(target_exp_q_vals, states) return target_exp_q_vals def _update_targets(self): self.target_critic.load_state_dict(self.critic.state_dict()) self.target_mixer.load_state_dict(self.mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.critic.cuda() self.mixer.cuda() self.target_critic.cuda() self.target_mixer.cuda() def save_models(self, path): self.mac.save_models(path) th.save(self.critic.state_dict(), "{}/critic.th".format(path)) th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.agent_optimiser.state_dict(), "{}/agent_opt.th".format(path)) th.save(self.critic_optimiser.state_dict(), "{}/critic_opt.th".format(path)) th.save(self.mixer_optimiser.state_dict(), "{}/mixer_opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) self.critic.load_state_dict( th.load("{}/critic.th".format(path), map_location=lambda storage, loc: storage)) self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) # Not quite right but I don't want to save target networks # self.target_critic.load_state_dict(self.critic.agent.state_dict()) self.target_mixer.load_state_dict(self.mixer.state_dict()) self.agent_optimiser.load_state_dict( th.load("{}/agent_opt.th".format(path), map_location=lambda storage, loc: storage)) self.critic_optimiser.load_state_dict( th.load("{}/critic_opt.th".format(path), map_location=lambda storage, loc: storage)) self.mixer_optimiser.load_state_dict( th.load("{}/mixer_opt.th".format(path), map_location=lambda storage, loc: storage))
class QLearner_3s_vs_4z: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.params = list(mac.parameters()) self.last_target_update_episode = 0 self.mixer = None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) self.params += list(self.mac.msg_rnn.parameters()) self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) self.target_mac = copy.deepcopy(mac) self.log_stats_t = -self.args.learner_log_interval - 1 self.loss_weight = [0.5, 1, 1.5] # its the beta in the Algorithm 1 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # Calculate estimated Q-Values mac_out = [] previous_msg_list = [] smooth_loss_list = [] regulariation_smooth = 1.5 regulariation_robust = 0.3 self.mac.init_hidden(batch.batch_size) smooth_loss = th.zeros((32 * 3)).cuda() for t in range(batch.max_seq_length): agent_local_outputs, input_hidden_states, vi = self.mac.forward( batch, t=t) input_hidden_states = input_hidden_states.view(-1, 64) self.mac.hidden_states_msg, dummy = self.mac.msg_rnn( self.mac.hidden_states_msg, input_hidden_states) ss = min(len(previous_msg_list), 3) #record the l2 difference in the window for i in range(ss): smooth_loss += self.loss_weight[i] * (( (dummy - previous_msg_list[i])**2).sum(dim=1)) / ( (ss * 32 * 3 * 10 * (dummy**2)).sum(dim=1)) previous_msg_list.append(dummy) if (len(previous_msg_list) > 3): previous_msg_list.pop(0) smooth_loss_reshape = smooth_loss.reshape(32, 3, 1).sum(1) #(32,1) smooth_loss_list.append(smooth_loss_reshape) # generate the message dummy_final = dummy.reshape(32, 3, 10) dummy0 = dummy_final[:, 0, :] dummy1 = dummy_final[:, 1, :] dummy2 = dummy_final[:, 2, :] agent0 = (dummy1 + dummy2) / 2.0 agent1 = (dummy0 + dummy2) / 2.0 agent2 = (dummy0 + dummy1) / 2.0 agent_global_outputs = th.cat((agent0.view( (32, 1, 10)), agent1.view((32, 1, 10)), agent2.view( (32, 1, 10))), 1) agent_outs = agent_local_outputs + agent_global_outputs mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1) # Concat over time ############compute the robustness loss################## robust_loss = th.topk(mac_out, 2)[0][:, :, :, 0] - th.topk( mac_out, 2)[0][:, :, :, 1] robust_loss = th.exp(-25.0 * robust_loss).sum( dim=2)[:, :-1].unsqueeze(2) / (32 * 6) #(32,38) # Pick the Q-Values for the actions taken by each agent chosen_action_qvals = th.gather(mac_out[:, :-1], dim=3, index=actions).squeeze( 3) # Remove the last dim # Calculate the Q-Values necessary for the target target_mac_out = [] self.target_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_local_outputs, target_input_hidden_states, tvi = self.target_mac.forward( batch, t=t) target_input_hidden_states = target_input_hidden_states.view( -1, 64) self.target_mac.hidden_states_msg, target_dummy = self.target_mac.msg_rnn( self.target_mac.hidden_states_msg, target_input_hidden_states) target_dummy_final = target_dummy.reshape(32, 3, 10) dummy0 = target_dummy_final[:, 0, :] dummy1 = target_dummy_final[:, 1, :] dummy2 = target_dummy_final[:, 2, :] target_agent0 = (dummy1 + dummy2) / 2.0 target_agent1 = (dummy0 + dummy2) / 2.0 target_agent2 = (dummy0 + dummy1) / 2.0 target_agent_global_outputs = th.cat((target_agent0.view( (32, 1, 10)), target_agent1.view( (32, 1, 10)), target_agent2.view((32, 1, 10))), 1) target_agent_outs = target_agent_local_outputs + target_agent_global_outputs target_mac_out.append(target_agent_outs) # We don't need the first timesteps Q-Value estimate for calculating targets target_mac_out = th.stack(target_mac_out[1:], dim=1) # Concat across time # Mask out unavailable actions target_mac_out[avail_actions[:, 1:] == 0] = -9999999 # Max over target Q-Values if self.args.double_q: # Get actions that maximise live Q (for double q-learning) mac_out[avail_actions == 0] = -9999999 cur_max_actions = mac_out[:, 1:].max(dim=3, keepdim=True)[1] target_max_qvals = th.gather(target_mac_out, 3, cur_max_actions).squeeze(3) else: target_max_qvals = target_mac_out.max(dim=3)[0] # Mix if self.mixer is not None: chosen_action_qvals = self.mixer(chosen_action_qvals, batch["state"][:, :-1]) target_max_qvals = self.target_mixer(target_max_qvals, batch["state"][:, 1:]) # Calculate 1-step Q-Learning targets targets = rewards + self.args.gamma * (1 - terminated) * target_max_qvals # Td-error td_error = (chosen_action_qvals - targets.detach()) mask = mask.expand_as(td_error) # 0-out the targets that came from padded data masked_td_error = td_error * mask ######compute the smooth_loss and robust_loss######### smooth_loss = th.stack(smooth_loss_list[0:-1], dim=1) smooth_loss = (smooth_loss * mask).sum() / mask.sum() robust_loss = (robust_loss * mask).sum() / mask.sum() # Normal L2 loss, take mean over actual data loss = (masked_td_error**2).sum() / mask.sum( ) + regulariation_smooth * smooth_loss + regulariation_robust * robust_loss # Optimise self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.optimiser.step() if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat( "td_error_abs", (masked_td_error.abs().sum().item() / mask_elems), t_env) self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class RODELearner: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.n_agents = args.n_agents self.params = list(mac.parameters()) self.last_target_update_episode = 0 self.mixer = None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) self.role_mixer = None if args.role_mixer is not None: if args.role_mixer == "vdn": self.role_mixer = VDNMixer() elif args.role_mixer == "qmix": self.role_mixer = QMixer(args) else: raise ValueError("Role Mixer {} not recognised.".format( args.role_mixer)) self.params += list(self.role_mixer.parameters()) self.target_role_mixer = copy.deepcopy(self.role_mixer) self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC self.target_mac = copy.deepcopy(mac) self.log_stats_t = -self.args.learner_log_interval - 1 self.role_interval = args.role_interval self.device = self.args.device self.role_action_spaces_updated = True # action encoder self.action_encoder_params = list(self.mac.action_encoder_params()) self.action_encoder_optimiser = RMSprop( params=self.action_encoder_params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # role_avail_actions = batch["role_avail_actions"] roles_shape_o = batch["roles"][:, :-1].shape role_at = int(np.ceil(roles_shape_o[1] / self.role_interval)) role_t = role_at * self.role_interval roles_shape = list(roles_shape_o) roles_shape[1] = role_t roles = th.zeros(roles_shape).to(self.device) roles[:, :roles_shape_o[1]] = batch["roles"][:, :-1] roles = roles.view(batch.batch_size, role_at, self.role_interval, self.n_agents, -1)[:, :, 0] # Calculate estimated Q-Values mac_out = [] role_out = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs, role_outs = self.mac.forward(batch, t=t) mac_out.append(agent_outs) if t % self.role_interval == 0 and t < batch.max_seq_length - 1: role_out.append(role_outs) mac_out = th.stack(mac_out, dim=1) # Concat over time role_out = th.stack(role_out, dim=1) # Concat over time # Pick the Q-Values for the actions taken by each agent chosen_action_qvals = th.gather(mac_out[:, :-1], dim=3, index=actions).squeeze( 3) # Remove the last dim chosen_role_qvals = th.gather(role_out, dim=3, index=roles.long()).squeeze(3) # Calculate the Q-Values necessary for the target target_mac_out = [] target_role_out = [] self.target_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_outs, target_role_outs = self.target_mac.forward( batch, t=t) target_mac_out.append(target_agent_outs) if t % self.role_interval == 0 and t < batch.max_seq_length - 1: target_role_out.append(target_role_outs) target_role_out.append( th.zeros(batch.batch_size, self.n_agents, self.mac.n_roles).to(self.device)) # We don't need the first timesteps Q-Value estimate for calculating targets target_mac_out = th.stack(target_mac_out[1:], dim=1) # Concat across time target_role_out = th.stack(target_role_out[1:], dim=1) # Mask out unavailable actions target_mac_out[avail_actions[:, 1:] == 0] = -9999999 # target_mac_out[role_avail_actions[:, 1:] == 0] = -9999999 # Max over target Q-Values if self.args.double_q: # Get actions that maximise live Q (for double q-learning) mac_out_detach = mac_out.clone().detach() mac_out_detach[avail_actions == 0] = -9999999 # mac_out_detach[role_avail_actions == 0] = -9999999 cur_max_actions = mac_out_detach[:, 1:].max(dim=3, keepdim=True)[1] target_max_qvals = th.gather(target_mac_out, 3, cur_max_actions).squeeze(3) role_out_detach = role_out.clone().detach() role_out_detach = th.cat( [role_out_detach[:, 1:], role_out_detach[:, 0:1]], dim=1) cur_max_roles = role_out_detach.max(dim=3, keepdim=True)[1] target_role_max_qvals = th.gather(target_role_out, 3, cur_max_roles).squeeze(3) else: target_max_qvals = target_mac_out.max(dim=3)[0] target_role_max_qvals = target_role_out.max(dim=3)[0] # Mix if self.mixer is not None: chosen_action_qvals = self.mixer(chosen_action_qvals, batch["state"][:, :-1]) target_max_qvals = self.target_mixer(target_max_qvals, batch["state"][:, 1:]) if self.role_mixer is not None: state_shape_o = batch["state"][:, :-1].shape state_shape = list(state_shape_o) state_shape[1] = role_t role_states = th.zeros(state_shape).to(self.device) role_states[:, :state_shape_o[1]] = batch["state"][:, :-1].detach( ).clone() role_states = role_states.view(batch.batch_size, role_at, self.role_interval, -1)[:, :, 0] chosen_role_qvals = self.role_mixer(chosen_role_qvals, role_states) role_states = th.cat([role_states[:, 1:], role_states[:, 0:1]], dim=1) target_role_max_qvals = self.target_role_mixer( target_role_max_qvals, role_states) # Calculate 1-step Q-Learning targets targets = rewards + self.args.gamma * (1 - terminated) * target_max_qvals rewards_shape = list(rewards.shape) rewards_shape[1] = role_t role_rewards = th.zeros(rewards_shape).to(self.device) role_rewards[:, :rewards.shape[1]] = rewards.detach().clone() role_rewards = role_rewards.view(batch.batch_size, role_at, self.role_interval).sum(dim=-1, keepdim=True) # role_terminated terminated_shape_o = terminated.shape terminated_shape = list(terminated_shape_o) terminated_shape[1] = role_t role_terminated = th.zeros(terminated_shape).to(self.device) role_terminated[:, :terminated_shape_o[1]] = terminated.detach().clone( ) role_terminated = role_terminated.view( batch.batch_size, role_at, self.role_interval).sum(dim=-1, keepdim=True) # role_terminated role_targets = role_rewards + self.args.gamma * ( 1 - role_terminated) * target_role_max_qvals # Td-error td_error = (chosen_action_qvals - targets.detach()) role_td_error = (chosen_role_qvals - role_targets.detach()) mask = mask.expand_as(td_error) mask_shape = list(mask.shape) mask_shape[1] = role_t role_mask = th.zeros(mask_shape).to(self.device) role_mask[:, :mask.shape[1]] = mask.detach().clone() role_mask = role_mask.view(batch.batch_size, role_at, self.role_interval, -1)[:, :, 0] # 0-out the targets that came from padded data masked_td_error = td_error * mask masked_role_td_error = role_td_error * role_mask # Normal L2 loss, take mean over actual data loss = (masked_td_error**2).sum() / mask.sum() role_loss = (masked_role_td_error**2).sum() / role_mask.sum() loss += role_loss # Optimise self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.optimiser.step() pred_obs_loss = None pred_r_loss = None pred_grad_norm = None if self.role_action_spaces_updated: # train action encoder no_pred = [] r_pred = [] for t in range(batch.max_seq_length): no_preds, r_preds = self.mac.action_repr_forward(batch, t=t) no_pred.append(no_preds) r_pred.append(r_preds) no_pred = th.stack(no_pred, dim=1)[:, :-1] # Concat over time r_pred = th.stack(r_pred, dim=1)[:, :-1] no = batch["obs"][:, 1:].detach().clone() repeated_rewards = batch["reward"][:, :-1].detach().clone( ).unsqueeze(2).repeat(1, 1, self.n_agents, 1) pred_obs_loss = th.sqrt(((no_pred - no)**2).sum(dim=-1)).mean() pred_r_loss = ((r_pred - repeated_rewards)**2).mean() pred_loss = pred_obs_loss + 10 * pred_r_loss self.action_encoder_optimiser.zero_grad() pred_loss.backward() pred_grad_norm = th.nn.utils.clip_grad_norm_( self.action_encoder_params, self.args.grad_norm_clip) self.action_encoder_optimiser.step() if t_env > self.args.role_action_spaces_update_start: self.mac.update_role_action_spaces() if 'noar' in self.args.mac: self.target_mac.role_selector.update_roles( self.mac.n_roles) self.role_action_spaces_updated = False self._update_targets() self.last_target_update_episode = episode_num if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", (loss - role_loss).item(), t_env) self.logger.log_stat("role_loss", role_loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) if pred_obs_loss is not None: self.logger.log_stat("pred_obs_loss", pred_obs_loss.item(), t_env) self.logger.log_stat("pred_r_loss", pred_r_loss.item(), t_env) self.logger.log_stat("action_encoder_grad_norm", pred_grad_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat( "td_error_abs", (masked_td_error.abs().sum().item() / mask_elems), t_env) self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("role_q_taken_mean", (chosen_role_qvals * role_mask).sum().item() / (role_mask.sum().item() * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) if self.role_mixer is not None: self.target_role_mixer.load_state_dict( self.role_mixer.state_dict()) self.target_mac.role_action_spaces_updated = self.role_action_spaces_updated self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() if self.role_mixer is not None: self.role_mixer.cuda() self.target_role_mixer.cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) if self.role_mixer is not None: th.save(self.role_mixer.state_dict(), "{}/role_mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) th.save(self.action_encoder_optimiser.state_dict(), "{}/action_repr_opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) if self.role_mixer is not None: self.role_mixer.load_state_dict( th.load("{}/role_mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage)) self.action_encoder_optimiser.load_state_dict( th.load("{}/action_repr_opt.th".format(path), map_location=lambda storage, loc: storage))
class LatentQLearner(QLearner): def __init__(self, mac, scheme, logger, args): super(LatentQLearner, self).__init__(mac, scheme, logger, args) self.args = args self.mac = mac self.logger = logger self.params = list(mac.parameters()) self.last_target_update_episode = 0 self.mixer = None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC self.target_mac = copy.deepcopy(mac) self.log_stats_t = -self.args.learner_log_interval - 1 self.role_save = 0 self.role_save_interval = 10 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # Calculate estimated Q-Values mac_out = [] self.mac.init_hidden(batch.batch_size) indicator, latent, latent_vae = self.mac.init_latent(batch.batch_size) reg_loss = 0 dis_loss = 0 ce_loss = 0 for t in range(batch.max_seq_length): agent_outs, loss_, dis_loss_, ce_loss_ = self.mac.forward( batch, t=t, t_glob=t_env, train_mode=True) # (bs,n,n_actions),(bs,n,latent_dim) reg_loss += loss_ dis_loss += dis_loss_ ce_loss += ce_loss_ # loss_cs=self.args.gamma*loss_cs + _loss mac_out.append(agent_outs) # [t,(bs,n,n_actions)] # mac_out_latent.append((agent_outs_latent)) #[t,(bs,n,latent_dim)] reg_loss /= batch.max_seq_length dis_loss /= batch.max_seq_length ce_loss /= batch.max_seq_length mac_out = th.stack(mac_out, dim=1) # Concat over time # (bs,t,n,n_actions), Q values of n_actions # mac_out_latent=th.stack(mac_out_latent,dim=1) # (bs,t,n,latent_dim) # mac_out_latent=mac_out_latent.reshape(-1,self.args.latent_dim) # Pick the Q-Values for the actions taken by each agent chosen_action_qvals = th.gather(mac_out[:, :-1], dim=3, index=actions).squeeze( 3) # Remove the last dim # (bs,t,n) Q value of an action # Calculate the Q-Values necessary for the target target_mac_out = [] self.target_mac.init_hidden(batch.batch_size) # (bs,n,hidden_size) self.target_mac.init_latent(batch.batch_size) # (bs,n,latent_size) for t in range(batch.max_seq_length): target_agent_outs, loss_cs_target, _, _ = self.target_mac.forward( batch, t=t) # (bs,n,n_actions), (bs,n,latent_dim) target_mac_out.append(target_agent_outs) # [t,(bs,n,n_actions)] # We don't need the first timesteps Q-Value estimate for calculating targets target_mac_out = th.stack( target_mac_out[1:], dim=1) # Concat across time, dim=1 is time index # (bs,t,n,n_actions) # Mask out unavailable actions target_mac_out[avail_actions[:, 1:] == 0] = -9999999 # Q values # Max over target Q-Values if self.args.double_q: # True for QMix # Get actions that maximise live Q (for double q-learning) mac_out_detach = mac_out.clone().detach( ) # return a new Tensor, detached from the current graph mac_out_detach[avail_actions == 0] = -9999999 # (bs,t,n,n_actions), discard t=0 cur_max_actions = mac_out_detach[:, 1:].max( dim=3, keepdim=True)[1] # indices instead of values # (bs,t,n,1) target_max_qvals = th.gather(target_mac_out, 3, cur_max_actions).squeeze(3) # (bs,t,n,n_actions) ==> (bs,t,n,1) ==> (bs,t,n) max target-Q else: target_max_qvals = target_mac_out.max(dim=3)[0] # Mix if self.mixer is not None: chosen_action_qvals = self.mixer(chosen_action_qvals, batch["state"][:, :-1]) target_max_qvals = self.target_mixer(target_max_qvals, batch["state"][:, 1:]) # (bs,t,1) # Calculate 1-step Q-Learning targets targets = rewards + self.args.gamma * (1 - terminated) * target_max_qvals # Td-error td_error = (chosen_action_qvals - targets.detach() ) # no gradient through target net # (bs,t,1) mask = mask.expand_as(td_error) # 0-out the targets that came from padded data masked_td_error = td_error * mask # Normal L2 loss, take mean over actual data loss = (masked_td_error**2).sum() / mask.sum() # entropy loss # mac_out_latent_norm=th.sqrt(th.sum(mac_out_latent*mac_out_latent,dim=1)) # mac_out_latent=mac_out_latent/mac_out_latent_norm[:,None] # loss+=(th.norm(mac_out_latent)/mac_out_latent.size(0))*self.args.entropy_loss_weight loss += reg_loss # Optimise self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_( self.params, self.args.grad_norm_clip) # max_norm self.optimiser.step() if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: # if self.role_save % self.role_save_interval == 0: # self.role_save = 0 # if self.args.latent_dim in [2, 3]: # fig = plt.figure() # ax = fig.add_subplot(111, projection='3d') # print(self.mac.agent.latent[:, :self.args.latent_dim], # self.mac.agent.latent[:, -self.args.latent_dim:]) # self.role_save += 1 self.logger.log_stat("loss", loss.item(), t_env) self.logger.log_stat("loss_reg", reg_loss.item(), t_env) self.logger.log_stat("loss_dis", dis_loss.item(), t_env) self.logger.log_stat("loss_ce", ce_loss.item(), t_env) #indicator=[var_mean,mi.max(),mi.min(),mi.mean(),mi.std(),di.max(),di.min(),di.mean(),di.std()] self.logger.log_stat("var_mean", indicator[0].item(), t_env) self.logger.log_stat("mi_max", indicator[1].item(), t_env) self.logger.log_stat("mi_min", indicator[2].item(), t_env) self.logger.log_stat("mi_mean", indicator[3].item(), t_env) self.logger.log_stat("mi_std", indicator[4].item(), t_env) self.logger.log_stat("di_max", indicator[5].item(), t_env) self.logger.log_stat("di_min", indicator[6].item(), t_env) self.logger.log_stat("di_mean", indicator[7].item(), t_env) self.logger.log_stat("di_std", indicator[8].item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat( "td_error_abs", (masked_td_error.abs().sum().item() / mask_elems), t_env) self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) if self.args.use_tensorboard: # log_vec(self,mat,metadata,label_img,global_step,tag) self.logger.log_vec(latent, list(range(self.args.n_agents)), t_env, "latent") self.logger.log_vec(latent_vae, list(range(self.args.n_agents)), t_env, "latent-VAE") self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class CateQLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.params = list(mac.parameters()) self.last_target_update_episode = 0 self.mixer = None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC self.target_mac = copy.deepcopy(mac) self.log_stats_t = -self.args.learner_log_interval - 1 self.s_mu = th.zeros(1) self.s_sigma = th.ones(1) def get_comm_beta(self, t_env): comm_beta = self.args.comm_beta if self.args.is_comm_beta_decay and t_env > self.args.comm_beta_start_decay: comm_beta += 1. * (self.args.comm_beta_target - self.args.comm_beta) / \ (self.args.comm_beta_end_decay - self.args.comm_beta_start_decay) * \ (t_env - self.args.comm_beta_start_decay) return comm_beta def get_comm_entropy_beta(self, t_env): comm_entropy_beta = self.args.comm_entropy_beta if self.args.is_comm_entropy_beta_decay and t_env > self.args.comm_entropy_beta_start_decay: comm_entropy_beta += 1. * (self.args.comm_entropy_beta_target - self.args.comm_entropy_beta) / \ (self.args.comm_entropy_beta_end_decay - self.args.comm_entropy_beta_start_decay) * \ (t_env - self.args.comm_entropy_beta_start_decay) return comm_entropy_beta def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # Calculate estimated Q-Values # shape = (bs, self.n_agents, -1) mac_out = [] mu_out = [] sigma_out = [] logits_out = [] m_sample_out = [] g_out = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): if self.args.comm and self.args.use_IB: agent_outs, (mu, sigma), logits, m_sample = self.mac.forward(batch, t=t) mu_out.append(mu) sigma_out.append(sigma) logits_out.append(logits) m_sample_out.append(m_sample) else: agent_outs = self.mac.forward(batch, t=t) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1) # Concat over time if self.args.use_IB: mu_out = th.stack(mu_out, dim=1)[:, :-1] # Concat over time sigma_out = th.stack(sigma_out, dim=1)[:, :-1] # Concat over time logits_out = th.stack(logits_out, dim=1)[:, :-1] m_sample_out = th.stack(m_sample_out, dim=1)[:, :-1] # Pick the Q-Values for the actions taken by each agent chosen_action_qvals = th.gather(mac_out[:, :-1], dim=3, index=actions).squeeze( 3) # Remove the last dim # I believe that code up to here is right... # Q values are right, the main issue is to calculate loss for message... # Calculate the Q-Values necessary for the target target_mac_out = [] self.target_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): if self.args.comm and self.args.use_IB: target_agent_outs, (target_mu, target_sigma), target_logits, target_m_sample = \ self.target_mac.forward(batch, t=t) else: target_agent_outs = self.target_mac.forward(batch, t=t) target_mac_out.append(target_agent_outs) # label label_target_max_out = th.stack(target_mac_out[:-1], dim=1) label_target_max_out[avail_actions[:, :-1] == 0] = -9999999 label_target_actions = label_target_max_out.max(dim=3, keepdim=True)[1] # We don't need the first timesteps Q-Value estimate for calculating targets target_mac_out = th.stack(target_mac_out[1:], dim=1) # Concat across time # Mask out unavailable actions target_mac_out[avail_actions[:, 1:] == 0] = -9999999 # Max over target Q-Values if self.args.double_q: # Get actions that maximise live Q (for double q-learning) mac_out[avail_actions == 0] = -9999999 cur_max_actions = mac_out[:, 1:].max(dim=3, keepdim=True)[1] target_max_qvals = th.gather(target_mac_out, 3, cur_max_actions).squeeze(3) else: target_max_qvals = target_mac_out.max(dim=3)[0] # Mix if self.mixer is not None: chosen_action_qvals = self.mixer(chosen_action_qvals, batch["state"][:, :-1]) target_max_qvals = self.target_mixer(target_max_qvals, batch["state"][:, 1:]) # Calculate 1-step Q-Learning targets targets = rewards + self.args.gamma * (1 - terminated) * target_max_qvals # Td-error td_error = (chosen_action_qvals - targets.detach()) mask = mask.expand_as(td_error) # 0-out the targets that came from padded data masked_td_error = td_error * mask # Normal L2 loss, take mean over actual data loss = (masked_td_error**2).sum() / mask.sum() if self.args.only_downstream or not self.args.use_IB: expressiveness_loss = th.Tensor([0.]) compactness_loss = th.Tensor([0.]) entropy_loss = th.Tensor([0.]) comm_loss = th.Tensor([0.]) comm_beta = th.Tensor([0.]) comm_entropy_beta = th.Tensor([0.]) else: # ### Optimize message # Message are controlled only by expressiveness and compactness loss. # Compute cross entropy with target q values of the same time step expressiveness_loss = 0 label_prob = th.gather(logits_out, 3, label_target_actions).squeeze(3) expressiveness_loss += ( -th.log(label_prob + 1e-6)).sum() / mask.sum() # Compute KL divergence compactness_loss = D.kl_divergence(D.Normal(mu_out, sigma_out), D.Normal(self.s_mu, self.s_sigma)).sum() / \ mask.sum() # Entropy loss entropy_loss = -D.Normal(self.s_mu, self.s_sigma).log_prob( m_sample_out).sum() / mask.sum() # Gate loss gate_loss = 0 # Total loss comm_beta = self.get_comm_beta(t_env) comm_entropy_beta = self.get_comm_entropy_beta(t_env) comm_loss = expressiveness_loss + comm_beta * compactness_loss + comm_entropy_beta * entropy_loss comm_loss *= self.args.c_beta loss += comm_loss comm_beta = th.Tensor([comm_beta]) comm_entropy_beta = th.Tensor([comm_entropy_beta]) # Optimise self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.optimiser.step() # Update target if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", loss.item(), t_env) self.logger.log_stat("comm_loss", comm_loss.item(), t_env) self.logger.log_stat("exp_loss", expressiveness_loss.item(), t_env) self.logger.log_stat("comp_loss", compactness_loss.item(), t_env) self.logger.log_stat("comm_beta", comm_beta.item(), t_env) self.logger.log_stat("entropy_loss", entropy_loss.item(), t_env) self.logger.log_stat("comm_beta", comm_beta.item(), t_env) self.logger.log_stat("comm_entropy_beta", comm_entropy_beta.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat( "td_error_abs", (masked_td_error.abs().sum().item() / mask_elems), t_env) self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) # self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() self.s_mu = self.s_mu.cuda() self.s_sigma = self.s_sigma.cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class FMACLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.n_agents = args.n_agents self.n_actions = args.n_actions self.logger = logger self.mac = mac self.target_mac = copy.deepcopy(self.mac) self.agent_params = list(mac.parameters()) self.critic = FMACCritic(scheme, args) self.target_critic = copy.deepcopy(self.critic) self.critic_params = list(self.critic.parameters()) self.mixer = None if args.mixer is not None and self.args.n_agents > 1: # if just 1 agent do not mix anything if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.critic_params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) print('Mixer Size: ') print(get_parameters_num(self.critic_params)) if getattr(self.args, "optimizer", "rmsprop") == "rmsprop": self.agent_optimiser = RMSprop(params=self.agent_params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) elif getattr(self.args, "optimizer", "rmsprop") == "adam": self.agent_optimiser = Adam(params=self.agent_params, lr=args.lr, eps=getattr(args, "optimizer_epsilon", 10E-8)) else: raise Exception("unknown optimizer {}".format( getattr(self.args, "optimizer", "rmsprop"))) if getattr(self.args, "optimizer", "rmsprop") == "rmsprop": self.critic_optimiser = RMSprop(params=self.critic_params, lr=args.critic_lr, alpha=args.optim_alpha, eps=args.optim_eps) elif getattr(self.args, "optimizer", "rmsprop") == "adam": self.critic_optimiser = Adam(params=self.critic_params, lr=args.critic_lr, eps=getattr(args, "optimizer_epsilon", 10E-8)) else: raise Exception("unknown optimizer {}".format( getattr(self.args, "optimizer", "rmsprop"))) self.log_stats_t = -self.args.learner_log_interval - 1 self.last_target_update_episode = 0 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int, off=False): # Get the relevant data rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] actions_onehot = batch["actions_onehot"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # Retrace Q target with th.no_grad(): q1, _ = self.target_critic(batch, batch["actions_onehot"].detach()) target_vals = self.target_mixer(q1, batch["state"]) lambd = 0 if off else self.args.lambd target_vals = build_td_lambda_targets(rewards, terminated, mask, target_vals, self.n_agents, self.args.gamma, lambd) # Train the critic # Current Q network forward q1, _ = self.critic(batch[:, :-1], actions_onehot.detach()) q_taken = self.mixer(q1, batch["state"][:, :-1]) critic_loss = 0.5 * ( (q_taken - target_vals.detach()) * mask).pow(2).sum() / mask.sum() self.critic_optimiser.zero_grad() critic_loss.backward() critic_grad_norm = th.nn.utils.clip_grad_norm_( self.critic_params, self.args.grad_norm_clip) self.critic_optimiser.step() # Train the actor if not off: pi = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length - 1): agent_outs = self.mac.forward(batch, t=t) pi.append(agent_outs) pi = th.stack(pi, dim=1) # Concat over time b, t, a, probs q1, _ = self.critic(batch[:, :-1], pi) q = self.mixer(q1, batch["state"][:, :-1]) pg_loss = -(q * mask).sum() / mask.sum() entropy_loss = categorical_entropy(pi).mean( -1, keepdim=True) # mean over agents entropy_loss[mask == 0] = 0 # fill nan entropy_loss = (entropy_loss * mask).sum() / mask.sum() loss = pg_loss - self.args.entropy_coef * entropy_loss / entropy_loss.item( ) self.agent_optimiser.zero_grad() loss.backward() agent_grad_norm = th.nn.utils.clip_grad_norm_( self.agent_params, self.args.grad_norm_clip) self.agent_optimiser.step() # target_update if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num # log if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("critic_loss", critic_loss.item(), t_env) self.logger.log_stat("critic_grad_norm", critic_grad_norm.item(), t_env) self.logger.log_stat("target_vals", (target_vals * mask).sum().item() / mask.sum().item(), t_env) if not off: self.logger.log_stat("pg_loss", pg_loss.item(), t_env) self.logger.log_stat("entropy_loss", entropy_loss.item(), t_env) self.logger.log_stat("agent_grad_norm", agent_grad_norm.item(), t_env) agent_mask = mask.repeat(1, 1, self.n_agents) self.logger.log_stat( "pi_max", (pi.max(dim=-1)[0] * agent_mask).sum().item() / agent_mask.sum().item(), t_env) self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) self.target_critic.load_state_dict(self.critic.state_dict()) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) # self.logger.console_logger.info("Updated all target networks") def cuda(self, device="cuda"): self.mac.cuda() self.target_mac.cuda() self.critic.to(device=device) self.target_critic.to(device=device) if self.mixer is not None: self.mixer.to(device=device) self.target_mixer.to(device=device) def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.agent_optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.agent_optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class MAXQLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.mac_params = list(mac.parameters()) self.params = list(self.mac.parameters()) self.last_target_update_episode = 0 self.mixer = None assert args.mixer is not None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.mixer_params = list(self.mixer.parameters()) self.params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC self.target_mac = copy.deepcopy(mac) # Central Q # TODO: Clean this mess up! self.central_mac = None if self.args.central_mixer in ["ff", "atten"]: if self.args.central_loss == 0: self.central_mixer = self.mixer self.central_mac = self.mac self.target_central_mac = self.target_mac else: if self.args.central_mixer == "ff": self.central_mixer = QMixerCentralFF( args ) # Feedforward network that takes state and agent utils as input # elif self.args.central_mixer == "atten": # self.central_mixer = QMixerCentralAtten(args) else: raise Exception("Error with central_mixer") assert args.central_mac == "basic_central_mac" self.central_mac = mac_REGISTRY[args.central_mac]( scheme, args ) # Groups aren't used in the CentralBasicController. Little hacky self.target_central_mac = copy.deepcopy(self.central_mac) self.params += list(self.central_mac.parameters()) else: raise Exception("Error with qCentral") self.params += list(self.central_mixer.parameters()) self.target_central_mixer = copy.deepcopy(self.central_mixer) print('Mixer Size: ') print( get_parameters_num( list(self.mixer.parameters()) + list(self.central_mixer.parameters()))) self.optimiser = Adam(params=self.params, lr=args.lr) self.log_stats_t = -self.args.learner_log_interval - 1 self.grad_norm = 1 self.mixer_norm = 1 self.mixer_norms = deque([1], maxlen=100) def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # Calculate estimated Q-Values mac_out = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.mac.forward(batch, t=t) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1) # Concat over time # Pick the Q-Values for the actions taken by each agent chosen_action_qvals_agents = th.gather(mac_out[:, :-1], dim=3, index=actions).squeeze( 3) # Remove the last dim chosen_action_qvals = chosen_action_qvals_agents # Calculate the Q-Values necessary for the target target_mac_out = [] self.target_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_outs = self.target_mac.forward(batch, t=t) target_mac_out.append(target_agent_outs) # We don't need the first timesteps Q-Value estimate for calculating targets target_mac_out = th.stack(target_mac_out[:], dim=1) # Concat across time # Mask out unavailable actions target_mac_out[avail_actions[:, :] == 0] = -9999999 # From OG deepmarl # Max over target Q-Values if self.args.double_q: # Get actions that maximise live Q (for double q-learning) mac_out_detach = mac_out.clone().detach() mac_out_detach[avail_actions == 0] = -9999999 cur_max_action_targets, cur_max_actions = mac_out_detach[:, :].max( dim=3, keepdim=True) target_max_agent_qvals = th.gather( target_mac_out[:, :], 3, cur_max_actions[:, :]).squeeze(3) else: raise Exception("Use double q") # Central MAC stuff central_mac_out = [] self.central_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.central_mac.forward(batch, t=t) central_mac_out.append(agent_outs) central_mac_out = th.stack(central_mac_out, dim=1) # Concat over time central_chosen_action_qvals_agents = th.gather( central_mac_out[:, :-1], dim=3, index=actions.unsqueeze(4).repeat( 1, 1, 1, 1, self.args.central_action_embed)).squeeze( 3) # Remove the last dim central_target_mac_out = [] self.target_central_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_outs = self.target_central_mac.forward(batch, t=t) central_target_mac_out.append(target_agent_outs) central_target_mac_out = th.stack(central_target_mac_out[:], dim=1) # Concat across time # Mask out unavailable actions central_target_mac_out[avail_actions[:, :] == 0] = -9999999 # From OG deepmarl # Use the Qmix max actions central_target_max_agent_qvals = th.gather( central_target_mac_out[:, :], 3, cur_max_actions[:, :].unsqueeze(4).repeat( 1, 1, 1, 1, self.args.central_action_embed)).squeeze(3) # --- # Mix chosen_action_qvals = self.mixer(chosen_action_qvals, batch["state"][:, :-1]) target_max_qvals = self.target_central_mixer( central_target_max_agent_qvals, batch["state"]) # We use the calculation function of sarsa lambda to approximate q star lambda targets = build_td_lambda_targets(rewards, terminated, mask, target_max_qvals, self.args.n_agents, self.args.gamma, self.args.td_lambda) # Td-error td_error = (chosen_action_qvals - (targets.detach())) mask = mask.expand_as(td_error) # 0-out the targets that came from padded data masked_td_error = td_error * mask # Training central Q central_chosen_action_qvals = self.central_mixer( central_chosen_action_qvals_agents, batch["state"][:, :-1]) central_td_error = (central_chosen_action_qvals - targets.detach()) central_mask = mask.expand_as(central_td_error) central_masked_td_error = central_td_error * central_mask central_loss = 0.5 * (central_masked_td_error**2).sum() / mask.sum() # QMIX loss with weighting ws = th.ones_like(td_error) * self.args.w if self.args.hysteretic_qmix: # OW-QMIX ws = th.where(td_error < 0, th.ones_like(td_error) * 1, ws) # Target is greater than current max w_to_use = ws.mean().item() # For logging else: # CW-QMIX is_max_action = (actions == cur_max_actions[:, :-1]).min(dim=2)[0] max_action_qtot = self.target_central_mixer( central_target_max_agent_qvals[:, :-1], batch["state"][:, :-1]) qtot_larger = targets > max_action_qtot ws = th.where(is_max_action | qtot_larger, th.ones_like(td_error) * 1, ws) # Target is greater than current max w_to_use = ws.mean().item() # Average of ws for logging qmix_loss = (ws.detach() * (masked_td_error**2)).sum() / mask.sum() # The weightings for the different losses aren't used (they are always set to 1) loss = self.args.qmix_loss * qmix_loss + self.args.central_loss * central_loss # Optimise self.optimiser.zero_grad() loss.backward() # Logging agent_norm = 0 for p in self.mac_params: param_norm = p.grad.data.norm(2) agent_norm += param_norm.item()**2 agent_norm = agent_norm**(1. / 2) mixer_norm = 0 for p in self.mixer_params: param_norm = p.grad.data.norm(2) mixer_norm += param_norm.item()**2 mixer_norm = mixer_norm**(1. / 2) self.mixer_norm = mixer_norm self.mixer_norms.append(mixer_norm) grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.grad_norm = grad_norm self.optimiser.step() if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", loss.item(), t_env) self.logger.log_stat("qmix_loss", qmix_loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) self.logger.log_stat("mixer_norm", mixer_norm, t_env) self.logger.log_stat("agent_norm", agent_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat( "td_error_abs", (masked_td_error.abs().sum().item() / mask_elems), t_env) self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("central_loss", central_loss.item(), t_env) self.logger.log_stat("w_to_use", w_to_use, t_env) self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) if self.central_mac is not None: self.target_central_mac.load_state(self.central_mac) self.target_central_mixer.load_state_dict( self.central_mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() if self.central_mac is not None: self.central_mac.cuda() self.target_central_mac.cuda() self.central_mixer.cuda() self.target_central_mixer.cuda() # TODO: Model saving/loading is out of date! def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class SACQLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.mac_params = list(mac.parameters()) self.params = list(self.mac.parameters()) self.last_target_update_episode = 0 self.mixer = None assert args.mixer is not None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.mixer_params = list(self.mixer.parameters()) self.params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC self.target_mac = copy.deepcopy(mac) # Central Q # TODO: Clean this mess up! self.central_mac = None assert self.args.central_mixer == "ff" self.central_mixer = QMixerCentralFF(args) assert args.central_mac == "basic_central_mac" self.central_mac = mac_REGISTRY[args.central_mac]( scheme, args ) # Groups aren't used in the CentralBasicController. Little hacky self.target_central_mac = copy.deepcopy(self.central_mac) self.params += list(self.central_mac.parameters()) self.params += list(self.central_mixer.parameters()) self.target_central_mixer = copy.deepcopy(self.central_mixer) self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) self.log_stats_t = -self.args.learner_log_interval - 1 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # Current policies mac_out = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.mac.forward(batch, t=t) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1) # Concat over time # Mask out unavailable actions, renormalise (as in action selection) mac_out[avail_actions == 0] = 0 mac_out = mac_out / mac_out.sum(dim=-1, keepdim=True) mac_out[avail_actions == 0] = 0 mac_out[(mac_out.sum(dim=-1, keepdim=True) == 0).expand_as( mac_out )] = 1 # Set any all 0 probability vectors to all 1s. They will be masked out later, but still need to be sampled. # Target policies target_mac_out = [] self.target_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_outs = self.target_mac.forward(batch, t=t) target_mac_out.append(target_agent_outs) target_mac_out = th.stack(target_mac_out, dim=1) # Concat across time # Mask out unavailable actions, renormalise (as in action selection) target_mac_out[avail_actions == 0] = 0 target_mac_out = target_mac_out / target_mac_out.sum(dim=-1, keepdim=True) target_mac_out[avail_actions == 0] = 0 target_mac_out[( target_mac_out.sum(dim=-1, keepdim=True) == 0 ).expand_as( target_mac_out )] = 1 # Set any all 0 probability vectors to all 1s. They will be masked out later, but still need to be sampled. # Sample actions sampled_actions = Categorical(mac_out).sample().long() sampled_target_actions = Categorical(target_mac_out).sample().long() # Central MAC stuff central_mac_out = [] self.central_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.central_mac.forward(batch, t=t) central_mac_out.append(agent_outs) central_mac_out = th.stack(central_mac_out, dim=1) # Concat over time # Actions chosen from replay buffer central_chosen_action_qvals_agents = th.gather( central_mac_out[:, :-1], dim=3, index=actions.unsqueeze(4).repeat( 1, 1, 1, 1, self.args.central_action_embed)).squeeze( 3) # Remove the last dim central_target_mac_out = [] self.target_central_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_outs = self.target_central_mac.forward(batch, t=t) central_target_mac_out.append(target_agent_outs) central_target_mac_out = th.stack(central_target_mac_out[:], dim=1) # Concat across time central_target_action_qvals_agents = th.gather(central_target_mac_out[:,:], 3, \ sampled_target_actions[:,:].unsqueeze(3).unsqueeze(4)\ .repeat(1,1,1,1,self.args.central_action_embed)).squeeze(3) # --- critic_bootstrap_qvals = self.target_central_mixer( central_target_action_qvals_agents[:, 1:], batch["state"][:, 1:]) target_chosen_action_probs = th.gather( target_mac_out, dim=3, index=sampled_target_actions.unsqueeze(3)).squeeze(dim=3) target_policy_logs = th.log(target_chosen_action_probs).sum( dim=2, keepdim=True) # Sum across agents # Calculate 1-step Q-Learning targets targets = rewards + self.args.gamma * (1 - terminated) * \ (critic_bootstrap_qvals - self.args.entropy_temp * target_policy_logs[:,1:]) # Training Critic central_chosen_action_qvals = self.central_mixer( central_chosen_action_qvals_agents, batch["state"][:, :-1]) central_td_error = (central_chosen_action_qvals - targets.detach()) central_mask = mask.expand_as(central_td_error) central_masked_td_error = central_td_error * central_mask central_loss = (central_masked_td_error**2).sum() / mask.sum() # Actor Loss central_sampled_action_qvals_agents = th.gather(central_mac_out[:, :-1], 3, \ sampled_actions[:, :-1].unsqueeze(3).unsqueeze(4) \ .repeat(1, 1, 1, 1, self.args.central_action_embed)).squeeze(3) central_sampled_action_qvals = self.central_mixer( central_sampled_action_qvals_agents, batch["state"][:, :-1]).repeat(1, 1, self.args.n_agents) sampled_action_probs = th.gather( mac_out, dim=3, index=sampled_actions.unsqueeze(3)).squeeze(3) policy_logs = th.log(sampled_action_probs)[:, :-1] actor_mask = mask.expand_as(policy_logs) actor_loss = ( (policy_logs * (self.args.entropy_temp * (policy_logs + 1) - central_sampled_action_qvals).detach()) * actor_mask).sum() / actor_mask.sum() loss = self.args.actor_loss * actor_loss + self.args.central_loss * central_loss # Optimise self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.grad_norm = grad_norm self.optimiser.step() if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", loss.item(), t_env) self.logger.log_stat("actor_loss", actor_loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat("target_mean", (targets * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("central_loss", central_loss.item(), t_env) ps = mac_out[:, :-1] * avail_actions[:, :-1] log_ps = th.log(mac_out[:, :-1] + 0.00001) * avail_actions[:, :-1] actor_entropy = -(( (ps * log_ps).sum(dim=3) * mask).sum() / mask.sum()) self.logger.log_stat("actor_entropy", actor_entropy.item(), t_env) self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) if self.central_mac is not None: self.target_central_mac.load_state(self.central_mac) self.target_central_mixer.load_state_dict( self.central_mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() if self.central_mac is not None: self.central_mac.cuda() self.target_central_mac.cuda() self.central_mixer.cuda() self.target_central_mixer.cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class PGLearner_v2: def __init__(self, mac, scheme, logger, args): self.args = args self.n_agents = args.n_agents self.n_actions = args.n_actions self.mac = mac self.logger = logger self.last_target_update_step = 0 self.critic_training_steps = 0 self.log_stats_t = -self.args.learner_log_interval - 1 self.target_mac = copy.deepcopy(mac) self.params = list(self.mac.parameters()) if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.params += list(self.mixer.parameters()) if self.args.optim == 'adam': self.optimiser = Adam(params=self.params, lr=args.lr) else: self.optimiser = RMSprop(params=self.params, lr=args.lr) def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities bs = batch.batch_size max_t = batch.max_seq_length rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"][:, :] critic_mask = mask.clone() mask = mask.repeat(1, 1, self.n_agents).view(-1) advantages, td_error, targets_taken, log_pi_taken, entropy = self._calculate_advs(batch, rewards, terminated, actions, avail_actions, critic_mask, bs, max_t) pg_loss = - ((advantages.detach() * log_pi_taken) * mask).sum() / mask.sum() vf_loss = ((td_error ** 2) * mask).sum() / mask.sum() entropy[mask == 0] = 0 entropy_loss = (entropy * mask).sum() / mask.sum() coma_loss = pg_loss + self.args.vf_coef * vf_loss if self.args.ent_coef: coma_loss -= self.args.ent_coef * entropy_loss # Optimise agents self.optimiser.zero_grad() coma_loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.optimiser.step() if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("critic_loss", ((td_error ** 2) * mask).sum().item() / mask.sum().item(), t_env) self.logger.log_stat("td_error_abs", (td_error.abs() * mask).sum().item() / mask.sum().item(), t_env) self.logger.log_stat("q_taken_mean", (targets_taken * mask).sum().item() / mask.sum().item(), t_env) self.logger.log_stat("target_mean", ((targets_taken + advantages) * mask).sum().item() / mask.sum().item(), t_env) self.logger.log_stat("pg_loss", - ((advantages.detach() * log_pi_taken) * mask).sum().item() / mask.sum().item(), t_env) self.logger.log_stat("advantage_mean", (advantages * mask).sum().item() / mask.sum().item(), t_env) self.logger.log_stat("coma_loss", coma_loss.item(), t_env) self.logger.log_stat("entropy_loss", entropy_loss.item(), t_env) self.logger.log_stat("agent_grad_norm", grad_norm, t_env) # self.logger.log_stat("pi_max", (pi.max(dim=1)[0] * mask).sum().item() / mask.sum().item(), t_env) self.log_stats_t = t_env def _calculate_advs(self, batch, rewards, terminated, actions, avail_actions, mask, bs, max_t): mac_out = [] q_outs = [] # Roll out experiences self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_out, q_out = self.mac.forward(batch, t=t) mac_out.append(agent_out) q_outs.append(q_out) mac_out = th.stack(mac_out, dim=1) # Concat over time q_outs = th.stack(q_outs, dim=1) # Concat over time # Mask out unavailable actions, renormalise (as in action selection) # mac_out[avail_actions == 0] = 0 # mac_out = mac_out/(mac_out.sum(dim=-1, keepdim=True) + 1e-5) # Calculated baseline pi = mac_out[:, :-1] #[bs, t, n_agents, n_actions] pi_taken = th.gather(pi, dim=-1, index=actions[:, :-1]).squeeze(-1) #[bs, t, n_agents] action_mask = mask.repeat(1, 1, self.n_agents) pi_taken[action_mask == 0] = 1.0 log_pi_taken = th.log(pi_taken).reshape(-1) # Calculate entropy entropy = categorical_entropy(pi).reshape(-1) #[bs, t, n_agents, 1] # Calculate q targets targets_taken = q_outs.squeeze(-1) #[bs, t, n_agents] if self.args.mixer: targets_taken = self.mixer(targets_taken, batch["state"][:, :]) #[bs, t, 1] # Calculate td-lambda targets targets = build_td_lambda_targets(rewards, terminated, mask, targets_taken, self.n_agents, self.args.gamma, self.args.td_lambda) advantages = targets - targets_taken[:, :-1] advantages = advantages.unsqueeze(2).repeat(1, 1, self.n_agents, 1).reshape(-1) td_error = targets_taken[:, :-1] - targets.detach() td_error = td_error.unsqueeze(2).repeat(1, 1, self.n_agents, 1).reshape(-1) return advantages, td_error, targets_taken[:, :-1].unsqueeze(2).repeat(1, 1, self.n_agents, 1).reshape(-1), log_pi_taken, entropy def cuda(self): self.mac.cuda() if self.args.mixer: self.mixer.cuda() def save_models(self, path): self.mac.save_models(path) if self.args.mixer: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) if self.args.mixer: self.mixer.load_state_dict(th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict(th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class SCLearner: def __init__(self, mac, scheme, logger, args): torch.autograd.set_detect_anomaly(True) self.args = args self.n_agents = args.n_agents self.n_actions = args.n_actions self.mac = mac self.logger = logger self.last_target_update_step = 0 self.critic_training_steps = 0 self.log_stats_t = -self.args.learner_log_interval - 1 self.control_critic = SCControlCritic(scheme, args) self.execution_critic = SCExecutionCritic(scheme, args) self.target_control_critic = copy.deepcopy(self.control_critic) self.target_execution_critic = copy.deepcopy(self.execution_critic) self.control_actor_params = list(self.mac.agent_dlstm_parameters()) + list(self.mac.latent_state_encoder_parameters()) self.execution_actor_params = list(self.mac.agent_lstm_parameters()) self.control_critic_params = list(self.control_critic.parameters()) self.execution_critic_params = list(self.execution_critic.parameters()) self.control_mixer = None if args.control_mixer is not None: if args.control_mixer == "vdn": self.control_mixer = VDNMixer() elif args.control_mixer == "qmix": self.control_mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.control_mixer)) self.target_control_mixer = copy.deepcopy(self.control_mixer) self.control_critic_params += list(self.control_mixer.parameters()) self.execution_mixer = None if args.execution_mixer is not None: if args.execution_mixer == "vdn": self.execution_mixer = VDNMixer() elif args.execution_mixer == "qmix": self.execution_mixer = DirMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.execution_mixer)) self.target_execution_mixer = copy.deepcopy(self.execution_mixer) self.execution_critic_params += list(self.execution_mixer.parameters()) self.control_actor_optimiser = RMSprop(params=self.control_actor_params, lr=args.control_actor_lr, alpha=args.optim_alpha, eps=args.optim_eps) self.control_critic_optimiser = RMSprop(params=self.control_critic_params, lr=args.control_critic_lr, alpha=args.optim_alpha, eps=args.optim_eps) self.execution_actor_optimiser = RMSprop(params=self.execution_actor_params, lr=args.execution_actor_lr, alpha=args.optim_alpha, eps=args.optim_eps) self.execution_critic_optimiser = RMSprop(params=self.execution_critic_params, lr=args.execution_critic_lr, alpha=args.optim_alpha, eps=args.optim_eps) # self.target_mac = copy.deepcopy(mac) # FIXME: implement double mac # control: # - decentralized actor pi(|): # - decentralized critic q(|) with mixer: estimate rewards # execution: # - decentralized actor pi(|): # - decentralized critic q(|) with mixer: estimate directional contributions def train(self, batch, t_env, epsidoe_num): bs = batch.batch_size max_t = batch.max_seq_length rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"][:, :-1] dirs_vals, execution_critic_train_stats = self._train_execution_critic(batch, terminated, mask) # [bs, seq_len, n_agents, latent_state_dim] q_vals, control_critic_train_stats = self._train_control_critic(batch, rewards, terminated, mask) # [bs, seq_len, n_agents] self.critic_training_steps += 1 actions = actions[:, :-1] lstm_out = [] dlstm_query_out = [] dlstm_key_out = [] dlstm_rule_out = [] lstm_r = [] self.mac.init_hidden(bs) self.mac.init_goals(bs) for t in range(batch["reward"].shape[1] - 1): # skip inferring latent state lstm_output, dlstm_output = self.mac.forward(batch, t) lstm_out.append(lstm_output) # [batch_num][n_agents][n_actions] dlstm_query_out.append(dlstm_output[0]) dlstm_key_out.append(dlstm_output[1]) dlstm_rule_out.append(dlstm_output[2]) intr_r = self._get_instrinsic_r(dirs_vals, t) lstm_r.append(self.args.instr_r_rate * intr_r + q_vals[:, t]) lstm_out = torch.stack(lstm_out, dim=1) # [batch_num][seq][n_agents][n_actions] dlstm_query_out = torch.stack(dlstm_query_out, dim=1) dlstm_key_out = torch.stack(dlstm_key_out, dim=1) dlstm_rule_out = torch.stack(dlstm_rule_out, dim=1) lstm_r = torch.stack(lstm_r, dim=1) dlstm_loss_partial = [] for t in range(batch["reward"].shape[1] - self.args.horizon-1): # FIXEME: can implement slice instead of t iterations dlstm_loss_partial_t = self._get_dlistm_partial(dirs_vals, dlstm_query_out, dlstm_key_out, dlstm_rule_out, t) dlstm_loss_partial.append(dlstm_loss_partial_t) dlstm_loss_partial = torch.stack(dlstm_loss_partial, dim=1) # [bs, seq_len-c, n_agents] # Mask out unavailable actions, renormalise (as in action selection) lstm_out[avail_actions == 0] = 0 # [batch_num][seq][n_agents][n_actions] lstm_out = lstm_out/lstm_out.sum(dim=-1, keepdim=True) lstm_out[avail_actions == 0] = 0 # FIXME: implement q baseline q_vals = q_vals[:, :-self.args.horizon] q_vals = q_vals.reshape(-1, 1).squeeze(1) lstm_r = lstm_r.reshape(-1, 1).squeeze(1) pi = lstm_out.reshape(-1, self.n_actions) mask_dlstm = mask.clone().repeat(1, 1, self.n_agents)[:, :-self.args.horizon].reshape(-1, 1).squeeze(1) mask = mask.repeat(1, 1, self.n_agents).view(-1) # print("mask_dlstm shape:{}".format(mask_dlstm.shape)) pi_taken = torch.gather(pi, dim=1, index=actions.reshape(-1, 1)).squeeze(1) pi_taken[mask == 0] = 1.0 log_pi_taken = torch.log(pi_taken) dlstm_loss_partial = dlstm_loss_partial.reshape(-1, 1).squeeze(1) dlstm_loss_partial[mask_dlstm == 0] = 0.0 dlstm_loss = ((dlstm_loss_partial * q_vals.detach()) * mask_dlstm).sum() / mask_dlstm.sum() dlstm_loss += self.mac.compute_lat_state_kl_div() lstm_loss = ((log_pi_taken * lstm_r) * mask).sum() / mask.sum() self.control_actor_optimiser.zero_grad() dlstm_loss.backward(retain_graph=True) dlstm_grad_norm = torch.nn.utils.clip_grad_norm_(self.control_actor_params, self.args.grad_norm_clip) self.control_actor_optimiser.step() self.execution_actor_optimiser.zero_grad() lstm_loss.backward() lstm_grad_norm = torch.nn.utils.clip_grad_norm_(self.execution_actor_params, self.args.grad_norm_clip) self.execution_actor_optimiser.step() if (self.critic_training_steps - self.last_target_update_step) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_step = self.critic_training_steps if t_env - self.log_stats_t >= self.args.learner_log_interval: ts_logged = len(control_critic_train_stats["control_critic_td_loss"]) for key in ["control_critic_td_loss", "control_critic_grad_norm"]: self.logger.log_stat(key, sum(control_critic_train_stats[key])/ts_logged, t_env) ts_logged = len(execution_critic_train_stats["execution_critic_td_loss"]) for key in ["execution_critic_td_loss", "execution_critic_grad_norm"]: self.logger.log_stat(key, sum(execution_critic_train_stats[key]) / ts_logged, t_env) self.logger.log_stat("control_actor_loss", dlstm_loss.item(), t_env) self.logger.log_stat("execution_actor_loss", lstm_loss.item(), t_env) self.logger.log_stat("control_grad_norm", dlstm_grad_norm, t_env) self.logger.log_stat("execution_grad_norm", lstm_grad_norm, t_env) self.log_stats_t = t_env def _get_instrinsic_r(self, dirs_val, t_ep): r = 0 for t in range(0, min(self.args.horizon, t_ep)): idx = t+1 dec_lat_states = dirs_val[:, t_ep-idx] # [bs][n_agents][lat_dim] goals = self.mac.get_prev_goal(idx) # [batch_num][n_agents][rule_dim] # calculate the cosine similarity between agent's estimated contribution and its goal r += F.cosine_similarity(dec_lat_states, goals, dim=2) # penalize if agent has no goal r -= self.args.cum_goal_zeros_penalty_rate * F.cosine_similarity(torch.zeros(goals.shape), goals, dim=2) return r/self.args.horizon # [bs][n_agents] def _get_dlistm_partial(self, dirs_vals, queries, keys, rules, t): # calculate dcos(zt+c−zt,gt(θ)) in multi-agent context # print("dirs_vals shape: {}".format(dirs_vals.shape)) dir_vals_t = dirs_vals[:, t] dir_vals_tc = dirs_vals[:, t+self.args.horizon] query = queries[:, t] key = keys[:, t] rule = rules[:, t] dlstm_partial = [] for i in range(self.n_agents): # calc decomposed latent's projection on rule at t p_i = [torch.einsum('ij,ij->i', dir_vals_t[:, j], rule[:, i]) for j in range(self.n_agents)] # rule is already normalized p_i = torch.stack(p_i, dim=1) # [bs, n_agents] p_t = [torch.einsum('i,ij->ij', p_i[:, j], rule[:, i]) for j in range(self.n_agents)] p_t = torch.stack(p_t, dim=1) # [bs, n_agents, latent_state_dim] # calc decomposed latent's projection on rule at t+c p_i_c = [torch.einsum('ij,ij->i', dir_vals_tc[:, j], rule[:, i]) for j in range(self.n_agents)] # rule is already normalized p_i_c = torch.stack(p_i_c, dim=1) # [bs, n_agents] p_tc = [torch.einsum('i,ij->ij', p_i_c[:, j], rule[:, i]) for j in range(self.n_agents)] p_tc = torch.stack(p_tc, dim=1) # [bs, n_agents, latent_state_dim] p_diff = p_t - p_tc # [bs, n_agents, latent_state_dim] qk_i = [torch.einsum('ij,ij->i', query[:, i], key[:, j]) # 2d torch.dot / self.args.attention_noramlization_squared_dk for j in range(self.n_agents)] qk_i_t = torch.stack(qk_i, dim=1) # [bs, n_agents] a_i = torch.nn.functional.softmax(qk_i_t, dim=1) # eq 2 in TarMac # FIXME: only save attention and use attention directly dlstm_partial_i = [torch.einsum("i,ij->ij", 1/a_i[:, j], p_diff[:, j]) for j in range(self.n_agents)] dlstm_partial_i = torch.stack(dlstm_partial_i, dim=1).sum(dim=1) # [bs, latent_state_dim] dlstm_partial_i = F.cosine_similarity(dlstm_partial_i, rule[:, i], dim=1) dlstm_partial.append(dlstm_partial_i) return torch.stack(dlstm_partial, dim=1) # [bs, n_agents] def _train_control_critic(self, batch, rewards, terminated, mask): # FIXME: terminated? bs = batch.batch_size qs_vals = [] running_log = { "control_critic_td_loss": [], "control_critic_grad_norm": [], } for t in range(batch["reward"].shape[1]-1): # print(mask.shape) mask_t = mask[:, t] if mask_t.sum() == 0: continue qs_t = self.control_critic(batch, t) # [bs, n_agents] qs_vals.append(qs_t.clone()) # [bs_size][t][lat_dim] qs_tot = self.control_mixer(qs_t.unsqueeze(1), batch["latent_state"][:, t].unsqueeze(1)) target_qs_t = self.target_control_critic(batch, t+1) target_qs_tot = self.target_control_mixer(target_qs_t.unsqueeze(1), batch["latent_state"][:, t+1].unsqueeze(1)) td_loss = qs_tot - (rewards[:, t] + self.args.control_discount*(1 - terminated[:,t])*target_qs_tot.detach()) td_loss = (td_loss**2).sum() self.control_critic_optimiser.zero_grad() td_loss.backward(retain_graph=True) grad_norm = torch.nn.utils.clip_grad_norm_(self.control_critic_params, self.args.grad_norm_clip) self.control_critic_optimiser.step() running_log["control_critic_td_loss"].append(td_loss.item()) running_log["control_critic_grad_norm"].append(grad_norm) return torch.stack(qs_vals, dim=1), running_log def _train_execution_critic(self, batch, terminated, mask): bs = batch.batch_size dirs_tot_vals = [] running_log = { "execution_critic_td_loss": [], "execution_critic_grad_norm": [], } for t in range(batch["reward"].shape[1]-1): # print(mask.shape) mask_t = mask[:, t] if mask_t.sum() == 0: continue # distance between latent states lat_state_target_dis = torch.sub(batch["latent_state"][:, t+1], batch["latent_state"][:, t]) # [bs_size, latent_state_dim] dirs_t = self.execution_critic(batch, t) # [bs, n_agents, latent_state_dim] dirs_tot_vals.append(dirs_t.clone()) # # [bs_size][t][lat_dim] dirs_tot = self.execution_mixer(dirs_t.unsqueeze(1), batch["latent_state"][:, t].unsqueeze(1)) target_dirs_t = self.target_execution_critic(batch, t+1) target_dirs_tot = self.target_execution_mixer(target_dirs_t.unsqueeze(1), batch["latent_state"][:, t+1].unsqueeze(1)) td_loss = dirs_tot - (lat_state_target_dis + self.args.execution_discount*(1 - terminated[:, t])*target_dirs_tot.detach()) td_loss = (td_loss ** 2).sum() self.execution_critic_optimiser.zero_grad() td_loss.backward(retain_graph=True) grad_norm = torch.nn.utils.clip_grad_norm_(self.control_critic_params, self.args.grad_norm_clip) self.execution_critic_optimiser.step() running_log["execution_critic_td_loss"].append(td_loss.item()) running_log["execution_critic_grad_norm"].append(grad_norm) return torch.stack(dirs_tot_vals, dim=1), running_log def _update_targets(self): self.target_control_critic.load_state_dict(self.control_critic.state_dict()) self.target_control_mixer.load_state_dict(self.control_mixer.state_dict()) self.target_execution_critic.load_state_dict(self.execution_critic.state_dict()) self.target_execution_mixer.load_state_dict(self.execution_mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.control_critic.cuda() self.execution_critic.cuda() self.target_control_critic.cuda() self.target_execution_critic.cuda() self.control_mixer.cuda() self.execution_mixer.cuda() self.target_control_mixer() self.target_execution_mixer() def save_models(self, path): self.mac.save_models(path) torch.save(self.control_critic.state_dict(), "{}/control_critic.th".format(path)) torch.save(self.execution_critic.state_dict(), "{}/execution_critic.th".format(path)) torch.save(self.control_mixer.state_dict(), "{}/control_mixer.th".format(path)) torch.save(self.execution_mixer.state_dict(), "{}/execution_mixer.th".format(path)) torch.save(self.control_critic_optimiser.state_dict(), "{}/control_critic_opt.th".format(path)) torch.save(self.control_actor_optimiser.state_dict(), "{}/control_actor_opt.th".format(path)) torch.save(self.execution_critic_optimiser.state_dict(), "{}/execution_critic_opt.th".format(path)) torch.save(self.execution_actor_optimiser.state_dict(), "{}/execution_actor_opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) self.control_critic.load_state_dict(torch.load( "{}/control_critic.th".format(path), map_location=lambda storage, loc: storage)) # Not quite right but I don't want to save target networks self.execution_critic.load_state_dict(torch.load( "{}/execution_critic.th".format(path), map_location=lambda storage, loc: storage)) self.control_mixer.load_state_dict( torch.load("{}/control_mixer.th".format(path), map_location=lambda storage, loc: storage)) self.execution_mixer.load_state_dict( torch.load("{}/execution_mixer.th".format(path), map_location=lambda storage, loc: storage)) self.control_actor_optimiser.load_state_dict( torch.load("{}/control_actor_opt.th".format(path), map_location=lambda storage, loc: storage)) self.execution_actor_optimiser.load_state_dict( torch.load("{}/execution_actor_opt.th".format(path), map_location=lambda storage, loc: storage)) self.control_critic_optimiser.load_state_dict( torch.load("{}/control_critic_opt.th".format(path), map_location=lambda storage, loc: storage)) self.execution_critic_optimiser.load_state_dict( torch.load("{}/execution_critic_opt.th".format(path), map_location=lambda storage, loc: storage))
class QLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.params = list(mac.parameters()) # agent들의 신경망을 등록합니다. self.last_target_update_episode = 0 self.mixer = None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.params += list( self.mixer.parameters()) # mixing network의 신경망을 추가 등록 합니다. self.target_mixer = copy.deepcopy(self.mixer) self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) self.target_mac = copy.deepcopy(mac) self.log_stats_t = -self.args.learner_log_interval - 1 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # 학습에 필요한 정보들을 정리 합니다. rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # agent 개개인의 Q 값을 산출 합니다. mac_out = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.mac.forward(batch, t=t) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1) # 시간 순에 따라 Concat 합니다. # agent가 수행한 action에 대한 Q 값을 가져 옵니다. chosen_action_qvals = th.gather( mac_out[:, :-1], dim=3, index=actions).squeeze(3) # 마지막 차원을 제거함. # agent 개개인의 target network의 Q 값을 산출 합니다 target_mac_out = [] self.target_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_outs = self.target_mac.forward(batch, t=t) target_mac_out.append(target_agent_outs) # 첫번째 step의 target network의 Q 값은 필요없습니다. target_mac_out = th.stack(target_mac_out[1:], dim=1) # 시간 순에 따라 Concat 합니다. # 수행 불가능한 action을 masking 합니다. target_mac_out[avail_actions[:, 1:] == 0] = -9999999 if self.args.double_q: mac_out_detach = mac_out.clone().detach() mac_out_detach[avail_actions == 0] = -9999999 cur_max_actions = mac_out_detach[:, 1:].max(dim=3, keepdim=True)[1] target_max_qvals = th.gather(target_mac_out, 3, cur_max_actions).squeeze(3) else: target_max_qvals = target_mac_out.max(dim=3)[0] # Q-tot과 target mixing network의 Q-tot를 산출합니다. if self.mixer is not None: chosen_action_qvals = self.mixer(chosen_action_qvals, batch["state"][:, :-1]) target_max_qvals = self.target_mixer(target_max_qvals, batch["state"][:, 1:]) # mixing network를 위한 loss를 산출 합니다. targets = rewards + self.args.gamma * (1 - terminated) * target_max_qvals td_error = (chosen_action_qvals - targets.detach()) mask = mask.expand_as(td_error) masked_td_error = td_error * mask loss = (masked_td_error**2).sum() / mask.sum() # 역전파를 수행하여 mixing network 및 agent 모두의 가중치들을 갱신합니다. self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.optimiser.step() # 일정 주기마다 target network를 업데이트 합니다. if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat( "td_error_abs", (masked_td_error.abs().sum().item() / mask_elems), t_env) self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class FacMADDPGLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.n_agents = args.n_agents self.n_actions = args.n_actions self.logger = logger self.mac = mac self.target_mac = copy.deepcopy(self.mac) self.agent_params = list(mac.parameters()) self.critic = FacMADDPGCritic(scheme, args) self.target_critic = copy.deepcopy(self.critic) self.critic_params = list(self.critic.parameters()) self.mixer = None if args.mixer is not None and self.args.n_agents > 1: # if just 1 agent do not mix anything if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "vdn-s": self.mixer = VDNState(args) elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.critic_params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) if getattr(self.args, "optimizer", "rmsprop") == "rmsprop": self.agent_optimiser = RMSprop(params=self.agent_params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) elif getattr(self.args, "optimizer", "rmsprop") == "adam": self.agent_optimiser = Adam(params=self.agent_params, lr=args.lr, eps=getattr(args, "optimizer_epsilon", 10E-8)) else: raise Exception("unknown optimizer {}".format( getattr(self.args, "optimizer", "rmsprop"))) if getattr(self.args, "optimizer", "rmsprop") == "rmsprop": self.critic_optimiser = RMSprop(params=self.critic_params, lr=args.critic_lr, alpha=args.optim_alpha, eps=args.optim_eps) elif getattr(self.args, "optimizer", "rmsprop") == "adam": self.critic_optimiser = Adam(params=self.critic_params, lr=args.critic_lr, eps=getattr(args, "optimizer_epsilon", 10E-8)) else: raise Exception("unknown optimizer {}".format( getattr(self.args, "optimizer", "rmsprop"))) self.log_stats_t = -self.args.learner_log_interval - 1 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = 1 - terminated # Train the critic inputs = self._build_inputs(batch, t=0) q_taken, _ = self.critic(inputs, actions.detach()) q_taken = q_taken.view(batch.batch_size, -1, 1) if self.mixer is not None: q_taken = self.mixer(q_taken, batch["state"][:, :-1]) # Use the target actor and target critic network to compute the target q target_actions = [] for t in range(1, batch.max_seq_length): agent_target_outs = self.target_mac.select_actions(batch, t_ep=t, t_env=None, test_mode=True) target_actions.append(agent_target_outs) target_actions = th.stack(target_actions, dim=1) # Concat over time target_inputs = self._build_inputs(batch, t=1) target_vals, _ = self.target_critic(target_inputs, target_actions.detach()) target_vals = target_vals.view(batch.batch_size, -1, 1) if self.mixer is not None: target_vals = self.target_mixer(target_vals, batch["state"][:, 1:]) targets = rewards.expand_as(target_vals) + self.args.gamma * ( 1 - terminated.expand_as(target_vals)) * target_vals td_error = (q_taken - targets.detach()) masked_td_error = td_error loss = (masked_td_error**2).mean() # Optimise the critic self.critic_optimiser.zero_grad() loss.backward() critic_grad_norm = th.nn.utils.clip_grad_norm_( self.critic_params, self.args.grad_norm_clip) self.critic_optimiser.step() # Train the actor pi = self.mac.forward(batch, t=0, select_actions=True)["actions"] q, _ = self.critic(self._build_inputs(batch, t=0), pi) q = q.view(batch.batch_size, -1, 1) # Use the joint Q to update the actor if self.mixer is not None: q = self.mixer(q, batch["state"][:, :-1]) pg_loss = -q.mean() + (pi**2).mean() * 1e-3 # Optimise the agents self.agent_optimiser.zero_grad() th.autograd.set_detect_anomaly(True) # DBG pg_loss.backward() agent_grad_norm = th.nn.utils.clip_grad_norm_(self.agent_params, self.args.grad_norm_clip) self.agent_optimiser.step() if getattr(self.args, "target_update_mode", "hard") == "hard": self._update_targets() elif getattr(self.args, "target_update_mode", "hard") in ["soft", "exponential_moving_average"]: self._update_targets_soft( tau=getattr(self.args, "target_update_tau", 0.001)) else: raise Exception("unknown target update mode: {}!".format( getattr(self.args, "target_update_mode", "hard"))) def _update_targets_soft(self, tau): for target_param, param in zip(self.target_mac.parameters(), self.mac.parameters()): target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau) for target_param, param in zip(self.target_critic.parameters(), self.critic.parameters()): target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau) if self.mixer is not None: for target_param, param in zip(self.target_mixer.parameters(), self.mixer.parameters()): target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau) if self.args.verbose: self.logger.console_logger.info( "Updated all target networks (soft update tau={})".format(tau)) def _build_inputs(self, batch, t): bs = batch.batch_size inputs = [] inputs.append(batch["obs"][:, t]) if self.args.obs_last_action: if t == 0: inputs.append(th.zeros_like(batch["actions"][:, t])) else: inputs.append(batch["actions"][:, t - 1]) if self.args.obs_agent_id: inputs.append( th.eye(self.n_agents, device=batch.device).unsqueeze(0).expand(bs, -1, -1)) inputs = th.cat([x.reshape(bs * self.n_agents, -1) for x in inputs], dim=1) return inputs def _update_targets(self): self.target_mac.load_state(self.mac) self.target_critic.load_state_dict(self.critic.state_dict()) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) self.logger.console_logger.info("Updated all target networks") def cuda(self, device="cuda:0"): self.mac.cuda(device=device) self.target_mac.cuda(device=device) self.critic.cuda(device=device) self.target_critic.cuda(device=device) if self.mixer is not None: self.mixer.cuda(device=device) self.target_mixer.cuda(device=device) def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.agent_optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.agent_optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class DDPGQLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.mac_params = list(mac.parameters()) self.params = list(self.mac.parameters()) self.last_target_update_episode = 0 self.mixer = None assert args.mixer is not None if args.mixer is not None: if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.mixer_params = list(self.mixer.parameters()) self.params += list(self.mixer.parameters()) self.target_mixer = copy.deepcopy(self.mixer) # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC self.target_mac = copy.deepcopy(mac) # Central Q # TODO: Clean this mess up! self.central_mac = None assert self.args.central_mixer == "ff" self.central_mixer = QMixerCentralFF(args) assert args.central_mac == "basic_central_mac" self.central_mac = mac_REGISTRY[args.central_mac]( scheme, args ) # Groups aren't used in the CentralBasicController. Little hacky self.target_central_mac = copy.deepcopy(self.central_mac) self.params += list(self.central_mac.parameters()) self.params += list(self.central_mixer.parameters()) self.target_central_mixer = copy.deepcopy(self.central_mixer) self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) self.log_stats_t = -self.args.learner_log_interval - 1 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, :-1] actions = batch["actions"][:, :-1] terminated = batch["terminated"][:, :-1].float() mask = batch["filled"][:, :-1].float() mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1]) avail_actions = batch["avail_actions"] # Calculate estimated Q-Values mac_out = [] self.mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.mac.forward(batch, t=t) mac_out.append(agent_outs) mac_out = th.stack(mac_out, dim=1) # Concat over time # Pick the Q-Values for the actions taken by each agent chosen_action_qvals_agents = th.gather(mac_out[:, :-1], dim=3, index=actions).squeeze( 3) # Remove the last dim chosen_action_qvals = chosen_action_qvals_agents # Calculate the Q-Values necessary for the target target_mac_out = [] self.target_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_outs = self.target_mac.forward(batch, t=t) target_mac_out.append(target_agent_outs) # We don't need the first timesteps Q-Value estimate for calculating targets target_mac_out = th.stack(target_mac_out[:], dim=1) # Concat across time # Max over target Q-Values if self.args.double_q: raise Exception("No double q for DDPG") else: mac_out_detach = mac_out.clone().detach() mac_out_detach[avail_actions == 0] = -9999999 _, cur_max_actions = mac_out_detach[:, :].max(dim=3, keepdim=True) target_mac_out_detach = target_mac_out.clone().detach() target_mac_out_detach[avail_actions == 0] = -9999999 _, tar_max_actions = target_mac_out_detach[:, :].max(dim=3, keepdim=True) # Central MAC stuff central_mac_out = [] self.central_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.central_mac.forward(batch, t=t) central_mac_out.append(agent_outs) central_mac_out = th.stack(central_mac_out, dim=1) # Concat over time central_chosen_action_qvals_agents = th.gather( central_mac_out[:, :-1], dim=3, index=actions.unsqueeze(4).repeat( 1, 1, 1, 1, self.args.central_action_embed)).squeeze( 3) # Remove the last dim central_target_mac_out = [] self.target_central_mac.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): target_agent_outs = self.target_central_mac.forward(batch, t=t) central_target_mac_out.append(target_agent_outs) central_target_mac_out = th.stack(central_target_mac_out[:], dim=1) # Concat across time # Use the Qmix max actions central_target_max_agent_qvals = th.gather( central_target_mac_out[:, :], 3, cur_max_actions[:, :].unsqueeze(4).repeat( 1, 1, 1, 1, self.args.central_action_embed)).squeeze(3) # central_target_max_agent_qvals = th.gather(central_target_mac_out[:,:], 3, tar_max_actions[:,:].unsqueeze(4).repeat(1,1,1,1,self.args.central_action_embed)).squeeze(3) # --- # Mix target_max_qvals = self.target_central_mixer( central_target_max_agent_qvals[:, 1:], batch["state"][:, 1:]) # Calculate 1-step Q-Learning targets targets = rewards + self.args.gamma * (1 - terminated) * target_max_qvals # Bad naming, its not a qmix_actor qmix_actor_loss = 0 for agent in range(self.args.n_agents): target_chosen_qvals = central_target_max_agent_qvals[:, :-1] chosen_utils = target_chosen_qvals.detach().clone() # For each agent compute Q(u_i, u_{-i}) for each u_i, keeping u_{-i} fixed qtots = [] for action in range(self.args.n_actions): chosen_utils[:, :, agent] = central_target_mac_out[:, :-1, agent, action] new_q_tot = self.target_central_mixer(chosen_utils, batch["state"][:, :-1]) qtots.append(new_q_tot) agent_q_tots = th.cat(qtots, dim=2) qs_to_use = agent_q_tots # Train via ST Gumbel Softmax log_agent_policy = F.gumbel_softmax(mac_out[:, :-1, agent], hard=True, dim=2, tau=self.args.policy_temp) # Train via expected policy gradient # log_agent_policy = F.softmax(mac_out[:, :-1, agent] / self.args.policy_temp, dim=2) agent_policy_loss = ( (log_agent_policy * qs_to_use.detach() * avail_actions[:, :-1, agent].float()).sum( dim=2, keepdim=True) * mask).sum() / mask.sum() qmix_actor_loss = qmix_actor_loss + agent_policy_loss # Logit entropy ps = F.softmax(mac_out[:, :-1], dim=3) * avail_actions[:, :-1] log_ps = F.log_softmax(mac_out[:, :-1], dim=3) * avail_actions[:, :-1] logit_entropy = -(((ps * log_ps).sum(dim=3) * mask).sum() / mask.sum()) # Training central Q central_chosen_action_qvals = self.central_mixer( central_chosen_action_qvals_agents, batch["state"][:, :-1]) central_td_error = (central_chosen_action_qvals - targets.detach()) central_mask = mask.expand_as(central_td_error) central_masked_td_error = central_td_error * central_mask central_loss = (central_masked_td_error**2).sum() / mask.sum() loss = -self.args.qmix_loss * qmix_actor_loss + self.args.central_loss * central_loss + -self.args.logit_entropy * logit_entropy # Optimise self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.grad_norm = grad_norm self.optimiser.step() if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", loss.item(), t_env) self.logger.log_stat("actor_loss", qmix_actor_loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) mask_elems = mask.sum().item() self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item() / (mask_elems * self.args.n_agents), t_env) self.logger.log_stat("central_loss", central_loss.item(), t_env) self.logger.log_stat("logit_entropy", logit_entropy.item(), t_env) self.log_stats_t = t_env def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) if self.central_mac is not None: self.target_central_mac.load_state(self.central_mac) self.target_central_mixer.load_state_dict( self.central_mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() if self.central_mac is not None: self.central_mac.cuda() self.target_central_mac.cuda() self.central_mixer.cuda() self.target_central_mixer.cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
class CQLearner: def __init__(self, mac, scheme, logger, args): self.args = args self.mac = mac self.logger = logger self.params = list(mac.parameters()) self.named_params = dict(mac.named_parameters()) self.last_target_update_episode = 0 self.mixer = None if args.mixer is not None and self.args.n_agents > 1: # if just 1 agent do not mix anything if args.mixer == "vdn": self.mixer = VDNMixer() elif args.mixer == 'vdn-s': self.mixer = VDNState(args) elif args.mixer == "qmix": self.mixer = QMixer(args) else: raise ValueError("Mixer {} not recognised.".format(args.mixer)) self.params += list(self.mixer.parameters()) self.named_params.update(dict(self.mixer.named_parameters())) self.target_mixer = copy.deepcopy(self.mixer) if getattr(self.args, "optimizer", "rmsprop") == "rmsprop": self.optimiser = RMSprop(params=self.params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps) elif getattr(self.args, "optimizer", "rmsprop") == "adam": self.optimiser = Adam(params=self.params, lr=args.lr, eps=getattr(args, "optimizer_epsilon", 10E-8)) else: raise Exception("unknown optimizer {}".format( getattr(self.args, "optimizer", "rmsprop"))) # a little wasteful to deepcopy (e.g. duplicates action selector), but should work for any MAC self.target_mac = copy.deepcopy(mac) self.log_stats_t = -self.args.learner_log_interval - 1 def train(self, batch: EpisodeBatch, t_env: int, episode_num: int): # Get the relevant quantities rewards = batch["reward"][:, 0].unsqueeze(-1) terminated = batch["terminated"][:, 0].float().unsqueeze(-1) mask = 1 - terminated t = 0 chosen_action_qvals, _ = self.mac.forward( batch, actions=batch["actions"][:, t:t + 1].detach(), t=t) t = 1 best_target_action = self.target_mac.select_actions(batch, t_ep=t, t_env=None, test_mode=True) target_max_qvals, _ = self.target_mac.forward( batch, t=t, actions=best_target_action.detach()) # Mix if self.mixer is not None: chosen_action_qvals = self.mixer(chosen_action_qvals, batch["state"][:, :-1]) target_max_qvals = self.target_mixer(target_max_qvals, batch["state"][:, 1:]) chosen_action_qvals = chosen_action_qvals.squeeze(-2) target_max_qvals = target_max_qvals.squeeze(-2) rewards = rewards.squeeze(-2) terminated = terminated.squeeze(-2) # Calculate 1-step Q-Learning targets targets = rewards.expand_as(target_max_qvals) + self.args.gamma * ( 1 - terminated.expand_as(target_max_qvals)) * target_max_qvals # Td-error td_error = (chosen_action_qvals - targets.detach()) # Normal L2 loss, take mean over actual data assert self.args.runner_scope == "transition", "Runner scope HAS to be transition!" loss = (td_error**2).mean() # Optimise self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.optimiser.step() if getattr(self.args, "target_update_mode", "hard") == "hard": if (episode_num - self.last_target_update_episode ) / self.args.target_update_interval >= 1.0: self._update_targets() self.last_target_update_episode = episode_num elif getattr(self.args, "target_update_mode", "hard") in ["soft", "exponential_moving_average"]: self._update_targets_soft( tau=getattr(self.args, "target_update_tau", 0.001)) else: raise Exception("unknown target update mode: {}!".format( getattr(self.args, "target_update_mode", "hard"))) if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("loss", loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) self.logger.log_stat("weight_norm", (th.sum( th.cat([th.sum(p**2).unsqueeze(0) for p in self.params]))**0.5).item(), t_env) self.logger.log_stat("q_taken_mean", (chosen_action_qvals * mask).sum().item() / (batch.batch_size * self.args.n_agents), t_env) self.logger.log_stat("target_mean", (targets * mask).sum().item() / (batch.batch_size * self.args.n_agents), t_env) self.log_stats_t = t_env def _update_targets_soft(self, tau): for target_param, param in zip(self.target_mac.parameters(), self.mac.parameters()): target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau) if self.mixer is not None: for target_param, param in zip(self.target_mixer.parameters(), self.mixer.parameters()): target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau) if self.args.verbose: self.logger.console_logger.info( "Updated target network (soft update tau={})".format(tau)) def _update_targets(self): self.target_mac.load_state(self.mac) if self.mixer is not None: self.target_mixer.load_state_dict(self.mixer.state_dict()) self.logger.console_logger.info("Updated target network") def cuda(self): self.mac.cuda() self.target_mac.cuda() if self.mixer is not None: self.mixer.cuda() self.target_mixer.cuda() def save_models(self, path): self.mac.save_models(path) if self.mixer is not None: th.save(self.mixer.state_dict(), "{}/mixer.th".format(path)) th.save(self.optimiser.state_dict(), "{}/opt.th".format(path)) def load_models(self, path): self.mac.load_models(path) # Not quite right but I don't want to save target networks self.target_mac.load_models(path) if self.mixer is not None: self.mixer.load_state_dict( th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage)) self.optimiser.load_state_dict( th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))