def ValueLoss(values, returns, value_loss_coeff):
  """Definition of the loss of the value function."""
  advantages = returns - values
  l2_value_loss = jnp.mean(advantages**2) * value_loss_coeff
  return l2_value_loss
Definition of the loss of the value function.
trax/rl/rl_layers.py