def ApproximateKLDivergence(dist_inputs, actions, old_log_probs, log_prob_fun):
  """Probability Ratio from the PPO algorithm."""
  new_log_probs = NewLogProbs(dist_inputs, actions, log_prob_fun)
  assert new_log_probs.shape == old_log_probs.shape, (
      f'new_log_probs.shape was {new_log_probs.shape} and'
      f'old_log_probs.shape was {old_log_probs.shape}')
  approximate_kl_divergence = 0.5 * \
      jnp.mean(new_log_probs - old_log_probs) ** 2
  return approximate_kl_divergence
Probability Ratio from the PPO algorithm.
trax/rl/rl_layers.py