Heuristics visualization

board_visualization.py

@@ -0,0 +1,58 @@
+from environment.config import *
+from environment.controls import Controller
+from environment.tetris import Tetris
+from environment.renderer import PyGameRenderer, PAPER_Tetris_Config
+from copy import copy
+
+board = [
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+]
+
+env = Tetris(10, 20)
+
+render_config = copy(PAPER_Tetris_Config)
+render_config.render_holes = True
+render_config.render_bumpiness = False
+render_config.render_max_height = False
+render_config.show_ghost_piece = False
+
+renderer = PyGameRenderer(30, render_config)
+
+key_actions = {
+    "left":     lambda: env.move(-1),
+    "right":    lambda: env.move(1),
+    "rotate_cw":env.rotate,
+    "soft_drop":env.soft_drop,
+    "hard_drop":env.hard_drop,
+    "hold":     env.hold,
+    "reset":    env.reset,
+    "down":     env.down,
+}
+
+if __name__ == '__main__':
+    try:
+        while True:
+            renderer.render(env)
+            controller = Controller(key_actions)
+            controller.handleEvents()
+            renderer.wait(1)
+    finally:
+        pass

environment/colors.py

@@ -9,5 +9,7 @@ class Color:
     CYAN   = (0,   240, 240)
     YELLOW = (240, 240, 0  )
     PURPLE = (160, 0,   240)
+    GRAY   = (240, 240, 240)
+    PINK   = (255, 20,  147)
 
     ALL = [BLACK, PURPLE, GREEN, RED, BLUE, ORANGE, CYAN, YELLOW]

environment/renderer.py

@@ -1,11 +1,29 @@
+import numpy as np
 import pygame
 from environment.colors import Color
 from environment.tetris import Tetris
+from utils.heuristics import Heuristics
+
+class RenderConfig():
+    def __init__(self, bg_color, grid_color, highlight_color, text_color, render_bumpiness=False, render_holes=False, render_max_height=False, show_ghost_piece=True) -> None:
+        self.bg_color = bg_color
+        self.grid_color = grid_color
+        self.highlight_color = highlight_color
+        self.render_bumpiness = render_bumpiness
+        self.render_holes = render_holes
+        self.render_max_height = render_max_height
+        self.show_ghost_piece = show_ghost_piece
+        self.text_color = text_color
+
+NES_Tetris_Config = RenderConfig(Color.BLACK, Color.BLACK, Color.WHITE, Color.WHITE)
+PAPER_Tetris_Config = RenderConfig(Color.GRAY, Color.WHITE, Color.PINK, Color.BLACK)
 
 class PyGameRenderer():
-    def __init__(self, cell_size):
+    def __init__(self, cell_size, config=NES_Tetris_Config):
         self.clock = pygame.time.Clock()
         self.cell_size = cell_size
+        self.heuristics = Heuristics()
+        self.config = config
 
     def render(self, env:Tetris):
         if not pygame.get_init():
@@ -18,21 +36,24 @@ def render(self, env:Tetris):
             self.font = pygame.font.Font(None, 36)
 
         self.surface.set_alpha(255)
-        self.screen.fill((0, 0, 0))
-        self.surface.fill((0, 0, 0))
+        self.screen.fill(self.config.bg_color)
+        self.surface.fill(self.config.bg_color)
 
         self.top_msg(f'Score: {env.score}')
-        self.draw_rect((255, 0, 0), 5, 1, env.cols, env.rows)
+        # self.draw_rect((255, 0, 0), 5, 1, env.cols, env.rows)
         for i, shape in enumerate(env.held_shapes):
             self.draw_matrix(shape, (1, i * 3 + 1))
         self.draw_matrix(env.board, (5, 1))
         self.draw_matrix(env.shape, (env.shape_x + 5, env.shape_y + 1))
-        self.draw_matrix(env.shape, (env.shape_x + 5, env._soft_drop(env.board, env.shape, (env.shape_x, env.shape_y)) + 1), 2)
+        if self.config.show_ghost_piece:
+            self.draw_matrix(env.shape, (env.shape_x + 5, env._soft_drop(env.board, env.shape, (env.shape_x, env.shape_y)) + 1), 2)
         for i, shape in enumerate(env.next_shapes):
             self.draw_matrix(shape, (env.board.shape[1] + 6, i * 3 + 1))
         self.draw_grid()
         self.draw_stats(env)
 
+        self.draw_heuristics(env)
+
         if env.done:
             self.add_backdrop()
             self.center_msg(f"Game Over!\nPress r to restart\nFinal score:{env.score}")
