docstr: Used for forward pass through loss computations.
  Works with batches of data, and optionally time_steps, but be consistent in usage: i.e. if using time_step,
  new_policy_probs, old_policy_probs, actions and advantages all must include a time_step dimension.
  
  NOTE: order of input arguments MUST NOT CHANGE because it matches the order
  parameters are passed in ppo_agent:train_network()
  
  :param (mx.nd or mx.sym) F: backend api (mx.sym if block has been hybridized).
  :param new_policy_probs: action probabilities predicted by DiscretePPOHead network,
      of shape (batch_size, num_actions) or
      of shape (batch_size, time_step, num_actions).
  :param old_policy_probs: action probabilities for previous policy,
      of shape (batch_size, num_actions) or
      of shape (batch_size, time_step, num_actions).
  :param actions: true actions taken during rollout,
      of shape (batch_size) or
      of shape (batch_size, time_step).
  :param clip_param_rescaler: scales epsilon to use for likelihood ratio clipping.
  :param advantages: change in state value after taking action (a.k.a advantage)
      of shape (batch_size) or
      of shape (batch_size, time_step).
  :param kl_coefficient: loss coefficient applied kl divergence loss (also see high_kl_penalty_coefficient).
  :return: loss, of shape (batch_size).
docstr_tok: ["Used","for","forward","pass","through","loss","computations",".","Works","with","batches","of","data",",","and","optionally","time_steps",",","but","be","consistent","in","usage",":","i.e",".","if","using","time_step",",","new_policy_probs",",","old_policy_probs",",","actions","and","advantages","all","must","include","a","time_step","dimension",".","NOTE",":","order","of","input","arguments","MUST","NOT","CHANGE","because","it","matches","the","order","parameters","are","passed","in","ppo_agent",":","train_network","(",")",":","param","(","mx.nd","or","mx.sym",")","F",":","backend","api","(","mx.sym","if","block","has","been","hybridized",")",".",":","param","new_policy_probs",":","action","probabilities","predicted","by","DiscretePPOHead","network",",","of","shape","(","batch_size",",","num_actions",")","or","of","shape","(","batch_size",",","time_step",",","num_actions",")",".",":","param","old_policy_probs",":","action","probabilities","for","previous","policy",",","of","shape","(","batch_size",",","num_actions",")","or","of","shape","(","batch_size",",","time_step",",","num_actions",")",".",":","param","actions",":","true","actions","taken","during","rollout",",","of","shape","(","batch_size",")","or","of","shape","(","batch_size",",","time_step",")",".",":","param","clip_param_rescaler",":","scales","epsilon","to","use","for","likelihood","ratio","clipping",".",":","param","advantages",":","change","in","state","value","after","taking","action","(","a.k.a","advantage",")","of","shape","(","batch_size",")","or","of","shape","(","batch_size",",","time_step",")",".",":","param","kl_coefficient",":","loss","coefficient","applied","kl","divergence","loss","(","also","see","high_kl_penalty_coefficient",")",".",":","return",":","loss",",","of","shape","(","batch_size",")","."]
code: def loss_forward(
      self,
      F: ModuleType,
      new_policy_probs: nd_sym_type,
      actions: nd_sym_type,
      old_policy_probs: nd_sym_type,
      clip_param_rescaler: nd_sym_type,
      advantages: nd_sym_type,
      kl_coefficient: nd_sym_type,
  ) -> List[Tuple[nd_sym_type, str]]:
      """
      Used for forward pass through loss computations.
      Works with batches of data, and optionally time_steps, but be consistent in usage: i.e. if using time_step,
      new_policy_probs, old_policy_probs, actions and advantages all must include a time_step dimension.
  
      NOTE: order of input arguments MUST NOT CHANGE because it matches the order
      parameters are passed in ppo_agent:train_network()
  
      :param (mx.nd or mx.sym) F: backend api (mx.sym if block has been hybridized).
      :param new_policy_probs: action probabilities predicted by DiscretePPOHead network,
          of shape (batch_size, num_actions) or
          of shape (batch_size, time_step, num_actions).
      :param old_policy_probs: action probabilities for previous policy,
          of shape (batch_size, num_actions) or
          of shape (batch_size, time_step, num_actions).
      :param actions: true actions taken during rollout,
          of shape (batch_size) or
          of shape (batch_size, time_step).
      :param clip_param_rescaler: scales epsilon to use for likelihood ratio clipping.
      :param advantages: change in state value after taking action (a.k.a advantage)
          of shape (batch_size) or
          of shape (batch_size, time_step).
      :param kl_coefficient: loss coefficient applied kl divergence loss (also see high_kl_penalty_coefficient).
      :return: loss, of shape (batch_size).
      """
  
      old_policy_dist = CategoricalDist(
          self.num_actions, old_policy_probs, F=F
      )
      action_probs_wrt_old_policy = old_policy_dist.log_prob(actions)
  
      new_policy_dist = CategoricalDist(
          self.num_actions, new_policy_probs, F=F
      )
      action_probs_wrt_new_policy = new_policy_dist.log_prob(actions)
  
      entropy_loss = -self.beta * new_policy_dist.entropy().mean()
  
      if self.use_kl_regularization:
          kl_div = old_policy_dist.kl_div(new_policy_dist).mean()
          weighted_kl_div = kl_coefficient * kl_div
          high_kl_div = (
              F.stack(F.zeros_like(kl_div), kl_div - self.kl_cutoff)
              .max()
              .square()
          )
          weighted_high_kl_div = (
              self.high_kl_penalty_coefficient * high_kl_div
          )
          kl_div_loss = weighted_kl_div + weighted_high_kl_div
      else:
          kl_div_loss = F.zeros(shape=(1,))
  
      # working with log probs, so minus first, then exponential (same as division)
      likelihood_ratio = (
          action_probs_wrt_new_policy - action_probs_wrt_old_policy
      ).exp()
  
      if self.clip_likelihood_ratio_using_epsilon is not None:
          # clipping of likelihood ratio
          min_value = (
              1
              - self.clip_likelihood_ratio_using_epsilon
              * clip_param_rescaler
          )
          max_value = (
              1
              + self.clip_likelihood_ratio_using_epsilon
              * clip_param_rescaler
          )
  
          # can't use F.clip (with variable clipping bounds), hence custom implementation
          clipped_likelihood_ratio = hybrid_clip(
              F,
              likelihood_ratio,
              clip_lower=min_value,
              clip_upper=max_value,
          )
  
          # lower bound of original, and clipped versions or each scaled advantage
          # element-wise min between the two ndarrays
          unclipped_scaled_advantages = likelihood_ratio * advantages
          clipped_scaled_advantages = (
              clipped_likelihood_ratio * advantages
          )
          scaled_advantages = F.stack(
              unclipped_scaled_advantages, clipped_scaled_advantages
          ).min(axis=0)
      else:
          scaled_advantages = likelihood_ratio * advantages
          clipped_likelihood_ratio = F.zeros_like(likelihood_ratio)
  
      # for each batch, calculate expectation of scaled_advantages across time steps,
      # but want code to work with data without time step too, so reshape to add timestep if doesn't exist.
      scaled_advantages_w_time = scaled_advantages.reshape(
          shape=(0, -1)
      )
      expected_scaled_advantages = scaled_advantages_w_time.mean(axis=1)
      # want to maximize expected_scaled_advantages, add minus so can minimize.
      surrogate_loss = (
          -expected_scaled_advantages * self.weight
      ).mean()
  
      return [
          (surrogate_loss, LOSS_OUT_TYPE_LOSS),
          (entropy_loss + kl_div_loss, LOSS_OUT_TYPE_REGULARIZATION),
          (kl_div_loss, LOSS_OUT_TYPE_KL),
          (entropy_loss, LOSS_OUT_TYPE_ENTROPY),
          (likelihood_ratio, LOSS_OUT_TYPE_LIKELIHOOD_RATIO),
          (
              clipped_likelihood_ratio,
              LOSS_OUT_TYPE_CLIPPED_LIKELIHOOD_RATIO,
          ),
      ]
  
imports: from typing import List, Tuple, Union
  from types import ModuleType
  import math
  import mxnet
  from mxnet.gluon import nn
  from rl_coach.base_parameters import AgentParameters
  from rl_coach.core_types import ActionProbabilities
  from rl_coach.spaces import SpacesDefinition, BoxActionSpace, DiscreteActionSpace
  from rl_coach.utils import eps
  from rl_coach.architectures.mxnet_components.heads.head import Head, HeadLoss, LossInputSchema, NormalizedRSSInitializer
  from rl_coach.architectures.mxnet_components.heads.head import LOSS_OUT_TYPE_LOSS, LOSS_OUT_TYPE_REGULARIZATION
  from rl_coach.architectures.mxnet_components.utils import hybrid_clip, broadcast_like
label_desc: Use image node_id 3 for calling the ClippedPPOLossDiscrete obj's underlying member method code with example usage: obj.loss_forward(F, new_policy_probs, actions, old_policy_probs, clip_param_rescaler, advantages, kl_coefficient) and returns: list