File size: 2,888 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import numpy as np
import torch
from gym import Env
from typing import Tuple, List
def get_vi_sequence(env: Env, observation: np.ndarray) -> Tuple[np.ndarray, List]:
"""
Overview:
Given an instance of the maze environment and the current observation, using Broad-First-Search (BFS) \
algorithm to plan an optimal path and record the result.
Arguments:
- env (:obj:`Env`): The instance of the maze environment.
- observation (:obj:`np.ndarray`): The current observation.
Returns:
- output (:obj:`Tuple[np.ndarray, List]`): The BFS result. ``output[0]`` contains the BFS map after each \
iteration and ``output[1]`` contains the optimal actions before reaching the finishing point.
"""
xy = np.where(observation[Ellipsis, -1] == 1)
start_x, start_y = xy[0][0], xy[1][0]
target_location = env.target_location
nav_map = env.nav_map
current_points = [target_location]
chosen_actions = {target_location: 0}
visited_points = {target_location: True}
vi_sequence = []
vi_map = np.full((env.size, env.size), fill_value=env.n_action, dtype=np.int32)
found_start = False
while current_points and not found_start:
next_points = []
for point_x, point_y in current_points:
for (action, (next_point_x, next_point_y)) in [(0, (point_x - 1, point_y)), (1, (point_x, point_y - 1)),
(2, (point_x + 1, point_y)), (3, (point_x, point_y + 1))]:
if (next_point_x, next_point_y) in visited_points:
continue
if not (0 <= next_point_x < len(nav_map) and 0 <= next_point_y < len(nav_map[next_point_x])):
continue
if nav_map[next_point_x][next_point_y] == 'x':
continue
next_points.append((next_point_x, next_point_y))
visited_points[(next_point_x, next_point_y)] = True
chosen_actions[(next_point_x, next_point_y)] = action
vi_map[next_point_x, next_point_y] = action
if next_point_x == start_x and next_point_y == start_y:
found_start = True
vi_sequence.append(vi_map.copy())
current_points = next_points
track_back = []
if found_start:
cur_x, cur_y = start_x, start_y
while cur_x != target_location[0] or cur_y != target_location[1]:
act = vi_sequence[-1][cur_x, cur_y]
track_back.append((torch.FloatTensor(env.process_states([cur_x, cur_y], env.get_maze_map())), act))
if act == 0:
cur_x += 1
elif act == 1:
cur_y += 1
elif act == 2:
cur_x -= 1
elif act == 3:
cur_y -= 1
return np.array(vi_sequence), track_back
|