Update app.py

app.py

@@ -7,9 +7,8 @@ import os
 import numpy as np
 import pygame
 import random
+import gymnasium as gym
 from stable_baselines3 import DQN
-from stable_baselines3.common.env_util import make_atari_env
-from stable_baselines3.common.vec_env import VecFrameStack
 from stable_baselines3.common.evaluation import evaluate_policy
 import gradio as gr
 
@@ -69,14 +68,12 @@ class Brick:
     def __init__(self, x, y):
         self.rect = pygame.Rect(x, y, BRICK_WIDTH, BRICK_HEIGHT)
 
-class ArkanoidEnv:
+class ArkanoidEnv(gym.Env):
     def __init__(self):
-        self.paddle = Paddle()
-        self.ball = Ball()
-        self.bricks = [Brick(x, y) for y in range(BRICK_HEIGHT, BRICK_HEIGHT * (BRICK_ROWS + 1), BRICK_HEIGHT) for x in range(BRICK_WIDTH, SCREEN_WIDTH - BRICK_WIDTH, BRICK_WIDTH)]
-        self.clock = pygame.time.Clock()
-        self.done = False
-        self.score = 0
+        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.reset()
 
     def reset(self):
         self.paddle = Paddle()
@@ -87,7 +84,13 @@ class ArkanoidEnv:
         return self._get_state()
 
     def step(self, action):
-        self.paddle.move(action)
+        if action == 0:
+            self.paddle.move(0)
+        elif action == 1:
+            self.paddle.move(-1)
+        elif action == 2:
+            self.paddle.move(1)
+
         self.ball.move()
 
         if self.ball.rect.colliderect(self.paddle.rect):
@@ -118,16 +121,20 @@ class ArkanoidEnv:
         ]
         for brick in self.bricks:
             state.extend([brick.rect.x, brick.rect.y])
+        state.extend([0, 0] * (BRICK_ROWS * BRICK_COLS - len(self.bricks)))  # Padding for missing bricks
         return np.array(state, dtype=np.float32)
 
-    def render(self):
+    def render(self, mode='human'):
         screen.fill(BLACK)
         pygame.draw.rect(screen, WHITE, self.paddle.rect)
         pygame.draw.ellipse(screen, WHITE, self.ball.rect)
         for brick in self.bricks:
             pygame.draw.rect(screen, RED, brick.rect)
         pygame.display.flip()
-        self.clock.tick(FPS)
+        pygame.time.Clock().tick(FPS)
+
+    def close(self):
+        pygame.quit()
 
 # Training function
 def train_model():
@@ -152,7 +159,7 @@ def play_game():
     frames = []
     while not done:
         action, _states = model.predict(obs, deterministic=True)
-        obs, rewards, done, info = env.step(action)
+        obs, reward, done, info = env.step(action)
         env.render()
         pygame.image.save(screen, "frame.png")
         frames.append(gr.Image(value="frame.png"))
@@ -190,6 +197,7 @@ if __name__ == "__main__":
 # - stable-baselines3
 # - torch
 # - gradio
+# - gymnasium
 #
 # You can install these dependencies using pip:
-# pip install pygame stable-baselines3 torch gradio
+# pip install pygame stable-baselines3 torch gradio gymnasium