GPU Optimization of k-armed Bandit Problem

I have the following code:

class k_armed_bandit:
    def __init__(self, epsilon, k, steps):
        self.batches = 4000
        self.epsilon = epsilon
        self.k = k
        self.steps = steps
        self.mean = 4 * torch.rand((self.batches, self.k), device=device) - 2
        self.var = torch.rand((self.batches, self.k), device=device) + 0.5
        self.estimates = 4 * torch.ones((self.batches, self.k), device=device)
        self.counts = torch.ones((self.batches, self.k), device=device)

    def run(self):
        rewards = torch.zeros(self.steps, device=device)

        for i in range(self.steps):
            pos = torch.where(
                torch.rand(self.batches, device=device) < self.epsilon,
                torch.randint(0, k, size=(self.batches, ), device=device),
                torch.argmax(self.estimates, dim=1)
            )
            pos_mask = F.one_hot(pos, num_classes=self.k)
            val = torch.normal(mean=torch.sum(pos_mask * self.mean, dim=1), std=torch.sum(pos_mask * self.var, dim=1))
            self.counts += pos_mask
            val = val[:, None] * pos_mask
            self.estimates += ((val - (self.estimates * pos_mask)) / self.counts) * pos_mask
            rewards[i] = torch.sum(val) / self.batches

        return rewards

If you're not familiar with the problem, the idea is that you have k options in front of you, and you are trying to maximize the sum of rewards received from the options across some number of turns. The reward from option i is sampled from N(mean[i], var[i]). You keep a running track of returns from each option (stored in self.estimates, which is updated in streaming fashion). With probability epsilon you pick a random option, and with probability 1 - epsilon you simply pick the option with the best return (the position selected is stored in pos). At the end, you will return the sequence of rewards you obtained from the options you selected.

Anyways, the code runs pretty fast, but I suspect that some of my naive mistakes make this slower than it should be. First of all, my torch.where requires that for every agent in the batch, both a random position and the argmax should be computed, even though only one is used. This can't be good. Next, pos_mask is a big problem. It allows me to update the averages each agent keeps track of, only if they were picked as the sampled position. It also allows me to update the number of times each option was selected. The issue is that pos_mask, and all related computations, are much larger than they need to be, as it a mask and I cannot seem to update the relevant indices in parallel. I tried to fiddle with torch's index related methods, but none of them seem to help.

Any clues how to speed this up? Or maybe this problem is just not nice to implement on GPU?