Repository: jonbinney/deep_rabbit_hole
Pull request #201: Use numba to speedup simple agent
Author: login: jonbinney
Created: 2025-05-15T21:23:50Z | Merged: 2025-05-17T13:02:05Z

The Quoridor game logic that was taking up the most computation time is re-implemented using numba, which just-in-time compiles it.

Good stuff:
- On my computer, the average running time of the Simple agent with default parameters goes from 24 ms to about 5 ms
- More speedup is possible. I think another 10x or so if the minimax implementation in the simple agent itself is numba-fied, and parallel processing is enabled

Meh stuff:
- To get much benefit from numba, we now do more batch processing in free functions. This blurs some of the abstraction provided by the Quoridor class (for example the numba grid functions in `qgrid.py` do action encoding/decoding themselves instead of using the ActionEncoder class)
- Compilation by numba takes a couple seconds each time `play.py` is run. This doesn't seem to get cached on disk and reused the next time the python interpreter is restarted.
- Numba works well for what it does, but has a lot of magic, doesn't feel particularly actively used or developed, and has limitations that would make we not want to use it in a more "serious" project.

A couple more notes:
- I've added tie-breaking of the minimax reward function using the number of walls remaining (suggested by @alejandromarcu when reviewing that last PR). This gives a nice bump to the results of the simple agent without slowing it down.
- We'd probably be better off removing the `Board` class and just calling `qgrid.py` functions directly from the Quoridor class. Three levels of abstraction feels like too much, but I didn't want to change too much in this PR so I left it for now.
- For the moment I'm using numba-fied BFS for both shortest path and for "does this wall block a player's path to goal". For blocking checks, we can use DFS instead, but it wasn't a bottleneck so I didn't do that yet.


Diff:
diff --git a/deep_quoridor/requirements.txt b/deep_quoridor/requirements.txt
index ea7acce8..a29120ff 100644
--- a/deep_quoridor/requirements.txt
+++ b/deep_quoridor/requirements.txt
@@ -1,4 +1,5 @@
 gymnasium
+numba
 numpy
 pettingzoo
 prettytable
diff --git a/deep_quoridor/src/agents/simple.py b/deep_quoridor/src/agents/simple.py
index fc722041..f4011e95 100644
--- a/deep_quoridor/src/agents/simple.py
+++ b/deep_quoridor/src/agents/simple.py
@@ -33,20 +33,16 @@ class SimpleParams(SubargsBase):
     discount_factor: float = 0.99
 
 
-def adjust_wall_position(position: tuple[int, int]) -> tuple[float, float]:
-    """
-    Compute a wall position that reflects where its center is relative to player positions.
-    """
-    return position[0] + 0.5, position[1] + 0.5
-
-
-def compute_wall_weight(wall: WallAction, position_1: tuple[int, int], position_2: tuple[int, int], sigma) -> float:
-    wall_position = np.array(adjust_wall_position(wall.position))
-    d1 = np.linalg.norm(np.array(position_1) - wall_position)
-    d2 = np.linalg.norm(np.array(position_2) - wall_position)
-    min_distance = min(d1, d2)
-    weight = np.exp(-0.5 * (min_distance / sigma) ** 2)
-    return weight
+def compute_wall_weights(
+    wall_actions: list[WallAction], position_1: tuple[int, int], position_2: tuple[int, int], sigma
+) -> float:
+    walls_array = np.array([wall.position for wall in wall_actions], dtype=np.float32)
+    walls_array += (0.5, 0.5)  # Adjust wall positions to reflect where their center is in player coordinates.
+    squared_distances_to_position_1 = ((walls_array - position_1) ** 2).sum(axis=1)
+    squared_distances_to_position_2 = ((walls_array - position_2) ** 2).sum(axis=1)
+    min_squared_distances = np.minimum(squared_distances_to_position_1, squared_distances_to_position_2)
+    weights = np.exp(-0.5 * min_squared_distances / sigma**2)
+    return weights
 
 
 def sample_actions(game: Quoridor, n: int, wall_sigma: float | None = None) -> list[Action]:
@@ -65,11 +61,13 @@ def sample_actions(game: Quoridor, n: int, wall_sigma: float | None = None) -> l
         wall_actions = game.get_valid_wall_actions()
         if len(wall_actions) > num_wall_actions:
             if wall_sigma is None:
+                # Use a uniform distribution to sample wall actions.
                 wall_actions = random.sample(wall_actions, num_wall_actions)
             else:
+                # Weight walls based on their distance to the players.
                 position_1 = game.board.get_player_position(Player.ONE)
                 position_2 = game.board.get_player_position(Player.TWO)
