def returns_mean(self):
    """The mean value of actions selected by the behavioral policy."""
    def f(values, index_max, returns, mask):
      del values, index_max
      return jnp.sum(returns) / jnp.sum(mask)
    return tl.Fn('ReturnsMean', f)
The mean value of actions selected by the behavioral policy.
trax/rl/training.py