def ClippedObjective(probs_ratio, advantages, epsilon):
  """Clipped Objective from the PPO algorithm."""
  assert probs_ratio.shape == advantages.shape, (
      f'probs_ratio.shape was {probs_ratio.shape} and'
      f'advantages.shape was {advantages.shape}')
  clipped_objective = jnp.clip(probs_ratio, 1 - epsilon,
                               1 + epsilon) * advantages
  assert probs_ratio.shape == clipped_objective.shape, (
      f'probs_ratio.shape was {probs_ratio.shape} and'
      f'clipped_objective.shape was {clipped_objective.shape}')
  return clipped_objective
Clipped Objective from the PPO algorithm.
trax/rl/rl_layers.py