def fix_inner_nodes(grid_map: GridTransitionMap, inner_node_pos: IntVector2D, rail_trans: RailEnvTransitions):
"""
Fix inner city nodes by connecting it to its neighbouring parallel track
:param grid_map:
:param inner_node_pos: inner city node to fix
:param rail_trans:
:return:
"""
corner_directions = []
for direction in range(4):
tmp_pos = get_new_position(inner_node_pos, direction)
if grid_map.grid[tmp_pos] > 0:
corner_directions.append(direction)
if len(corner_directions) == 2:
transition = 0
transition = rail_trans.set_transition(transition, mirror(corner_directions[0]), corner_directions[1], 1)
transition = rail_trans.set_transition(transition, mirror(corner_directions[1]), corner_directions[0], 1)
grid_map.grid[inner_node_pos] = transition
tmp_pos = get_new_position(inner_node_pos, corner_directions[0])
transition = grid_map.grid[tmp_pos]
transition = rail_trans.set_transition(transition, corner_directions[0], mirror(corner_directions[0]), 1)
grid_map.grid[tmp_pos] = transition
tmp_pos = get_new_position(inner_node_pos, corner_directions[1])
transition = grid_map.grid[tmp_pos]
transition = rail_trans.set_transition(transition, corner_directions[1], mirror(corner_directions[1]),
1)
grid_map.grid[tmp_pos] = transition
return
def mod_rail_2cells(self, lrcCells, bAddRemove=True, iCellToMod=0, bPop=False):
"""
Add transitions for rail between two cells
lrcCells -- list of two rc cells
bAddRemove -- whether to add (True) or remove (False) the transition
iCellToMod -- the index of the cell to modify: either 0 or 1
"""
rc2Cells = array(lrcCells[:2]) # the 2 cells
rcMod = rc2Cells[iCellToMod] # the cell which we will update
# get the row, col delta between the 2 cells, eg [-1,0] = North
rc1Trans = np.diff(rc2Cells, axis=0)
# get the direction index for the transition
liTrans = []
for rcTrans in rc1Trans:
iTrans = np.argwhere(np.all(self.gRCTrans - rcTrans == 0, axis=1))
if len(iTrans) > 0:
iTrans = iTrans[0][0]
liTrans.append(iTrans)
# check that we have one transition
if len(liTrans) == 1:
# Set the transition as a deadend
# The transition is going from cell 0 to cell 1.
if iCellToMod == 0:
# if 0, reverse the transition, we need to be entering cell 0
self.env.rail.set_transition((*rcMod, mirror(liTrans[0])), liTrans[0], bAddRemove)
else:
# if 1, the transition is entering cell 1
self.env.rail.set_transition((*rcMod, liTrans[0]), mirror(liTrans[0]), bAddRemove)
if bPop:
lrcCells.pop(0)
def mod_rail_3cells(self, lrcStroke, bAddRemove=True, bPop=True):
"""
Add transitions for rail spanning three cells.
lrcStroke -- list containing "stroke" of cells across grid
bAddRemove -- whether to add (True) or remove (False) the transition
The transition is added to or removed from the 2nd cell, consistent with
entering from the 1st cell, and exiting into the 3rd.
Both the forward and backward transitions are added,
eg rcCells [(3,4), (2,4), (2,5)] would result in the transitions
N->E and W->S in cell (2,4).
"""
rc3Cells = array(lrcStroke[:3]) # the 3 cells
rcMiddle = rc3Cells[1] # the middle cell which we will update
bDeadend = np.all(
lrcStroke[0] == lrcStroke[2]) # deadend means cell 0 == cell 2
# get the 2 row, col deltas between the 3 cells, eg [[-1,0],[0,1]] = North, East
rc2Trans = np.diff(rc3Cells, axis=0)
# get the direction index for the 2 transitions
liTrans = []
for rcTrans in rc2Trans:
# gRCTrans - rcTrans gives an array of vector differences between our rcTrans
# and the 4 directions stored in gRCTrans.
# Where the vector difference is zero, we have a match...
# np.all detects where the whole row,col vector is zero.
# argwhere gives the index of the zero vector, ie the direction index
iTrans = np.argwhere(np.all(self.gRCTrans - rcTrans == 0, axis=1))
if len(iTrans) > 0:
iTrans = iTrans[0][0]
liTrans.append(iTrans)
# check that we have two transitions
if len(liTrans) == 2:
# Set the transition
# If this transition spans 3 cells, it is not a deadend, so remove any deadends.
# The user will need to resolve any conflicts.
self.env.rail.set_transition((*rcMiddle, liTrans[0]),
liTrans[1],
bAddRemove,
remove_deadends=not bDeadend)
# Also set the reverse transition
# use the reversed outbound transition for inbound
# and the reversed inbound transition for outbound
self.env.rail.set_transition((*rcMiddle, mirror(liTrans[1])),
mirror(liTrans[0]),
bAddRemove,
remove_deadends=not bDeadend)
if bPop:
lrcStroke.pop(0) # remove the first cell in the stroke
def connect_straight_line_in_grid_map_(
grid_map: GridTransitionMap, path: IntVector2DArray,
rail_trans: RailEnvTransitions) -> IntVector2DArray:
"""
Generates a straight rail line from start cell to end cell.
Diagonal lines are not allowed
:param rail_trans:
:param grid_map:
:param start: Cell coordinates for start of line
:param end: Cell coordinates for end of line
:return: A list of all cells in the path
"""
plen = len(path)
to_return = path.copy()
if plen < 2:
return []
if path[0] == path[1]:
path = path[1:]
to_return.remove(path[0])
current_dir = get_direction(path[0], path[1])
end_pos = path[-1]
for index in range(len(path) - 1):
current_pos = path[index]
new_pos = path[index + 1]
if current_pos == new_pos:
to_return.remove(current_pos)
continue
new_dir = get_direction(current_pos, new_pos)
try:
new_trans = grid_map.grid[current_pos]
except:
import pdb
pdb.set_trace()
new_trans = rail_trans.set_transition(new_trans, current_dir, new_dir,
1)
new_trans = rail_trans.set_transition(new_trans, mirror(new_dir),
mirror(current_dir), 1)
grid_map.grid[current_pos] = new_trans
if new_pos == end_pos:
new_trans_e = grid_map.grid[end_pos]
new_trans_e = rail_trans.set_transition(new_trans_e, new_dir,
new_dir, 1)
grid_map.grid[end_pos] = new_trans_e
current_dir = new_dir
return to_return
def connect_straight_line_in_grid_map(
grid_map: GridTransitionMap, start: IntVector2D, end: IntVector2D,
rail_trans: RailEnvTransitions) -> IntVector2DArray:
"""
Generates a straight rail line from start cell to end cell.
Diagonal lines are not allowed
:param rail_trans:
:param grid_map:
:param start: Cell coordinates for start of line
:param end: Cell coordinates for end of line
:return: A list of all cells in the path
"""
if not (start[0] == end[0] or start[1] == end[1]):
print("No straight line possible!")
return []
direction = direction_to_point(start, end)
if direction is Grid4TransitionsEnum.NORTH or direction is Grid4TransitionsEnum.SOUTH:
start_row = min(start[0], end[0])
end_row = max(start[0], end[0]) + 1
rows = np.arange(start_row, end_row)
length = np.abs(end[0] - start[0]) + 1
cols = np.repeat(start[1], length)
else: # Grid4TransitionsEnum.EAST or Grid4TransitionsEnum.WEST
start_col = min(start[1], end[1])
end_col = max(start[1], end[1]) + 1
cols = np.arange(start_col, end_col)
length = np.abs(end[1] - start[1]) + 1
rows = np.repeat(start[0], length)
path = list(zip(rows, cols))
for cell in path:
transition = grid_map.grid[cell]
transition = rail_trans.set_transition(transition, direction,
direction, 1)
transition = rail_trans.set_transition(transition, mirror(direction),
mirror(direction), 1)
grid_map.grid[cell] = transition
return path
def connect_rail_in_grid_map(grid_map: GridTransitionMap, start: IntVector2D, end: IntVector2D,
rail_trans: RailEnvTransitions,
a_star_distance_function: IntVector2DDistance = Vec2d.get_manhattan_distance,
flip_start_node_trans: bool = False, flip_end_node_trans: bool = False,
respect_transition_validity: bool = True, forbidden_cells: IntVector2DArray = None,
avoid_rail=False) -> IntVector2DArray:
"""
Creates a new path [start,end] in `grid_map.grid`, based on rail_trans, and
returns the path created as a list of positions.
:param avoid_rail:
:param rail_trans: basic rail transition object
:param grid_map: grid map
:param start: start position of rail
:param end: end position of rail
:param flip_start_node_trans: make valid start position by adding dead-end, empty start if False
:param flip_end_node_trans: make valid end position by adding dead-end, empty end if False
:param respect_transition_validity: Only draw rail maps if legal rail elements can be use, False, draw line without
respecting rail transitions.
:param a_star_distance_function: Define what distance function a-star should use
:param forbidden_cells: cells to avoid when drawing rail. Rail cannot go through this list of cells
:return: List of cells in the path
"""
path: IntVector2DArray = a_star(grid_map, start, end, a_star_distance_function, avoid_rail,
respect_transition_validity,
forbidden_cells)
if len(path) < 2:
return []
current_dir = get_direction(path[0], path[1])
end_pos = path[-1]
for index in range(len(path) - 1):
current_pos = path[index]
new_pos = path[index + 1]
new_dir = get_direction(current_pos, new_pos)
new_trans = grid_map.grid[current_pos]
if index == 0:
if new_trans == 0:
# end-point
if flip_start_node_trans:
# need to flip direction because of how end points are defined
new_trans = rail_trans.set_transition(new_trans, mirror(current_dir), new_dir, 1)
else:
new_trans = 0
else:
# into existing rail
new_trans = rail_trans.set_transition(new_trans, current_dir, new_dir, 1)
else:
# set the forward path
new_trans = rail_trans.set_transition(new_trans, current_dir, new_dir, 1)
# set the backwards path
new_trans = rail_trans.set_transition(new_trans, mirror(new_dir), mirror(current_dir), 1)
grid_map.grid[current_pos] = new_trans
if new_pos == end_pos:
# setup end pos setup
new_trans_e = grid_map.grid[end_pos]
if new_trans_e == 0:
# end-point
if flip_end_node_trans:
new_trans_e = rail_trans.set_transition(new_trans_e, new_dir, mirror(new_dir), 1)
else:
new_trans_e = 0
else:
# into existing rail
new_trans_e = rail_trans.set_transition(new_trans_e, new_dir, new_dir, 1)
grid_map.grid[end_pos] = new_trans_e
current_dir = new_dir
return path
def generator(width: int,
height: int,
num_agents: int,
num_resets: int = 0,
np_random: RandomState = None) -> RailGenerator:
if num_agents > nr_start_goal:
num_agents = nr_start_goal
print(
"complex_rail_generator: num_agents > nr_start_goal, changing num_agents"
)
grid_map = GridTransitionMap(width=width,
height=height,
transitions=RailEnvTransitions())
rail_array = grid_map.grid
rail_array.fill(0)
# generate rail array
# step 1:
# - generate a start and goal position
# - validate min/max distance allowed
# - validate that start/goals are not placed too close to other start/goals
# - draw a rail from [start,goal]
# - if rail crosses existing rail then validate new connection
# - possibility that this fails to create a path to goal
# - on failure generate new start/goal
#
# step 2:
# - add more rails to map randomly between cells that have rails
# - validate all new rails, on failure don't add new rails
#
# step 3:
# - return transition map + list of [start_pos, start_dir, goal_pos] points
#
rail_trans = grid_map.transitions
start_goal = []
start_dir = []
nr_created = 0
created_sanity = 0
sanity_max = 9000
while nr_created < nr_start_goal and created_sanity < sanity_max:
all_ok = False
for _ in range(sanity_max):
start = (np_random.randint(0, height),
np_random.randint(0, width))
goal = (np_random.randint(0,
height), np_random.randint(0, width))
# check to make sure start,goal pos is empty?
if rail_array[goal] != 0 or rail_array[start] != 0:
continue
# check min/max distance
dist_sg = distance_on_rail(start, goal)
if dist_sg < min_dist:
continue
if dist_sg > max_dist:
continue
# check distance to existing points
sg_new = [start, goal]
def check_all_dist(sg_new):
"""
Function to check the distance betweens start and goal
:param sg_new: start and goal tuple
:return: True if distance is larger than 2, False otherwise
"""
for sg in start_goal:
for i in range(2):
for j in range(2):
dist = distance_on_rail(sg_new[i], sg[j])
if dist < 2:
return False
return True
if check_all_dist(sg_new):
all_ok = True
break
if not all_ok:
# we might as well give up at this point
break
new_path = connect_rail_in_grid_map(
grid_map,
start,
goal,
rail_trans,
Vec2d.get_chebyshev_distance,
flip_start_node_trans=True,
flip_end_node_trans=True,
respect_transition_validity=True,
forbidden_cells=None)
if len(new_path) >= 2:
nr_created += 1
start_goal.append([start, goal])
start_dir.append(
mirror(get_direction(new_path[0], new_path[1])))
else:
# after too many failures we will give up
created_sanity += 1
# add extra connections between existing rail
created_sanity = 0
nr_created = 0
while nr_created < nr_extra and created_sanity < sanity_max:
all_ok = False
for _ in range(sanity_max):
start = (np_random.randint(0, height),
np_random.randint(0, width))
goal = (np_random.randint(0,
height), np_random.randint(0, width))
# check to make sure start,goal pos are not empty
if rail_array[goal] == 0 or rail_array[start] == 0:
continue
else:
all_ok = True
break
if not all_ok:
break
new_path = connect_rail_in_grid_map(
grid_map,
start,
goal,
rail_trans,
Vec2d.get_chebyshev_distance,
flip_start_node_trans=True,
flip_end_node_trans=True,
respect_transition_validity=True,
forbidden_cells=None)
if len(new_path) >= 2:
nr_created += 1
else:
# after too many failures we will give up
created_sanity += 1
return grid_map, {
'agents_hints': {
'start_goal': start_goal,
'start_dir': start_dir
}
}
def generator() -> RailGenerator:
"""
Arguments are ignored and taken directly from the curriculum except the np_random: RandomState which is the last
argument (args[-1])
"""
if curriculum.get("n_agents") > curriculum.get("n_cities"):
raise Exception("complex_rail_generator: n_agents > n_cities!")
grid_map = GridTransitionMap(width=curriculum.get("x_dim"),
height=curriculum.get("y_size"),
transitions=RailEnvTransitions())
rail_array = grid_map.grid
rail_array.fill(0)
# generate rail array
# step 1:
# - generate a start and goal position
# - validate min/max distance allowed
# - validate that start/goals are not placed too close to other start/goals
# - draw a rail from [start,goal]
# - if rail crosses existing rail then validate new connection
# - possibility that this fails to create a path to goal
# - on failure generate new start/goal
#
# step 2:
# - add more rails to map randomly between cells that have rails
# - validate all new rails, on failure don't add new rails
#
# step 3:
# - return transition map + list of [start_pos, start_dir, goal_pos] points
#
rail_trans = grid_map.transitions
start_goal = []
start_dir = []
nr_created = 0
created_sanity = 0
sanity_max = 9000
free_cells = set([(r, c) for r, row in enumerate(rail_array)
for c, col in enumerate(row) if col == 0])
while nr_created < curriculum.get(
"n_cities") and created_sanity < sanity_max:
all_ok = False
if len(free_cells) == 0:
break
for _ in range(sanity_max):
start = random.sample(free_cells, 1)[0]
goal = random.sample(free_cells, 1)[0]
# check min/max distance
dist_sg = distance_on_rail(start, goal)
if dist_sg < curriculum.get("min_dist"):
continue
if dist_sg > curriculum.get("max_dist"):
continue
# check distance to existing points
sg_new = [start, goal]
def check_all_dist():
"""
Function to check the distance betweens start and goal
:param sg_new: start and goal tuple
:return: True if distance is larger than 2, False otherwise
"""
for sg in start_goal:
for i in range(2):
for j in range(2):
dist = distance_on_rail(sg_new[i], sg[j])
if dist < 2:
return False
return True
if check_all_dist():
all_ok = True
free_cells.remove(start)
free_cells.remove(goal)
break
if not all_ok:
# we might as well give up at this point
break
new_path = connect_rail_in_grid_map(
grid_map,
start,
goal,
rail_trans,
Vec2d.get_chebyshev_distance,
flip_start_node_trans=True,
flip_end_node_trans=True,
respect_transition_validity=True,
forbidden_cells=None)
if len(new_path) >= 2:
nr_created += 1
start_goal.append([start, goal])
start_dir.append(
mirror(get_direction(new_path[0], new_path[1])))
else:
# after too many failures we will give up
created_sanity += 1
# add extra connections between existing rail
created_sanity = 0
nr_created = 0
while nr_created < curriculum.get(
"n_extra") and created_sanity < sanity_max:
if len(free_cells) == 0:
break
for _ in range(sanity_max):
start = random.sample(free_cells, 1)[0]
goal = random.sample(free_cells, 1)[0]
new_path = connect_rail_in_grid_map(
grid_map,
start,
goal,
rail_trans,
Vec2d.get_chebyshev_distance,
flip_start_node_trans=True,
flip_end_node_trans=True,
respect_transition_validity=True,
forbidden_cells=None)
if len(new_path) >= 2:
nr_created += 1
else:
# after too many failures we will give up
created_sanity += 1
return grid_map, {
'agents_hints': {
'start_goal': start_goal,
'start_dir': start_dir
}
}
