Update app.py

app.py

@@ -106,14 +106,18 @@ class ArkanoidEnv(gym.Env):
                 self.bricks.remove(brick)
                 self.ball.velocity[1] = -self.ball.velocity[1]
                 self.score += 1
+                reward = 1
+                if not self.bricks:
+                    reward += 10  # Bonus reward for breaking all bricks
+                    self.done = True
+                return self._get_state(), reward, self.done, {}
 
         if self.ball.rect.bottom >= SCREEN_HEIGHT:
             self.done = True
+            reward = -1
+        else:
+            reward = 0
 
-        if not self.bricks:
-            self.done = True
-
-        reward = 1 if self.score > 0 else -1
         return self._get_state(), reward, self.done, {}
 
     def _get_state(self):
@@ -142,16 +146,14 @@ class ArkanoidEnv(gym.Env):
         pygame.quit()
 
 # Training function
-def train_model():
-    env = ArkanoidEnv()
+def train_model(env, total_timesteps=10000):
     model = DQN('MlpPolicy', env, verbose=1)
-    model.learn(total_timesteps=10000)
+    model.learn(total_timesteps=total_timesteps)
     model.save("arkanoid_model")
     return model
 
 # Evaluation function
-def evaluate_model(model):
-    env = ArkanoidEnv()
+def evaluate_model(model, env):
     mean_reward, _ = evaluate_policy(model, env, n_eval_episodes=10, render=False)
     return mean_reward
 
@@ -170,22 +172,33 @@ def play_game():
         frames.append(gr.Image(value="frame.png"))
     return frames
 
-# Main function
-def main():
-    if not os.path.exists("arkanoid_model.zip"):
-        print("Training model...")
-        train_model()
-    else:
-        print("Model already trained.")
+# Real-time training function
+def train_and_play():
+    env = ArkanoidEnv()
+    model = DQN('MlpPolicy', env, verbose=1)
+    total_timesteps = 10000
+    timesteps_per_update = 1000
+    frames = []
 
-    print("Evaluating model...")
-    model = DQN.load("arkanoid_model")
-    mean_reward = evaluate_model(model)
-    print(f"Mean reward: {mean_reward}")
+    for i in range(0, total_timesteps, timesteps_per_update):
+        model.learn(total_timesteps=timesteps_per_update)
+        obs = env.reset()[0]
+        done = False
+        episode_frames = []
+        while not done:
+            action, _states = model.predict(obs, deterministic=True)
+            obs, reward, done, info = env.step(action)
+            env.render()
+            pygame.image.save(screen, "frame.png")
+            episode_frames.append(gr.Image(value="frame.png"))
+        frames.extend(episode_frames)
+        yield frames
 
+# Main function
+def main():
     # Gradio interface
     iface = gr.Interface(
-        fn=play_game,
+        fn=train_and_play,
         inputs=None,
         outputs="image",
         live=True