-                wall_weights = [compute_wall_weight(wall, position_1, position_2, wall_sigma) for wall in wall_actions]
+                wall_weights = compute_wall_weights(wall_actions, position_1, position_2, wall_sigma)
                 wall_actions = np.random.choice(
                     wall_actions,
                     size=num_wall_actions,
@@ -101,32 +99,24 @@ def choose_action(
     elif game.check_win(opponent):
         return None, -WINNING_REWARD
     elif max_depth == 0:
-        return None, game.player_distance_to_target(opponent) - game.player_distance_to_target(player)
+        distance_reward = game.player_distance_to_target(opponent) - game.player_distance_to_target(player)
+        wall_reward = game.board.get_walls_remaining(player) - game.board.get_walls_remaining(opponent)
+        return None, distance_reward + wall_reward / 100
 
     actions = sample_actions(game, branching_factor, wall_sigma)
     assert len(actions) > 0, "There are no valid actions for the current player."
 
-    chosen_action = None
-    chosen_value = -float("inf") if game.get_current_player() == player else float("inf")
+    values = []
+    current_player_before_action = game.get_current_player()
+    position_before_action = game.board.get_player_position(current_player_before_action)
     for action in actions:
-        current_player_before_action = game.get_current_player()
-        position_before_action = game.board.get_player_position(current_player_before_action)
-
         # Apply the action to the game.
         game.step(action)
 
         # Recursively call this function to choose the next player's action.
         _, value = choose_action(game, player, opponent, max_depth - 1, branching_factor, wall_sigma, discount_factor)
         value = discount_factor * value
-
-        if current_player_before_action == player:
-            if value >= chosen_value:
-                chosen_value = value
-                chosen_action = action
-        else:
-            if value <= chosen_value:
-                chosen_value = value
-                chosen_action = action
+        values.append(value)
 
         # Undo the action to put the game back to its original state.
         if isinstance(action, WallAction):
@@ -135,7 +125,18 @@ def choose_action(
             game.board.move_player(current_player_before_action, position_before_action)
         game.set_current_player(current_player_before_action)
 
-    return chosen_action, chosen_value
+    # "player" tries to maximize the reward, while "opponent" tries to minimize it.
+    values = np.array(values)
+    if current_player_before_action == player:
+        best_value = values.max()
+    else:
+        best_value = values.min()
+
+    best_action_indices = np.flatnonzero(values == best_value)
+    # There may be multiple actions with the same value, so we randomly choose one of them.
+    chosen_action_index = np.random.choice(best_action_indices)
+
+    return actions[chosen_action_index], values[chosen_action_index]
 
 
 class SimpleAgent(Agent):
diff --git a/deep_quoridor/src/qgrid.py b/deep_quoridor/src/qgrid.py
new file mode 100644
index 00000000..cbae7d64
--- /dev/null
+++ b/deep_quoridor/src/qgrid.py
@@ -0,0 +1,458 @@
+"""
+Fast implementation of quoridor functions using Numba.
+
+The core data structure in this implemtation is a 2D grid that represents both the quoridor board and walls. Here is
+an ASCII representation of a 5x5 grid in an example game state:
+
+#####################
+#####################
+##.   .   .   .   .##
+##                 ##
+##.   .   .   . # .##
+##        # # # #  ##
+##.   . # .   . # 2##
+##      #          ##
+##.   . # .   .   .##
+##        # # #    ##
+##.   .   1   .   .##
+##    # # #   # # ###
+##.   .   .   .   .##
+#####################
+#####################
+
+The hashes represent walls, the dots represent positions that players can be in, and the numbers represent the players.
+Note that we add a border of two wall cells around the grid. This reduces the number of special cases we need to handle
+when checking for valid moves and wall placements.  The cells in the grid (other than the boarder) alternate between
+player cells and wall cells. Because of the border and the alternating wall and player cells, for an NxN Quoridor
+board, the corresponding grid will be (2N+2)x(2N+2) in size.
+
+Terminology:
+- row or col - the row or column on the quoridor board (each between 0 and 8 on a normal board)
+- i or j - the indices into the combined grid used to represent the board and walls (more info below)
+- cell - an element in the combined grid.
+
+In addition to the grid, there are a few arrays used to represent the game state:
+- player_positions (np.ndarray): 2D array of containing [[p1_row, p1_col], [p2_row, p2_col]]
+- walls_remaining (np.ndarray): 1D array of containing [p1_walls_remaining, p2_walls_remaining]
+- goal_rows (np.ndarray): 1D array of containing [p1_goal_row, p2_goal_row]
+- current_player (scalar): 0 or 1
+"""
+
+import numpy as np
+from numba import njit
+
+WALL_ORIENTATION_VERTICAL = 0
+WALL_ORIENTATION_HORIZONTAL = 1
+
+# Possible values for each cell in the grid.
+CELL_FREE = -1
+# The player numbers start at 0 so that they can also be used as indices into the player positions array.
+CELL_PLAYER1 = 0
+CELL_PLAYER2 = 1
+CELL_WALL = 10
+
+
+@njit
+def distance_to_row(grid: np.ndarray, start_row: int, start_col: int, target_row: int) -> int:
+    """
+    Args:
+        row (int): The current row of the pawn
+        col (int): The current column of the pawn
+        target_row (int): The target row to reach
+        visited (numpy.array): A 2D boolean array with the same shape as the board,
+            indicating which positions have been visited
+
+    Returns:
+        int: Number of steps to reach the target or -1 if it's unreachable
+    """
+    grid_height = grid.shape[0]
+    grid_width = grid.shape[1]
+
+    start_i = start_row * 2 + 2
+    start_j = start_col * 2 + 2
+    target_i = target_row * 2 + 2
+
+    if target_i == start_i:
+        return 0
+
+    visited = np.zeros(grid.shape, dtype="bool")
+    visited[start_i, start_j] = True
+
+    queue = [(start_i, start_j, 0)]
+
+    while queue:
+        i, j, steps = queue.pop(0)
+        # Iterate in the 4 directions, and if we can move to that position and we haven't already visited it, add it to the queue.
+        # This was done in a for loop before, but making everything explicit makes it significantly faster, and this method is called
+        # very often.
+
+        # Down
+        new_i = i + 2
+        wall_i = i + 1
+        if new_i < grid_height and not visited[new_i, j] and grid[wall_i, j] != CELL_WALL:
+            visited[new_i, j] = True
+            if target_i == new_i:
+                return steps + 1
+            queue.append((new_i, j, steps + 1))
+
+        # Up
+        new_i = i - 2
+        wall_i = i - 1
+        if new_i >= 0 and not visited[new_i, j] and grid[wall_i, j] != CELL_WALL:
+            visited[new_i, j] = True
+            if target_i == new_i:
+                return steps + 1
+            queue.append((new_i, j, steps + 1))
+
+        # Right
+        new_j = j + 2
+        wall_j = j + 1
+        if new_j < grid_width and not visited[i, new_j] and grid[i, wall_j] != CELL_WALL:
+            visited[i, new_j] = True
+            if target_i == i:
+                return steps + 1
+            queue.append((i, new_j, steps + 1))
+
+        # Left
+        new_j = j - 2
+        wall_j = j - 1
+        if new_j >= 0 and not visited[i, new_j] and grid[i, wall_j] != CELL_WALL:
+            visited[i, new_j] = True
+            if target_i == i:
+                return steps + 1
+            queue.append((i, new_j, steps + 1))
+
+    return -1
+
+
+@njit
+def are_wall_cells_free(grid: np.ndarray, wall_row: int, wall_col: int, wall_orientation: int) -> bool:
+    grid_height = grid.shape[0]
+    grid_width = grid.shape[1]
+
+    if wall_orientation == WALL_ORIENTATION_VERTICAL:
+        start_i = wall_row * 2 + 2
+        start_j = wall_col * 2 + 3
+        assert start_i >= 0 and start_i + 2 < grid_height and start_j >= 0 and start_j < grid_width
+
+        return (
+            grid[start_i, start_j] == CELL_FREE
+            and grid[start_i + 1, start_j] == CELL_FREE
+            and grid[start_i + 2, start_j] == CELL_FREE
+        )
+    else:
+        start_i = wall_row * 2 + 3
+        start_j = wall_col * 2 + 2
+        assert start_i >= 0 and start_i < grid_height and start_j >= 0 and start_j + 2 < grid_width
+
+        return (
+            grid[start_i, start_j] == CELL_FREE
+            and grid[start_i, start_j + 1] == CELL_FREE
+            and grid[start_i, start_j + 2] == CELL_FREE
+        )
+
+
+@njit
+def is_wall_potential_block(grid, wall_row, wall_col, wall_orientation):
+    grid_height = grid.shape[0]
+    grid_width = grid.shape[1]
+
+    if wall_orientation == WALL_ORIENTATION_VERTICAL:
+        start_i = wall_row * 2 + 2
+        start_j = wall_col * 2 + 3
+        assert start_i >= 0 and start_i + 2 < grid_height and start_j >= 0 and start_j < grid_width
+
+        touches = 0
+        if (
+            grid[start_i - 1, start_j - 1] == CELL_WALL
+            or grid[start_i - 2, start_j] == CELL_WALL
+            or grid[start_i - 1, start_j + 1] == CELL_WALL
+        ):
+            touches += 1
+        if grid[start_i + 1, start_j - 1] == CELL_WALL or grid[start_i + 1, start_j + 1] == CELL_WALL:
+            touches += 1
+        if (
+            grid[start_i + 3, start_j - 1] == CELL_WALL
+            or grid[start_i + 4, start_j] == CELL_WALL
+            or grid[start_i + 3, start_j + 1] == CELL_WALL
+        ):
+            touches += 1
+        return touches >= 2
+
+    else:  # HORIZONTAL
+        start_i = wall_row * 2 + 3
+        start_j = wall_col * 2 + 2
+        assert start_i >= 0 and start_i < grid_height and start_j >= 0 and start_j + 2 < grid_width
+
+        touches = 0
+        if (
+            grid[start_i - 1, start_j - 1] == CELL_WALL
+            or grid[start_i, start_j - 2] == CELL_WALL
+            or grid[start_i + 1, start_j - 1] == CELL_WALL
+        ):
+            touches += 1
+        if grid[start_i - 1, start_j + 1] == CELL_WALL or grid[start_i + 1, start_j + 1] == CELL_WALL:
+            touches += 1
+        if (
+            grid[start_i - 1, start_j + 3] == CELL_WALL
+            or grid[start_i, start_j + 4] == CELL_WALL
+            or grid[start_i + 1, start_j + 3] == CELL_WALL
+        ):
+            touches += 1
+        return touches >= 2
+
+
+@njit
+def set_wall_cells(grid, wall_row, wall_col, wall_orientation, cell_value):
+    grid_height = grid.shape[0]
+    grid_width = grid.shape[1]
+
+    if wall_orientation == WALL_ORIENTATION_VERTICAL:
+        start_i = wall_row * 2 + 2
+        start_j = wall_col * 2 + 3
+        assert start_i >= 0 and start_i + 2 < grid_height and start_j >= 0 and start_j < grid_width
+
+        grid[start_i, start_j] = cell_value
+        grid[start_i + 1, start_j] = cell_value
+        grid[start_i + 2, start_j] = cell_value
+    else:
+        start_i = wall_row * 2 + 3
+        start_j = wall_col * 2 + 2
+        assert start_i >= 0 and start_i < grid_height and start_j >= 0 and start_j + 2 < grid_width
+
+        grid[start_i, start_j] = cell_value
+        grid[start_i, start_j + 1] = cell_value
+        grid[start_i, start_j + 2] = cell_value
+
+
+@njit
+def check_wall_cells(grid, wall_row, wall_col, wall_orientation, cell_value):
+    """
+    Return True iff all the grid cells for the wall equal the given value.
+    """
+    grid_height = grid.shape[0]
+    grid_width = grid.shape[1]
+
+    if wall_orientation == WALL_ORIENTATION_VERTICAL:
+        start_i = wall_row * 2 + 2
+        start_j = wall_col * 2 + 3
+        assert start_i >= 0 and start_i + 2 < grid_height and start_j >= 0 and start_j < grid_width
+        return (
+            grid[start_i, start_j] == cell_value
+            and grid[start_i + 1, start_j] == cell_value
+            and grid[start_i + 2, start_j] == cell_value
+        )
+    else:
+        start_i = wall_row * 2 + 3
+        start_j = wall_col * 2 + 2
+        assert start_i >= 0 and start_i < grid_height and start_j >= 0 and start_j + 2 < grid_width
+
+        return (
+            grid[start_i, start_j] == cell_value
+            and grid[start_i, start_j + 1] == cell_value
+            and grid[start_i, start_j + 2] == cell_value
+        )
+
+
+@njit
+def is_move_action_valid(grid, player_positions, current_player, destination_row, destination_col):
+    grid_height = grid.shape[0]
+    grid_width = grid.shape[1]
+
+    player_i = player_positions[current_player][0] * 2 + 2
+    player_j = player_positions[current_player][1] * 2 + 2
+    destination_i = destination_row * 2 + 2
+    destination_j = destination_col * 2 + 2
+    opponent = 1 - current_player
+
+    # This suffices for bounds cheking, since the grid cells we check for a move are always between the player its destination.
+    assert player_i >= 0 and player_i < grid_height and player_j >= 0 and player_j < grid_width
+    if destination_i < 0 or destination_i >= grid_height or destination_j < 0 or destination_j >= grid_width:
+        return False
+
+    if grid[destination_i, destination_j] != CELL_FREE:
+        return False
+
+    delta_i = destination_i - player_i
+    delta_j = destination_j - player_j
+    if delta_i == 2 and delta_j == 0:
+        return grid[player_i + 1, player_j] == CELL_FREE
+    elif delta_i == -2 and delta_j == 0:
+        return grid[player_i - 1, player_j] == CELL_FREE
+    elif delta_i == 0 and delta_j == 2:
+        return grid[player_i, player_j + 1] == CELL_FREE
+    elif delta_i == 0 and delta_j == -2:
+        return grid[player_i, player_j - 1] == CELL_FREE
+    elif delta_i == 4 and delta_j == 0:
+        return (
+            grid[player_i + 1, player_j] == CELL_FREE
+            and grid[player_i + 2, player_j] == opponent
+            and grid[player_i + 3, player_j] == CELL_FREE
+        )
+    elif delta_i == -4 and delta_j == 0:
+        return (
+            grid[player_i - 1, player_j] == CELL_FREE
+            and grid[player_i - 2, player_j] == opponent
+            and grid[player_i - 3, player_j] == CELL_FREE
+        )
+    elif delta_i == 0 and delta_j == 4:
+        pass
+        return (
+            grid[player_i, player_j + 1] == CELL_FREE
+            and grid[player_i, player_j + 2] == opponent
+            and grid[player_i, player_j + 3] == CELL_FREE
+        )
+    elif delta_i == 0 and delta_j == -4:
+        return (
+            grid[player_i, player_j - 1] == CELL_FREE
+            and grid[player_i, player_j - 2] == opponent
+            and grid[player_i, player_j - 3] == CELL_FREE
+        )
+    elif delta_i == 2 and delta_j == 2:
+        return (
+            grid[player_i + 1, player_j] == CELL_FREE
+            and grid[player_i + 2, player_j] == opponent
+            and grid[player_i + 2, player_j + 1] == CELL_FREE
+            and grid[player_i + 3, player_j] == CELL_WALL  # Wall behind opponent
+        ) or (
+            grid[player_i, player_j + 1] == CELL_FREE
+            and grid[player_i, player_j + 2] == opponent
+            and grid[player_i + 1, player_j + 2] == CELL_FREE
+            and grid[player_i, player_j + 3] == CELL_WALL  # Wall behind opponent
+        )
+    elif delta_i == -2 and delta_j == 2:
+        return (
+            grid[player_i - 1, player_j] == CELL_FREE
+            and grid[player_i - 2, player_j] == opponent
+            and grid[player_i - 2, player_j + 1] == CELL_FREE
+            and grid[player_i - 3, player_j] == CELL_WALL  # Wall behind opponent
+        ) or (
+            grid[player_i, player_j + 1] == CELL_FREE
+            and grid[player_i, player_j + 2] == opponent
+            and grid[player_i - 1, player_j + 2] == CELL_FREE
+            and grid[player_i, player_j + 3] == CELL_WALL  # Wall behind opponent
+        )
+    elif delta_i == 2 and delta_j == -2:
+        return (
+            grid[player_i + 1, player_j] == CELL_FREE
+            and grid[player_i + 2, player_j] == opponent
+            and grid[player_i + 2, player_j - 1] == CELL_FREE
+            and grid[player_i + 3, player_j] == CELL_WALL  # Wall behind opponent
+        ) or (
+            grid[player_i, player_j - 1] == CELL_FREE
+            and grid[player_i, player_j - 2] == opponent
+            and grid[player_i + 1, player_j - 2] == CELL_FREE
+            and grid[player_i, player_j - 3] == CELL_WALL  # Wall behind opponent
+        )
+    elif delta_i == -2 and delta_j == -2:
+        return (
+            grid[player_i - 1, player_j] == CELL_FREE
+            and grid[player_i - 2, player_j] == opponent
+            and grid[player_i - 2, player_j - 1] == CELL_FREE
+            and grid[player_i - 3, player_j] == CELL_WALL  # Wall behind opponent
+        ) or (
+            grid[player_i, player_j - 1] == CELL_FREE
+            and grid[player_i, player_j - 2] == opponent
+            and grid[player_i - 1, player_j - 2] == CELL_FREE
+            and grid[player_i, player_j - 3] == CELL_WALL  # Wall behind opponent
+        )
+    else:
+        return False
+
+
+@njit
+def is_wall_action_valid(
+    grid, player_positions, walls_remaining, goal_rows, current_player, wall_row, wall_col, wall_orientation
+):
+    is_valid = True
+    if walls_remaining[current_player] <= 0:
+        is_valid = False
+    elif not check_wall_cells(grid, wall_row, wall_col, wall_orientation, CELL_FREE):
+        is_valid = False
+    elif is_wall_potential_block(grid, wall_row, wall_col, wall_orientation):
+        # Temprorarily add the wall.
+        set_wall_cells(grid, wall_row, wall_col, wall_orientation, CELL_WALL)
+
+        # Make sure all players still have some path to their goal row.
+        for i in range(len(player_positions)):
+            if distance_to_row(grid, player_positions[i][0], player_positions[i][1], goal_rows[i]) == -1:
+                is_valid = False
+                break
+
+        # Restore the grid to its previous state.
+        set_wall_cells(grid, wall_row, wall_col, wall_orientation, CELL_FREE)
+
+    return is_valid
+
+
+@njit
+def compute_move_action_mask(
+    grid: np.ndarray,
+    player_positions: np.ndarray,
+    current_player: int,
+    action_mask: np.ndarray,
+) -> np.ndarray:
+    grid_width = grid.shape[1]
+    board_size = (grid_width - 4) // 2 + 1
+
+    assert action_mask.shape == (board_size**2,)
+
+    for k in range(len(action_mask)):
+        action_mask[k] = 0
+
+    # Check all possible moves
+    for delta_row in np.arange(-2, 3):
+        for delta_col in np.arange(-2, 3):
+            destination_row = player_positions[current_player][0] + delta_row
+            destination_col = player_positions[current_player][1] + delta_col
+            if is_move_action_valid(grid, player_positions, current_player, destination_row, destination_col):
+                action_mask[destination_row * board_size + destination_col] = 1
+
+    return action_mask
+
+
+@njit
+def compute_wall_action_mask(
+    grid: np.ndarray,
+    player_positions: np.ndarray,
+    walls_remaining: np.ndarray,
+    goal_rows: np.ndarray,
+    current_player: int,
+    action_mask: np.ndarray,
+) -> np.ndarray:
+    grid_width = grid.shape[1]
+    board_size = (grid_width - 4) // 2 + 1
+    wall_size = board_size - 1
+
+    assert action_mask.shape == (2 * wall_size**2,)
+
+    for k in range(len(action_mask)):
+        action_mask[k] = 0
+
+    for wall_row in range(wall_size):
+        for wall_col in range(wall_size):
+            if is_wall_action_valid(
+                grid,
+                player_positions,
+                walls_remaining,
+                goal_rows,
+                current_player,
+                wall_row,
+                wall_col,
+                WALL_ORIENTATION_VERTICAL,
+            ):
+                action_mask[wall_row * wall_size + wall_col] = 1
+
+            if is_wall_action_valid(
+                grid,
+                player_positions,
+                walls_remaining,
+                goal_rows,
+                current_player,
+                wall_row,
+                wall_col,
+                WALL_ORIENTATION_HORIZONTAL,
+            ):
+                action_mask[wall_size**2 + wall_row * wall_size + wall_col] = 1
+
+    return action_mask
diff --git a/deep_quoridor/src/quoridor.py b/deep_quoridor/src/quoridor.py
index acb954a5..953cf8f0 100644
--- a/deep_quoridor/src/quoridor.py
+++ b/deep_quoridor/src/quoridor.py
@@ -1,9 +1,11 @@
 import copy
 from dataclasses import dataclass
 from enum import IntEnum, unique
+from functools import cache
 from typing import Optional, Sequence
 
 import numpy as np
+import qgrid
 
 
 @unique
@@ -14,8 +16,8 @@ class Player(IntEnum):
 
 @unique
 class WallOrientation(IntEnum):
-    VERTICAL = 0
-    HORIZONTAL = 1
+    VERTICAL = qgrid.WALL_ORIENTATION_VERTICAL
+    HORIZONTAL = qgrid.WALL_ORIENTATION_HORIZONTAL
 
 
 class Action:
@@ -35,34 +37,35 @@ class WallAction(Action):
 
 class ActionEncoder:
     def __init__(self, board_size: int):
-        self._board_size = board_size
-        self._wall_size = board_size - 1
+        self.board_size = board_size
+        self.wall_size = board_size - 1
 
     def action_to_index(self, action) -> int:
         """
         Converts an action object to an action index
         """
         if isinstance(action, MoveAction):
-            return action.destination[0] * self._board_size + action.destination[1]
+            return action.destination[0] * self.board_size + action.destination[1]
         elif isinstance(action, WallAction) and action.orientation == WallOrientation.VERTICAL:
-            return self._board_size**2 + action.position[0] * self._wall_size + action.position[1]
+            return self.board_size**2 + action.position[0] * self.wall_size + action.position[1]
         elif isinstance(action, WallAction) and action.orientation == WallOrientation.HORIZONTAL:
-            return self._board_size**2 + self._wall_size**2 + action.position[0] * self._wall_size + action.position[1]
+            return self.board_size**2 + self.wall_size**2 + action.position[0] * self.wall_size + action.position[1]
         else:
             raise ValueError(f"Invalid action type: {action}")
 
+    @cache
     def index_to_action(self, idx) -> Action:
         """
         Converts an action index to an action object.
         """
         action = None
-        if idx < self._board_size**2:  # Pawn movement
-            action = MoveAction(divmod(idx, self._board_size))
-        elif idx < self._board_size**2 + self._wall_size**2:
-            action = WallAction(divmod(idx - self._board_size**2, self._wall_size), WallOrientation.VERTICAL)
-        elif idx < self._board_size**2 + (self._wall_size**2) * 2:
+        if idx < self.board_size**2:  # Pawn movement
+            action = MoveAction(divmod(idx, self.board_size))
+        elif idx < self.board_size**2 + self.wall_size**2:
+            action = WallAction(divmod(idx - self.board_size**2, self.wall_size), WallOrientation.VERTICAL)
+        elif idx < self.board_size**2 + (self.wall_size**2) * 2:
             action = WallAction(
-                divmod(idx - self._board_size**2 - self._wall_size**2, self._wall_size),
+                divmod(idx - self.board_size**2 - self.wall_size**2, self.wall_size),
                 WallOrientation.HORIZONTAL,
             )
         else:
@@ -73,31 +76,11 @@ def index_to_action(self, idx) -> Action:
 
 class Board:
     # Possible values for each cell in the grid.
-    FREE = -1
+    FREE = qgrid.CELL_FREE
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
-    PLAYER1 = 0
-    PLAYER2 = 1
-    WALL = 10
-
-    # When we check whether a new wall could potentially block all routes to the goal for a player,
-    # we first check to see whether it spans between two walls already on the board, or between
-    # one wall on the board and the border of the board. There are three possible places a new wall
-    # could touch an existing wall (or the border). At the start of the new wall, in the middle of the
-    # new wall, or at the end of the new wall. To simplify these checks, we create this list of wall
-    # grid cells that touch each of those three places on a new wall. THe positions here are relative
-    # to the starting position of the new wall.
-    _potential_wall_neighbors = {
-        WallOrientation.VERTICAL: [
-            np.array([(-1, -1), (-2, 0), (-1, 1)]),
-            np.array([(1, -1), (1, 1)]),
-            np.array([(3, -1), (4, 0), (3, 1)]),
-        ],
-        WallOrientation.HORIZONTAL: [
-            np.array([(-1, -1), (0, -2), (1, -1)]),
-            np.array([(-1, 1), (1, 1)]),
-            np.array([(-1, 3), (0, 4), (1, 3)]),
-        ],
-    }
+    PLAYER1 = qgrid.CELL_PLAYER1
+    PLAYER2 = qgrid.CELL_PLAYER2
+    WALL = qgrid.CELL_WALL
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
@@ -116,8 +99,8 @@ def __init__(self, board_size: int = 9, max_walls: int = 10):
         self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
 
         self._players = [Player.ONE, Player.TWO]
-        self._player_positions = [(0, self.board_size // 2), (self.board_size - 1, self.board_size // 2)]
-        self._walls_remaining = [self.max_walls, self.max_walls]
+        self._player_positions = np.array([(0, self.board_size // 2), (self.board_size - 1, self.board_size // 2)])
+        self._walls_remaining = np.array([self.max_walls, self.max_walls])
         for player, position in zip(self._players, self._player_positions):
             self.set_player_cell(position, player)
 
@@ -125,7 +108,7 @@ def get_player_position(self, player: Player) -> tuple[int, int]:
         """
         Get the position of the player's pawn.
         """
-        return self._player_positions[player]
+        return tuple(self._player_positions[player])
 
     def move_player(self, player: Player, new_position: tuple[int, int]):
         """
@@ -190,7 +173,7 @@ def can_place_wall(self, player: Player, position: tuple[int, int], orientation:
         if self._walls_remaining[player] < 1:
             return False
 
-        return (self._grid[self._get_wall_slice(position, orientation)] == Board.FREE).all()
+        return qgrid.are_wall_cells_free(self._grid, position[0], position[1], int(orientation))
 
     def _are_adjacent(self, position1: tuple[int, int], position2: tuple[int, int]) -> bool:
         dr = position1[0] - position2[0]
@@ -241,114 +224,6 @@ def _get_wall_slice(self, position: tuple[int, int], orientation: WallOrientatio
 
         return wall_slice
 
-    def _is_wall_potential_block(self, position, orientation):
-        wall_start_index = self._wall_position_to_grid_index(position, orientation)
-
-        touches = 0
-        for neighbor_offsets in Board._potential_wall_neighbors[orientation]:
-            neighbors = wall_start_index + neighbor_offsets
-            if (self._grid[neighbors[:, 0], neighbors[:, 1]] == Board.WALL).any():
-                touches += 1
-
-        return touches > 1
-
-    def _bfs(self, start_position, target_row, visited):
-        """
-        Performs a breadth-first search to find the shortest path to the target row.
-
-
-        Args:
-            row (int): The current row of the pawn
-            col (int): The current column of the pawn
-            target_row (int): The target row to reach
-            visited (numpy.array): A 2D boolean array with the same shape as the board,
-                indicating which positions have been visited
-
-        Returns:
-            int: Number of steps to reach the target or -1 if it's unreachable
-        """
-        row, col = start_position
-        if target_row == row:
-            return 0
-
-        queue = [(row, col, 0)]
-        visited[row, col] = True
-        N = self.board_size
-
-        while queue:
-            curr_row, curr_col, steps = queue.pop(0)
-            # Iterate in the 4 directions, and if we can move to that position and we haven't already visited it, add it to the queue.
-            # This was done in a for loop before, but making everything explicit makes it significantly faster, and this method is called
-            # very often.
-
-            # Down
-            wall_row = curr_row * 2 + 2
-            wall_col = curr_col * 2 + 2
-            new_row = curr_row + 1
-            if new_row < N and not visited[new_row, curr_col] and self._grid[wall_row + 1, wall_col] != Board.WALL:
-                visited[new_row, curr_col] = True
-                if target_row == new_row:
-                    return steps + 1
-                queue.append((new_row, curr_col, steps + 1))
-
-            # Up
-            new_row = curr_row - 1
-            if new_row >= 0 and not visited[new_row, curr_col] and self._grid[wall_row - 1, wall_col] != Board.WALL:
-                visited[new_row, curr_col] = True
-                if target_row == new_row:
-                    return steps + 1
-                queue.append((new_row, curr_col, steps + 1))
-
-            # Right
-            new_col = curr_col + 1
-            if new_col < N and not visited[curr_row, new_col] and self._grid[wall_row, wall_col + 1] != Board.WALL:
-                visited[curr_row, new_col] = True
-                queue.append((curr_row, new_col, steps + 1))
-
-            # Left
-            new_col = curr_col - 1
-            if new_col >= 0 and not visited[curr_row, new_col] and self._grid[wall_row, wall_col - 1] != Board.WALL:
-                visited[curr_row, new_col] = True
-                queue.append((curr_row, new_col, steps + 1))
-
-        return -1
-
-    def _dfs(self, start_position, target_row, visited):
-        """
-        Performs a depth-first search to find whether the pawn can reach the target row.
-
-        Args:
-            row (int): The current row of the pawn
-            col (int): The current column of the pawn
-            target_row (int): The target row to reach
-            visited (numpy.array): A 2D boolean array with the same shape as the board,
-                indicating which positions have been visited
-
-        Returns:
-            int: Number of steps to reach the target or -1 if it's unreachable
-        """
-        row, col = start_position
-        if row == target_row:
-            return 0
-
-        visited[row, col] = True
-
-        # Find out the forward direction to try it first and maybe get to the target faster
-        fwd = 1 if target_row > row else -1
-
-        moves = [(row + fwd, col), (row, col - 1), (row, col + 1), (row - fwd, col)]
-        for new_row, new_col in moves:
-            if (
-                self.is_position_on_board((new_row, new_col))
-                and not visited[new_row, new_col]
-                and not self.is_wall_between((row, col), (new_row, new_col))
-            ):
-                dfs = self._dfs((new_row, new_col), target_row, visited)
-                if dfs != -1:
-                    return dfs + 1
-
-        return -1
-
     def __str__(self):
         """
         Return a pretty-printed string representation of the grid.
@@ -372,9 +247,9 @@ def __init__(self, board: Board, current_player: Player = Player.ONE):
         """
         self.board = board
         self.current_player = current_player
+        self.action_encoder = ActionEncoder(board.board_size)
 
-        self._goal_rows = {Player.ONE: self.board.board_size - 1, Player.TWO: 0}
-        self._jump_checks = create_jump_checks()
+        self._goal_rows = np.array([self.board.board_size - 1, 0])
 
     def step(self, action: Action, validate: bool = True):
         """
@@ -401,59 +276,29 @@ def is_action_valid(self, action: Action, for_player: Optional[Player] = None) -
         Check whether the given action is valid given the current game state.
         """
         if for_player is None:
-            player = self.get_current_player()
-        else:
-            player = for_player
-        opponent = Player(1 - player)
+            for_player = self.get_current_player()
 
         is_valid = True
         if isinstance(action, MoveAction):
-            current_position = self.board.get_player_position(player)
-            opponent_position = self.board.get_player_position(opponent)
-            opponent_offset = (opponent_position[0] - current_position[0], opponent_position[1] - current_position[1])
-            position_delta = (action.destination[0] - current_position[0], action.destination[1] - current_position[1])
-
-            # Destination cell must be free
-            if self.board.get_player_cell(action.destination) != Board.FREE:
-                is_valid = False
-
-            elif np.sum(np.abs(position_delta)) == 1:
-                # Moving to an adjacent cell, just make sure no wall is in the way.
-                if self.board.is_wall_between(current_position, action.destination):
-                    is_valid = False
-
-            elif (opponent_offset, position_delta) in self._jump_checks:
-                jump_checks = self._jump_checks[(opponent_offset, position_delta)]
-
-                for check_delta_1, check_delta_2 in jump_checks["wall"]:
-                    if not self.board.is_wall_between(
-                        (current_position[0] + check_delta_1[0], current_position[1] + check_delta_1[1]),
-                        (current_position[0] + check_delta_2[0], current_position[1] + check_delta_2[1]),
-                    ):
-                        is_valid = False
-                for check_delta_1, check_delta_2 in jump_checks["nowall"]:
-                    if self.board.is_wall_between(
-                        (current_position[0] + check_delta_1[0], current_position[1] + check_delta_1[1]),
-                        (current_position[0] + check_delta_2[0], current_position[1] + check_delta_2[1]),
-                    ):
-                        is_valid = False
-            else:
-                # This isn't a move by 1, and it isn't a jump, so it's invalid.
-                is_valid = False
+            return qgrid.is_move_action_valid(
+                self.board._grid,
+                self.board._player_positions,
+                int(for_player),
+                action.destination[0],
+                action.destination[1],
+            )
 
         elif isinstance(action, WallAction):
-            if self.board.can_place_wall(player, action.position, action.orientation):
-                if self.board._is_wall_potential_block(action.position, action.orientation):
-                    # Temporarily place the wall so that we can check player paths to their goals.
-                    self.board.add_wall(self.current_player, action.position, action.orientation, check_if_valid=False)
-                    for p in [Player.ONE, Player.TWO]:
-                        if not self.can_reach(self.board.get_player_position(p), self.get_goal_row(p)):
-                            is_valid = False
-                            break
-                    # Restore the board to its previous state
-                    self.board.remove_wall(self.current_player, action.position, action.orientation)
-            else:
-                is_valid = False
+            is_valid = qgrid.is_wall_action_valid(
+                self.board._grid,
+                self.board._player_positions,
+                self.board._walls_remaining,
+                self._goal_rows,
+                int(for_player),
+                action.position[0],
+                action.position[1],
+                int(action.orientation),
+            )
         else:
             raise ValueError("Invalid action type")
 
@@ -463,30 +308,34 @@ def get_valid_move_actions(self, player: Optional[Player] = None) -> list[MoveAc
         if player is None:
             player = self.get_current_player()
 
-        position = self.board.get_player_position(player)
-
-        valid_move_actions = list()
-        for delta_row in range(-2, 3):
-            for delta_col in range(-2, 3):
-                destination = (position[0] + delta_row, position[1] + delta_col)
-                if self.board.is_position_on_board(destination):
-                    move_action = MoveAction(destination)
-                    if self.is_action_valid(move_action, player):
-                        valid_move_actions.append(move_action)
-        return valid_move_actions
+        action_mask = np.zeros(self.action_encoder.board_size**2, dtype=bool)
+        qgrid.compute_move_action_mask(
+            self.board._grid,
+            self.board._player_positions,
+            int(player),
+            action_mask,
+        )
+        actions = [self.action_encoder.index_to_action(action_i) for action_i in np.flatnonzero(action_mask)]
+        return actions
 
     def get_valid_wall_actions(self, player: Optional[Player] = None) -> list[WallAction]:
         if player is None:
             player = self.get_current_player()
 
-        valid_wall_actions = list()
-        for row in range(self.board.board_size - 1):
-            for col in range(self.board.board_size - 1):
-                for orientation in [WallOrientation.VERTICAL, WallOrientation.HORIZONTAL]:
-                    wall_action = WallAction((row, col), orientation)
-                    if self.is_action_valid(wall_action, player):
-                        valid_wall_actions.append(wall_action)
-        return valid_wall_actions
+        action_mask = np.zeros(2 * self.action_encoder.wall_size**2, dtype=bool)
+        qgrid.compute_wall_action_mask(
+            self.board._grid,
+            self.board._player_positions,
+            self.board._walls_remaining,
+            self._goal_rows,
+            int(player),
+            action_mask,
+        )
+        actions = [
+            self.action_encoder.index_to_action(self.action_encoder.board_size**2 + action_i)
+            for action_i in np.flatnonzero(action_mask)
+        ]
+        return actions
 
     def get_valid_actions(self, player: Optional[Player] = None) -> Sequence[Action]:
         return self.get_valid_move_actions(player) + self.get_valid_wall_actions(player)
@@ -508,15 +357,14 @@ def check_win(self, player):
         return row == self.get_goal_row(player)
 
     def can_reach(self, start_position: tuple[int, int], target_row: int):
-        visited = np.zeros((self.board.board_size, self.board.board_size), dtype="bool")
-        return self.board._dfs(start_position, target_row, visited) != -1
+        d = qgrid.distance_to_row(self.board._grid, start_position[0], start_position[1], target_row)
+        return d != -1
 
     def distance_to_target(self, start_position, target_row):
         """
         Returns the minimum number of moves it takes to reach the target row, or -1 if it's not reachable.
         """
-        visited = np.zeros((self.board.board_size, self.board.board_size), dtype="bool")
-        return self.board._bfs(start_position, target_row, visited)
+        return qgrid.distance_to_row(self.board._grid, start_position[0], start_position[1], target_row)
 
     def player_distance_to_target(self, player: Player):
         start_position = self.board.get_player_position(player)
@@ -530,53 +378,6 @@ def __str__(self):
         return str(self.board)
 
 
-def create_jump_checks():
-    """
-    Pre-generate a sequence of checks for possible jumps.
-
-    Pre-generating the checks saves some computation time checking validity of move actions later.
-
-    There are 2 types of checks:
-    - nowall: There must not be a wall between each of these pairs of positions.
-    - wall: There must be a wall between each of these pairs of positions.
-
-    Positions are relative to the player's current position.
-    """
-    # Checks for the case where we are moving to (row + 1, col) or jumping over
-    # an opponent at (row + 1, col)
-    checks_for_one_direction = {
-        (2, 0): {
-            "wall": [],
-            "nowall": [((0, 0), (1, 0)), ((1, 0), (2, 0))],
-        },
-        (1, 1): {
-            "wall": [((1, 0), (2, 0))],
-            "nowall": [((0, 0), (1, 0)), ((1, 0), (1, 1))],
-        },
-        (1, -1): {
-            "wall": [((1, 0), (2, 0))],
-            "nowall": [((0, 0), (1, 0)), ((1, 0), (1, -1))],
-        },
-    }
-    # Rotate the checks for all 4 possible directions
-    checks = {}
-    for opponent_offset in np.array([(1, 0), (0, 1), (-1, 0), (0, -1)]):
-        rotation_matrix = np.array((opponent_offset, -opponent_offset[::-1])).T
-        for move in checks_for_one_direction:
-            rotated_move = np.dot(rotation_matrix, move)
-            lookup_key = (tuple(opponent_offset), tuple(rotated_move))
-            checks[lookup_key] = {}
-            for check_type in checks_for_one_direction[move]:
-                checks[lookup_key][check_type] = []
-                for wall_start, wall_end in checks_for_one_direction[move][check_type]:
-                    rotated_wall = (
-                        tuple(np.dot(rotation_matrix, wall_start)),
-                        tuple(np.dot(rotation_matrix, wall_end)),
-                    )
-                    checks[lookup_key][check_type].append(rotated_wall)
-    return checks
-
-
 def construct_game_from_observation(observation: dict, player_id: str) -> tuple[Quoridor, Player, Player]:
     if player_id == "player_0":
         player = Player.ONE


Reviews:
user: login: alejandromarcu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-16T05:19:54Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-16T13:57:23Z

Review comments:
diff_hunk: @@ -135,7 +125,19 @@ def choose_action(
             game.board.move_player(current_player_before_action, position_before_action)
         game.set_current_player(current_player_before_action)
 
-    return chosen_action, chosen_value
+    # "player" tries to maximize the reward, while "opponent" tries to minimize it.
+    values = np.array(values)
+    if current_player_before_action == player:
+        best_value = values.max()
+    else:
+        best_value = values.min()
+
+    best_action_indices = np.flatnonzero(values == best_value)
+    # There may be multiple actions with the same value, so we randomly choose one of them.
+    # chosen_action_index = np.random.choice(best_action_indices)
path: deep_quoridor/src/agents/simple.py
user: login: alejandromarcu
body: Any reason to use the last rather than picking a random one?
created_at: 2025-05-16T05:19:54Z
author_association: COLLABORATOR
diff_hunk: @@ -135,7 +125,19 @@ def choose_action(
             game.board.move_player(current_player_before_action, position_before_action)
         game.set_current_player(current_player_before_action)
 
-    return chosen_action, chosen_value
+    # "player" tries to maximize the reward, while "opponent" tries to minimize it.
+    values = np.array(values)
+    if current_player_before_action == player:
+        best_value = values.max()
+    else:
+        best_value = values.min()
+
+    best_action_indices = np.flatnonzero(values == best_value)
+    # There may be multiple actions with the same value, so we randomly choose one of them.
+    # chosen_action_index = np.random.choice(best_action_indices)
path: deep_quoridor/src/agents/simple.py
user: login: jonbinney
body: Ah, great catch! Didn't mean to leave that in; fixed now. Interestingly fixing that makes simple a bit better as P2 but a bit worse as P1.....
created_at: 2025-05-16T13:57:23Z
author_association: OWNER

PR comments:
user: login: alejandromarcu
created_at: 2025-05-16T05:18:08Z
body: Wow, very nice!
In my computer I saw Simple's average time coming down from 85 to 7.7, so more than 10x!
One thing that I noticed is that the functions that you numba-fied don't appear anymore in the profiler, so it's hard to continue profiling.   One easy solution if you're in Python 3.12 is that when you want to profile those functions, you do `export NUMBA_ENABLE_SYS_MONITORING=1` and then they appear.

One thing that I noticed when profiling is that `index_to_action` appears in the 4th position, which seemed to much for such a simple function:
```
   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    20901    0.611    0.000    0.846    0.000 ffi.py:195(__call__)
625363/104832    0.337    0.000    0.530    0.000 ir.py:321(_rec_list_vars)
3982182/3982176    0.307    0.000    0.371    0.000 {built-in method builtins.isinstance}
   490365    0.254    0.000    0.431    0.000 quoridor.py:57(index_to_action)
82811/283    0.198    0.000    3.542    0.013 simple.py:83(choose_action)
```
The thing is that every time you call it, it is instantiating an object, which is probably the slow part.  Notice that the objects are immutable and there's a limited quantity of possible objects, so a very easy solution that I tried is to add `@cache` to index_to_action (you need to do `from functools import cache`), and that gave me a 10% bump in speed for simple (now at 7ms). With this we'll just create about 200 objects rather than 490k (it was in just 10 games), which not sure if will have any impact in garbage collection or memory but won't hurt. 

> * Compilation by numba takes a couple seconds each time `play.py` is run. This doesn't seem to get cached on disk and reused the next time the python interpreter is restarted.

I think usually is not a big deal.  I saw there are options to cache the compilation but didn't read much.


> * I've added tie-breaking of the minimax reward function using the number of walls remaining (suggested by @alejandromarcu when reviewing that last PR). This gives a nice bump to the results of the simple agent without slowing it down.

Yey!

> * For the moment I'm using numba-fied BFS for both shortest path and for "does this wall block a player's path to goal". For blocking checks, we can use DFS instead, but it wasn't a bottleneck so I didn't do that yet.

DFS may be a bit faster, but not sure if it would give a significant gain to make it worth implementing as well, since now everything is much faster

author_association: COLLABORATOR
user: login: jonbinney
created_at: 2025-05-16T23:41:11Z
body: Added `@cache` to `index_to_action` and it did help a bit, thanks!

@alejandromarcu are you ok with merging this? I'll have another PR doing a bit more optimization in a few days.
author_association: OWNER