Hierarchical, Symbolic, Language, and Multi-Agent Planning › Worked example: options with explicit initiation and failure

from collections import deque
from dataclasses import dataclass

import numpy as np

GRID = np.array(
    [
        [0, 0, 0, 1, 0, 0],
        [0, 1, 0, 1, 0, 1],
        [0, 1, 2, 1, 0, 0],
        [0, 0, 0, 0, 3, 0],
        [1, 1, 1, 1, 0, 0],
        [0, 0, 0, 0, 4, 0],
    ]
)

START = (0, 0)
KEY = tuple(map(int, np.argwhere(GRID == 2)[0]))
DOOR = tuple(map(int, np.argwhere(GRID == 3)[0]))
DOOR_APPROACH = (DOOR[0], DOOR[1] - 1)
GOAL = tuple(map(int, np.argwhere(GRID == 4)[0]))
MOVES = [(-1, 0), (1, 0), (0, -1), (0, 1)]


@dataclass(frozen=True)
class EnvState:
    cell: tuple
    has_key: bool
    door_open: bool


@dataclass(frozen=True)
class OptionResult:
    state: EnvState
    cells: tuple
    steps: int
    succeeded: bool
    reason: str


def step(state, action):
    """Apply a movement or unlock action and return the next environment state."""
    if action == "unlock":
        if state.cell != DOOR_APPROACH or not state.has_key:
            return state
        return EnvState(state.cell, True, True)

    dr, dc = action
    nr, nc = state.cell[0] + dr, state.cell[1] + dc
    if not (0 <= nr < GRID.shape[0] and 0 <= nc < GRID.shape[1]):
        return state
    if GRID[nr, nc] == 1 or ((nr, nc) == DOOR and not state.door_open):
        return state
    cell = (int(nr), int(nc))
    return EnvState(cell, state.has_key or cell == KEY, state.door_open)

def bfs_policy(target, door_open):
    """Map each reachable cell to one shortest-path movement toward target."""
    frontier = deque([target])
    distance = {target: 0}
    action_from = {}
    while frontier:
        cell = frontier.popleft()
        for action in MOVES:
            nr, nc = cell[0] - action[0], cell[1] - action[1]
            previous = (nr, nc)
            if not (0 <= nr < GRID.shape[0] and 0 <= nc < GRID.shape[1]):
                continue
            if GRID[nr, nc] == 1 or (previous == DOOR and not door_open):
                continue
            if previous in distance:
                continue
            distance[previous] = distance[cell] + 1
            action_from[previous] = action
            frontier.append(previous)
    return action_from


NAV_POLICIES = {
    "go_to_key": (KEY, bfs_policy(KEY, door_open=False)),
    "go_to_door": (DOOR_APPROACH, bfs_policy(DOOR_APPROACH, door_open=False)),
    "go_to_goal": (GOAL, bfs_policy(GOAL, door_open=True)),
}


def option_is_initiable(name, state):
    if name == "go_to_key":
        return not state.has_key and state.cell in NAV_POLICIES[name][1]
    if name == "go_to_door":
        return (
            state.has_key
            and not state.door_open
            and state.cell in NAV_POLICIES[name][1]
        )
    if name == "unlock_door":
        return (
            state.has_key
            and not state.door_open
            and state.cell == DOOR_APPROACH
        )
    if name == "go_to_goal":
        return (
            state.has_key
            and state.door_open
            and state.cell in NAV_POLICIES[name][1]
        )
    return False


def run_option(name, state, step_cap=64):
    if not option_is_initiable(name, state):
        return OptionResult(
            state, (state.cell,), 0, False, "initiation condition failed"
        )
    if name == "unlock_door":
        next_state = step(state, "unlock")
        return OptionResult(
            next_state,
            (state.cell,),
            1,
            next_state.door_open,
            "terminated" if next_state.door_open else "unlock failed",
        )

    target, policy = NAV_POLICIES[name]
    cells = [state.cell]
    for used in range(1, step_cap + 1):
        if state.cell == target:
            return OptionResult(
                state, tuple(cells), used - 1, True, "terminated"
            )
        action = policy.get(state.cell)
        if action is None:
            return OptionResult(
                state, tuple(cells), used - 1, False, "no policy action"
            )
        next_state = step(state, action)
        if next_state.cell == state.cell:
            return OptionResult(
                state, tuple(cells), used, False, "blocked transition"
            )
        state = next_state
        cells.append(state.cell)
        if state.cell == target:
            return OptionResult(state, tuple(cells), used, True, "terminated")
    return OptionResult(
        state, tuple(cells), step_cap, False, "step cap exhausted"
    )


def execute_option_sequence(names, initial):
    state = initial
    trajectory = [state.cell]
    total_steps = 0
    log = []
    for name in names:
        result = run_option(name, state)
        log.append((name, result.succeeded, result.reason, result.steps))
        state = result.state
        trajectory.extend(result.cells[1:])
        total_steps += result.steps
        if not result.succeeded:
            return state, trajectory, total_steps, log
    return state, trajectory, total_steps, log