Best-first, A-star and Beam Search

#algorithms/a-star-search#algorithms/beam-search#algorithms/best-first-search

1. Best-first search algorithm

(1) Introduction of the theory

The best-first search algorithm^[1], or as said, $B^{\ast }$ algorithm, In this search algorithm, we choose the

However, in real applications, we do not know the distance of every point to the target, so the distance is often a heuristic value, that not consider some other cases, for example, for a robotic search with some blocks in the path, we just take :

$$\begin{array}{}\text{(1.1.1)}& d=\sqrt{{|x_t-x_i|}^2+{|y_t-y_i|}^2}\end{array}$$

In every step, it greedily takes the position with the smallest distance.

(2) Code Example

1. Initialization

Consider a robotic way planning problem, that has some blocks in the way to search for the optimal path within the limited time. An example can be made by following initializations :

import numpy as np
from dataclasses import dataclass
from typing import Dict, List, Tuple

Position = Tuple[int, int]

# A small warehouse: 0 = free floor, 1 = shelf/wall.
WAREHOUSE_MAP = np.array([
    [0, 0, 0, 1, 0, 0, 0],
    [1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 0, 0, 1, 0],
    [0, 1, 1, 1, 0, 0, 0],
    [0, 0, 0, 0, 0, 1, 0],
])

START: Position = (0, 0)
TARGET: Position = (6, 4)

MOVES: List[Position] = [
    (0, -1),  # north
    (1, 0),   # east
    (0, 1),   # south
    (-1, 0),  # west
]


@dataclass
class SearchState:
    grid: List[List[int]]
    start: Position
    target: Position
    visited: set[Position]
    queue: List[Position]    # the queue of positions to explore next
    parent: Dict[Position, Position | None]


def manhattan_distance(point: Position, target: Position) -> int:
    """Cheap estimate of remaining walking distance, ignoring shelves."""
    return abs(point[0] - target[0]) + abs(point[1] - target[1])


initial_state = SearchState(
    grid=WAREHOUSE_MAP,
    start=START,
    target=TARGET,
    visited={START},
    queue=[START],
    parent={START: None},
)

# Example heuristic values available to the planner:
start_estimate = manhattan_distance(START, TARGET)  # 6

2. Algorithm

def best_first_search(state: SearchState) -> List[Position] | None:
    """
    Perform a best-first search to find a path from start to target.
    """
    while state.queue:
        # Sort the queue based on the heuristic (Manhattan distance) to the target
        state.queue.sort(key=lambda pos: manhattan_distance(pos, state.target))
        current = state.queue.pop(0)  # use the position with the lowest heuristic value 
        
        if current == state.target:
            # Reconstruct the path from start to target
            path = []
            while current is not None:
                path.append(current)
                current = state.parent[current]
            return path[::-1]  # reverse the path to get it from start to target 
        else:
            
            state.visited.add(current)

        # explore all valid neighboring positions 
        for move in MOVES:
            neighbor = (current[0] + move[0], current[1] + move[1])
            is_valid_move = (neighbor[0] >= 0 and neighbor[0] < state.grid.shape[1]) and \
                            (neighbor[1] >= 0 and neighbor[1] < state.grid.shape[0])
            if not is_valid_move:
                continue

            is_possible_move = (state.grid[neighbor[1], neighbor[0]] == 0)
            if not is_possible_move:
                continue

            # attention : marking a node visited only when popped allows 
            #   the same neighbor to be added to queue multiple times before its first turn
            # we also note here, if the path has length, we need to update the minimal length 
            if neighbor not in state.visited and  neighbor not in state.queue:
                state.queue.append(neighbor)
                state.parent[neighbor] = current

    return None  # No path found


if __name__ == "__main__":
    path = best_first_search(initial_state)
    if path:
        print("Path found:", path)
    else:
        print("No path found.")

(3) A-star search algorithm

A-star algorithm^[2][3] is a pathfinding algorithm in the computer science and robotics traveling. The thought is completely same as the Best-first algorithm, with the only difference being the evaluation function.

For $A^{\ast }$ search algorithm, we set depth of shallowest solution as $d$ and branching factor $b$, which is the maximum number of successor. The thought is, if the decision obviously can't be on an optimal path, it's wasting effort.

According to the essay^[3:1], We firstly suppose that some function $f(n)$ can be calculated for every node $n$,

1. Mark $s$ as open and compute $f(s)$
2. Select the node $n$ with smallest $f$, Resolve ties arbitrarily,
3. if $n\in T$ (is the target), mark the $n$ as "closed" and terminate the algorithm.
4. otherwise, still mark $n$ closed, and apply the successor operator $\mathrm{\Gamma }$ to $n$, calculate $\hat{f}$ for ==each successor== of $n$,

We note the successor here is the sub-branch. But since possible solution is $\{n_j,c_{ij}\}$, where $c_{ij}$ is the cost, we only take the first several choice with the lowest cost.

Also, we note the $f(s)=h(s)$ is the cost of an unconstrained optimal path. so we use following :

$$\begin{array}{}\text{(1.3.1)}& f(n)=g(n)+h(n)\end{array}$$

where $g(n)$ is the actual cost of an optimal path from $s$ to $n$, like the following, and $h(n)$ is actual cost from $n$ to a preferred goal node

$$\begin{array}{}\text{(1.3.2)}& \hat{f}(n)=\hat{g}(n)+\hat{h}(n)\end{array}$$

where $\hat{h}(n)$ is the proven lower bound to reach the solution.

To give the code of it, we modify the initial state :

