def get_waypoint_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
"""
trainrun = self.trainrun_dict[agent_id]
entry_time_step = trainrun[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 = trainrun[-1].scheduled_at
if step >= exit_time_step:
# agent loses position as soon as target cell is reached
return Waypoint(position=None, direction=trainrun[-1].waypoint.direction)
waypoint = None
for trainrun_waypoint in trainrun:
if step < trainrun_waypoint.scheduled_at:
return waypoint
if step >= trainrun_waypoint.scheduled_at:
waypoint = trainrun_waypoint.waypoint
assert waypoint is not None
return waypoint
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,agent_pos,agent_dir):
if agent_pos is None :
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
else :
position = agent_pos
direction = agent_dir
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_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
# N.B. use OrderedSet to make result deterministic!
heap: OrderedSet[Tuple[Waypoint]] = OrderedSet()
# 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_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)
def test_get_shortest_paths():
#env = load_flatland_environment_from_file('test_002.mpk', 'env_data.tests')
env, env_dict = RailEnvPersister.load_new("test_002.mpk", "env_data.tests")
#print("env len(agents): ", len(env.agents))
#print(env.distance_map)
#print("env number_of_agents:", env.number_of_agents)
#print("env agents:", env.agents)
#env.distance_map.reset(env.agents, env.rail)
#actual = get_shortest_paths(env.distance_map)
#print("shortest paths:", actual)
#print(env.distance_map)
#print("Dist map agents:", env.distance_map.agents)
#print("\nenv reset()")
env.reset()
actual = get_shortest_paths(env.distance_map)
#print("env agents: ", len(env.agents))
#print("env number_of_agents: ", env.number_of_agents)
assert len(
actual) == 2, "get_shortest_paths should return a dict of length 2"
expected = {
0: [
Waypoint(position=(1, 1), direction=1),
Waypoint(position=(1, 2), direction=1),
Waypoint(position=(1, 3), direction=1),
Waypoint(position=(2, 3), direction=2),
Waypoint(position=(2, 4), direction=1),
Waypoint(position=(2, 5), direction=1),
Waypoint(position=(2, 6), direction=1),
Waypoint(position=(2, 7), direction=1),
Waypoint(position=(2, 8), direction=1),
Waypoint(position=(2, 9), direction=1),
Waypoint(position=(2, 10), direction=1),
Waypoint(position=(2, 11), direction=1),
Waypoint(position=(2, 12), direction=1),
Waypoint(position=(2, 13), direction=1),
Waypoint(position=(2, 14), direction=1),
Waypoint(position=(2, 15), direction=1),
Waypoint(position=(2, 16), direction=1),
Waypoint(position=(2, 17), direction=1),
Waypoint(position=(2, 18), direction=1)
],
1: [
Waypoint(position=(3, 18), direction=3),
Waypoint(position=(3, 17), direction=3),
Waypoint(position=(3, 16), direction=3),
Waypoint(position=(2, 16), direction=0),
Waypoint(position=(2, 15), direction=3),
Waypoint(position=(2, 14), direction=3),
Waypoint(position=(2, 13), direction=3),
Waypoint(position=(2, 12), direction=3),
Waypoint(position=(2, 11), direction=3),
Waypoint(position=(2, 10), direction=3),
Waypoint(position=(2, 9), direction=3),
Waypoint(position=(2, 8), direction=3),
Waypoint(position=(2, 7), direction=3),
Waypoint(position=(2, 6), direction=3),
Waypoint(position=(2, 5), direction=3),
Waypoint(position=(2, 4), direction=3),
Waypoint(position=(2, 3), direction=3),
Waypoint(position=(2, 2), direction=3),
Waypoint(position=(2, 1), 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 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, _ = RailEnvPersister.load_new("Level_distance_map_shortest_path.mpk",
"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])