def sample_reward(env, policy, t_max=100):
    """
    Interact with an environment, return sum of all rewards.
    If game doesn't end on t_max (e.g. agent walks into a wall),
    force end the game and return whatever reward you got so far.
    Tip: see signature of env.step(...) method above.
    """
    s = env.reset()
    total_reward = 0

    for _ in range(t_max):
        action = policy[s]
        s, reward, done, info = env.step(action)
        total_reward += reward
        if done:
            break
    return total_reward