def loss_fn(policy: Policy, model: ModelV2, dist_class: TorchDistributionWrapper, sample_batch: SampleBatch): max_seq_len = sample_batch['seq_lens'].max().item() mask = sequence_mask(sample_batch['seq_lens'], max_seq_len, time_major=model.is_time_major()).view((-1, 1)) mean_reg = sample_batch['seq_lens'].sum() * model.nbr_agents actions = sample_batch['actions'].view( (sample_batch['actions'].shape[0], model.nbr_agents, -1))[:, :, :1].to(torch.long) actions = add_time_dimension(actions, max_seq_len=max_seq_len, framework='torch', time_major=True).reshape_as(actions) logits_pi, _ = model(sample_batch, [ sample_batch['state_in_0'], ], sample_batch['seq_lens']) logits_pi = logits_pi.view((logits_pi.shape[0], model.nbr_agents, -1)) logits_pi_action = logits_pi[:, :, :model.nbr_actions] log_pi_action = nn.functional.log_softmax(logits_pi_action, dim=-1) pi_action = torch.exp(log_pi_action) log_pi_action_selected = torch.gather(log_pi_action, -1, actions).squeeze(-1) q_values = model.q_values(sample_batch, target=False) q_values = add_time_dimension(q_values, max_seq_len=max_seq_len, framework="torch", time_major=True).reshape_as(q_values) q_values_selected = torch.gather(q_values, -1, actions).squeeze(-1) q_values_target = sample_batch[Postprocessing.VALUE_TARGETS] q_values_target = add_time_dimension( q_values_target, max_seq_len=max_seq_len, framework="torch", time_major=True).reshape_as(q_values_target) td_error = q_values_selected - q_values_target with torch.no_grad(): coma_avg = q_values_selected - (pi_action * q_values).sum(-1) entropy = -(log_pi_action * pi_action).sum(-1) critic_loss = torch.pow(mask * td_error, 2.0) actor_loss = mask * coma_avg * log_pi_action_selected entropy = mask * entropy policy.actor_loss = -actor_loss.sum() / mean_reg policy.critic_loss = critic_loss.sum() / mean_reg policy.entropy = entropy.sum() / mean_reg pi_loss = policy.actor_loss - policy.config[ 'entropy_coeff'] * policy.entropy return pi_loss, policy.critic_loss
def ppo_surrogate_loss( policy: Policy, model: ModelV2, dist_class: Type[TorchDistributionWrapper], train_batch: SampleBatch) -> Union[TensorType, List[TensorType]]: """Constructs the loss for Proximal Policy Objective. Args: policy (Policy): The Policy to calculate the loss for. model (ModelV2): The Model to calculate the loss for. dist_class (Type[ActionDistribution]: The action distr. class. train_batch (SampleBatch): The training data. Returns: Union[TensorType, List[TensorType]]: A single loss tensor or a list of loss tensors. """ logits, state = model.from_batch(train_batch, is_training=True) curr_action_dist = dist_class(logits, model) # RNN case: Mask away 0-padded chunks at end of time axis. if state: B = len(train_batch["seq_lens"]) max_seq_len = logits.shape[0] // B mask = sequence_mask(train_batch["seq_lens"], max_seq_len, time_major=model.is_time_major()) mask = torch.reshape(mask, [-1]) num_valid = torch.sum(mask) def reduce_mean_valid(t): return torch.sum(t[mask]) / num_valid # non-RNN case: No masking. else: mask = None reduce_mean_valid = torch.mean prev_action_dist = dist_class(train_batch[SampleBatch.ACTION_DIST_INPUTS], model) logp_ratio = torch.exp( curr_action_dist.logp(train_batch[SampleBatch.ACTIONS]) - train_batch[SampleBatch.ACTION_LOGP]) action_kl = prev_action_dist.kl(curr_action_dist) mean_kl = reduce_mean_valid(action_kl) curr_entropy = curr_action_dist.entropy() mean_entropy = reduce_mean_valid(curr_entropy) surrogate_loss = torch.min( train_batch[Postprocessing.ADVANTAGES] * logp_ratio, train_batch[Postprocessing.ADVANTAGES] * torch.clamp(logp_ratio, 1 - policy.config["clip_param"], 1 + policy.config["clip_param"])) mean_policy_loss = reduce_mean_valid(-surrogate_loss) if policy.config["use_gae"]: prev_value_fn_out = train_batch[SampleBatch.VF_PREDS] value_fn_out = model.value_function() vf_loss1 = torch.pow( value_fn_out - train_batch[Postprocessing.VALUE_TARGETS], 2.0) vf_clipped = prev_value_fn_out + torch.clamp( value_fn_out - prev_value_fn_out, -policy.config["vf_clip_param"], policy.config["vf_clip_param"]) vf_loss2 = torch.pow( vf_clipped - train_batch[Postprocessing.VALUE_TARGETS], 2.0) vf_loss = torch.max(vf_loss1, vf_loss2) mean_vf_loss = reduce_mean_valid(vf_loss) total_loss = reduce_mean_valid(-surrogate_loss + policy.kl_coeff * action_kl + policy.config["vf_loss_coeff"] * vf_loss - policy.entropy_coeff * curr_entropy) else: mean_vf_loss = 0.0 total_loss = reduce_mean_valid(-surrogate_loss + policy.kl_coeff * action_kl - policy.entropy_coeff * curr_entropy) # Store stats in policy for stats_fn. policy._total_loss = total_loss policy._mean_policy_loss = mean_policy_loss policy._mean_vf_loss = mean_vf_loss policy._vf_explained_var = explained_variance( train_batch[Postprocessing.VALUE_TARGETS], policy.model.value_function()) policy._mean_entropy = mean_entropy policy._mean_kl = mean_kl return total_loss
def loss(self, model: ModelV2, dist_class: Type[ActionDistribution], train_batch: SampleBatch) -> Union[TensorType, List[TensorType]]: """Constructs the loss for Proximal Policy Objective. Args: model: The Model to calculate the loss for. dist_class: The action distr. class. train_batch: The training data. Returns: The PPO loss tensor given the input batch. """ logits, state = model(train_batch) curr_action_dist = dist_class(logits, model) # RNN case: Mask away 0-padded chunks at end of time axis. if state: B = len(train_batch[SampleBatch.SEQ_LENS]) max_seq_len = logits.shape[0] // B mask = sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len, time_major=model.is_time_major()) mask = torch.reshape(mask, [-1]) num_valid = torch.sum(mask) def reduce_mean_valid(t): return torch.sum(t[mask]) / num_valid # non-RNN case: No masking. else: mask = None reduce_mean_valid = torch.mean prev_action_dist = dist_class( train_batch[SampleBatch.ACTION_DIST_INPUTS], model) logp_ratio = torch.exp( curr_action_dist.logp(train_batch[SampleBatch.ACTIONS]) - train_batch[SampleBatch.ACTION_LOGP]) # Only calculate kl loss if necessary (kl-coeff > 0.0). if self.config["kl_coeff"] > 0.0: action_kl = prev_action_dist.kl(curr_action_dist) mean_kl_loss = reduce_mean_valid(action_kl) else: mean_kl_loss = torch.tensor(0.0, device=logp_ratio.device) curr_entropy = curr_action_dist.entropy() mean_entropy = reduce_mean_valid(curr_entropy) surrogate_loss = torch.min( train_batch[Postprocessing.ADVANTAGES] * logp_ratio, train_batch[Postprocessing.ADVANTAGES] * torch.clamp(logp_ratio, 1 - self.config["clip_param"], 1 + self.config["clip_param"])) mean_policy_loss = reduce_mean_valid(-surrogate_loss) # Compute a value function loss. if self.config["use_critic"]: prev_value_fn_out = train_batch[SampleBatch.VF_PREDS] value_fn_out = model.value_function() vf_loss1 = torch.pow( value_fn_out - train_batch[Postprocessing.VALUE_TARGETS], 2.0) vf_clipped = prev_value_fn_out + torch.clamp( value_fn_out - prev_value_fn_out, -self.config["vf_clip_param"], self.config["vf_clip_param"]) vf_loss2 = torch.pow( vf_clipped - train_batch[Postprocessing.VALUE_TARGETS], 2.0) vf_loss = torch.max(vf_loss1, vf_loss2) mean_vf_loss = reduce_mean_valid(vf_loss) # Ignore the value function. else: vf_loss = mean_vf_loss = 0.0 total_loss = reduce_mean_valid(-surrogate_loss + self.config["vf_loss_coeff"] * vf_loss - self.entropy_coeff * curr_entropy) # Add mean_kl_loss (already processed through `reduce_mean_valid`), # if necessary. if self.config["kl_coeff"] > 0.0: total_loss += self.kl_coeff * mean_kl_loss # Store values for stats function in model (tower), such that for # multi-GPU, we do not override them during the parallel loss phase. model.tower_stats["total_loss"] = total_loss model.tower_stats["mean_policy_loss"] = mean_policy_loss model.tower_stats["mean_vf_loss"] = mean_vf_loss model.tower_stats["vf_explained_var"] = explained_variance( train_batch[Postprocessing.VALUE_TARGETS], model.value_function()) model.tower_stats["mean_entropy"] = mean_entropy model.tower_stats["mean_kl_loss"] = mean_kl_loss return total_loss