def test_get_shortest_paths_max_depth():
env = load_flatland_environment_from_file('test_002.pkl', 'env_data.tests')
env.reset()
actual = get_shortest_paths(env.distance_map, max_depth=2)
expected = {
0: [
WayPoint(position=(1, 1), direction=1),
WayPoint(position=(1, 2), direction=1)
],
1: [
WayPoint(position=(3, 18), direction=3),
WayPoint(position=(3, 17), direction=3),
]
}
for agent_handle in expected:
assert np.array_equal(actual[agent_handle], expected[agent_handle]), \
"[{}] actual={},expected={}".format(agent_handle, actual[agent_handle], expected[agent_handle])
def _shortest_path_for_agent(agent):
if agent.status == RailAgentStatus.READY_TO_DEPART:
position = agent.initial_position
elif agent.status == RailAgentStatus.ACTIVE:
position = agent.position
elif agent.status == RailAgentStatus.DONE:
position = agent.target
else:
shortest_paths[agent.handle] = None
return
direction = agent.direction
shortest_paths[agent.handle] = []
distance = math.inf
depth = 0
while (position != agent.target
and (max_depth is None or depth < max_depth)):
next_actions = get_valid_move_actions_(direction, position,
distance_map.rail)
best_next_action = None
for next_action in next_actions:
next_action_distance = distance_map.get()[
agent.handle, next_action.next_position[0],
next_action.next_position[1], next_action.next_direction]
if next_action_distance < distance:
best_next_action = next_action
distance = next_action_distance
shortest_paths[agent.handle].append(WayPoint(position, direction))
depth += 1
# if there is no way to continue, the rail must be disconnected!
# (or distance map is incorrect)
if best_next_action is None:
shortest_paths[agent.handle] = None
return
position = best_next_action.next_position
direction = best_next_action.next_direction
if max_depth is None or depth < max_depth:
shortest_paths[agent.handle].append(WayPoint(position, direction))
def get_way_point_before_or_at_step(self, agent_id: int,
step: int) -> WayPoint:
"""
Get the way point point from which the current position can be extracted.
Parameters
----------
agent_id
step
Returns
-------
WalkingElement
"""
train_run = self.train_run_dict[agent_id]
entry_time_step = train_run[0].scheduled_at
# the agent has no position before and at choosing to enter the grid (one tick elapses before the agent enters the grid)
if step <= entry_time_step:
return WayPoint(
position=None,
direction=self.env.agents[agent_id].initial_direction)
# the agent has no position as soon as the target is reached
exit_time_step = train_run[-1].scheduled_at
if step >= exit_time_step:
# agent loses position as soon as target cell is reached
return WayPoint(position=None,
direction=train_run[-1].way_point.direction)
way_point = None
for train_run_way_point in train_run:
if step < train_run_way_point.scheduled_at:
return way_point
if step >= train_run_way_point.scheduled_at:
way_point = train_run_way_point.way_point
assert way_point is not None
return way_point
def test_get_k_shortest_paths(rendering=False):
rail, rail_map = make_simple_rail_with_alternatives()
env = RailEnv(
width=rail_map.shape[1],
height=rail_map.shape[0],
rail_generator=rail_from_grid_transition_map(rail),
schedule_generator=random_schedule_generator(),
number_of_agents=1,
obs_builder_object=GlobalObsForRailEnv(),
)
env.reset()
initial_position = (3, 1) # west dead-end
initial_direction = Grid4TransitionsEnum.WEST # west
target_position = (3, 9) # east
# set the initial position
agent = env.agents[0]
agent.position = initial_position
agent.initial_position = initial_position
agent.direction = initial_direction
agent.target = target_position # east dead-end
agent.moving = True
env.reset(False, False)
if rendering:
renderer = RenderTool(env, gl="PILSVG")
renderer.render_env(show=True, show_observations=False)
input()
actual = set(
get_k_shortest_paths(
env=env,
source_position=initial_position, # west dead-end
source_direction=int(initial_direction), # east
target_position=target_position,
k=10))
expected = set([
(WayPoint(position=(3, 1),
direction=3), WayPoint(position=(3, 0), direction=3),
WayPoint(position=(3, 1),
direction=1), WayPoint(position=(3, 2), direction=1),
WayPoint(position=(3, 3),
direction=1), WayPoint(position=(2, 3), direction=0),
WayPoint(position=(1, 3),
direction=0), WayPoint(position=(0, 3), direction=0),
WayPoint(position=(0, 4),
direction=1), WayPoint(position=(0, 5), direction=1),
WayPoint(position=(0, 6),
direction=1), WayPoint(position=(0, 7), direction=1),
WayPoint(position=(0, 8),
direction=1), WayPoint(position=(0, 9), direction=1),
WayPoint(position=(1, 9),
direction=2), WayPoint(position=(2, 9), direction=2),
WayPoint(position=(3, 9), direction=2)),
(WayPoint(position=(3, 1),
direction=3), WayPoint(position=(3, 0), direction=3),
WayPoint(position=(3, 1),
direction=1), WayPoint(position=(3, 2), direction=1),
WayPoint(position=(3, 3),
direction=1), WayPoint(position=(3, 4), direction=1),
WayPoint(position=(3, 5),
direction=1), WayPoint(position=(3, 6), direction=1),
WayPoint(position=(4, 6),
direction=2), WayPoint(position=(5, 6), direction=2),
WayPoint(position=(6, 6),
direction=2), WayPoint(position=(5, 6), direction=0),
WayPoint(position=(4, 6),
direction=0), WayPoint(position=(4, 7), direction=1),
WayPoint(position=(4, 8),
direction=1), WayPoint(position=(4, 9), direction=1),
WayPoint(position=(3, 9), direction=0))
])
assert actual == expected, "actual={},expected={}".format(actual, expected)
def test_get_shortest_paths_agent_handle():
env = load_flatland_environment_from_file(
'Level_distance_map_shortest_path.pkl', 'env_data.tests')
env.reset()
actual = get_shortest_paths(env.distance_map, agent_handle=6)
print(actual, file=sys.stderr)
expected = {
6: [
WayPoint(position=(5, 5), direction=0),
WayPoint(position=(4, 5), direction=0),
WayPoint(position=(3, 5), direction=0),
WayPoint(position=(2, 5), direction=0),
WayPoint(position=(1, 5), direction=0),
WayPoint(position=(0, 5), direction=0),
WayPoint(position=(0, 6), direction=1),
WayPoint(position=(0, 7), direction=1),
WayPoint(position=(0, 8), direction=1),
WayPoint(position=(0, 9), direction=1),
WayPoint(position=(0, 10), direction=1),
WayPoint(position=(1, 10), direction=2),
WayPoint(position=(2, 10), direction=2),
WayPoint(position=(3, 10), direction=2),
WayPoint(position=(4, 10), direction=2),
WayPoint(position=(5, 10), direction=2),
WayPoint(position=(6, 10), direction=2),
WayPoint(position=(7, 10), direction=2),
WayPoint(position=(8, 10), direction=2),
WayPoint(position=(9, 10), direction=2),
WayPoint(position=(10, 10), direction=2),
WayPoint(position=(11, 10), direction=2),
WayPoint(position=(12, 10), direction=2),
WayPoint(position=(13, 10), direction=2),
WayPoint(position=(14, 10), direction=2),
WayPoint(position=(15, 10), direction=2),
WayPoint(position=(16, 10), direction=2),
WayPoint(position=(17, 10), direction=2),
WayPoint(position=(18, 10), direction=2),
WayPoint(position=(19, 10), direction=2),
WayPoint(position=(20, 10), direction=2),
WayPoint(position=(20, 9), direction=3),
WayPoint(position=(20, 8), direction=3),
WayPoint(position=(21, 8), direction=2),
WayPoint(position=(21, 7), direction=3),
WayPoint(position=(21, 6), direction=3),
WayPoint(position=(21, 5), direction=3)
]
}
for agent_handle in expected:
assert np.array_equal(actual[agent_handle], expected[agent_handle]), \
"[{}] actual={},expected={}".format(agent_handle, actual[agent_handle], expected[agent_handle])
def get_k_shortest_paths(env: RailEnv,
source_position: Tuple[int, int],
source_direction: int,
target_position=Tuple[int, int],
k: int = 1,
debug=False) -> List[Tuple[WayPoint]]:
"""
Computes the k shortest paths using modified Dijkstra
following pseudo-code https://en.wikipedia.org/wiki/K_shortest_path_routing
In contrast to the pseudo-code in wikipedia, we do not a allow for loopy paths.
Parameters
----------
env : RailEnv
source_position: Tuple[int,int]
source_direction: int
target_position: Tuple[int,int]
k : int
max number of shortest paths
debug: bool
print debug statements
Returns
-------
List[Tuple[WalkingElement]]
We use tuples since we need the path elements to be hashable.
We use a list of paths in order to keep the order of length.
"""
# P: set of shortest paths from s to t
# P =empty,
shortest_paths: List[Tuple[WayPoint]] = []
# countu: number of shortest paths found to node u
# countu = 0, for all u in V
count = {(r, c, d): 0
for r in range(env.height) for c in range(env.width)
for d in range(4)}
# B is a heap data structure containing paths
heap: Set[Tuple[WayPoint]] = set()
# insert path Ps = {s} into B with cost 0
heap.add((WayPoint(source_position, source_direction), ))
# while B is not empty and countt < K:
while len(heap) > 0 and len(shortest_paths) < k:
if debug:
print("iteration heap={}, shortest_paths={}".format(
heap, shortest_paths))
# – let Pu be the shortest cost path in B with cost C
cost = np.inf
pu = None
for path in heap:
if len(path) < cost:
pu = path
cost = len(path)
u: WayPoint = pu[-1]
if debug:
print(" looking at pu={}".format(pu))
# – B = B − {Pu }
heap.remove(pu)
# – countu = countu + 1
urcd = (*u.position, u.direction)
count[urcd] += 1
# – if u = t then P = P U {Pu}
if u.position == target_position:
if debug:
print(" found of length {} {}".format(len(pu), pu))
shortest_paths.append(pu)
# – if countu ≤ K then
# CAVEAT: do not allow for loopy paths
elif count[urcd] <= k:
possible_transitions = env.rail.get_transitions(*urcd)
if debug:
print(" looking at neighbors of u={}, transitions are {}".
format(u, possible_transitions))
# for each vertex v adjacent to u:
for new_direction in range(4):
if debug:
print(" looking at new_direction={}".format(
new_direction))
if possible_transitions[new_direction]:
new_position = get_new_position(u.position, new_direction)
if debug:
print(" looking at neighbor v={}".format(
(*new_position, new_direction)))
v = WayPoint(position=new_position,
direction=new_direction)
# CAVEAT: do not allow for loopy paths
if v in pu:
continue
# – let Pv be a new path with cost C + w(u, v) formed by concatenating edge (u, v) to path Pu
pv = pu + (v, )
# – insert Pv into B
heap.add(pv)
# return P
return shortest_paths
def test_action_plan(rendering: bool = False):
"""Tests ActionPlanReplayer: does action plan generation and replay work as expected."""
rail, rail_map = make_simple_rail()
env = RailEnv(width=rail_map.shape[1],
height=rail_map.shape[0],
rail_generator=rail_from_grid_transition_map(rail),
schedule_generator=random_schedule_generator(seed=77),
number_of_agents=2,
obs_builder_object=GlobalObsForRailEnv(),
remove_agents_at_target=True)
env.reset()
env.agents[0].initial_position = (3, 0)
env.agents[0].target = (3, 8)
env.agents[0].initial_direction = Grid4TransitionsEnum.WEST
env.agents[1].initial_position = (3, 8)
env.agents[1].initial_direction = Grid4TransitionsEnum.WEST
env.agents[1].target = (0, 3)
env.agents[1].speed_data['speed'] = 0.5 # two
env.reset(False, False, False)
for handle, agent in enumerate(env.agents):
print("[{}] {} -> {}".format(handle, agent.initial_position,
agent.target))
chosen_path_dict = {
0: [
TrainRunWayPoint(scheduled_at=0,
way_point=WayPoint(position=(3, 0), direction=3)),
TrainRunWayPoint(scheduled_at=2,
way_point=WayPoint(position=(3, 1), direction=1)),
TrainRunWayPoint(scheduled_at=3,
way_point=WayPoint(position=(3, 2), direction=1)),
TrainRunWayPoint(scheduled_at=14,
way_point=WayPoint(position=(3, 3), direction=1)),
TrainRunWayPoint(scheduled_at=15,
way_point=WayPoint(position=(3, 4), direction=1)),
TrainRunWayPoint(scheduled_at=16,
way_point=WayPoint(position=(3, 5), direction=1)),
TrainRunWayPoint(scheduled_at=17,
way_point=WayPoint(position=(3, 6), direction=1)),
TrainRunWayPoint(scheduled_at=18,
way_point=WayPoint(position=(3, 7), direction=1)),
TrainRunWayPoint(scheduled_at=19,
way_point=WayPoint(position=(3, 8), direction=1)),
TrainRunWayPoint(scheduled_at=20,
way_point=WayPoint(position=(3, 8), direction=5))
],
1: [
TrainRunWayPoint(scheduled_at=0,
way_point=WayPoint(position=(3, 8), direction=3)),
TrainRunWayPoint(scheduled_at=3,
way_point=WayPoint(position=(3, 7), direction=3)),
TrainRunWayPoint(scheduled_at=5,
way_point=WayPoint(position=(3, 6), direction=3)),
TrainRunWayPoint(scheduled_at=7,
way_point=WayPoint(position=(3, 5), direction=3)),
TrainRunWayPoint(scheduled_at=9,
way_point=WayPoint(position=(3, 4), direction=3)),
TrainRunWayPoint(scheduled_at=11,
way_point=WayPoint(position=(3, 3), direction=3)),
TrainRunWayPoint(scheduled_at=13,
way_point=WayPoint(position=(2, 3), direction=0)),
TrainRunWayPoint(scheduled_at=15,
way_point=WayPoint(position=(1, 3), direction=0)),
TrainRunWayPoint(scheduled_at=17,
way_point=WayPoint(position=(0, 3), direction=0))
]
}
expected_action_plan = [
[
# take action to enter the grid
ActionPlanElement(0, RailEnvActions.MOVE_FORWARD),
# take action to enter the cell properly
ActionPlanElement(1, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(2, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(3, RailEnvActions.STOP_MOVING),
ActionPlanElement(13, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(14, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(15, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(16, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(17, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(18, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(19, RailEnvActions.STOP_MOVING)
],
[
ActionPlanElement(0, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(1, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(3, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(5, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(7, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(9, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(11, RailEnvActions.MOVE_RIGHT),
ActionPlanElement(13, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(15, RailEnvActions.MOVE_FORWARD),
ActionPlanElement(17, RailEnvActions.STOP_MOVING),
]
]
MAX_EPISODE_STEPS = 50
deterministic_controller = ControllerFromTrainRuns(env, chosen_path_dict)
deterministic_controller.print_action_plan()
ControllerFromTrainRuns.assert_actions_plans_equal(
expected_action_plan, deterministic_controller.action_plan)
ControllerFromTrainRunsReplayer.replay_verify(MAX_EPISODE_STEPS,
deterministic_controller,
env, rendering)
def test_shortest_path_predictor(rendering=False):
rail, rail_map = make_simple_rail()
env = RailEnv(
width=rail_map.shape[1],
height=rail_map.shape[0],
rail_generator=rail_from_grid_transition_map(rail),
schedule_generator=random_schedule_generator(),
number_of_agents=1,
obs_builder_object=TreeObsForRailEnv(
max_depth=2, predictor=ShortestPathPredictorForRailEnv()),
)
env.reset()
# set the initial position
agent = env.agents[0]
agent.initial_position = (5, 6) # south dead-end
agent.position = (5, 6) # south dead-end
agent.direction = 0 # north
agent.initial_direction = 0 # north
agent.target = (3, 9) # east dead-end
agent.moving = True
agent.status = RailAgentStatus.ACTIVE
env.reset(False, False)
if rendering:
renderer = RenderTool(env, gl="PILSVG")
renderer.render_env(show=True, show_observations=False)
input("Continue?")
# compute the observations and predictions
distance_map = env.distance_map.get()
assert distance_map[0, agent.initial_position[0], agent.initial_position[1], agent.direction] == 5.0, \
"found {} instead of {}".format(
distance_map[agent.handle, agent.initial_position[0], agent.position[1], agent.direction], 5.0)
paths = get_shortest_paths(env.distance_map)[0]
assert paths == [
WayPoint((5, 6), 0),
WayPoint((4, 6), 0),
WayPoint((3, 6), 0),
WayPoint((3, 7), 1),
WayPoint((3, 8), 1),
WayPoint((3, 9), 1)
]
# extract the data
predictions = env.obs_builder.predictions
positions = np.array(
list(map(lambda prediction: [*prediction[1:3]], predictions[0])))
directions = np.array(
list(map(lambda prediction: [prediction[3]], predictions[0])))
time_offsets = np.array(
list(map(lambda prediction: [prediction[0]], predictions[0])))
# test if data meets expectations
expected_positions = [
[5, 6],
[4, 6],
[3, 6],
[3, 7],
[3, 8],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
[3, 9],
]
expected_directions = [
[Grid4TransitionsEnum.NORTH], # next is [5,6] heading north
[Grid4TransitionsEnum.NORTH], # next is [4,6] heading north
[Grid4TransitionsEnum.NORTH], # next is [3,6] heading north
[Grid4TransitionsEnum.EAST], # next is [3,7] heading east
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
[Grid4TransitionsEnum.EAST],
]
expected_time_offsets = np.array([
[0.],
[1.],
[2.],
[3.],
[4.],
[5.],
[6.],
[7.],
[8.],
[9.],
[10.],
[11.],
[12.],
[13.],
[14.],
[15.],
[16.],
[17.],
[18.],
[19.],
[20.],
])
assert np.array_equal(time_offsets, expected_time_offsets), \
"time_offsets {}, expected {}".format(time_offsets, expected_time_offsets)
assert np.array_equal(positions, expected_positions), \
"positions {}, expected {}".format(positions, expected_positions)
assert np.array_equal(directions, expected_directions), \
"directions {}, expected {}".format(directions, expected_directions)