flappy_bird.py

# -*- coding: utf-8 -*-
"""flappy_bird.ipynb

Author: Fu-Lin Hsu
Last Update Date: 04/11/2021

Before you start:
	1. Be sure you upload game assets in the directory order './assets/sprites/*.png'
	(you can set the BASE_PATH as the prefix of assets/sprites/*.png, such as a mounted drive in Colab. Remember to end it with a slash '/')
	2. Set up the directory you want to save checkpoints

To Continue Training:
	1. Set LOAD = True
	2. Put the NUMBER of checkpoints files to the ckpt list in order to populate Replay Memory before training
	3. Put the latest checkpoint file to the path within the last .load_model()

"""

#from google.colab import drive
#drive.mount('/content/drive')

import pygame
from pygame.locals import *
from itertools import cycle
import random
import numpy as np
import cv2
import sys
import os
os.environ['SDL_VIDEODRIVER'] = 'dummy' # Run Headless Pygame environment

"""## Load Game Resources"""

def getHitmask(image):
    """returns a hitmask using an image's alpha."""
    mask = []
    for x in range(image.get_width()):
        mask.append([])
        for y in range(image.get_height()):
            mask[x].append(bool(image.get_at((x,y))[3]))
    return mask

def load(BASE_PATH = './'):
    # path of player with different states
    PLAYER_PATH = (
            BASE_PATH + 'assets/sprites/redbird-upflap.png',
            BASE_PATH + 'assets/sprites/redbird-midflap.png',
            BASE_PATH + 'assets/sprites/redbird-downflap.png'
    )

    # path of background
    BACKGROUND_PATH = BASE_PATH + 'assets/sprites/background-black.png'

    # path of pipe
    PIPE_PATH = BASE_PATH + 'assets/sprites/pipe-green.png'

    IMAGES, HITMASKS = {}, {}

    # numbers sprites for score display
    IMAGES['numbers'] = (
        pygame.image.load(BASE_PATH + 'assets/sprites/0.png').convert_alpha(),
        pygame.image.load(BASE_PATH + 'assets/sprites/1.png').convert_alpha(),
        pygame.image.load(BASE_PATH + 'assets/sprites/2.png').convert_alpha(),
        pygame.image.load(BASE_PATH + 'assets/sprites/3.png').convert_alpha(),
        pygame.image.load(BASE_PATH + 'assets/sprites/4.png').convert_alpha(),
        pygame.image.load(BASE_PATH + 'assets/sprites/5.png').convert_alpha(),
        pygame.image.load(BASE_PATH + 'assets/sprites/6.png').convert_alpha(),
        pygame.image.load(BASE_PATH + 'assets/sprites/7.png').convert_alpha(),
        pygame.image.load(BASE_PATH + 'assets/sprites/8.png').convert_alpha(),
        pygame.image.load(BASE_PATH + 'assets/sprites/9.png').convert_alpha()
    )

    # base (ground) sprite
    IMAGES['base'] = pygame.image.load(BASE_PATH + 'assets/sprites/base.png').convert_alpha()

    # select random background sprites
    IMAGES['background'] = pygame.image.load(BACKGROUND_PATH).convert()

    # select random player sprites
    IMAGES['player'] = (
        pygame.image.load(PLAYER_PATH[0]).convert_alpha(),
        pygame.image.load(PLAYER_PATH[1]).convert_alpha(),
        pygame.image.load(PLAYER_PATH[2]).convert_alpha(),
    )

    # select random pipe sprites
    IMAGES['pipe'] = (
        pygame.transform.rotate(
            pygame.image.load(PIPE_PATH).convert_alpha(), 180),
        pygame.image.load(PIPE_PATH).convert_alpha(),
    )

    # hismask for pipes
    HITMASKS['pipe'] = (
        getHitmask(IMAGES['pipe'][0]),
        getHitmask(IMAGES['pipe'][1]),
    )

    # hitmask for player
    HITMASKS['player'] = (
        getHitmask(IMAGES['player'][0]),
        getHitmask(IMAGES['player'][1]),
        getHitmask(IMAGES['player'][2]),
    )
    return IMAGES, HITMASKS

"""## Game Parameters Setting"""

FPS = 30
SCREENWIDTH = 288
SCREENHEIGHT = 512

pygame.init()
FPSCLOCK = pygame.time.Clock()
SCREEN = pygame.display.set_mode((SCREENWIDTH, SCREENHEIGHT))
pygame.display.set_caption('Flappy Bird')

IMAGES, HITMASKS = load()
PIPEGAPSIZE = 100 # gap between upper and lower part of pipe
BASEY = SCREENHEIGHT * 0.79

