## Method definition
def save_video_of_model(model, env_name, filename='agent.mp4', preprocess=lambda x : x):
    display = Display(visible=0, size=(40, 30))
    display.start()

    env = gym.make(env_name)
    obs = preprocess(env.reset())
    shape = env.render(mode='rgb_array').shape[0:2]

    out = skvideo.io.FFmpegWriter(filename)

    done = False
    while not done:
        frame = env.render(mode='rgb_array')
        out.writeFrame(frame)
        action = model(tf.convert_to_tensor(obs.reshape((1,-1)), tf.float32)).numpy().argmax()
        obs, reward, done, info = env.step(action)
        obs = preprocess(obs)
    out.close()
    print ("Successfully saved into {}!".format(filename))


## Call
save_video_of_model(pong_model, "Pong-v0", filename='pong_agent.mp4', preprocess=pre_process)