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