def value_mean(self):
    """The mean value of the value function."""
    # TODO(henrykm): A better solution would take into account the masks
    def f(values):
      return jnp.mean(values)
    return tl.Fn('ValueMean', f)
The mean value of the value function.
trax/rl/actor_critic.py