snake_game_AI.py

import pygame
import random
from collections import namedtuple
from enum import Enum
import numpy as np
pygame.init()


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

Point = namedtuple('Point','x , y')
BLOCK_SIZE = 20
SPEED = 200
#rgb colors
BLACK = (0,0,0)
WHITE = (255,255,255)
RED = (0,0,255)
BLUE1 = (0,100,255)
BLUE2 = (0,100,255)

#font - pygame.font.Font('arial.ttf',25) 
font  = pygame.font.SysFont('arial',25)

#reset
#reward
#play(action)->direction
#game_iteration
#is_collision




class SnakeGameAI:
    
    def __init__(self,w=640,h=480):
        self.w = w
        self.h = h

        #init display
        self.display = pygame.display.set_mode((self.w,self.h))
        pygame.display.set_caption('AI_That_Plays_Snake - SANDEEP PADHI - IIIT LUCKNOW')
        self.clock = pygame.time.Clock()
        self.reset()
        self.count =0 
        
    
    def reset(self):
        self.frame_iteration = 0
        #init 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

    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, action):
        self.frame_iteration += 1
        #1.collect user input
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()
        

        #2.move
        self._move(action) #update the head 
        self.snake.insert(0,self.head)

        #check if game over
        reward = 0
        game_over=False
        if self.is_collision() or self.frame_iteration > 100*len(self.snake):
            game_over=True
            reward = self.count-10
            self.count = 0
            return reward ,game_over,self.score


        #4.place new food or just move
        if self.head == self.food:
            self.score+=1
            self._place_food()
            reward = self.count + 10
            self.count = 0
        else:
            self.snake.pop()
            self.count=0
            
        #.update Ui and clock
        self._update_ui()
        self.clock.tick(SPEED)

        #6.return game and score
        return reward ,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)
        
        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,BLUE1,pygame.Rect(pt.x+4 ,pt.y+4 ,12,12))

        pygame.draw.rect(self.display,RED,pygame.Rect(self.food.x,self.food.y,BLOCK_SIZE,BLOCK_SIZE))

        text = font.render("Score:"+str(self.score),True,WHITE)
        self.display.blit(text,[0,0])
        pygame.display.flip()

    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 r->d->l->u
        else: #[0,0,1]
            next_idx = (idx - 1)%4
            new_dir = clock_wise[next_idx]
        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.UP:
            y-=BLOCK_SIZE
        elif self.direction == Direction.DOWN:
            y+=BLOCK_SIZE
        
        self.head= Point(x,y)