s = torch.tensor([[0, 0], [-1, 0], [0, 1], [1, 0], [0, -1]], device=self.device)
b = a.clone()
b[:, 1:] = b[:, 1:] + s[actions[b[:, 0]]]
s = torch.tensor([[0, 0], [-1, 0], [0, 1], [1, 0], [0, -1]], device=self.device)
b = a.clone()
b[:, 1:] = b[:, 1:] + s[actions[b[:, 0]]]