AI Revisited: Breaking Down BFS

So we're approaching the end of the year, and of all the topics that I've tended to focus on in my writings, there's one that I haven't really written about in probably over a year, and this is AI and Machine Learning. I've still been doing some work behind the scenes, as you'll know if you keep following the blog for the next few weeks. But I figured I'd spend the last few weeks of the year with some AI related topics. This week, I'll go over an algorithm that is really useful to understand when it comes to writing simple AI programs, and this is Breadth-First-Search.

All the code for the next few weeks can be found in this GitHub repository! For this week, all the code can be found in the BFS module.

The Algorithm

To frame this problem in a more concrete way, let's imagine we have a 2-dimensional grid. Some spaces are free, other spaces are "walls". We want to use breadth first search to find a path from a start point to a destination point.

a___
_xx_
_xb_
____

So our algorithm will take two locations, and return a path from location A to Location B, or an empty list if no path can be found.

The key data structure when executing a breadth-first-search is a queue. Our basic approach is this: we will place our starting location in the queue. Then, we'll go through a loop as long as the queue is not empty. We'll pull an item off, and then add each of the empty neighbors on to the back of the queue, as long as they haven't been added yet. If we dequeue the destination, we're done! But if we reach an empty queue, then we don't have a valid path.

The last tricky part is that we to track the "parent" of each location. That is, which of its neighbors placed it on the queue? This will allow us to reconstruct the path we need to take to get from point a to point b.

So let's imagine we have a simple graph like in the ASCII art above. We start at (0,0). Our queue will operate like this.

It contains (0,0). We'll then enqueue (0, 1) and (1, 0), since those are the neighbors of (0, 0).

(0, 0) <-- Current
(0, 1)
(1, 0)

Then we're done with (0, 0). So we dequeue (0, 1). This its only neighbor is (0, 2), so that gets placed on the end of the queue.

(0, 1) <-- Current
(1, 0)
(0, 2)

And then we repeat the process with (1, 0), placing (0, 2).

(1, 0) <-- Current
(0, 2)
(2, 0)

We keep doing this until we navigate around to our destination at (2,2).

Types First

How do we translate this to Haskell? My favorite approach to problems like this is to use a top-down, type driven, compile-first method of writing the algorithm. Because before we can really get started in earnest, we have to define our data structures and our types. First, let's alias an integer tuple as a "Location":

type Location = (Int, Int)

Now, we're going to imagine we're navigating a 2D grid, and we'll represent this with an array where the indices are tuples which represent locations, and each value is either "empty" or "wall". We can move through empty spaces, but we cannot move through walls.

data Cell = Empty | Wall
  deriving (Eq)
type Grid = A.Array Location Cell

Now we're ready to define the type signature for our function. This takes the grid as an input, as well as the start and end location:

bfsSearch :: Grid -> Location -> Location -> [Location]

We'll need one more type to help frame the problem. This algorithm will use the State monad, because there's a lot of state we need to track here. First off, we need the queue itself. We represent this with the Sequence type in Haskell. Then, we need our set of visited locations. Each time we enqueue a location, we'll save it here. Last, we need our "parents" map. This will help us determine the path at the very end.

data BFSState = BFSState
  { queue :: S.Seq Location
  , visited :: Set.Set Location
  , parents :: M.Map Location Location
  }

A Stateful Skeleton

With these types, we can start framing the problem a bit more. First, we want to construct our initial state. Everything is empty except our queue has the starting location on it.

bfsSearch :: Grid -> Location -> Location -> [Location]
bfsSearch grid start finish = ...
  where
    initialState = BFSState (S.singleton start) Set.empty M.empty

Now we want to pass this function to a stateful computation that returns our list. So we'll imagine we have a helper in the State monad which returns our location. We'll call this bfsSearch'. We can then fill in our original function with evalState.