PLAYER_WIDTH = IMAGES['player'][0].get_width()
PLAYER_HEIGHT = IMAGES['player'][0].get_height()
PIPE_WIDTH = IMAGES['pipe'][0].get_width()
PIPE_HEIGHT = IMAGES['pipe'][0].get_height()
BACKGROUND_WIDTH = IMAGES['background'].get_width()

PLAYER_INDEX_GEN = cycle([0, 1, 2, 1])


class GameState:
    def __init__(self):
        self.score = self.playerIndex = self.loopIter = 0
        self.playerx = int(SCREENWIDTH * 0.2)
        self.playery = int((SCREENHEIGHT - PLAYER_HEIGHT) / 2)
        self.basex = 0
        self.baseShift = IMAGES['base'].get_width() - BACKGROUND_WIDTH

        newPipe1 = getRandomPipe()
        newPipe2 = getRandomPipe()
        self.upperPipes = [
            {'x': SCREENWIDTH, 'y': newPipe1[0]['y']},
            {'x': SCREENWIDTH + (SCREENWIDTH / 2), 'y': newPipe2[0]['y']},
        ]
        self.lowerPipes = [
            {'x': SCREENWIDTH, 'y': newPipe1[1]['y']},
            {'x': SCREENWIDTH + (SCREENWIDTH / 2), 'y': newPipe2[1]['y']},
        ]

        # player velocity, max velocity, downward accleration, accleration on flap
        self.pipeVelX = -4
        self.playerVelY    =  0    # player's velocity along Y, default same as playerFlapped
        self.playerMaxVelY =  10   # max vel along Y, max descend speed
        self.playerMinVelY =  -8   # min vel along Y, max ascend speed
        self.playerAccY    =   1   # players downward accleration
        self.playerFlapAcc =  -9   # players speed on flapping
        self.playerFlapped = False # True when player flaps

    def frame_step(self, input_actions):
        pygame.event.pump()

        reward = 0.1
        terminal = False

        if sum(input_actions) != 1:
            raise ValueError('Multiple input actions!')

        # input_actions[0] == 1: do nothing
        # input_actions[1] == 1: flap the bird
        if input_actions[1] == 1:
            if self.playery > -2 * PLAYER_HEIGHT:
                self.playerVelY = self.playerFlapAcc
                self.playerFlapped = True

        # check for score
        playerMidPos = self.playerx + PLAYER_WIDTH / 2
        for pipe in self.upperPipes:
            pipeMidPos = pipe['x'] + PIPE_WIDTH / 2
            if pipeMidPos <= playerMidPos < pipeMidPos + 4:
                self.score += 1
                reward = 1

        # playerIndex basex change
        if (self.loopIter + 1) % 3 == 0:
            self.playerIndex = next(PLAYER_INDEX_GEN)
        self.loopIter = (self.loopIter + 1) % 30
        self.basex = -((-self.basex + 100) % self.baseShift)

        # player's movement
        if self.playerVelY < self.playerMaxVelY and not self.playerFlapped:
            self.playerVelY += self.playerAccY
        if self.playerFlapped:
            self.playerFlapped = False
        self.playery += min(self.playerVelY, BASEY - self.playery - PLAYER_HEIGHT)
        if self.playery < 0:
            self.playery = 0

        # move pipes to left
        for uPipe, lPipe in zip(self.upperPipes, self.lowerPipes):
            uPipe['x'] += self.pipeVelX
            lPipe['x'] += self.pipeVelX

        # add new pipe when first pipe is about to touch left of screen
        if 0 < self.upperPipes[0]['x'] < 5:
            newPipe = getRandomPipe()
            self.upperPipes.append(newPipe[0])
            self.lowerPipes.append(newPipe[1])

        # remove first pipe if its out of the screen
        if self.upperPipes[0]['x'] < -PIPE_WIDTH:
            self.upperPipes.pop(0)
            self.lowerPipes.pop(0)

        # check if crash here
        isCrash= checkCrash({'x': self.playerx, 'y': self.playery,
                             'index': self.playerIndex},
                            self.upperPipes, self.lowerPipes)
        if isCrash:
            terminal = True
            #self.__init__()
            reward = -1

        # draw sprites
        SCREEN.blit(IMAGES['background'], (0,0))

        for uPipe, lPipe in zip(self.upperPipes, self.lowerPipes):
            SCREEN.blit(IMAGES['pipe'][0], (uPipe['x'], uPipe['y']))
            SCREEN.blit(IMAGES['pipe'][1], (lPipe['x'], lPipe['y']))

        SCREEN.blit(IMAGES['base'], (self.basex, BASEY))
        # print score so player overlaps the score
        # showScore(self.score)
        SCREEN.blit(IMAGES['player'][self.playerIndex],
                    (self.playerx, self.playery))

        image_data = pygame.surfarray.array3d(pygame.display.get_surface())
        pygame.display.update()
        FPSCLOCK.tick(FPS)
        return image_data, reward, terminal


