diff --git a/flatland/envs/generators.py b/flatland/envs/generators.py
index e3c978e192b315f2bbeb26cd776b822505ca9987..525db36e8c5a09a451c0c59c1d03f1352f66e827 100644
--- a/flatland/envs/generators.py
+++ b/flatland/envs/generators.py
@@ -1,5 +1,4 @@
import warnings
-from enum import IntEnum
import msgpack
import numpy as np
@@ -543,414 +542,6 @@ def random_rail_generator(cell_type_relative_proportion=[1.0] * 11):
return generator
-def realistic_rail_generator(nr_start_goal=1, seed=0, add_max_dead_end=4, goals_only_in_dead_end=False,
- two_track_back_bone=True):
- """
- Parameters
- -------
- width : int
- The width (number of cells) of the grid to generate.
- height : int
- The height (number of cells) of the grid to generate.
-
- Returns
- -------
- numpy.ndarray of type numpy.uint16
- The matrix with the correct 16-bit bitmaps for each cell.
-
-
-
- transition_list = [int('0000000000000000', 2), # empty cell - Case 0
- int('1000000000100000', 2), # Case 1 - straight
- int('1001001000100000', 2), # Case 2 - simple switch
- int('1000010000100001', 2), # Case 3 - diamond drossing
- int('1001011000100001', 2), # Case 4 - single slip
- int('1100110000110011', 2), # Case 5 - double slip
- int('0101001000000010', 2), # Case 6 - symmetrical
- int('0010000000000000', 2), # Case 7 - dead end
- int('0100000000000010', 2), # Case 1b (8) - simple turn right
- int('0001001000000000', 2), # Case 1c (9) - simple turn left
- int('1100000000100010', 2)] # Case 2b (10) - simple switch mirrored
-
- """
-
- def min_max_cut(min_v, max_v, v):
- return max(min_v, min(max_v, v))
-
- def add_rail(width, height, grid_map, pt_from, pt_via, pt_to, bAddRemove=True):
- gRCTrans = np.array([[-1, 0], [0, 1], [1, 0], [0, -1]]) # NESW in RC
-
- lrcStroke = [[min_max_cut(0, height - 1, pt_from[0]),
- min_max_cut(0, width - 1, pt_from[1])],
- [min_max_cut(0, height - 1, pt_via[0]),
- min_max_cut(0, width - 1, pt_via[1])],
- [min_max_cut(0, height - 1, pt_to[0]),
- min_max_cut(0, width - 1, pt_to[1])]]
-
- rc3Cells = np.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(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
- # 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.
- grid_map.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
- grid_map.set_transition((*rcMiddle, mirror(liTrans[1])),
- mirror(liTrans[0]), bAddRemove, remove_deadends=not bDeadend)
-
- def make_switch_w_e(width, height, grid_map, center):
- # e -> w
- start = (center[0] + 1, center[1] - 1)
- via = (center[0], center[1] - 1)
- goal = (center[0], center[1])
- add_rail(width, height, grid_map, start, via, goal)
- start = (center[0], center[1] - 1)
- via = (center[0] + 1, center[1] - 1)
- goal = (center[0] + 1, center[1] - 2)
- add_rail(width, height, grid_map, start, via, goal)
-
- def make_switch_e_w(width, height, grid_map, center):
- # e -> w
- start = (center[0] + 1, center[1])
- via = (center[0] + 1, center[1] - 1)
- goal = (center[0], center[1] - 1)
- add_rail(width, height, grid_map, start, via, goal)
- start = (center[0] + 1, center[1] - 1)
- via = (center[0], center[1] - 1)
- goal = (center[0], center[1] - 2)
- add_rail(width, height, grid_map, start, via, goal)
-
- class Grid4TransitionsEnum(IntEnum):
- NORTH = 0
- EAST = 1
- SOUTH = 2
- WEST = 3
-
- @staticmethod
- def to_char(int: int):
- return {0: 'N',
- 1: 'E',
- 2: 'S',
- 3: 'W'}[int]
-
- def generator(width, height, num_agents, num_resets=0):
-
- if num_agents > nr_start_goal:
- num_agents = nr_start_goal
- print("complex_rail_generator: num_agents > nr_start_goal, changing num_agents")
- rail_trans = RailEnvTransitions()
- grid_map = GridTransitionMap(width=width, height=height, transitions=rail_trans)
- rail_array = grid_map.grid
- rail_array.fill(0)
-
- np.random.seed(seed + num_resets)
-
- max_n_track_seg = np.random.choice(np.arange(3, int(height / 2))) + int(two_track_back_bone)
- x_offsets = np.arange(0, height, max_n_track_seg).astype(int)
-
- agents_positions = []
- agents_directions = []
- agents_targets = []
-
- for off_set_loop in range(len(x_offsets)):
- off_set = x_offsets[off_set_loop]
- # second track
- data = np.arange(4, width - 4)
- n_track_seg = np.random.choice([1, 2, 3])
-
- track_2 = False
- if two_track_back_bone:
- if off_set + 1 < height:
- start_track = (off_set + 1, int((off_set_loop) % 2) * int(two_track_back_bone))
- goal_track = (off_set + 1, width - 1 - int((off_set_loop + 1) % 2) * int(two_track_back_bone))
- new_path = connect_rail(rail_trans, rail_array, start_track, goal_track)
- if len(new_path):
- track_2 = True
-
- start_track = (off_set, int((off_set_loop + 1) % 2) * int(two_track_back_bone) * int(track_2))
- goal_track = (off_set, width - 1 - int((off_set_loop) % 2) * int(two_track_back_bone) * int(track_2))
- new_path = connect_rail(rail_trans, rail_array, start_track, goal_track)
-
- if track_2:
- if np.random.random() < 0.75:
- c = (off_set, 3)
- if np.random.random() < 0.5:
- make_switch_e_w(width, height, grid_map, c)
- else:
- make_switch_w_e(width, height, grid_map, c)
- if np.random.random() < 0.5:
- c = (off_set, width - 3)
- if np.random.random() < 0.5:
- make_switch_e_w(width, height, grid_map, c)
- else:
- make_switch_w_e(width, height, grid_map, c)
-
- # track one (full track : left right)
- for two_track_back_bone_loop in range(1 + int(track_2) * int(two_track_back_bone)):
- if off_set_loop > 0:
- if off_set_loop % 2 == 1:
- start_track = (
- x_offsets[off_set_loop - 1] + 1 + int(two_track_back_bone_loop),
- width - 1 - int(two_track_back_bone_loop))
- goal_track = (x_offsets[off_set_loop] - 1 + int(two_track_back_bone) * int(track_2) - int(
- two_track_back_bone_loop),
- width - 1 - int(
- two_track_back_bone_loop))
- new_path = connect_rail(rail_trans, rail_array, start_track, goal_track)
-
- if (goal_track[1] - start_track[1]) > 1:
- add_pos = (
- int((start_track[0] + goal_track[0]) / 2), int((start_track[1] + goal_track[1]) / 2))
- agents_positions.append(add_pos)
- agents_directions.append(([1, 3][off_set_loop % 2]))
- if not goals_only_in_dead_end:
- agents_targets.append(add_pos)
-
- add_rail(width, height, grid_map,
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop),
- width - 2 - int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop),
- width - 1 - int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop) + 1,
- width - 1 - int(two_track_back_bone_loop)))
- add_rail(width, height, grid_map,
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2),
- width - 2 - int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2),
- width - 1 - int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2) - 1,
- width - 1 - int(two_track_back_bone_loop)))
- add_rail(width, height, grid_map,
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop),
- width - 1 - int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop) + 1,
- width - 1 - int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop) + 2,
- width - 1 - int(two_track_back_bone_loop)))
- add_rail(width, height, grid_map,
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2),
- width - 1 - int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2) - 1,
- width - 1 - int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2) - 2,
- width - 1 - int(two_track_back_bone_loop)))
-
- else:
- start_track = (
- x_offsets[off_set_loop - 1] + 1 + int(two_track_back_bone_loop),
- int(two_track_back_bone_loop))
- goal_track = (x_offsets[off_set_loop] - 1 + int(two_track_back_bone) * int(track_2) - int(
- two_track_back_bone_loop),
- int(two_track_back_bone_loop))
- new_path = connect_rail(rail_trans, rail_array, start_track, goal_track)
-
- if (goal_track[1] - start_track[1]) > 1:
- add_pos = (
- int((start_track[0] + goal_track[0]) / 2), int((start_track[1] + goal_track[1]) / 2))
- agents_positions.append(add_pos)
- agents_directions.append(([1, 3][off_set_loop % 2]))
- if not goals_only_in_dead_end:
- agents_targets.append(add_pos)
-
- add_rail(width, height, grid_map,
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop),
- 1 + int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop),
- 0 + int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop) + 1,
- 0 + int(two_track_back_bone_loop)))
- add_rail(width, height, grid_map,
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2),
- 1 + int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2),
- 0 + int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2) - 1,
- 0 + int(two_track_back_bone_loop)))
- add_rail(width, height, grid_map,
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop),
- 0 + int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop) + 1,
- 0 + int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop - 1] + int(two_track_back_bone_loop) + 2,
- 0 + int(two_track_back_bone_loop)))
- add_rail(width, height, grid_map,
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2),
- 0 + int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2) - 1,
- 0 + int(two_track_back_bone_loop)),
- (x_offsets[off_set_loop] - int(two_track_back_bone_loop) + int(
- two_track_back_bone) * int(track_2) - 2,
- 0 + int(two_track_back_bone_loop)))
-
- for nbr_track_loop in range(max_n_track_seg - 1):
- n_track_seg = 1
- if len(data) < 2 * n_track_seg + 1:
- break
- x = np.sort(np.random.choice(data, 2 * n_track_seg, False)).astype(int)
- data = []
- for x_loop in range(int(len(x) / 2)):
- start = (
- max(0, min(off_set + nbr_track_loop + 1, height - 1)), max(0, min(x[2 * x_loop], width - 1)))
- goal = (
- max(0, min(off_set + nbr_track_loop + 1, height - 1)),
- max(0, min(x[2 * x_loop + 1], width - 1)))
-
- d = np.arange(x[2 * x_loop] + 1, x[2 * x_loop + 1] - 1)
- data.extend(d)
-
- new_path = connect_rail(rail_trans, rail_array, start, goal)
- if len(new_path) > 0:
- c = (off_set + nbr_track_loop, x[2 * x_loop] + 1)
- make_switch_e_w(width, height, grid_map, c)
- c = (off_set + nbr_track_loop, x[2 * x_loop + 1] + 1)
- make_switch_w_e(width, height, grid_map, c)
-
- add_pos = (int((start[0] + goal[0]) / 2), int((start[1] + goal[1]) / 2))
- agents_positions.append(add_pos)
- agents_directions.append(([1, 3][off_set_loop % 2]))
- add_pos = (int((start[0] + goal[0]) / 2), int((2 * start[1] + goal[1]) / 3))
- if not goals_only_in_dead_end:
- agents_targets.append(add_pos)
-
- for off_set_loop in range(len(x_offsets)):
- off_set = x_offsets[off_set_loop]
- pos_ys = np.random.choice(np.arange(width - 7) + 4, min(width - 7, add_max_dead_end), False)
- for pos_y in pos_ys:
- pos_x = off_set + 1 + int(two_track_back_bone)
- if pos_x < height - 1:
- ok = True
- for k in range(5):
- if two_track_back_bone:
- c = (pos_x - 1, pos_y - k + 2)
- ok &= grid_map.grid[c[0]][c[1]] == 1025
- c = (pos_x, pos_y - k + 2)
- ok &= grid_map.grid[c[0]][c[1]] == 0
- if ok:
- if np.random.random() < 0.5:
- start_track = (pos_x, pos_y)
- goal_track = (pos_x, pos_y - 2)
- new_path = connect_rail(rail_trans, rail_array, start_track, goal_track)
- if len(new_path) > 0:
- c = (pos_x - 1, pos_y - 1)
- make_switch_e_w(width, height, grid_map, c)
- add_pos = (
- int((goal_track[0] + start_track[0]) / 2),
- int((goal_track[1] + start_track[1]) / 2))
- agents_positions.append(add_pos)
- agents_directions.append(3)
- add_pos = (
- int((goal_track[0] + start_track[0]) / 2),
- int((goal_track[1] + start_track[1]) / 2))
- agents_targets.append(add_pos)
- else:
- start_track = (pos_x, pos_y)
- goal_track = (pos_x, pos_y - 2)
- new_path = connect_rail(rail_trans, rail_array, start_track, goal_track)
- if len(new_path) > 0:
- c = (pos_x - 1, pos_y + 1)
- make_switch_w_e(width, height, grid_map, c)
- add_pos = (
- int((goal_track[0] + start_track[0]) / 2),
- int((goal_track[1] + start_track[1]) / 2))
- agents_positions.append(add_pos)
- agents_directions.append(1)
- add_pos = (
- int((goal_track[0] + start_track[0]) / 2),
- int((goal_track[1] + start_track[1]) / 2))
- agents_targets.append(add_pos)
-
- agents_position = []
- agents_target = []
- agents_direction = []
-
- remove_a = []
- for a in range(len(agents_positions)):
- cell_transitions = grid_map.get_transitions(agents_positions[a][0], agents_positions[a][1],
- agents_directions[a])
- if np.sum(cell_transitions) == 0:
- for i in range(4):
- agents_directions[a] = i
- cell_transitions = grid_map.get_transitions(agents_positions[a][0], agents_positions[a][1],
- agents_directions[a])
- if np.sum(cell_transitions) != 0:
- break
- if np.sum(cell_transitions):
- remove_a.extend([a])
- for i in range(len(remove_a)):
- agents_positions.pop(i)
- agents_directions.pop(i)
-
- for a in range(min(len(agents_targets), num_agents)):
- t = np.random.choice(range(len(agents_targets)))
- tp = agents_targets[t]
- agents_targets.pop(t)
- agents_target.append((tp[0], tp[1]))
-
- if len(agents_positions) == 0:
- print("no more position left")
- break
-
- sel = np.random.choice(range(len(agents_positions)))
- # backward
- p = agents_positions[sel]
- d = agents_directions[sel]
- cnt = 0
- while (p[0] == tp[0] and p[1] == tp[1]):
- sel = np.random.choice(range(len(agents_positions)))
- # backward
- p = agents_positions[sel]
- d = agents_directions[sel]
- cnt += 1
- if cnt > 10:
- print("target postion == agent postion !")
- break
- agents_positions.pop(sel)
- agents_directions.pop(sel)
- agents_position.append((p[0], p[1]))
- agents_direction.append(d)
-
- return grid_map, agents_position, agents_direction, agents_target, [1.0] * len(agents_position)
-
- return generator
-
-
def sparse_rail_generator(num_cities=5, num_intersections=4, num_trainstations=2, min_node_dist=20, node_radius=2,
num_neighb=3, realistic_mode=False, enhance_intersection=False, seed=0):
"""
diff --git a/tests/test_flatland_env_sparse_rail_generator.py b/tests/test_flatland_env_sparse_rail_generator.py
index 1dbb788cf20b2a82e7fd55d3d94408f5ad29ac30..d59e684575e9410b2859bb011ecb835a267b1c36 100644
--- a/tests/test_flatland_env_sparse_rail_generator.py
+++ b/tests/test_flatland_env_sparse_rail_generator.py
@@ -1,35 +1,7 @@
-import os
-
-import numpy as np
-
-from flatland.envs.generators import sparse_rail_generator, realistic_rail_generator
+from flatland.envs.generators import sparse_rail_generator
from flatland.envs.observations import GlobalObsForRailEnv
from flatland.envs.rail_env import RailEnv
-from flatland.utils.rendertools import RenderTool, AgentRenderVariant
-
-
-def test_realistic_rail_generator(vizualization_folder_name=None):
- num_agents = np.random.randint(10, 30)
- env = RailEnv(width=np.random.randint(40, 80),
- height=np.random.randint(10, 20),
- rail_generator=realistic_rail_generator(nr_start_goal=num_agents + 1,
- seed=1,
- add_max_dead_end=4,
- two_track_back_bone=1 % 2 == 0),
- number_of_agents=num_agents,
- obs_builder_object=GlobalObsForRailEnv())
- # reset to initialize agents_static
- env_renderer = RenderTool(env, gl="PILSVG", agent_render_variant=AgentRenderVariant.ONE_STEP_BEHIND,
- screen_height=600,
- screen_width=800)
- env_renderer.render_env(show=True, show_observations=True, show_predictions=False)
- if vizualization_folder_name is not None:
- env_renderer.gl.save_image(
- os.path.join(
- vizualization_folder_name,
- "flatland_frame_{:04d}.png".format(0)
- ))
- env_renderer.close_window()
+from flatland.utils.rendertools import RenderTool
def test_sparse_rail_generator():