Computed distance from main path for the retargeting data

procgen_tools/maze.py

@@ -945,6 +945,20 @@ def get_path_to_corner(inner_grid, graph, start_node=(0, 0)):
     return nx.shortest_path(graph, start_node, corner_node)
 
 
+def distance_to_tr_path(
+    grid: np.ndarray, target: Square, tr_path: Optional[List[Square]] = None
+) -> int:
+    """Computes the distance between a target and the path from the
+    origin to the top-right square. If no `tr_path`
+    is provided, then computes the path."""
+    if tr_path is None:
+        tr_path = pathfind(
+            grid, (0, 0), (grid.shape[0] - 1, grid.shape[1] - 1)
+        )
+    target_path: List[Square] = pathfind(grid, (0, 0), target)
+    return len(set(target_path) - set(tr_path))
+
+
 def pathfind(grid: np.ndarray, start, end):
     _, came_from, extra = shortest_path(
         grid, start, stop_condition=lambda _, sq: sq == end

procgen_tools/visualization.py

@@ -621,22 +621,6 @@ def show_grid_heatmap(
                 extra_adjustment=False,
             )
 
-            # Get the center square of size size in the pixel slices
-            # center = (
-            #     (pixel_slices[0].stop - pixel_slices[0].start) // 2,
-            #     (pixel_slices[1].stop - pixel_slices[1].start) // 2,
-            # )
-            # pixel_slices = (
-            #     slice(
-            #         center[0] - int(size * center[0]),
-            #         center[0] + int(size * center[0]),
-            #     ),
-            #     slice(
-            #         center[1] - int(size * center[1]),
-            #         center[1] + int(size * center[1]),
-            #     ),
-            # )
-
             heatmap_overlay[pixel_slices[0], pixel_slices[1]] = cm(
                 heatmap[coord]
             )  # Flip because the rows are upside-down by default