class_docstr: ## Trainer
class_docstr_tok: ["#","#","Trainer"]
docstr: ### Calculate total loss
docstr_tok: ["#","#","#","Calculate","total","loss"]
code: def _calc_loss(
      self, samples: Dict[str, torch.Tensor]
  ) -> torch.Tensor:
      """
      ### Calculate total loss
      """
  
      # $R_t$ returns sampled from $\pi_{\theta_{OLD}}$
      sampled_return = samples["values"] + samples["advantages"]
  
      # $\bar{A_t} = \frac{\hat{A_t} - \mu(\hat{A_t})}{\sigma(\hat{A_t})}$,
      # where $\hat{A_t}$ is advantages sampled from $\pi_{\theta_{OLD}}$.
      # Refer to sampling function in [Main class](#main) below
      #  for the calculation of $\hat{A}_t$.
      sampled_normalized_advantage = self._normalize(
          samples["advantages"]
      )
  
      # Sampled observations are fed into the model to get $\pi_\theta(a_t|s_t)$ and $V^{\pi_\theta}(s_t)$;
      #  we are treating observations as state
      pi, value = self.model(samples["obs"])
  
      # $-\log \pi_\theta (a_t|s_t)$, $a_t$ are actions sampled from $\pi_{\theta_{OLD}}$
      log_pi = pi.log_prob(samples["actions"])
  
      # Calculate policy loss
      policy_loss = self.ppo_loss(
          log_pi,
          samples["log_pis"],
          sampled_normalized_advantage,
          self.clip_range(),
      )
  
      # Calculate Entropy Bonus
      #
      # $\mathcal{L}^{EB}(\theta) =
      #  \mathbb{E}\Bigl[ S\bigl[\pi_\theta\bigr] (s_t) \Bigr]$
      entropy_bonus = pi.entropy()
      entropy_bonus = entropy_bonus.mean()
  
      # Calculate value function loss
      value_loss = self.value_loss(
          value, samples["values"], sampled_return, self.clip_range()
      )
  
      # $\mathcal{L}^{CLIP+VF+EB} (\theta) =
      #  \mathcal{L}^{CLIP} (\theta) +
      #  c_1 \mathcal{L}^{VF} (\theta) - c_2 \mathcal{L}^{EB}(\theta)$
      loss = (
          policy_loss
          + self.value_loss_coef() * value_loss
          - self.entropy_bonus_coef() * entropy_bonus
      )
  
      # for monitoring
      approx_kl_divergence = (
          0.5 * ((samples["log_pis"] - log_pi) ** 2).mean()
      )
  
      # Add to tracker
      tracker.add(
          {
              "policy_reward": -policy_loss,
              "value_loss": value_loss,
              "entropy_bonus": entropy_bonus,
              "kl_div": approx_kl_divergence,
              "clip_fraction": self.ppo_loss.clip_fraction,
          }
      )
  
      return loss
  
imports: from typing import Dict
  import numpy
  import torch
  from torch import nn
  from torch import optim
  from torch.distributions import Categorical
  from labml import monit, tracker, logger, experiment
  from labml.configs import FloatDynamicHyperParam, IntDynamicHyperParam
  from labml_helpers.module import Module
  from labml_nn.rl.game import Worker
  from labml_nn.rl.ppo import ClippedPPOLoss, ClippedValueFunctionLoss
  from labml_nn.rl.ppo.gae import GAE
label_desc: Use image node_id 5 for calling the Trainer obj's underlying member method code with example usage: obj._calc_loss(samples) and returns: loss