Create Snake_EnvAndAgent.py

Snake_EnvAndAgent.py

@@ -0,0 +1,251 @@
+import gymnasium as gym
+from gymnasium import spaces
+import random
+import pygame
+import numpy as np
+import collections
+from collections import deque 
+from Environment_Constants import (
+    GRID_SIZE, CELL_SIZE, SCREEN_WIDTH, SCREEN_HEIGHT,
+    WHITE, BLACK, GREEN, RED, BLUE,
+    UP, DOWN, LEFT, RIGHT,
+    REWARD_FOOD, REWARD_COLLISION, REWARD_STEP,
+    FPS
+)
+
+class SnakeGameEnv(gym.Env):
+    metadata = {'render_modes': ['human', 'rgb_array'], 'render_fps': FPS}
+
+    def __init__(self, render_mode=None):
+        super().__init__()
+        self.grid_size = GRID_SIZE
+        self.cell_size = CELL_SIZE
+        self.screen_width = SCREEN_WIDTH
+        self.screen_height = SCREEN_HEIGHT
+
+        self.action_space = spaces.Discrete(3)
+
+        self.observation_space = spaces.Box(low=0, high=1, shape=(11,), dtype=np.float32)
+
+        self.render_mode = render_mode
+        self.window = None
+        self.clock = None
+
+        self._init_game_state()
+
+    def _init_game_state(self):
+        self.snake = deque()
+        self.head = (self.grid_size // 2, self.grid_size // 2) 
+        self.snake.append(self.head)
+    
+        for _ in range(2):
+            self.snake.append((self.head[0], self.head[1] + (_ + 1))) 
+
+        self.direction = UP
+        self.score = 0
+        self.food = self._place_food()
+        self.game_over = False
+        self.steps_since_food = 0
+
+    def _place_food(self):
+        while True:
+            x = random.randrange(self.grid_size)
+            y = random.randrange(self.grid_size)
+            food_pos = (x, y)
+            if food_pos not in self.snake:
+                return food_pos
+
+    def _get_observation(self):
+      
+        obs = np.zeros(11, dtype=np.float32)
+
+        hx, hy = self.head
+
+        if self.direction == UP:
+            dir_straight = UP
+            dir_right = RIGHT
+            dir_left = LEFT
+        elif self.direction == DOWN:
+            dir_straight = DOWN
+            dir_right = LEFT
+            dir_left = RIGHT
+        elif self.direction == LEFT:
+            dir_straight = LEFT
+            dir_right = UP
+            dir_left = DOWN
+        elif self.direction == RIGHT:
+            dir_straight = RIGHT
+            dir_right = DOWN
+            dir_left = UP
+
+        check_pos_straight = (hx + dir_straight[0], hy + dir_straight[1])
+        check_pos_right = (hx + dir_right[0], hy + dir_right[1])
+        check_pos_left = (hx + dir_left[0], hy + dir_left[1])
+
+        def is_danger(pos):
+            px, py = pos
+
+            if not (0 <= px < self.grid_size and 0 <= py < self.grid_size):
+                return True
+            
+            if pos in list(self.snake)[1:]:
+                return True
+            return False
+
+        obs[0] = 1 if is_danger(check_pos_straight) else 0 
+        obs[1] = 1 if is_danger(check_pos_right) else 0  
+        obs[2] = 1 if is_danger(check_pos_left) else 0 
+
+        
+        fx, fy = self.food
+        if fy < hy: obs[3] = 1
+        if fy > hy: obs[4] = 1 
+        if fx < hx: obs[5] = 1 
+        if fx > hx: obs[6] = 1 
+
+        if self.direction == UP: obs[7] = 1
+        elif self.direction == DOWN: obs[8] = 1
+        elif self.direction == LEFT: obs[9] = 1
+        elif self.direction == RIGHT: obs[10] = 1
+
+        return obs
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed)
+        self._init_game_state()
+
+        observation = self._get_observation()
+        info = self._get_info()
+        
+        if self.render_mode == 'human':
+            self._render_frame()
+
+        return observation, info
+
+    def _get_info(self):
+        return {"score": self.score, "snake_length": len(self.snake)}
+
+    def step(self, action):
+      
+        current_dir_idx = [UP, DOWN, LEFT, RIGHT].index(self.direction)
+        
+        if action == 0: 
+            new_direction = self.direction
+        elif action == 1:
+            if self.direction == UP: new_direction = RIGHT
+            elif self.direction == DOWN: new_direction = LEFT
+            elif self.direction == LEFT: new_direction = UP
+            elif self.direction == RIGHT: new_direction = DOWN
+        elif action == 2:
+            if self.direction == UP: new_direction = LEFT
+            elif self.direction == DOWN: new_direction = RIGHT
+            elif self.direction == LEFT: new_direction = DOWN
+            elif self.direction == RIGHT: new_direction = UP
+        else:
+            raise ValueError(f"Received invalid action={action} which is not part of the action space.")
+
+        self.direction = new_direction
+
+        hx, hy = self.head
+        dx, dy = self.direction
+        new_head = (hx + dx, hy + dy)
+
+        reward = REWARD_STEP
+
+        terminated = False
+       
+        if not (0 <= new_head[0] < self.grid_size and 0 <= new_head[1] < self.grid_size):
+            terminated = True
+            reward = REWARD_COLLISION
+        
+        elif new_head in list(self.snake):
+            terminated = True
+            reward = REWARD_COLLISION
+
+        if terminated:
+            self.game_over = True
+        else:
+            self.snake.appendleft(new_head) 
+            self.head = new_head 
+
+            if new_head == self.food:
+                self.score += 1
+                reward = REWARD_FOOD
+                self.food = self._place_food()
+                self.steps_since_food = 0 
+            else:
+                self.snake.pop()
+                self.steps_since_food += 1
+
+            if self.steps_since_food > self.grid_size * self.grid_size * 2:
+                 terminated = True
+                 reward = REWARD_COLLISION 
+
+        observation = self._get_observation()
+        info = self._get_info()
+        truncated = False 
+
+        if self.render_mode == 'human':
+            self._render_frame()
+
+        return observation, reward, terminated, truncated, info
+
+    def _render_frame(self):
+        if self.window is None and self.render_mode == 'human':
+            pygame.init()
+            pygame.display.init()
+            self.window = pygame.display.set_mode((self.screen_width, self.screen_height))
+            pygame.display.set_caption("Snake AI Training")
+        if self.clock is None and self.render_mode == 'human':
+            self.clock = pygame.time.Clock()
+
+        if self.render_mode == 'human':
+        
+            self.window.fill(BLACK)
+
+            pygame.draw.rect(self.window, RED, (self.food[0] * self.cell_size,
+                                               self.food[1] * self.cell_size,
+                                               self.cell_size, self.cell_size))
+
+            for i, segment in enumerate(self.snake):
+                color = BLUE if i == 0 else GREEN 
+                pygame.draw.rect(self.window, color, (segment[0] * self.cell_size,
+                                                      segment[1] * self.cell_size,
+                                                      self.cell_size, self.cell_size))
+
+            for x in range(0, self.screen_width, self.cell_size):
+                pygame.draw.line(self.window, WHITE, (x, 0), (x, self.screen_height))
+            for y in range(0, self.screen_height, self.cell_size):
+                pygame.draw.line(self.window, WHITE, (0, y), (self.screen_width, y))
+
+            font = pygame.font.Font(None, 25)
+            text = font.render(f"Score: {self.score}", True, WHITE)
+            self.window.blit(text, (5, 5))
+
+            pygame.event.pump() 
+            pygame.display.flip() 
+
+            self.clock.tick(self.metadata["render_fps"])
+        elif self.render_mode == "rgb_array":
+
+            surf = pygame.Surface((self.screen_width, self.screen_height))
+            surf.fill(BLACK)
+
+            pygame.draw.rect(surf, RED, (self.food[0] * self.cell_size,
+                                         self.food[1] * self.cell_size,
+                                         self.cell_size, self.cell_size))
+            
+            for i, segment in enumerate(self.snake):
+                color = BLUE if i == 0 else GREEN
+                pygame.draw.rect(surf, color, (segment[0] * self.cell_size,
+                                               segment[1] * self.cell_size,
+                                               self.cell_size, self.cell_size))
+
+            return np.transpose(np.array(pygame.surfarray.pixels3d(surf)), axes=(1, 0, 2))
+
+    def close(self):
+        if self.window is not None:
+            pygame.display.quit()
+            pygame.quit()
+            self.window = None
+            self.clock = None