def getRandomPipe():
    """returns a randomly generated pipe"""
    # y of gap between upper and lower pipe
    gapYs = [20, 30, 40, 50, 60, 70, 80, 90]
    index = random.randint(0, len(gapYs)-1)
    gapY = gapYs[index]

    gapY += int(BASEY * 0.2)
    pipeX = SCREENWIDTH + 10

    return [
        {'x': pipeX, 'y': gapY - PIPE_HEIGHT},  # upper pipe
        {'x': pipeX, 'y': gapY + PIPEGAPSIZE},  # lower pipe
    ]


def checkCrash(player, upperPipes, lowerPipes):
    """returns True if player collders with base or pipes."""
    pi = player['index']
    player['w'] = IMAGES['player'][0].get_width()
    player['h'] = IMAGES['player'][0].get_height()

    # if player crashes into ground
    if player['y'] + player['h'] >= BASEY - 1:
        return True
    else:

        playerRect = pygame.Rect(player['x'], player['y'],
                      player['w'], player['h'])

        for uPipe, lPipe in zip(upperPipes, lowerPipes):
            # upper and lower pipe rects
            uPipeRect = pygame.Rect(uPipe['x'], uPipe['y'], PIPE_WIDTH, PIPE_HEIGHT)
            lPipeRect = pygame.Rect(lPipe['x'], lPipe['y'], PIPE_WIDTH, PIPE_HEIGHT)

            # player and upper/lower pipe hitmasks
            pHitMask = HITMASKS['player'][pi]
            uHitmask = HITMASKS['pipe'][0]
            lHitmask = HITMASKS['pipe'][1]

            # if bird collided with upipe or lpipe
            uCollide = pixelCollision(playerRect, uPipeRect, pHitMask, uHitmask)
            lCollide = pixelCollision(playerRect, lPipeRect, pHitMask, lHitmask)

            if uCollide or lCollide:
                return True

    return False

def pixelCollision(rect1, rect2, hitmask1, hitmask2):
    """Checks if two objects collide and not just their rects"""
    rect = rect1.clip(rect2)

    if rect.width == 0 or rect.height == 0:
        return False

    x1, y1 = rect.x - rect1.x, rect.y - rect1.y
    x2, y2 = rect.x - rect2.x, rect.y - rect2.y

    for x in range(rect.width):
        for y in range(rect.height):
            if hitmask1[x1+x][y1+y] and hitmask2[x2+x][y2+y]:
                return True
    return False

"""# DQN Model"""

import torch
import torch.nn.functional as F
device = 'cuda' if torch.cuda.is_available() else 'cpu'

def weights_init(layer):
  if isinstance(layer, torch.nn.Conv2d) or isinstance(layer, torch.nn.Linear):
    torch.nn.init.normal_(layer.weight, mean = 0., std = 0.01)
    layer.bias.data.fill_(0.01)

class DQN_net(torch.nn.Module):
  def __init__(self, in_channels = 4, out_actions = 2):
    super(DQN_net, self).__init__()
    self.conv1 = torch.nn.Conv2d(in_channels, 32, kernel_size = 8, stride = 4, padding = 2)
    self.conv2 = torch.nn.Conv2d(32, 64, kernel_size = 4, stride = 2, padding = 1)
    self.conv3 = torch.nn.Conv2d(64, 64, kernel_size = 3, stride = 1, padding = 1)
    self.fc1 = torch.nn.Linear(10*10*64, 512)
    self.fc2 = torch.nn.Linear(512, out_actions)

  def forward(self, x):
    # (84, 84, 4)
    x = F.relu(self.conv1(x))  # (20, 20, 32)
    x = F.relu(self.conv2(x)) # (10, 10, 64)
    x = F.relu(self.conv3(x)) # (10, 10, 64)
    x = x.reshape(-1, 10*10*64)  # (1, 6400)
    x = F.relu(self.fc1(x)) # (1, 512)
    x = self.fc2(x) # (1, 2)
    return x

"""# Replay Memory"""

class ReplayMemory:
  def __init__(self, capacity):
    self.capacity = capacity
    self.container = []
  def store(self, transition):
    self.container.append(transition)
    if len(self.container) > self.capacity:
      del self.container[0]
  def sample(self, batch_size):
    return random.sample(self.container, batch_size)
  def __len__(self):
    return len(self.container)

