Update app.py

app.py

@@ -11,6 +11,8 @@ import gymnasium as gym
 from stable_baselines3 import DQN
 from stable_baselines3.common.evaluation import evaluate_policy
 import gradio as gr
+import cv2
+import imageio
 
 # Constants
 SCREEN_WIDTH = 640
@@ -182,6 +184,7 @@ def train_and_play():
     total_timesteps = 10000
     timesteps_per_update = 1000
     frames = []
+    video_frames = []
 
     for i in range(0, total_timesteps, timesteps_per_update):
         model.learn(total_timesteps=timesteps_per_update)
@@ -193,18 +196,32 @@ def train_and_play():
             action, _states = model.predict(obs, deterministic=True)
             obs, reward, done, truncated, info = env.step(action)
             env.render()
-            pygame.image.save(screen, "frame.png")
+            # Capture the current frame
+            frame = pygame.surfarray.array3d(pygame.display.get_surface())
+            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+            video_frames.append(frame)
             episode_frames.append(gr.Image(value="frame.png"))
         frames.extend(episode_frames)
         yield frames
 
+    # Save the video
+    video_path = "arkanoid_training.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video_writer = cv2.VideoWriter(video_path, fourcc, FPS, (SCREEN_WIDTH, SCREEN_HEIGHT))
+    for frame in video_frames:
+        video_writer.write(frame)
+    video_writer.release()
+
+    # Return the video path
+    return gr.Video(video_path)
+
 # Main function
 def main():
     # Gradio interface
     iface = gr.Interface(
         fn=train_and_play,
         inputs=None,
-        outputs="image",
+        outputs="video",
         live=True
     )
     iface.launch()
@@ -220,6 +237,8 @@ if __name__ == "__main__":
 # - torch
 # - gradio
 # - gymnasium
+# - opencv-python
+# - imageio
 #
 # You can install these dependencies using pip:
-# pip install pygame stable-baselines3 torch gradio gymnasium
+# pip install pygame stable-baselines3 torch gradio gymnasium opencv-python imageio