game.py

import pygame
import random
from enum import Enum
from collections import namedtuple
import numpy as np

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


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 = 40


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("Snake Game - Varun Patkar")
		self.clock = pygame.time.Clock()
		self.reset()

	def reset(self):
		# init game state

		# initial direction is random
		self.direction = random.choice(list(Direction))

		self.head = Point(self.w/2, self.h/2)
		if self.direction == Direction.RIGHT:
			self.snake = [self.head, Point(
				self.head.x-BLOCK_SIZE, self.head.y), Point(self.head.x-(2*BLOCK_SIZE), self.head.y)]
		if self.direction == Direction.LEFT:
			self.snake = [self.head, Point(
				self.head.x+BLOCK_SIZE, self.head.y), Point(self.head.x+(2*BLOCK_SIZE), self.head.y)]
		if self.direction == Direction.UP:
			self.snake = [self.head, Point(
				self.head.x, self.head.y+BLOCK_SIZE), Point(self.head.x, self.head.y+(2*BLOCK_SIZE))]
		if self.direction == Direction.DOWN:
			self.snake = [self.head, Point(
				self.head.x, self.head.y-BLOCK_SIZE), Point(self.head.x, self.head.y-(2*BLOCK_SIZE))]

		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)
		while self.food in self.snake:
			self._place_food()

	def play_step(self, action, n_games, high_score):
		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 head
		self.snake.insert(0, self.head)
		# 3. check if game over and quit if yes
		reward = 0
		game_over = False
		if self.is_collision() or self.frame_iteration > 100*len(self.snake):
			game_over = True
			reward = -10
			return reward, game_over, self.score
		# 4. place new food if self.head==self.food else move
		if self.head == self.food:
			self.score += 1
			reward = 10
			self._place_food()
		else:
			self.snake.pop()
		# 5. update UI and clock
		self._update_ui(n_games,high_score)
		self.clock.tick(SPEED)
		# 6. return game over state and score
		game_over = False
		return reward, game_over, self.score

	def _update_ui(self, n_games, high_score):
		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, BLUE2, 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)+"  Game No.:"+str(n_games)+"  High Score:"+str(high_score), True, WHITE)
		self.display.blit(text, (0, 0))
		pygame.display.flip()

	def _move(self, action):
		# [straight,right,left]
		clockwise = [Direction.RIGHT, Direction.DOWN,
					 Direction.LEFT, Direction.UP]
		idx = clockwise.index(self.direction)
		if np.array_equal(action, [1, 0, 0]):
			new_dir = clockwise[idx]  # no change
		elif np.array_equal(action, [0, 1, 0]):
			new_dir = clockwise[(idx+1) % 4]  # right turn
		elif np.array_equal(action, [0, 0, 1]):
			new_dir = clockwise[(idx-1) % 4]  # 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.UP:
			y -= BLOCK_SIZE
		elif self.direction == Direction.DOWN:
			y += BLOCK_SIZE
		self.head = Point(x, y)

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