10-fold computation speed-up with numba and other optimizations

+ documentation

Author: mikhail-vlasenko

README.md

@@ -35,6 +35,6 @@ This is only checked when a new piece appears, so you need to hold the key.
     2 - no hard drop
     3 - let the piece land on its own, the bot is always scared
     Number of computing paths for the next piece:
-    z, x, c - 1, 3, 5
+    z, x, c - 1, 4, 8
     n - try to clean the field
     m - disable cleaning mode (focus on getting tetrises)

config.py

@@ -1,7 +1,9 @@
-# pieces are encoded as
-# 0 - line, 1 - square, 2 - T(flip), 3 - |__, 4 - __|, 5 - -|_,6 - _|-
+import numpy as np
 from src.display_consts import DisplayConsts
 
+
+# pieces are encoded as
+# 0 - line, 1 - square, 2 - T(flip), 3 - |__, 4 - __|, 5 - -|_,6 - _|-
 PIECE_NAMES = ['line', 'square', 'T(flip)', '|__', '__|', '-|_', '_|-']
 
 
@@ -10,7 +12,7 @@ def name_piece(piece: int) -> str:
 
 
 # in BGR
-original_colors = [(0, 0, 0) for _ in range(7)]
+original_colors = np.zeros((7, 3), np.int)
 original_colors[0] = (230, 228, 180)
 original_colors[1] = (182, 228, 247)
 original_colors[2] = (177, 99, 140)
@@ -20,7 +22,7 @@ def name_piece(piece: int) -> str:
 original_colors[6] = (171, 240, 177)
 
 # tetr.io colors in RGB
-tetrio_colors = [(0, 0, 0) for _ in range(7)]
+tetrio_colors = np.zeros((7, 3), np.int)
 tetrio_colors[0] = (36, 214, 150)
 tetrio_colors[1] = (210, 171, 42)
 tetrio_colors[2] = (212, 67, 195)
@@ -50,9 +52,9 @@ def name_piece(piece: int) -> str:
     'debug status': 1,  # greater is more information, 0 is none
     'key press delay': 0.02,  # increase if facing misclicks, decrease to go faster
     'tetrio garbage': True,
-    'starting choices for 2nd': 3,
+    'starting choices for 2nd': 8,
     # if true, looks at another frame to check for correct piece placement
-    # reduces speed, improves robustness
+    # reduces speed, increases robustness
     'confirm placement': True,
     'play safe': False,  # ai is even more robust
     'play for survival': False,  # if true, starts in 'cleaning' mode
@@ -75,7 +77,7 @@ def name_piece(piece: int) -> str:
 
 
 def configure_fast():
-    CONFIG['starting choices for 2nd'] = 4
+    CONFIG['starting choices for 2nd'] = 8
     CONFIG['confirm placement'] = False
     CONFIG['play for survival'] = True
     CONFIG['override delay'] = True

src/AI_main.py

@@ -1,5 +1,7 @@
 from typing import List
 
+from numba import jit
+
 from find_landings import all_landings
 import numpy as np
 from direct_keys import *
@@ -49,6 +51,7 @@ def clear_line(field):
         return field, full_cnt
 
     @staticmethod
