def UnclippedObjective(probs_ratio, advantages):
  """Unclipped 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}')
  unclipped_objective = probs_ratio * advantages
  return unclipped_objective
Unclipped Objective from the PPO algorithm.
trax/rl/rl_layers.py