data.py

"""
BFS based expert which can be used to generate (state, action) training pairs automatically.

"""

from collections import deque
from snakenv import SnakeEnv, GRID, TURN_LEFT, TURN_RIGHT

def _bfs(snake_set: set, head: tuple, food: tuple, grid: int):
    """
    BFS from head to food. snake_set are walls
    """
    queue = deque([(head, None)]) #(pos, first_step)
    visited = {head}

    while queue:
        pos, first_step = queue.popleft()

        if pos==food:
            return first_step
        
        for d in [(0, -1), (0, 1), (-1, 0), (1, 0)]:
            npos = (pos[0]+d[0], pos[1]+d[1])
            if (0 <= npos[0] < grid and 0 <= npos[1] < grid and npos not in snake_set and npos not in visited):
                visited.add(npos)
                queue.append((npos, d if first_step is None else first_step))

    return None

def get_best_action(env: SnakeEnv) -> int:
    """returns the best action (0, 1, 2 for st,l,r respectively) fr the curr env state.
    """

    head = env.snake[0]
    snake_set = set(env.snake)

    best_abs_dir = _bfs(snake_set, head, env.food, env.grid)

    if best_abs_dir is None:
        for d in [(0, -1), (0, 1), (-1, 0), (1, 0)]:
            npos = (head[0]+d[0], head[1]+d[1])
            if (0 <= npos[0] < env.grid and 0 <= npos[1] < env.grid and npos not in snake_set):
                best_abs_dir = d
                break

    if best_abs_dir is None:
        return 0
    
    if best_abs_dir == env.direction:
        return 0
    elif best_abs_dir == TURN_RIGHT[env.direction]:
        return 1
    elif best_abs_dir == TURN_LEFT[env.direction]:
        return 2
    else:
        right = TURN_RIGHT[env.direction]
        npos = (head[0]+right[0], head[1]+right[1])
        if(0 <= npos[0] < env.grid and 0<=npos[1] <env.grid and npos not in snake_set):
            return 1
        return 2