web_game_rl.py

import pygame
import random
import numpy as np
import base64
import io
import torch
from enum import Enum
from collections import namedtuple, deque
from PIL import Image
from model import Linear_QNet, QTrainer

# Initialize pygame
pygame.init()
font = pygame.font.Font('arial.ttf', 25)

# Constants from original project
MAX_MEMORY = 100_000
BATCH_SIZE = 1000
LR = 0.001

class Direction(Enum):
    RIGHT = 1
    LEFT = 2
    UP = 3
    DOWN = 4

Point = namedtuple('Point', 'x, y')

# RGB colors
WHITE = (255, 255, 255)
RED = (200, 0, 0)
BLUE1 = (0, 0, 255)
BLUE2 = (0, 100, 255)
BLACK = (0, 0, 0)

BLOCK_SIZE = 20
SPEED = 20  # Increased speed for faster visualization

class Agent:
    def __init__(self):
        self.n_games = 0
        self.epsilon = 0  # randomness
        self.gamma = 0.9  # discount rate
        self.memory = deque(maxlen=MAX_MEMORY)  # popleft()
        self.model = Linear_QNet(11, 256, 3)
        
        # Try to load the model if it exists
        try:
            model_folder_path = './model'
            file_name = 'model.pth'
            model_path = f"{model_folder_path}/{file_name}"
            self.model.load_state_dict(torch.load(model_path))
            self.model.eval()
            print("Loaded existing model")
        except:
            print("No existing model found, starting fresh")
        
        self.trainer = QTrainer(self.model, lr=LR, gamma=self.gamma)

    def get_state(self, game):
        head = game.snake[0]
        point_l = Point(head.x - BLOCK_SIZE, head.y)
        point_r = Point(head.x + BLOCK_SIZE, head.y)
        point_u = Point(head.x, head.y - BLOCK_SIZE)
        point_d = Point(head.x, head.y + BLOCK_SIZE)
        
        dir_l = game.direction == Direction.LEFT
        dir_r = game.direction == Direction.RIGHT
        dir_u = game.direction == Direction.UP
        dir_d = game.direction == Direction.DOWN

        state = [
            # Danger straight
            (dir_r and game.is_collision(point_r)) or 
            (dir_l and game.is_collision(point_l)) or 
            (dir_u and game.is_collision(point_u)) or 
            (dir_d and game.is_collision(point_d)),

            # Danger right
            (dir_u and game.is_collision(point_r)) or 
            (dir_d and game.is_collision(point_l)) or 
            (dir_l and game.is_collision(point_u)) or 
            (dir_r and game.is_collision(point_d)),

            # Danger left
            (dir_d and game.is_collision(point_r)) or 
            (dir_u and game.is_collision(point_l)) or 
            (dir_r and game.is_collision(point_u)) or 
            (dir_l and game.is_collision(point_d)),
            
            # Move direction
            dir_l,
            dir_r,
            dir_u,
            dir_d,
            
            # Food location 
            game.food.x < game.head.x,  # food left
            game.food.x > game.head.x,  # food right
            game.food.y < game.head.y,  # food up
            game.food.y > game.head.y   # food down
        ]

        return np.array(state, dtype=int)

    def remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))  # popleft if MAX_MEMORY is reached

    def train_long_memory(self):
        if len(self.memory) > BATCH_SIZE:
            mini_sample = random.sample(self.memory, BATCH_SIZE)  # list of tuples
        else:
            mini_sample = self.memory

        states, actions, rewards, next_states, dones = zip(*mini_sample)
        self.trainer.train_step(states, actions, rewards, next_states, dones)

    def train_short_memory(self, state, action, reward, next_state, done):
        self.trainer.train_step(state, action, reward, next_state, done)

    def get_action(self, state):
        # random moves: tradeoff exploration / exploitation
        self.epsilon = 80 - self.n_games
        final_move = [0, 0, 0]
        
        # For web visualization, reduce randomness
        if random.randint(0, 200) < self.epsilon and self.n_games < 80:
            move = random.randint(0, 2)
            final_move[move] = 1
        else:
            state0 = torch.tensor(state, dtype=torch.float)
            prediction = self.model(state0)
            move = torch.argmax(prediction).item()
            final_move[move] = 1

        return final_move

