Environment for grid based Reinforcement learning game

GridEnvironment.py

import numpy as np
from scipy.spatial import distance
from PIL import Image, ImageDraw, ImageFont
import cv2
import time

class BaseBlob:
    game_blobs = []
    wall_blobs = []

    def __init__(self, size=20, bgr=(255,255,255)):
        self.size = size
        self.x = np.random.randint(0, self.size)
        self.y = np.random.randint(0, self.size)

        self.bgr = bgr

    def __str__(self):
        return f"{self.x}, {self.y}"

    def __sub__(self, rhs_blob):
        return (self.x-rhs_blob.x, self.y-rhs_blob.y)

    def __eq__(self, rhs_blob):
        return (self - rhs_blob) == (0,0)

    def action(self, choice):
        raise Exception("No Action was defined")

    def move(self, x,y):
        raise Exception("No Move was defined")

class WallBlob(BaseBlob):

    def __init__(self, size=20, bgr=(192,192,192)):
        super().__init__(size, bgr)
        while (self.x, self.y) in BaseBlob.wall_blobs:
            self.x = np.random.randint(0, self.size)
            self.y = np.random.randint(0, self.size)

        BaseBlob.wall_blobs.append((self.x,self.y))

class Blob(BaseBlob):

    def __init__(self, size=20, bgr=(255,255,255)):
        super().__init__(size, bgr)
        # make sure that no blob overlaps with another
        while (self.x, self.y) in BaseBlob.game_blobs or (self.x, self.y) in BaseBlob.wall_blobs:
            self.x = np.random.randint(0, self.size)
            self.y = np.random.randint(0, self.size)

        BaseBlob.game_blobs.append((self.x, self.y))


    def action(self, choice):
        """
        Gives us 8 total movement options:
            - 0,1,2,3
            - 0: upper right
            - 1: lower left
            - 2: upper left
            - 3: lower right

            - 4: up
            - 5: down
            - 6: left
            - 7: right

        :param choice:
        :return:
        """
        if choice == 0: # upper right
            self.move(x=1, y=1)
        elif choice == 1: # lower left
            self.move(x=-1, y=-1)
        elif choice == 2: # upper left
            self.move(x=-1, y=1)
        elif choice == 3: # lower right
            self.move(x=1, y=-1)
        elif choice == 4: # up
            self.move(x=0, y=1)
        elif choice == 5: # down
            self.move(x=0, y=-1)
        elif choice == 6: # left
            self.move(x=-1, y=0)
        elif choice == 7: # right
            self.move(x=1, y=0)


    def move(self, x=-99, y=-99):
        """

        :param x: amount in the x direction to move
        :param y: amount in the y direction to move
        :return: True - the move was possible, False - could not move either out of bounds or hit a wall
        """
        # if no value for x, move random
        original_x = self.x
        original_y = self.y
        result = True # assume we can perform the action

        if x == -99:
            self.x += np.random.randint(-1,2)
        else:
            self.x += x
        # if no value for y, move random
        if y == -99:
            self.y += np.random.randint(-1,2)
        else:
            self.y += y

        # check to make sure we did not move to a wall
        for wall_blob in BaseBlob.wall_blobs:
            if wall_blob == (self.x, self.y):
                # if we moved onto a way, go back to where we were before hitting the wall
                self.x = original_x
                self.y = original_y
                result = False
                break

        # check for out of bounds
        if self.x < 0:
            self.x = 0
            result = False
        elif self.x > self.size-1:
            self.x = self.size - 1
            result = False

        if self.y < 0:
            self.y = 0
            result = False
        elif self.y > self.size - 1:
            self.y = self.size - 1
            result = False

        return result

