Update app.py

app.py

@@ -10,6 +10,8 @@ from stable_baselines3.common.vec_env import VecFrameStack
 
 from stable_baselines3.common.env_util import make_atari_env
 
+max_steps = 5000 # Let's try with 5000 steps.
+
 # Loading functions were taken from Edward Beeching code
 def load_env(env_name):
     env = make_atari_env(env_name, n_envs=1)
@@ -32,19 +34,23 @@ def load_model(env_name):
 
     return model
 
-def replay(env_name, time_sleep): #, num_episodes):
+def replay(env_name, max_steps, time_sleep):
   env = load_env(env_name)
   model = load_model(env_name)
   #for i in range(num_episodes):
   obs = env.reset()
   done = False
+  i = 0
   while not done:
-      frame = env.render(mode="rgb_array")
-      action, _states = model.predict(obs)
-      obs, reward, done, info = env.step([action])
-      time.sleep(time_sleep)
-      yield frame
-
+      i++
+      if i < max_steps:
+          frame = env.render(mode="rgb_array")
+          action, _states = model.predict(obs)
+          obs, reward, done, info = env.step([action])
+          time.sleep(time_sleep)
+          yield frame
+     else:
+         break
 
 demo = gr.Interface(
     replay,
@@ -53,6 +59,7 @@ demo = gr.Interface(
         "SeaquestNoFrameskip-v4",
         "QbertNoFrameskip-v4",
         ]),
+     gr.Slider(100, 10000, value=500),
      gr.Slider(0.01, 1, value=0.05),
      #gr.Slider(1, 20, value=5)
      ],