def ExplainedVariance(values, returns):
  """Definition of explained variance - an approach from OpenAI baselines."""
  assert returns.shape == values.shape, (
      f'returns.shape was {returns.shape} and values.shape was {values.shape}')
  # TODO(henrykm): it would be good to explain the relation with the time dim.
  returns_variance = jnp.var(returns)
  explained_variance = 1 - jnp.var(returns-values)/returns_variance
  return explained_variance
Definition of explained variance - an approach from OpenAI baselines.
trax/rl/rl_layers.py