code: def sample(self, logits, argmax_sampling=False):
          
          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)
docs: Sample from a probability space of all actions