breadth first search (flood fill) in torch

bfs.py

class BFS(nn.Module):
    def __init__(self):
        super().__init__()
        # convolutional kernel to connect with neighbors on the grid
        self.step = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False)
        self.step.weight.data = 0.3 * torch.tensor([[1, 1, 1],
                                                    [1, 2, 1],
                                                    [1, 1, 1]]).view(1,1,3,3)
        self.steps = 4

    def forward(
        self,
        grid, # float zero grid with 1 where agent is
        closed # bool zero grid with 1 where obstacles are
    ):
        # propagate signal from starting cell
        for _ in range(self.steps):
            grid = self.step(grid)
            print(grid, 'pre', _)
            # activation: squash and restore obstacles
            grid = -0.01 * closed + grid.tanh() * (1 - closed.float())
            print(grid, _)

        # take only reached cells
        grid = (grid>0).float()
        print(grid, 'signed and restored')

        # connect nearby obstacles
        grid = self.step(grid)
        print(grid)

        return grid>0