def value_loss(self):
    """Value loss computed using smooth L1 loss or L2 loss."""
    def f(values, actions, returns, mask):
      ind_0, ind_1 = np.indices(actions.shape)
      # We calculate length using the shape of returns
      # and adequatly remove a superflous slice of values.
      # An analogous operation is done in value_batches_stream.
      length = returns.shape[1]
      values = values[:, :length, :]
      selected_values = values[ind_0, ind_1, actions]
      shapes.assert_same_shape(selected_values, returns)
      shapes.assert_same_shape(selected_values, mask)
      if self._smoothl1loss:
        return tl.SmoothL1Loss().forward((selected_values, returns, mask))
      else:
        return tl.L2Loss().forward((selected_values, returns, mask))
    return tl.Fn('ValueLoss', f)
Value loss computed using smooth L1 loss or L2 loss.
trax/rl/training.py