# A small warehouse: 0 = free floor, 1 = shelf/wall.
WAREHOUSE_MAP = np.array([
    [0, 0, 0, 1, 0, 0, 0],
    [1, 1, 0, 1, 0, 1, 0],
    [0, 0, 0, 0, 0, 1, 0],
    [0, 1, 1, 1, 0, 1, 0],
    [0, 0, 0, 0, 0, 1, 0],
])

Then the code of $A^{\ast }$ can be reached by slightly modifying the best first search :

def a_star_search(state: SearchState) -> tuple[dict[Position, int], List[Position] | None]:
    best_dist_sofar = {state.start: 0}  # distance from start to each position 
    
    while state.queue:
        # Sort the queue based on the heuristic (Manhattan distance) to the target
        state.queue.sort(key=lambda pos: best_dist_sofar.get(pos, float('inf')) + manhattan_distance(pos, state.target))
        current = state.queue.pop(0)  # use the position with the lowest heuristic value 
        
        if current == state.target:
            # Reconstruct the path from start to target
            path = []
            while current is not None:
                path.append(current)
                current = state.parent[current]
            return best_dist_sofar, path[::-1]  # reverse the path to get it from start to target 
        else:
            state.visited.add(current)
            
        # explore all valid neighboring positions 
        for move in MOVES:
            neighbor = (current[0] + move[0], current[1] + move[1])
            is_valid_move = (neighbor[0] >= 0 and neighbor[0] < state.grid.shape[1]) and \
                            (neighbor[1] >= 0 and neighbor[1] < state.grid.shape[0])
            if not is_valid_move:
                continue

            is_possible_move = (state.grid[neighbor[1], neighbor[0]] == 0)
            if not is_possible_move:
                continue

            # attention : marking a node visited only when popped allows 
            #   the same neighbor to be added to queue multiple times before its first turn
            # we also note here, if the path has length, we need to update the minimal length 
            if neighbor not in state.visited and  neighbor not in state.queue:
                state.queue.append(neighbor)
                state.parent[neighbor] = current
                if current not in best_dist_sofar:
                    raise ValueError(f"current {current} not in best_dist_sofar")
                best_dist_sofar[neighbor] = best_dist_sofar[current] + 1

    return best_dist_sofar, None  # No path found

2. Beam-Search Algorithm

(1) Introduction

The beam-search algorithm ^[4] allows to explore multiple possible paths simultaneously. It is a modification of the Best-first search algorithm. It is also a heuristic search method. Such a search algorithm is useful for the planning of the next state. For example, planning for the next move. In beam search, only a predetermined number of best partial solutions are kept as candidates.

In RL, a value estimate is usually not guaranteed to be accurate or admissible.

The beam search ==only keeps the best B partial cube states at every depth==, which can reduce the memory requirements for the search.

Property	A*	Beam search
Finds a solution	Yes, in a finite state space	Not guaranteed
Finds shortest solution	Yes, with admissible h	No
Memory use	Can be large	Fixed by beam width
Uses RL policy well	As ordering/tie-breaker	Very naturally
Risk	Slow or memory-heavy	May discard the only path to solution

The beam search is often used by LLM inference ^[5], behavior prediction and the action taking process. This is the reason why we often get multiple responses in LLM models.

(2) Code Implementation

We note the beam-search is not simply apply a queue clipping into the best-first-search algorithm, the standard beam search is level-based :

1. Expand every state in the current frontier, up to beam_width.
2. Combine all their children.
3. Keep the best beam_width children as the next frontier.
4. Repeat at the next depth.

def beam_search(state: SearchState, beam_width: int = 3) -> List[Position] | None:
    """
    Perform a beam search to find a path from start to target.
    """
    beams = [state.start]  # Initialize the beam with the start position

    # we use the queue as the child node set (stop until it becomes empty)
    while beams:
        state.queue.clear()  # Clear the queue for the next layer of child nodes  
        # search a layer under the current beam. 
        beams.sort(key=lambda pos: manhattan_distance(pos, state.target))
        beams = beams[:beam_width]  # Keep only the top beam_width positions
        
        # Keep only the top beam_width positions
        for beam in beams:
            if beam == state.target:
                # Reconstruct the path from start to target
                path = []
                current = beam
                while current is not None:
                    path.append(current)
                    current = state.parent[current]
                return path[::-1]  # reverse the path to get it from start to target -> here we only return the first beam  
            else:
                state.visited.add(beam)

            # explore all valid neighboring positions 
            for move in MOVES:
                neighbor = (beam[0] + move[0], beam[1] + move[1])
                is_valid_move = (neighbor[0] >= 0 and neighbor[0] < state.grid.shape[1]) and \
                                (neighbor[1] >= 0 and neighbor[1] < state.grid.shape[0])
                if not is_valid_move:
                    continue

                is_possible_move = (state.grid[neighbor[1], neighbor[0]] == 0)
                if not is_possible_move:
                    continue
                
                # attention : marking a node visited only when popped allows 
                #   the same neighbor to be added to queue multiple times before its first turn
                # we also note here, if the path has length, we need to update the minimal length 
                if neighbor not in state.visited and neighbor not in state.queue:
                    state.visited.add(neighbor)
                    state.parent[neighbor] = beam
                    state.queue.append(neighbor)   # child nodes 
        # after exploring all beams, we update the beams to be the next layer of child nodes 
        beams = state.queue.copy() # Move to the next layer of child nodes 
    return None  # No path found

Reference

---

1. https://en.wikipedia.org/wiki/Best-first_search ↩︎

2. https://en.wikipedia.org/wiki/A*_search_algorithm ↩︎

3. https://www.cs.auckland.ac.nz/courses/compsci709s2c/resources/Mike.d/astarNilsson.pdf ↩︎ ↩︎

4. https://en.wikipedia.org/wiki/Beam_search ↩︎

5. https://huggingface.co/learn/llm-course/chapter1/8 ↩︎