Dynamic update and show gym env in Notebook, done but the next cell goes up and down. How to deal with it?

Question

I have created an openai gym env and would like to render it in vscode notebook. Following previous answers, How to dynamically update a plot in a loop in IPython notebook (within one cell), my code is as follows:

import gym
from IPython import display
from IPython.display import Image
%matplotlib inline

from matplotlib import pyplot as plt
# plt.ion()

env = gym.make('CartPole-v0')
env.reset()

img = plt.imshow(env.render(mode='rgb_array'))

done = False
while not done:
    action = env.action_space.sample()
    observation, reward, done, _ = env.step(action)

    img.set_data(env.render(mode='rgb_array'))
    display.display(plt.gcf())
    display.clear_output(wait=True)

env.close()

This solution did make it but the next cell keep going up and down while clearing output and re-creating images. Does anyone could help me to solve it? Thank you in advance.