+    @jit(nopython=True)
     def find_roofs(field: np.array) -> (int, int, np.array, int):
         """
         finds blank squares under landed pieces
@@ -70,6 +73,7 @@ def find_roofs(field: np.array) -> (int, int, np.array, int):
         return blank_cnt, int(np.max(tops[:, 0])), tops[:, 0], blank_depth
 
     @staticmethod
+    @jit(nopython=True)
     def almost_full_line(field):
         score = 0
         for i in range(len(field)):
@@ -81,6 +85,7 @@ def almost_full_line(field):
         return score
 
     @staticmethod
+    @jit(nopython=True)
     def find_pit(field, tops):
         gap_idx = []
         for i in range(len(field)):
@@ -104,6 +109,7 @@ def find_pit(field, tops):
         return max_pit_h
 
     @staticmethod
+    @jit(nopython=True)
     def find_hole(tops):
         cnt_hole = 0
         tops = np.insert(tops, 0, 20)
@@ -223,6 +229,14 @@ def choose_action(self, field: np.array, piece_idx, can_hold) -> Position:
         return result
 
     def choose_action_depth2(self, field: np.array, piece_idx: int, next_piece: int, can_hold: bool) -> Position:
+        """
+        finds best action considering the next piece as well
+        :param field:
+        :param piece_idx:
+        :param next_piece:
+        :param can_hold:
+        :return:
+        """
         if self.choices_for_2nd == 1:
             # can simplify
             return self.choose_action(field, piece_idx, can_hold)
@@ -249,14 +263,18 @@ def choose_action_depth2(self, field: np.array, piece_idx: int, next_piece: int,
                 sub_score_hold = self.calc_best(results[i].field, piece_idx)[0].score
                 results[i].next_score = max(sub_score, sub_score_hold)
 
-        # sort by total score, prioritizing tetris on current turn
-        # (not another move, and then tetris)
-        results.sort(key=lambda x: x.next_score + x.score + 1000 * x.expect_tetris, reverse=True)
-        if results[0].piece == self.held_piece and self.held_piece != piece_idx:
+        # take best by total score, prioritizing tetris on current turn
+        # (instead of another move, and then tetris)
+        optimal = max(results, key=lambda x: x.next_score + x.score + 1000 * x.expect_tetris)
+        if optimal.piece == self.held_piece and self.held_piece != piece_idx:
             self.hold_piece(piece_idx)
-        return results[0]
+        return optimal
 
     def place_piece(self, piece: int, rotation: int, x_pos: int, height: int, rot_now=0, x_pos_now=3, depth=0):
+        """
+        puts the piece into correct position (before lowering)
+        optionally verifies placement
+        """
         if depth == 3:
             if CONFIG['debug status'] >= 1:
                 print('depth 3 reached in place_piece')
@@ -322,7 +340,7 @@ def runtime_tuning(self):
         2 - medium speed, no hard placing, turn on at level 6
         3 - for the late game, always scared
         control number of paths for the next piece:
-        z, x, c - 1, 3, 5
+        z, x, c - 1, 4, 8
         n - try to clean the field
         m - disable cleaning mode (try to get tetrises)
 
@@ -342,9 +360,9 @@ def runtime_tuning(self):
         if keyboard.is_pressed('z'):
             self.choices_for_2nd = 1
         elif keyboard.is_pressed('x'):
-            self.choices_for_2nd = 3
+            self.choices_for_2nd = 4
         elif keyboard.is_pressed('c'):
-            self.choices_for_2nd = 5
+            self.choices_for_2nd = 8
 
         if keyboard.is_pressed('n'):
             self.clearing = True

src/find_landings.py

@@ -2,14 +2,17 @@
 from typing import List
 
 import numpy as np
+from numba import jit
 
 from position import Position
 from figures import array_of_figures as pieces
 
 
+@jit(nopython=True)
 def check_collision(field, piece, piece_pos, piece_idx):
     r = 4
     if piece_idx != 0:
+        # only the line takes 4x4 grid, so for other pieces only check 3x3
         r -= 1
     for i in range(r):
         for j in range(r):
@@ -20,18 +23,28 @@ def check_collision(field, piece, piece_pos, piece_idx):
     return False
 
 
-def land(field: np.array, piece: np.array, pos_now: List[int], piece_idx: int) -> np.array:
-    res = deepcopy(field)
-    while not check_collision(res, piece, pos_now, piece_idx):
+@jit(nopython=True)
+def land(field: np.array, piece: np.array, x_pos: int, piece_idx: int) -> np.array:
+    """
+    simulates a piece falling from a set position onto the field
+    helps predict the outcome of an action
+    :param field: np.array, WILL BE MODIFIED
+    :param piece: array that tells the shape of the piece (with rotation applied)
+    :param x_pos: horizontal position of the piece
+    :param piece_idx: piece index
+    :return: resulting field
+    """
+    pos_now = [0, x_pos]
+    while not check_collision(field, piece, pos_now, piece_idx):
         pos_now[0] += 1
     if pos_now[0] == 0:
         return None
     pos_now[0] -= 1
     for i in range(4):
         for j in range(4):
             if i + pos_now[0] < len(field) and j + pos_now[1] < len(field[0]):
-                res[i + pos_now[0]][j + pos_now[1]] += piece[i][j]
-    return res
+                field[i + pos_now[0]][j + pos_now[1]] += piece[i][j]
+    return field
 
 
 def all_landings(field: np.array, piece_index: int) -> List[Position]:
@@ -44,7 +57,7 @@ def all_landings(field: np.array, piece_index: int) -> List[Position]:
     results = []
     for rotation in range(len(pieces[piece_index])):
         for x_pos in range(-3, 10):
-            res = land(field, pieces[piece_index][rotation], [0, x_pos], piece_index)
+            res = land(deepcopy(field), pieces[piece_index][rotation], x_pos, piece_index)
             if res is not None:
                 results.append(Position(res, rotation, x_pos, piece_index))
     return results

src/main.py

@@ -11,7 +11,12 @@ def main():
     can_hold_flag = True
     expected_rwd = 0
     ai = AI()
-    placement = None
+    position = None
+
+    # call jit-compiling functions to compile them
+    ai.calc_best(np.zeros((20, 10), dtype=np.int), 0)
+    field, _ = get_field()
+    print("Compilation complete")
 
     # infinite playing cycle
     while True:
@@ -30,14 +35,14 @@ def main():
         if ai.held_piece == -1:
             ai.hold_piece(piece_idx)
             can_hold_flag = False
-            time.sleep(0.2)
             continue
+
         # shenanigans for better parsing of the original game
         if CONFIG['game'] == 'original':
-            if placement is not None and placement.expect_tetris:
+            if position is not None and position.expect_tetris:
                 # hoping that it was not a misclick, not taking a screenshot because TETRIS blocks the view
                 field = np.zeros((3, 10), dtype=np.int)
-                field = np.concatenate((field, ai.clear_line(placement.field)[0]))
+                field = np.concatenate((field, ai.clear_line(position.field)[0]))
                 time.sleep(0.2)
             elif not ai.scared:
                 field, next_piece = get_field()
@@ -53,27 +58,29 @@ def main():
             if CONFIG['debug status'] >= 2:
                 print(field)
             print(f'current score {actual_score}')
+
+        # next piece is not recognized
         if next_piece == -1:
             if CONFIG['debug status'] >= 1:
                 print("unknown next")
             next_piece = 1  # assume square as it is the most neutral one
 
         calc_start_time = time.time()
         # compute best outcome
-        placement = ai.choose_action_depth2(field[3:], piece_idx, next_piece, can_hold_flag)
+        position = ai.choose_action_depth2(field[3:], piece_idx, next_piece, can_hold_flag)
 
         if CONFIG['debug status'] >= 1:
-            # useful info
+            # print useful info
             print('calculation took', time.time() - calc_start_time)
-            print(f'chosen placement for {name_piece(placement.piece)}: '
-                  f'({placement.rotation}, {placement.x_pos}) with score {placement.score}')
-            print(f'next figure {name_piece(next_piece)} should give {placement.next_score}')
-            if placement.expect_tetris:
+            print(f'chosen placement for {name_piece(position.piece)}: '
+                  f'({position.rotation}, {position.x_pos}) with score {position.score}')
+            print(f'next figure {name_piece(next_piece)} should give {position.next_score}')
+            if position.expect_tetris:
                 print('expecting TETRIS')
 
-        expected_rwd = ai.get_score(ai.clear_line(placement.field)[0])[0]
+        expected_rwd = ai.get_score(ai.clear_line(position.field)[0])[0]
         # emulate key presses to place the piece
-        ai.place_piece(placement.piece, placement.rotation, placement.x_pos, ai.find_roofs(placement.field)[1])
+        ai.place_piece(position.piece, position.rotation, position.x_pos, ai.find_roofs(position.field)[1])
         # wait for everything to settle down
         ai.place_piece_delay()
 
@@ -83,5 +90,5 @@ def main():
 
 
 if __name__ == '__main__':
-    time.sleep(2)
+    time.sleep(1)
     main()

src/scan_field.py

@@ -1,10 +1,13 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from mss import mss
+from numba import jit
+
 from config import CONFIG
 
 screen_capture = mss()
 monitor = {"left": 0, "top": 0, "width": CONFIG['screen width'], "height": CONFIG['screen height']}
+piece_colors = CONFIG['piece colors']
 
 
 def print_image(arr, figure_size=10):
@@ -45,11 +48,13 @@ def simplified(pixels: np.array) -> np.array:
     return field
 
 
+@jit(nopython=True)
 def cmp_pixel(p1, p2):
     return abs(p1[0] - p2[0]) + abs(p1[1] - p2[1]) + abs(p1[2] - p2[2])
 
 
-def get_figure_by_color(screen):
+@jit(nopython=True)
+def get_figure_by_color(screen: np.array):
     """
     finds next piece class based on color
     shape recognition is difficult because the piece is not aligned with a grid
@@ -59,8 +64,8 @@ def get_figure_by_color(screen):
     for i in range(len(screen)):
         for j in range(len(screen[0])):
             pixel = screen[i, j][:3]
-            for piece_idx in range(len(CONFIG['piece colors'])):
-                distance = cmp_pixel(CONFIG['piece colors'][piece_idx][::-1], pixel)
+            for piece_idx in range(len(piece_colors)):
+                distance = cmp_pixel(piece_colors[piece_idx][::-1], pixel)
                 if distance < 10:
                     return piece_idx
     return -1