def entropy(self, params):
        """
        Categorical distribution entropy calculation - sum probs * log(probs).
        In case of diagonal gaussian distribution - 1/2 log(2 pi e sigma^2)
        """
        log_std = params[:, :, 1]
        return (log_std + 0.5 * (self.LOG2PI + 1)).sum(dim=-1)