code: def A2CObjective(
      dist_inputs,
      values,
      returns,
      dones,
      rewards,
      actions,
      mask,
      log_prob_fun,
      normalize_advantages,
  ):
      """Definition of the Advantage Actor Critic (A2C) loss."""
      # dist_inputs of the shape float32[128,1,18]
      # values of the shape float32[128,1,1]
      # returns of the shape float32[128,1,1]
      # dones of the shape int32[128,1,1]
      # actions of the shape int32[128,1]
      # and mask of the shape float32[128,1]
      # We have to squeeze values and returns, because we
      # are planning to compute (return - values) * new_log_probs * mask
      # and all of them should be of the same dimension
      values = values.squeeze(axis=2)
      returns = returns.squeeze(axis=2)
      dones = dones.squeeze(axis=2)
      rewards = rewards.squeeze(axis=2)
      assert (
          rewards.shape == dones.shape
      ), f"rewards.shape was {rewards.shape} and dones.shape was {dones.shape}"
      assert (
          dones.shape == values.shape
      ), f"dones.shape was {dones.shape} and values.shape was {values.shape}"
      assert (
          returns.shape == values.shape
      ), f"returns.shape was {returns.shape} and values.shape was {values.shape}"
      assert (
          values.shape == mask.shape
      ), f"values.shape was {values.shape} and mask.shape was {mask.shape}"
      assert returns.shape[0] == dist_inputs.shape[0], (
          f"returns.shape[0] was {returns.shape[0]} and dist_inputs.shape[0] was "
          f"{dist_inputs.shape[0]}"
      )
  
      new_log_probs = NewLogProbs(dist_inputs, actions, log_prob_fun)
      assert new_log_probs.shape == mask.shape, (
          f"new_log_probs.shape was {new_log_probs.shape} and mask.shape was "
          f"{mask.shape}"
      )
  
      # jaxified versions of
      # returns[dones] = rewards[dones]
      # values[dones] = 0
      returns = jnp.where(dones, rewards, returns)
      values = jnp.where(dones, jnp.zeros_like(values), values)
      advantages = returns - values
      if normalize_advantages:
          advantages = advantages - jnp.mean(advantages)
          advantages /= jnp.std(advantages) + 1e-8
      assert new_log_probs.shape == advantages.shape, (
          f"new_log_probs.shape was {new_log_probs.shape} and advantages.shape was "
          f"{advantages.shape}"
      )
  
      # One of the motivation to the squeezes and assertions is to
      # avoid [128,1] * [128,1,1] * [128] multiplications in the definition
      # of the a2c objective - we insist on the same shapes
      a2c_objective = -jnp.sum(
          new_log_probs * advantages * mask
      ) / jnp.sum(mask)
      return a2c_objective
  
file_path: trax/trax/rl/rl_layers.py
imports: from trax.fastmath import numpy
label_desc: Use image node_id 11 for calling a global function with example usage: A2CObjective(dist_inputs, values, returns, dones, rewards, actions, mask, log_prob_fun, normalize_advantages) and returns: a2c_objective