"""# DQN Training Object"""

class DQN:
  STACK_FRAMES = 4
  def __init__(self, memory_capacity, batch_size, epsilon, explore, replace_period, alpha, gamma, num_frames, num_actions):
    # Hyper-parameters
    self.replace_period = replace_period
    self.replace_counter = 1
    self.epsilon = epsilon
    self.epsilon_step = (0.1 - 0.0001) / explore
    self.alpha = alpha
    self.gamma = gamma

    # NN, loss, optimizer
    self.policy_net = DQN_net(num_frames, num_actions).to(device)
    self.target_net = DQN_net(num_frames, num_actions).to(device)
    self.policy_net.apply(weights_init)
    self.target_net.load_state_dict(self.policy_net.state_dict())
    self.loss_function = torch.nn.MSELoss().to(device)
    self.optimizer = torch.optim.Adam(self.policy_net.parameters(), lr = self.alpha)
    # Replay Memory
    self.replay_memory = ReplayMemory(memory_capacity)
    self.batch_size = batch_size

  def train(self):
    # Sample transition
    batch = self.replay_memory.sample(self.batch_size)
    state, action, reward, state_, terminal = zip(*batch)
    state = torch.tensor(state, dtype = torch.float32, requires_grad = True, device = device).reshape(self.batch_size, STACK_FRAMES, 80, 80)
    action = torch.cat(action).to(device)
    reward = torch.tensor(reward, dtype = torch.float32, requires_grad = False, device = device).reshape(self.batch_size, 1)
    state_ = torch.tensor(state_, dtype = torch.float32, requires_grad = False, device = device).reshape(self.batch_size, STACK_FRAMES, 80, 80)
    # (R + gamma * Q_) - Q
    Q = self.policy_net(state).gather(dim = 1, index = action.view(-1, 1))
    Q_ = self.target_net(state_).max(dim = 1)[0].view(-1, 1)
    TD_target = torch.zeros(self.batch_size, 1).to(device)
    # G = reward + self.gamma * Q_
    for i in range(self.batch_size):
      if not terminal[i]:
        TD_target[i, 0] = reward[i, 0] + self.gamma * Q_[i, 0]
      else:
        TD_target[i, 0] = reward[i, 0]
    # loss
    loss = self.loss_function(Q, TD_target)
    # Optimize
    self.optimizer.zero_grad()
    loss.backward()
    self.optimizer.step()

    if self.replace_counter % self.replace_period == 0:
      self.update_target_net()
      self.replace_counter = 1
    self.replace_counter += 1

  def choose_action(self, obs, is_train = True):
    if is_train:
      if random.random() > self.epsilon:
        return self.policy_net(obs).max(dim = 1)[1]
      else:
        return torch.tensor([random.randint(0, 1)], dtype = torch.int64, device = device)
    else:
      return self.policy_net(obs).max(dim = 1)[1]

  def memory_store(self, transition):
    self.replay_memory.store(transition)

  def update_epsilon(self):
    if self.epsilon > 0.0001:
      self.epsilon -= self.epsilon_step

  def update_target_net(self):
    print('Update Target Net')
    self.target_net.load_state_dict(self.policy_net.state_dict())

  def load_model(self, PATH):
    checkpoint = torch.load(PATH)
    self.policy_net.load_state_dict(checkpoint['policy_net_state_dict'])
    self.target_net.load_state_dict(checkpoint['target_net_state_dict'])
    self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    self.epsilon = checkpoint['epsilon']
    return checkpoint['episode'], checkpoint['iterations']

  def save_model(self, episode, iterations):
    torch.save({
        'episode': episode,
        'iterations': iterations,
        'policy_net_state_dict': self.policy_net.state_dict(),
        'target_net_state_dict': self.target_net.state_dict(),
        'optimizer_state_dict': self.optimizer.state_dict(),
        'epsilon': self.epsilon
    }, './checkpoint' + str(episode) + '.tar')


"""# DQN Process

"""

OBSERVE = 10000
EXPLORE = 3000000.
EPISODE = 1000000
ACTION_IDLE = 1
SAVE_ITER = 1000
STACK_FRAMES = 4
LOAD = False

# Initialize Game
game = GameState()
episode = 1
iterations = 0
# Initialize Model
if not LOAD:
  dqn = DQN(memory_capacity = 50000,
            batch_size = 32,
            epsilon = 0.1,
            explore = EXPLORE,
            replace_period = 5000,
            alpha = 1e-6,
            gamma = 0.99,
            num_frames = STACK_FRAMES,
            num_actions = 2)
