Update app.py

app.py

@@ -65,11 +65,13 @@ class Brick:
         self.rect = pygame.Rect(x, y, BRICK_WIDTH - 5, BRICK_HEIGHT - 5)
 
 class ArkanoidEnv(gym.Env):
-    def __init__(self):
+    def __init__(self, reward_size=1, penalty_size=-1):
         super(ArkanoidEnv, self).__init__()
         self.action_space = gym.spaces.Discrete(3)  # 0: stay, 1: move left, 2: move right
         self.observation_space = gym.spaces.Box(low=0, high=SCREEN_WIDTH, shape=(5 + BRICK_ROWS * BRICK_COLS * 2,), dtype=np.float32)
         self.seed_value = None
+        self.reward_size = reward_size
+        self.penalty_size = penalty_size
         self.reset()
 
     def reset(self, seed=None, options=None):
@@ -103,16 +105,16 @@ class ArkanoidEnv(gym.Env):
                 self.bricks.remove(brick)
                 self.ball.velocity[1] = -self.ball.velocity[1]
                 self.score += 1
-                reward = 1
+                reward = self.reward_size
                 if not self.bricks:
-                    reward += 10  # Bonus reward for breaking all bricks
+                    reward += self.reward_size * 10  # Bonus reward for breaking all bricks
                     self.done = True
                     truncated = False
                     return self._get_state(), reward, self.done, truncated, {}
 
         if self.ball.rect.bottom >= SCREEN_HEIGHT:
             self.done = True
-            reward = -1
+            reward = self.penalty_size
             truncated = False
         else:
             reward = 0
@@ -145,27 +147,14 @@ class ArkanoidEnv(gym.Env):
     def close(self):
         pygame.quit()
 
-# Training function
-def train_model(env, total_timesteps=10000):
-    model = DQN('MlpPolicy', env, verbose=1)
-    model.learn(total_timesteps=total_timesteps)
-    model.save("arkanoid_model")
-    return model
-
-# Evaluation function
-def evaluate_model(model, env):
-    mean_reward, _ = evaluate_policy(model, env, n_eval_episodes=10, render=False)
-    return mean_reward
-
-# Real-time training function
-def train_and_play():
-    env = ArkanoidEnv()
+# Training and playing with custom parameters
+def train_and_play(reward_size, penalty_size, iterations):
+    env = ArkanoidEnv(reward_size=reward_size, penalty_size=penalty_size)
     model = DQN('MlpPolicy', env, verbose=1)
-    total_timesteps = 10000
     timesteps_per_update = 1000
     video_frames = []
 
-    for i in range(0, total_timesteps, timesteps_per_update):
+    for i in range(0, iterations, timesteps_per_update):
         model.learn(total_timesteps=timesteps_per_update)
         obs, _ = env.reset()
         done = False
@@ -192,10 +181,14 @@ def train_and_play():
 
 # Main function
 def main():
-    # Gradio interface
+    # Gradio interface with parameters
     iface = gr.Interface(
         fn=train_and_play,
-        inputs=None,
+        inputs=[
+            gr.Number(label="Reward Size", value=1),
+            gr.Number(label="Penalty Size", value=-1),
+            gr.Slider(label="Iterations", minimum=10, maximum=100000, step=10, value=10000)
+        ],
         outputs="video",
         live=True
     )