def _policy_inputs_to_advantages(self, preprocess):
    """A layer that computes advantages from policy inputs."""
    def fn(dist_inputs, actions, q_values, act_log_probs, mask):
      del dist_inputs, actions, mask
      q_values = jnp.swapaxes(q_values, 0, 1)
      act_log_probs = jnp.swapaxes(act_log_probs, 0, 1)
      if self._sample_all_discrete_actions:
        values = jnp.sum(q_values * jnp.exp(act_log_probs), axis=0)
      else:
        values = jnp.mean(q_values, axis=0)
      advantages = q_values - values  # Broadcasting values over n_samples
      if preprocess:
        advantages = self._preprocess_advantages(advantages)
      return advantages
    return tl.Fn('PolicyInputsToAdvantages', fn)
A layer that computes advantages from policy inputs.
trax/rl/actor_critic.py