class _ActionSpace:
    number_of_actions = 8
    """
        Gives us 8 total movement options:
        - 0,1,2,3
        - 0: upper right
        - 1: lower left
        - 2: upper left
        - 3: lower right

        - 4: up
        - 5: down
        - 6: left
        - 7: right

    """
    def __init__(self):
        self.n = _ActionSpace.number_of_actions
        self.actionspace = [0,1,2,3,4,5,6,7]

    def action(self, choice, blob):
        """
        Gives us 8 total movement options:
            - 0,1,2,3
            - 0: upper right
            - 1: lower left
            - 2: upper left
            - 3: lower right

            - 4: up
            - 5: down
            - 6: left
            - 7: right

        :param choice:
        :return: True - if action was successful, False if not
        """
        result = True
        if choice == 0: # upper right
            result = blob.move(x=1, y=1)
        elif choice == 1: # lower left
            result = blob.move(x=-1, y=-1)
        elif choice == 2: # upper left
            result = blob.move(x=-1, y=1)
        elif choice == 3: # lower right
            result = blob.move(x=1, y=-1)
        elif choice == 4: # up
            result = blob.move(x=0, y=1)
        elif choice == 5: # down
            result = blob.move(x=0, y=-1)
        elif choice == 6: # left
            result = blob.move(x=-1, y=0)
        elif choice == 7: # right
            result = blob.move(x=1, y=0)

        return result

class _ObservationSpace:
    """
    The observation space, or state space is:


    for an NxN grid:

    Distance calculation from player to win square, [above,below,inline] to win square, [left,right,inline] to win square, distance calculation from player to lose square, [above,below,inline] to lose square, [left,right,inline] to lose square, actions

    """

    def __init__(self, size, number_of_blobs=2):
        """

        :param size:
        :param number_of_blobs: typically 2, 1-win and 1-lose blob
        """
        self.size = size
        self.max_distance_indexes = int(np.sqrt(self.size ** 2 + self.size ** 2))
        total_space = ()
        for _ in range(number_of_blobs):
            total_space = total_space + (self.max_distance_indexes, 3, 3,)

        total_space = total_space + (_ActionSpace.number_of_actions, )

        self.dqn_observation_space = (size,size,3)
        self.observation_space = total_space

    def _distance(self, from_blob, to_blob):
        return int(distance.euclidean((to_blob.x, to_blob.y), (from_blob.x, from_blob.y)))

    def state(self, player_blob, win_blob, lose_blobs):
        d = self._distance(player_blob, win_blob)

        above_below = 0 # is this blob above, below or inline with the to_blob
        left_right = 0 # is this blob to the left, right or inline with the to_blob
        if win_blob.y > player_blob.y:
            above_below = 2 # win above player
        elif win_blob.y < player_blob.y:
            above_below = 1 # win below player

        if win_blob.x > player_blob.x:
            left_right = 2 # win blob to right
        elif win_blob.x < player_blob.x:
            left_right = 1 # win blob to the left

        state = (d, above_below, left_right)

        for lose_blob in lose_blobs:
            d = self._distance(player_blob, lose_blob)

            above_below = 0  # is this blob above, below or inline with the to_blob
            left_right = 0  # is this blob to the left, right or inline with the to_blob
            if lose_blob.y > player_blob.y:
                above_below = 2
            elif lose_blob.y < player_blob.y:
                above_below = 1

            if lose_blob.x > player_blob.x:
                left_right = 2
            elif lose_blob.x < player_blob.x:
                left_right = 1

            state = state + (d, above_below, left_right, )

        return state




