def f(dist_inputs, values, returns, dones, rewards, actions, old_log_probs):
      """Clipped objective from the PPO algorithm."""
      ppo_objective = rl_layers.PPOObjective(
          dist_inputs, values, returns, dones, rewards, actions, old_log_probs,
          log_prob_fun=self._policy_dist.log_prob,
          epsilon=self._epsilon,
          normalize_advantages=self._normalize_advantages)
      return jnp.mean(ppo_objective)
Clipped objective from the PPO algorithm.
trax/rl/actor_critic_joint.py