def clip_fraction(self):
    """Joint policy and value loss layer."""
    def ClipFraction(dist_inputs, actions, old_log_probs):
      """Probability Ratio Mean from the PPO algorithm."""
      probs_ratio = rl_layers.ProbsRatio(
          dist_inputs, actions, old_log_probs,
          log_prob_fun=self._policy_dist.log_prob)
      return jnp.mean(jnp.abs(probs_ratio - 1) > self._epsilon)

    def f(dist_inputs, values, returns, dones, rewards, actions, old_log_probs):
      del values, returns, dones, rewards
      return ClipFraction(dist_inputs, actions, old_log_probs)
    return tl.Fn('ClipFraction', f)
Joint policy and value loss layer.
trax/rl/actor_critic_joint.py