def test(agent, env, policy):
    """
    Tests an agent given an environment and the agent's policy
    :param agent: the agent to test
    :param env: the environment to test in
    :param policy: the policy to follow
    :return: agent's total reward as float
    """
    done = False
    obs = env.reset()
    total_reward = 0.0
    while not done:
        action = policy[agent.discretize(obs)]
        next_obs, reward, done, info = env.step(action)
        obs = next_obs
        total_reward += reward
    return total_reward