import pickle


def load_data(log_dir):
    """
    Returns the average rewards along training, along with the number of arms.

    Parameters
    ----------
    log_dir : str
        The directory where agentData.pk is stored.

    Returns
    -------
    The data in the agentData.pk file.

    """
    with open(f"{log_dir}/agentData.pk", 'rb') as f:
        data = pickle.load(f)
    return data