elif LOAD:
  dqn = DQN(memory_capacity = 50000,
            batch_size = 32,
            epsilon = 0.1,
            explore = EXPLORE,
            replace_period = 5000,
            alpha = 1e-6,
            gamma = 0.99,
            num_frames = STACK_FRAMES,
            num_actions = 2)
  # Populate
  ckpts = [] # Put the NUMBER of checkpoints files here, will use them to populate the Replay Memory
  for ckpt in ckpts:
    dqn.load_model('./checkpoint' + str(ckpt) + '.tar')
    print('Start Populating by Model: ', ckpt)
    for ep in range(20):
      game.__init__()
      R = 0
      obs, reward, terminal = game.frame_step(np.array([1, 0]))
      obs = cv2.cvtColor(cv2.resize(obs, (80, 80)), cv2.COLOR_BGR2GRAY)
      _, obs = cv2.threshold(obs, 1, 255, cv2.THRESH_BINARY)
      obs = np.reshape(obs, (1, 80, 80))
      obs = np.concatenate([obs] * STACK_FRAMES, axis = 0)
      while not terminal:
        # Choose actions
        if iterations % ACTION_IDLE == 0:
          obs_tmp = torch.tensor(obs, dtype = torch.float32, device = device).reshape(1, STACK_FRAMES, 80, 80)
          action = dqn.choose_action(obs_tmp, True)
        else:
          action = torch.tensor(0, dtype = torch.int64, device = device)
        # Get next state
        if action.cpu().numpy()[0] == 0:
          act = np.array([1, 0])
        elif action.cpu().numpy()[0] == 1:
          act = np.array([0, 1])
        obs_, reward, terminal = game.frame_step(act)
        obs_ = cv2.cvtColor(cv2.resize(obs_, (80, 80)), cv2.COLOR_BGR2GRAY)
        _, obs_ = cv2.threshold(obs_, 1, 255, cv2.THRESH_BINARY)
        obs_ = np.reshape(obs_, (1, 80, 80))
        obs_ = np.concatenate([obs_, obs[:3, ...]], axis = 0)
        # Push transition to replay memory
        transition = [obs, action, reward, obs_, terminal]
        dqn.memory_store(transition)  
        # Update
        obs = obs_
        R += reward
      print('Episode: {}, Total Reward: {}'.format(ep, R))
  print('Start training by Model: ', 0) # Put the LAST checkpoint file here to continue training
  episode, iterations = dqn.load_model('./checkpoint' + str(0) + '.tar')


while episode <= EPISODE:
  game.__init__()
  R = 0
  # Get the first frame and stack it 4 times
  obs, reward, terminal = game.frame_step(np.array([1, 0]))
  obs = cv2.cvtColor(cv2.resize(obs, (80, 80)), cv2.COLOR_BGR2GRAY)
  _, obs = cv2.threshold(obs, 1, 255, cv2.THRESH_BINARY)
  obs = np.reshape(obs, (1, 80, 80))
  obs = np.concatenate([obs] * STACK_FRAMES, axis = 0)
  while not terminal:
    # Choose actions
    if iterations % ACTION_IDLE == 0:
      obs_tmp = torch.tensor(obs, dtype = torch.float32, device = device).reshape(1, STACK_FRAMES, 80, 80)
      action = dqn.choose_action(obs_tmp)
    else:
      action = torch.tensor(0, dtype = torch.int64, device = device)

    # Get next state
    if action.cpu().numpy()[0] == 0:
      act = np.array([1, 0])
    elif action.cpu().numpy()[0] == 1:
      act = np.array([0, 1])
    obs_, reward, terminal = game.frame_step(act)
    obs_ = cv2.cvtColor(cv2.resize(obs_, (80, 80)), cv2.COLOR_BGR2GRAY)
    _, obs_ = cv2.threshold(obs_, 1, 255, cv2.THRESH_BINARY)
    obs_ = np.reshape(obs_, (1, 80, 80))
    obs_ = np.concatenate([obs_, obs[:3, ...]], axis = 0)

    # Push transition to replay memory
    transition = [obs, action, reward, obs_, terminal]
    dqn.memory_store(transition)

    # Train
    if iterations > OBSERVE:
      dqn.update_epsilon()
      dqn.train()
    
    # Update
    obs = obs_
    iterations += 1
    R += reward
    if iterations % SAVE_ITER  == 0 and iterations > OBSERVE:
      dqn.save_model(episode, iterations)
  
  print('Episode: {}, Total Reward: {}, Iterations: {}, Epsilon: {}, Memory Size: {}'.format(episode, R, iterations, dqn.epsilon, len(dqn.replay_memory)))
  episode += 1