Sample from a probability space of all actions

def sample(self, params, argmax_sampling=False):
        """ Sample from a probability space of all actions """
        means = params[:, :, 0]
        log_std = params[:, :, 1]

        if argmax_sampling:
            return means
        else:
            return torch.randn_like(means) * torch.exp(log_std) + means