class GridEnvironment:

    def __init__(self, size=20, lose_blob_count=1, number_of_wall_blobs=0):
        self.size = size
        self.actionspace = _ActionSpace()
        self.observationspace = _ObservationSpace(self.size, number_of_blobs=lose_blob_count+1)
        self.q_table = None
        self.win_blob = None
        self.lose_blobs = []
        self.lose_blob_count = lose_blob_count
        self.player_blob = None
        self.MOVE_PENALTY = -1
        self.LOSE_PENALTY = -300
        self.WIN_REWARD = 25
        self.NON_MOVE_PENALTY = self.MOVE_PENALTY * 4
        self.draw_font = ImageFont.truetype("/Library/Fonts/Arial Narrow.ttf", 20)
        self.wall_blobs = []
        if number_of_wall_blobs > 0:
            for wall_blob in range(0, number_of_wall_blobs):
                wall_blob = WallBlob(size=size)
                self.wall_blobs.append(wall_blob)


    def reset(self):
        BaseBlob.game_blobs = []
        self.win_blob = Blob(size=self.size, bgr=(0,255,0))
        self.player_blob = Blob(size=self.size, bgr=(255,175,0))
        self.lose_blobs = []
        for _ in range(self.lose_blob_count):
            self.lose_blobs.append(Blob(size=self.size, bgr=(0,0,255)))

        not_done = True
        while not_done:
            blobs_moved = []
            for i, lose_blob in enumerate(self.lose_blobs):
                if self.win_blob - lose_blob == (0, 0):
                    blobs_moved.append(i)
            if len(blobs_moved) > 0:
                for x in blobs_moved:
                    self.lose_blobs[x] = Blob(size=self.size, bgr=(0,0,255))
            else:
                not_done = False

        initial_state = self.observationspace.state(self.player_blob, self.win_blob, self.lose_blobs)
        return initial_state

    def create_q_table(self):
        self.q_table = np.random.uniform(size=self.observationspace.observation_space)
        return self.q_table

    def step(self, action_choice):
        """
        Step the player
        :param action_choice: one of: [0,1,2,3,4,5,6,7]
        :return: next_state, reward
        """
        action_result = self.actionspace.action(action_choice, self.player_blob)
        new_state = self.observationspace.state(self.player_blob, self.win_blob, self.lose_blobs)

        # Calculate Reward
        done = False
        for lose_blob in self.lose_blobs:
            if self.player_blob - lose_blob == (0,0):
                reward = self.LOSE_PENALTY
                done = True
                break

        if not done:
            if self.player_blob - self.win_blob == (0, 0):
                reward = self.WIN_REWARD
                done = True
            else:
                if action_result:
                    # then we were able to move
                    reward = self.MOVE_PENALTY
                else:
                    # we could not perform the action.  perhaps at the boundary or we hit a wall
                    # we want to penalize these slightly more
                    reward = self.NON_MOVE_PENALTY


        return new_state, reward, done

    def _draw_grid(self, image):

        step_count = self.size

        # Draw a grid
        draw = ImageDraw.Draw(image)
        y_start = 0
        y_end = image.height
        step_size = int(image.width / step_count)

        for x in range(0, image.width, step_size):
            line = ((x, y_start), (x, y_end))
            draw.line(line, fill=255)

        x_start = 0
        x_end = image.width

        for y in range(0, image.height, step_size):
            line = ((x_start, y), (x_end, y))
            draw.line(line, fill=255)

    def render(self, show=True, slow_render=0, done=False, episode_number = 0, move_number = 0, show_grid=False, total_episodes = None, episode_reward=None):
        """

        :param show: True - display the game
        :param slow_render delay to slow the rendering down. 0 - do not slow down rendering
        :return: boolean:
        """
        if show:
            game_canvas = np.zeros((self.size, self.size, 3), dtype=np.uint8)  # starts an rbg of our size
            if self.win_blob - self.player_blob != (0,0):
                # if the win blob and player blog are in different places draw the
                # winning blob
                game_canvas[self.win_blob.x][self.win_blob.y] = self.win_blob.bgr  # sets the food location tile to green color
            for lose_blob in self.lose_blobs:
                game_canvas[lose_blob.x][lose_blob.y] = lose_blob.bgr  # sets the enemy location to red
            game_canvas[self.player_blob.x][self.player_blob.y] = self.player_blob.bgr  # sets the player tile to blue

            for wall_blob in self.wall_blobs:
                game_canvas[wall_blob.x][wall_blob.y] = wall_blob.bgr  # sets the wall location

            img = Image.fromarray(game_canvas, 'RGB')  # reading to rgb. Apparently. Even tho color definitions are bgr. ???
            img = img.resize((400, 400))  # resizing so we can see our agent in all its glory.

            if show_grid:
                self._draw_grid(img)

            if episode_number >= 0 and move_number >= 0:
                draw = ImageDraw.Draw(img)
                episode_str = f"Episode: {episode_number}"
                if total_episodes:
                    episode_str += f" / {total_episodes}"

                episode_reward_str = ""
                if episode_reward is not None:
                    episode_reward_str = f"Reward: {episode_reward}"

                summary_text = f"{episode_str} Move: {move_number} {episode_reward_str}"
                draw.text((0, 0), summary_text, font=self.draw_font)
                # print(summary_text)

            cv2.imshow("GridEnvironment", np.array(img))  # show it!
            if cv2.waitKey(1) & 0xFF == ord('q'):
                cv2.destroyAllWindows()

            if slow_render:
                time.sleep(slow_render)

            if done:
                time.sleep(1.0)
    def destroy(self):
        cv2.destroyAllWindows()