text: #!/usr/bin/env python3
  
  # Copyright (c) Facebook, Inc. and its affiliates.
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import numpy as np
  import torch
  import torch.distributed as distrib
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  
  from habitat_baselines.rl.aux_losses import AuxLosses
  from habitat_baselines.rl.ddppo.policy.running_mean_and_var import RunningMeanAndVar
  
  EPS_PPO = 1e-5
  
  
  class PPO(nn.Module):
      def __init__(
          self,
          actor_critic,
          clip_param,
          ppo_epoch,
          num_mini_batch,
          value_loss_coef,
          entropy_coef,
          lr=None,
          eps=None,
          max_grad_norm=None,
          use_clipped_value_loss=True,
          use_normalized_advantage=True,
      ):
  
          super().__init__()
  
          self.actor_critic = actor_critic
  
          self.clip_param = clip_param
          self.ppo_epoch = ppo_epoch
          self.num_mini_batch = num_mini_batch
  
          self.value_loss_coef = value_loss_coef
          self.entropy_coef = entropy_coef
  
          self.max_grad_norm = max_grad_norm
          self.use_clipped_value_loss = use_clipped_value_loss
  
          self.optimizer = optim.Adam(
              list(filter(lambda p: p.requires_grad, actor_critic.parameters())),
              lr=lr,
              eps=eps,
          )
          self.device = next(actor_critic.parameters()).device
          self.use_normalized_advantage = use_normalized_advantage
  
          self.reward_whitten = RunningMeanAndVar(shape=(1,))
          self.reward_whitten.to(self.device)
  
      def forward(self, *x):
          raise NotImplementedError
  
      def get_advantages(self, rollouts):
          advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
          if not self.use_normalized_advantage:
              return advantages
  
          return (advantages - advantages.mean()) / (advantages.std() + EPS_PPO)
  
      def update(self, rollouts):
          advantages = self.get_advantages(rollouts)
  
          value_loss_epoch = 0
          action_loss_epoch = 0
          aux_losses_epoch = 0
          dist_entropy_epoch = 0
  
          AuxLosses.activate()
          for e in range(self.ppo_epoch):
              data_generator = rollouts.recurrent_generator(
                  advantages, self.num_mini_batch
              )
  
              for sample in data_generator:
                  (
                      obs_batch,
                      recurrent_hidden_states_batch,
                      actions_batch,
                      prev_actions_batch,
                      value_preds_batch,
                      return_batch,
                      masks_batch,
                      old_action_log_probs_batch,
                      adv_targ,
                  ) = sample
  
                  AuxLosses.clear()
                  # Reshape to do in a single forward pass for all steps
                  (
                      values,
                      action_log_probs,
                      dist_entropy,
                  ) = self.actor_critic.evaluate_actions(
                      obs_batch,
                      recurrent_hidden_states_batch,
                      prev_actions_batch,
                      masks_batch,
                      actions_batch,
                  )
  
                  ratio = torch.exp(action_log_probs - old_action_log_probs_batch)
                  surr1 = ratio * adv_targ
                  surr2 = (
                      torch.clamp(ratio, 1.0 - self.clip_param, 1.0 + self.clip_param)
                      * adv_targ
                  )
                  action_loss = -torch.min(surr1, surr2).mean()
  
                  if self.use_clipped_value_loss:
                      value_pred_clipped = value_preds_batch + (
                          values - value_preds_batch
                      ).clamp(-self.clip_param, self.clip_param)
                      value_losses = (values - return_batch).pow(2)
                      value_losses_clipped = (value_pred_clipped - return_batch).pow(2)
                      value_loss = (
                          0.5 * torch.max(value_losses, value_losses_clipped).mean()
                      )
                  else:
                      value_loss = 0.5 * (return_batch - values).pow(2).mean()
  
                  self.optimizer.zero_grad()
                  total_loss = (
                      value_loss * self.value_loss_coef
                      + action_loss
                      - dist_entropy * self.entropy_coef
                  )
  
                  use_aux_loss = True
  
                  aux_losses = AuxLosses.reduce()
                  aux_losses = (1.0 if use_aux_loss else 0.0) * aux_losses
                  aux_losses_epoch += aux_losses.item()
  
                  total_loss = total_loss + aux_losses
  
                  self.before_backward(total_loss)
                  total_loss.backward()
                  self.after_backward(total_loss)
  
                  self.before_step()
                  self.optimizer.step()
                  self.after_step()
  
                  value_loss_epoch += value_loss.item()
                  action_loss_epoch += action_loss.item()
                  dist_entropy_epoch += dist_entropy.item()
  
          num_updates = self.ppo_epoch * self.num_mini_batch
  
          value_loss_epoch /= num_updates
          action_loss_epoch /= num_updates
          dist_entropy_epoch /= num_updates
          aux_losses_epoch /= num_updates
  
          AuxLosses.deactivate()
  
          return (
              value_loss_epoch,
              action_loss_epoch,
              dist_entropy_epoch,
              aux_losses_epoch,
          )
  
      def before_backward(self, loss):
          pass
  
      def after_backward(self, loss):
          pass
  
      def before_step(self):
          nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
  
      def after_step(self):
          pass