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) mac_out[avail_actions == 0] = 0 # 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 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 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"][:, :-1] old_probs = batch["probs"][:, :-1] old_probs[avail_actions == 0] = 1e-10 old_logprob = th.log(th.gather(old_probs, dim=3, index=actions)).detach() mask_agent = mask.unsqueeze(2).repeat(1, 1, self.n_agents, 1) # targets and advantages with th.no_grad(): old_values = [] self.critic.init_hidden(batch.batch_size) for t in range(batch.max_seq_length): agent_outs = self.critic.forward(batch, t=t) old_values.append(agent_outs) old_values = th.stack(old_values, dim=1) if self.use_value_norm: value_shape = old_values.shape values = self.value_norm.denormalize( old_values.view(-1)).view(value_shape) advantages, targets = build_gae_targets( rewards.unsqueeze(2).repeat(1, 1, self.n_agents, 1), mask_agent, values, self.args.gamma, self.args.gae_lambda) if self.use_value_norm: targets_shape = targets.shape targets = targets.reshape(-1) self.value_norm.update(targets) targets = self.value_norm.normalize(targets).view( targets_shape) advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-6) # PPO Loss for _ in range(self.args.mini_epochs): # Critic values = [] self.critic.init_hidden(batch.batch_size) for t in range(batch.max_seq_length - 1): agent_outs = self.critic.forward(batch, t=t) values.append(agent_outs) values = th.stack(values, dim=1) # value clip values_clipped = old_values[:, :-1] + ( values - old_values[:, :-1]).clamp(-self.args.eps_clip, self.args.eps_clip) # 0-out the targets that came from padded data td_error = th.max((values - targets.detach())**2, (values_clipped - targets.detach())**2) masked_td_error = td_error * mask_agent critic_loss = 0.5 * masked_td_error.sum() / mask_agent.sum() # Actor 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 pi[avail_actions == 0] = 1e-10 pi_taken = th.gather(pi, dim=3, index=actions) log_pi_taken = th.log(pi_taken) ratios = th.exp(log_pi_taken - old_logprob) surr1 = ratios * advantages surr2 = th.clamp(ratios, 1 - self.args.eps_clip, 1 + self.args.eps_clip) * advantages actor_loss = -(th.min(surr1, surr2) * mask_agent).sum() / mask_agent.sum() # entropy 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 = actor_loss + self.args.critic_coef * critic_loss - self.args.entropy * entropy_loss / entropy_loss.item( ) # Optimise agents 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 t_env - self.log_stats_t >= self.args.learner_log_interval: mask_elems = mask_agent.sum().item() self.logger.log_stat("advantage_mean", (advantages * mask_agent).sum().item() / mask_elems, t_env) self.logger.log_stat("actor_loss", actor_loss.item(), t_env) self.logger.log_stat("entropy_loss", entropy_loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) self.logger.log_stat("lr", self.last_lr, t_env) self.logger.log_stat("critic_loss", critic_loss.item(), t_env) self.logger.log_stat("target_mean", (targets * mask_agent).sum().item() / mask_elems, t_env) self.log_stats_t = t_env
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"][:, :-1] old_probs = batch["probs"][:, :-1] old_probs[avail_actions == 0] = 1e-10 old_logprob = th.log(th.gather(old_probs, dim=3, index=actions)).detach() mask_agent = mask.unsqueeze(2).repeat(1, 1, self.n_agents, 1) for _ in range(self.args.mini_epochs): # Critic values = self.critic(batch) advantages, targets = build_gae_targets( rewards, mask, values, self.args.gamma, self.args.gae_lambda) # 0-out the targets that came from padded data td_error = (values[:, :-1] - targets.detach()) masked_td_error = td_error * mask critic_loss = 0.5 * (masked_td_error ** 2).sum() / mask.sum() # Actor advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) advantages = advantages.unsqueeze(2).repeat(1, 1, self.n_agents, 1) 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 pi[avail_actions == 0] = 1e-10 pi_taken = th.gather(pi, dim=3, index=actions) log_pi_taken = th.log(pi_taken) ratios = th.exp(log_pi_taken - old_logprob) surr1 = ratios * advantages surr2 = th.clamp(ratios, 1-self.args.eps_clip, 1+self.args.eps_clip) * advantages actor_loss = -(th.min(surr1, surr2) * mask_agent).sum() / mask_agent.sum() # entropy 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 = actor_loss + self.args.critic_coef * critic_loss - self.args.entropy * entropy_loss / entropy_loss.item() # Optimise agents self.optimiser.zero_grad() loss.backward() grad_norm = th.nn.utils.clip_grad_norm_(self.params, self.args.grad_norm_clip) self.optimiser.step() # Dynamic LR if self.args.lr_threshold: with th.no_grad(): kl_dist = 0.5 * ((old_logprob - log_pi_taken) ** 2) kl_dist = (kl_dist * mask).sum() / mask.sum() kl_dist = kl_dist.item() if kl_dist > (2.0 * self.args.lr_threshold): self.last_lr = max(self.last_lr / 1.5, 1e-6) if kl_dist < (0.5 * self.args.lr_threshold): self.last_lr = min(self.last_lr * 1.5, 1e-2) for param_group in self.optimiser.param_groups: param_group['lr'] = self.last_lr if t_env - self.log_stats_t >= self.args.learner_log_interval: self.logger.log_stat("advantage_mean", (advantages * mask_agent).sum().item() / mask_agent.sum().item(), t_env) self.logger.log_stat("actor_loss", actor_loss.item(), t_env) self.logger.log_stat("entropy_loss", entropy_loss.item(), t_env) self.logger.log_stat("grad_norm", grad_norm, t_env) self.logger.log_stat("lr", self.last_lr, t_env) self.logger.log_stat("critic_loss", critic_loss.item(), 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("target_mean", (targets * mask).sum().item() / mask_elems, t_env) self.log_stats_t = t_env