def run_episode(env, policy, **kwargs):
    """
    Run a single episode for some time steps

    :param env: object
        Initialized OpenAI's gym environment.
    :param policy: ndarray
        Policy to be followed by the agent.
    :param kwargs:
     :T int default 1k
        Time steps which this episode lasts
     :render boolean default False
        Display the game play or not
     :discount float default None
        (gamma) discount factor
    :return: total_reward:
        Total reward for this episode
    """
    # kwargs
    T = kwargs.get('T', 1000)
    render = kwargs.get('render', False)
    discount = kwargs.get('discount', None)

    total_reward = 0
    state = env.reset()
    for t in range(T):
        if render:
            env.render()
        action = policy[state]
        state, reward, done, _ = env.step(action)
        total_reward += pow(discount, t) * reward if discount else reward
        if done:
            break
    return total_reward