def value_loss(self):
    """Value loss - so far generic for all A2C."""
    def f(dist_inputs, values, returns):
      del dist_inputs
      return rl_layers.ValueLoss(values, returns, self._value_loss_coeff)
    return tl.Fn('ValueLoss', f)
Value loss - so far generic for all A2C.
trax/rl/actor_critic_joint.py