def LossInput(dist_inputs, actions, advantages, old_dist_inputs):  # pylint: disable=invalid-name
      """Calculates action log probabilities and normalizes advantages."""
      advantages = self._preprocess_advantages(advantages)
      log_probs = self._policy_dist.log_prob(dist_inputs, actions)
      old_log_probs = self._policy_dist.log_prob(old_dist_inputs, actions)
      return (log_probs, advantages, old_log_probs)
Calculates action log probabilities and normalizes advantages.
trax/rl/actor_critic.py