Sample from a probability space of all actions

def sample(self, logits, argmax_sampling=False):
        """ Sample from a probability space of all actions """
        if argmax_sampling:
            return torch.argmax(logits, dim=-1)
        else:
            u = torch.rand_like(logits)
            return torch.argmax(logits - torch.log(-torch.log(u)), dim=-1)