import torch
import random


def epsilon_greedy(q_values, epsilon):
  """Returns the index of the highest q value with prob 1 - epsilon,
    otherwise uniformly at random with prob epsilon.

    Args:
        q_values (Variable[FloatTensor]): (batch_size, num_actions)
        epsilon (float)

    Returns:
        list[int]: actions
    """
  batch_size, num_actions = q_values.size()
  _, max_indices = torch.max(q_values, 1)
  max_indices = max_indices.cpu().data.numpy()
  actions = []
  for i in range(batch_size):
    if random.random() > epsilon:
      actions.append(max_indices[i])
    else:
      actions.append(random.randint(0, num_actions - 1))
  return actions