import torch


def get_reward_mab(a_t, rewards):
    """define the reward function at time t

    Parameters
    ----------
    a_t : int
        action
    rewards : list
        rewards across arms

    Returns
    -------
    torch.FloatTensor, scalar
        immediate reward at time t

    """
    r_t = rewards[a_t]
    return r_t.type(torch.FloatTensor).data