minor changes.

app.py

@@ -54,8 +54,9 @@ def play():
                 img = env.render()
                 plt.imshow( img )
                 plt.show()
-                time.sleep(0.1)
-                yield (img.copy())
+                #time.sleep(0.1)
+                images.append(img.copy())
+                yield (img.copy()), None
 
             # Run agent on the state
             with torch.no_grad():
@@ -72,24 +73,38 @@ def play():
                 break
             
         if info["flag_get"]:
-            break
-    
+            #break
+            imageio.mimsave('movie_new.gif', images)
+            time.sleep(5)
+            return (img.copy()), 'movie_new.gif'
+
+def refresh_playback():
+    return 'movie.gif'
+
+
 with gr.Blocks() as demo:
     gr.HTML("""<h1 align="center">Mario AI</h1>""")
     gr.HTML("""<h1 align="center">(May take a few re-plays to pass full scenario. )</h1>""")
     session_data = gr.State([])
 
     with gr.Row():
+        #with gr.Column(scale=1):
+        with gr.Column(scale=1):
+            play_mario = gr.Button("Let AI Play")
+            mario_image = gr.Image(height=400,width=400, label="New play.") 
         with gr.Column(scale=1):
-            play_mario = gr.Button("Play")
-        with gr.Column(scale=3):
-            mario_image = gr.Image(height=400,width=400) 
+            refresh = gr.Button("Refresh")
+            mario_gif = gr.Image(height=400,width=400, value='movie.gif', label="Playback previous AI run.")
    
-
+    refresh.click(
+        refresh_playback,
+        [],
+        [mario_gif]
+    )
     play_mario.click(
         play,
         [],
-        [mario_image],
+        [mario_image, mario_gif],
         #show_progress=True,
 
     )