agent.py

import torch
import random
import numpy as np
from collections import deque
from game import SnakeGameAI, Direction, Point
from model import Linear_QNet, QTrainer

MAX_MEMORY = 100_000 
BATCH_SIZE = 1000
LR = 0.001

class Agent:

    def __init__(self):
        self.n_games = 0
        self.epsilon = 0  # parameter to control randomness
        self. gamma =  0.9  # discount rate must be smaller than 1
        self.memory = deque(maxlen=MAX_MEMORY) # popleft()
        self.model = Linear_QNet(11, 256, 3)   # 11 values size of state, hidden value size, output value size
        self.trainer = QTrainer(self.model, lr=LR, gamma=self.gamma)
        

    def get_state(self, game):
    
    # Create head
    head = game.snake[0]

    # Create points next to the head in all directions to check if it hits boundary for danger
    point_l = Point(head.x - 20, head.y)
    point_r = Point(head.x + 20, head.y)
    point_u = Point(head.x, head.y - 20)
    point_d = Point(head.x, head.y + 20)

    # Check if the current game direction == direction
    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 is straight ahead
        (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 is to the 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 is to the 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 directions
        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
        ]

    # data type = int will convert bool to 0 or 1
    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_MEOMORY is reached

    def train_long_memory(self):
        if len(self.memory) > BATCH_SIZE:
            mini_sample = random.sample(self.memory, BATCH_SIZE) # return a 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 between exploration / exploitation 
        # the more games the less the randomness(epsilon)
        self.epsilon = 80 - self.n_games
        final_move = [0,0,0] 
        # smaller epsilon means less the random moves
        if random.randint(0, 200) < self.epsilon:
            move = random.randint(0, 2)
            final_move[move] = 1
        else:
        # else there will be a move based off the model 
            state0 = torch.tensor(state, dtype=torch.float)
            prediction = self.model(state0)
            move = torch.argmax(prediction).item()
            final_move[move] = 1

        return final_move

    def train():
        plot_scores = []
        plot_mean_scores = []
        total_score = 0
        record = 0
        agent = Agent()
        game = SnakeGameAI()
        while True:
            # get old state 
            state_old = agent.get_state(game)

            # get move 
            final_move = agent.get_action(state_old)

            # perform move and get new state
            reward, done, score = game.play_step(final_move)
            state_new = agent.get_state(game)

            # train the short memory of agent
            agent.train_short_memory(state_old, final_move, reward, state_new, done)

            # remember 
            agent.remember(state_old, final_move, reward, state_new, done)

            if done:
                # train the long memory(experience replay)
                game.reset
                agent.n_games += 1
                agent.train_long_memory()

                if score > record:
                    record = score
                    agent.model.save()

                print('Game', agent.n_games, 'Score', score, 'Record:', record)

                

    if __name__ == '__main__':
        train()