This repository has been archived by the owner on Nov 29, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
testAI.py
158 lines (123 loc) · 3.87 KB
/
testAI.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
import pygame
import numpy
import random
import neat
import os
import pickle
from resources.reference import *
from resources.util import *
from gamesrc.grid import Grid
from gamesrc.snake import Snake
from gamesrc.food import Food
def has_failed(snake):
"""
This function checks to see if the snake has met any of the fail conditions.
Since the grid is a standard 30 by 30 grid, the grid object is not required
to check extremes.
Arguments:
snake {Snake} -- The snake whose status is to be found.
Returns:
boolean -- true if it has met a fail condition and False otherwise.
"""
x, y = snake.coords[0][0], snake.coords[0][1]
# Did it hit its own body?
if ((x, y) in snake.coords[1:]):
return True
elif x < 0 or x > 29 or y < 0 or y > 29: # Did it go outside the grid?
return True
else:
return False
def draw(window, snake, food, score):
"""
This function draws and updates the pygame window with the given information.
Arguments:
window {Surface} -- The active PyGame window
snake {Snake} -- The snake in the game
food {Food} -- The food in the game
score {int} -- The current score
"""
global ANIMATION_TICK
ANIMATION_TICK -= 1 # Used to change food color
window.fill((0, 0, 51))
# Food color changing and drawing is handled here
if ANIMATION_TICK == 0:
global FOOD_RGB
FOOD_RGB = (
random.randrange(
50, 255), random.randrange(
50, 255), random.randrange(
50, 255))
FOOD_IMG.fill(FOOD_RGB)
ANIMATION_TICK = 25
else:
window.blit(FOOD_IMG, (food.x * 15, food.y * 15))
# Draw snake
for i, coord in enumerate(snake.coords):
x = coord[0] * 15
y = coord[1] * 15
if i == 0:
head = pygame.Surface((15, 15))
head.fill((255, 255, 255))
WINDOW.blit(head, (x, y))
else:
WINDOW.blit(SNAKE_IMG, (x, y))
# Draw Score
score_txt = STAT_FONT.render("Score: " + str(score), 1, (255, 255, 255))
window.blit(
score_txt,
(WIN_WIDTH - 10 - score_txt.get_width(), 10)) # top right of screen
pygame.display.update()
def run_model(nn):
"""
This function plays the game of snake using the decisions made by the nn
passed to it.
Arguments:
nn -- The neural network
"""
global FOOD_RGB
grid = Grid()
snake = Snake()
food = generate_food(grid, snake)
score = 0
isRunning = True
while isRunning:
game_clock.tick(17)
# Check for quitting
for event in pygame.event.get():
if event.type == pygame.QUIT:
isRunning = False
pygame.quit()
quit()
# Make a decision at every tick or move forward by default
decision = make_decision(nn, grid, snake, food)
max_val = max(decision)
if max_val == decision[0]:
snake.move(grid, "L")
elif max_val == decision[1]:
snake.move(grid, "R")
else:
snake.tick(grid)
# Check if snake collided with food as a result
if snake.collide(food):
score += 1
snake.elongate(grid)
# Fancy color changing for snake :)
SNAKE_IMG.fill(FOOD_RGB)
food = generate_food(grid, snake)
# If it fails, then stop
if has_failed(snake):
isRunning = False
break
draw(WINDOW, snake, food, score)
def main():
"""
This function is called to run the program with the stored neural net.
"""
# Load stored NN
nn_file = open("best_model.pickle", "rb")
neural_net = pickle.load(nn_file)
nn_file.close()
# Use NN to run Flappy Bird
run_model(neural_net)
if __name__ == "__main__":
main()