class SnakeGameRL:
    def __init__(self, w=480, h=360):
        self.w = w
        self.h = h
        # Initialize display (hidden)
        pygame.display.init()
        self.display = pygame.Surface((self.w, self.h))
        self.clock = pygame.time.Clock()
        
        # Initialize agent
        self.agent = Agent()
        
        # Initialize game state
        self.reset()
        
        # Training data
        self.plot_scores = []
        self.plot_mean_scores = []
        self.total_score = 0
        self.record = 0
        
        # Auto-restart timer
        self.restart_timer = 0
        self.restart_delay = 20  # Frames to wait before restarting

    def reset(self):
        # Initialize game state
        self.direction = Direction.RIGHT
        self.head = Point(self.w/2, self.h/2)
        self.snake = [self.head,
                      Point(self.head.x-BLOCK_SIZE, self.head.y),
                      Point(self.head.x-(2*BLOCK_SIZE), self.head.y)]
        self.score = 0
        self.food = None
        self._place_food()
        self.frame_iteration = 0
        self.game_over = False
        
        # Get initial state
        self.state = self.agent.get_state(self)

    def _place_food(self):
        x = random.randint(0, (self.w-BLOCK_SIZE)//BLOCK_SIZE)*BLOCK_SIZE
        y = random.randint(0, (self.h-BLOCK_SIZE)//BLOCK_SIZE)*BLOCK_SIZE
        self.food = Point(x, y)
        if self.food in self.snake:
            self._place_food()

    def play_step(self):
        # If game is over, handle restart logic
        if self.game_over:
            self.restart_timer += 1
            if self.restart_timer >= self.restart_delay:
                self.reset()
                self.restart_timer = 0
            return self.game_over, self.score
            
        self.frame_iteration += 1
        
        # Get action from agent
        action = self.agent.get_action(self.state)
        
        # Move
        self._move(action)
        self.snake.insert(0, self.head)
        
        # Check if game over
        reward = 0
        game_over = False
        
        # Check collision or timeout
        if self.is_collision() or self.frame_iteration > 100*len(self.snake):
            game_over = True
            reward = -10
            self.game_over = True
            
            # Update training data
            self.agent.n_games += 1
            self.agent.train_short_memory(self.state, action, reward, self.state, True)
            self.agent.remember(self.state, action, reward, self.state, True)
            
            if self.score > self.record:
                self.record = self.score
                # Save model
                self.agent.model.save()
            
            # Add score to plot data
            self.plot_scores.append(self.score)
            self.total_score += self.score
            mean_score = self.total_score / self.agent.n_games
            self.plot_mean_scores.append(mean_score)
            
            # Print game info
            print('Game', self.agent.n_games, 'Score', self.score, 'Record:', self.record)
            
            return game_over, self.score
        
        # Place new food or just move
        if self.head == self.food:
            self.score += 1
            reward = 10
            self._place_food()
        else:
            self.snake.pop()
        
        # Update UI
        self._update_ui()
        self.clock.tick(SPEED)
        
        # Get new state and train
        new_state = self.agent.get_state(self)
        
        # Train short memory
        self.agent.train_short_memory(self.state, action, reward, new_state, game_over)
        
        # Remember
        self.agent.remember(self.state, action, reward, new_state, game_over)
        
        # Update state
        self.state = new_state
        
        return game_over, self.score

    def is_collision(self, pt=None):
        if pt is None:
            pt = self.head
        # Hits boundary
        if pt.x > self.w - BLOCK_SIZE or pt.x < 0 or pt.y > self.h - BLOCK_SIZE or pt.y < 0:
            return True
        # Hits itself
        if pt in self.snake[1:]:
            return True
        return False

    def _update_ui(self):
        self.display.fill(BLACK)
        
        # Draw snake
        for pt in self.snake:
            pygame.draw.rect(self.display, BLUE1, pygame.Rect(pt.x, pt.y, BLOCK_SIZE, BLOCK_SIZE))
            pygame.draw.rect(self.display, BLUE2, pygame.Rect(pt.x+4, pt.y+4, 12, 12))
        
        # Draw food
        pygame.draw.rect(self.display, RED, pygame.Rect(self.food.x, self.food.y, BLOCK_SIZE, BLOCK_SIZE))
        
        # Draw score
        text = font.render("Score: " + str(self.score), True, WHITE)
        self.display.blit(text, [0, 0])
        
        # Draw games played
        games_text = font.render("Games: " + str(self.agent.n_games), True, WHITE)
        self.display.blit(games_text, [0, 30])
        
        # Draw "Game Over" text if game is over
        if self.game_over:
            game_over_font = pygame.font.Font('arial.ttf', 40)
            game_over_text = game_over_font.render("GAME OVER", True, (255, 0, 0))
            text_rect = game_over_text.get_rect(center=(self.w/2, self.h/2))
            self.display.blit(game_over_text, text_rect)
            
            restart_font = pygame.font.Font('arial.ttf', 20)
            restart_text = restart_font.render("Restarting...", True, WHITE)
            restart_rect = restart_text.get_rect(center=(self.w/2, self.h/2 + 50))
            self.display.blit(restart_text, restart_rect)

    def _move(self, action):
        # [straight, right, left]
        clock_wise = [Direction.RIGHT, Direction.DOWN, Direction.LEFT, Direction.UP]
        idx = clock_wise.index(self.direction)
        
        if np.array_equal(action, [1, 0, 0]):
            new_dir = clock_wise[idx]  # No change
        elif np.array_equal(action, [0, 1, 0]):
            next_idx = (idx + 1) % 4
            new_dir = clock_wise[next_idx]  # Right turn
        else:  # [0, 0, 1]
            next_idx = (idx - 1) % 4
            new_dir = clock_wise[next_idx]  # Left turn
        
        self.direction = new_dir
        
        x = self.head.x
        y = self.head.y
        if self.direction == Direction.RIGHT:
            x += BLOCK_SIZE
        elif self.direction == Direction.LEFT:
            x -= BLOCK_SIZE
        elif self.direction == Direction.DOWN:
            y += BLOCK_SIZE
        elif self.direction == Direction.UP:
            y -= BLOCK_SIZE
        
        self.head = Point(x, y)

    def train_long_memory(self):
        if len(self.agent.memory) > 0:
            self.agent.train_long_memory()

    def get_frame(self):
        """Convert the pygame surface to a PIL Image"""
        # Get the pygame surface data
        raw_str = pygame.image.tostring(self.display, 'RGB')
        # Create a PIL image
        pil_img = Image.frombytes('RGB', (self.w, self.h), raw_str)
        return pil_img

    def get_base64_frame(self):
        """Get the current frame as a base64 encoded string"""
        img = self.get_frame()
        buffered = io.BytesIO()
        img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        return img_str