@@ -63,9 +84,9 @@ def draw_rect(self, color, left, top, width, height):
 
     def draw_grid(self):
         for i in range(self.width // self.cell_size):
-            pygame.draw.line(self.surface, (0, 0, 0), (i * self.cell_size, 0), (i * self.cell_size, self.height * self.cell_size))
+            pygame.draw.line(self.surface, self.config.grid_color, (i * self.cell_size, 0), (i * self.cell_size, self.height * self.cell_size))
         for i in range(self.height // self.cell_size):
-            pygame.draw.line(self.surface, (0, 0, 0), (0, i * self.cell_size), (self.width * self.cell_size, i * self.cell_size))
+            pygame.draw.line(self.surface, self.config.grid_color, (0, i * self.cell_size), (self.width * self.cell_size, i * self.cell_size))
 
     def wait(self, ms):
         self.clock.tick(1000 / ms)
@@ -78,7 +99,7 @@ def top_msg(self, msg):
         self.surface.blit(msg_image, (self.width // 2-msgim_center_x, 0))
 
     def font_img(self, text):
-        return pygame.font.Font(pygame.font.get_default_font(), 12).render(text, False, (255,255,255), (0,0,0))
+        return pygame.font.Font(pygame.font.get_default_font(), 12).render(text, False, self.config.text_color, self.config.bg_color)
 
     def center_msg(self, msg):
         for i, line in enumerate(msg.splitlines()):
@@ -98,3 +119,46 @@ def draw_matrix(self, matrix:list[list[int]], offset:tuple[int, int], width:int=
                 color = Color.ALL[val] if val < len(Color.ALL) else Color.GRAY
                 rect = pygame.Rect((off_x+x) * self.cell_size, (off_y+y) * self.cell_size, self.cell_size, self.cell_size)
                 pygame.draw.rect(self.surface, (*color, opacity), rect, width)
+
+    def draw_heuristics(self, env):
+        if self.config.render_bumpiness:
+            self.draw_bumpiness(env.board)
+
+        if self.config.render_holes:
+            self.draw_holes(env.board)
+
+        if self.config.render_max_height:
+            self.draw_max_height(env.board)
+
+    def draw_bumpiness(self, board):
+        heights = self.heuristics._get_heights(board)
+        for i, height in enumerate(heights):
+            rows, _ = board.shape
+            start = ((i + 5) * self.cell_size, (rows - height + 1) * self.cell_size)
+            end = ((i + 6) * self.cell_size, (rows - height + 1) * self.cell_size)
+            pygame.draw.line(self.surface, self.config.highlight_color, start, end, 5)
+            if i < len(heights) - 1:
+                start = ((i + 6) * self.cell_size, (rows - heights[i + 1] + 1) * self.cell_size)
+                pygame.draw.line(self.surface, self.config.highlight_color, end, start, 5)
+
+    def draw_holes(self, board):
+        bridge_mask = board != 0
+        rows, cols = board.shape
+
+        for col in range(cols):
+            if np.any(bridge_mask[:-1, col]):
+                bridge_row = np.argmax(bridge_mask[:, col])
+                for row in range(bridge_row + 1, rows):
+                    if not bridge_mask[row, col]:
+                        x = (col + 5) * self.cell_size
+                        y = (row + 1 ) * self.cell_size
+                        pygame.draw.rect(self.surface, self.config.highlight_color, (x, y, self.cell_size, self.cell_size), 4)
+
+    def draw_max_height(self, board):
+        max_heights = np.max(self.heuristics._get_heights(board))
+        rows, cols = board.shape
+
+        start = (5 * self.cell_size, (rows - max_heights + 1) * self.cell_size)
+        end = ((cols + 6) * self.cell_size, (rows - max_heights + 1) * self.cell_size)
+        pygame.draw.line(self.surface, self.config.highlight_color, start , end, 5)
+

utils/heuristics.py

@@ -17,7 +17,7 @@ def _count_bridges(self, board:np.ndarray):
         return empty_count
 
     def _calculate_max_height_and_bumpiness(self, board:np.ndarray):
-        heights = np.max(np.where(board != 0, len(board) - np.arange(len(board))[:, None], 0), axis=0)
+        heights = self._get_heights(board)
         max_height = np.max(heights)
         bumpiness = np.sum(np.abs(np.diff(heights)))
 
@@ -52,4 +52,7 @@ def get_heuristics(self, board:np.ndarray):
         holes = self._count_bridges(board)
         bumpiness, height = self._calculate_max_height_and_bumpiness(board)
         check_pillar = self._check_for_pillar(board)
-        return np.array([cleared_lines, holes, bumpiness, height,check_pillar])
+        return np.array([cleared_lines, holes, bumpiness, height,check_pillar])
+
+    def _get_heights(self, board):
+        return np.max(np.where(board != 0, len(board) - np.arange(len(board))[:, None], 0), axis=0)