bfsSearch :: Grid -> Location -> Location -> [Location]
bfsSearch grid start finish = evalState (bfsSearch' grid finish) initialState
  where
    initialState = BFSState (S.singleton start) Set.empty M.empty

bfsSearch' :: Grid -> Location -> State BFSState [Location]
...

Base Case

Now within our stateful helper, we can recognize that this will be a recursive function. We dequeue an element, enqueue its neighbors, and then repeat the process. So let's handle the base cases first. We'll retrieve our sequence from the state and check if it's empty or not. If it's empty, we return the empty list. This means that we couldn't find a path.

bfsSearch' :: Grid -> Location -> State BFSState [Location]
bfsSearch' grid finish = do
  (BFSState q v p) <- get
  case S.viewl q of
    (top S.:< rest) -> ...
    _ -> return []

Now another base case is where the top of our queue is the destination. In this case, we're ready to "unwind" the path from that destination in our stateful map. Let's imagine we have a function to handle this unwinding process. We'll fill it in later.

bfsSearch' :: Grid -> Location -> State BFSState [Location]
bfsSearch' grid finish = do
  (BFSState q v p) <- get
  case S.viewl q of
    (top S.:< rest) -> if top == finish
      then return (unwindPath p [finish])
      else ...
    _ -> return []

unwindPath :: M.Map Location Location -> [Location] -> [Location]

The General Case

Now let's write out the steps for our general case.

1. Get the neighbors of the top element on the queue
2. Append these to the "rest" of the queue (discarding the top element).
3. Insert this top element into our "visited" set v.
4. For each new location, insert it into our "parents" map with the current top as its "parent".
5. Update our final state and recurse!

Each of these statements is 1-2 lines in our function, except we'll want to make a helper for the first line. Let's imagine we have a function that can give us the unvisited neighbors of a space in our grid. This will require passing the location, the grid, and the visited set.

let valid adjacent = getValidNeighbors top grid v
...

getValidNeighbors ::
  Location -> Grid -> Set.Set Location -> [Location]

The next lines involve data structure manipulation, with a couple tricky folds. First, appending the new elements into the queue.

let newQueue = foldr (flip (S.|>)) rest validAdjacent

Next, inserting the top into the visited set. This one's easy.

let newVisited = Set.insert top v

Now, insert each new neighbor into the parents map. The new location is the "key", and the current top is the value.

let newParentsMap = foldr (\loc -> M.insert loc top) p validAdjacent

Last of all, we replace the state and recurse!

put (BFSState newQueue newVisited newParentsMap)
bfsSearch' grid finish

Here's our complete function!

bfsSearch' :: Grid -> Location -> State BFSState [Location]
bfsSearch' grid finish = do
  (BFSState q v p) <- get
  case S.viewl q of
    (top S.:< rest) -> if top == finish
      then return (unwindPath p [finish])
      else do
        let validAdjacent = getValidNeighbors top grid v
        let newQueue = foldr (flip (S.|>)) rest validAdjacent
        let newVisited = Set.insert top v
        let newParentsMap = foldr (\loc -> M.insert loc top) p validAdjacent
        put (BFSState newQueue newVisited newParentsMap)
        bfsSearch' grid finish
    _ -> return []

Filling in Helpers

Now we just need to fill in our helper functions. Unwinding the map is a fairly straightforward tail-recursive problem. We get the parent of the current element, and keep an accumulating list of the places we've gone:

unwindPath :: M.Map Location Location -> [Location] -> [Location]
unwindPath parentsMap currentPath = case M.lookup (head currentPath) parentsMap of
  Nothing -> tail currentPath
  Just parent -> unwindPath parentsMap (parent : currentPath)

Finding the neighbors is slightly tricker. For each direction (right, down, left, and right), we have to consider if the "neighbor" cell is in bounds. Then we have to consider if it's empty. Finally, we need to know if it is still "unvisited". As long as all three of these conditions hold, we can potentially add it. Here's what this process looks like for finding the "right" neighbor.

getValidNeighbors :: Location -> Grid -> Set.Set Location -> [Location]
getValidNeighbors (r, c) grid v = ...
  where
    (rowMax, colMax) = snd . A.bounds $ grid
    right = (r, c + 1)
    right' = if c + 1 <= colMax && grid A.! right == Empty && not (Set.member right v)
      then Just right
      else Nothing

We do this in every direction, and we'll use catMaybes so we only get the correct ones in the end!

getValidNeighbors :: Location -> Grid -> Set.Set Location -> [Location]
getValidNeighbors (r, c) grid v = catMaybes [right', down', left', up']
  where
    (rowMax, colMax) = snd . A.bounds $ grid
    right = (r, c + 1)
    right' = if c + 1 <= colMax && grid A.! right == Empty && not (Set.member right v)
      then Just right
      else Nothing
    down = (r + 1, c)
    down' = if r + 1 <= rowMax && grid A.! down == Empty && not (Set.member down v)
      then Just down
      else Nothing
    left = (r, c - 1)
    left' = if c - 1 >= 0 && grid A.! left == Empty && not (Set.member left v)
      then Just left
      else Nothing
    up = (r - 1, c)
    up' = if r - 1 >= 0 && grid A.! up == Empty && not (Set.member up v)
      then Just up
      else Nothing

Conclusion

This basic structure can also be adapted to use depth-first search as well! The main difference is that you must treat the Sequence as a stack instead of a queue, appending new items to the left side of the sequence. Both of these algorithms are guaranteed to find a path if it exists. But BFS will find the shortest path in this kind of scenario, whereas DFS probably won't!

Next week, we'll continue a basic AI exploration by putting this algorithm to work in a game environment with Gloss!