diff --git a/flatland/envs/rail_generators.py b/flatland/envs/rail_generators.py index c100e3db88dcc8f2a42396baf8e6565cb14dcc48..b69d80a065670ebafb60da2972d31efb05add047 100644 --- a/flatland/envs/rail_generators.py +++ b/flatland/envs/rail_generators.py @@ -678,8 +678,8 @@ def sparse_rail_generator(num_cities=5, num_intersections=4, num_trainstations=2 if num_cities > 1: train_stations = [[] for i in range(num_cities)] built_num_trainstation = 0 - spot_found = True for station in range(num_trainstations): + spot_found = True trainstation_node = int(station / num_trainstations * num_cities) station_x = np.clip(node_positions[trainstation_node][0] + np.random.randint(-node_radius, node_radius), @@ -705,7 +705,7 @@ def sparse_rail_generator(num_cities=5, num_intersections=4, num_trainstations=2 if tries > 100: warnings.warn("Could not set trainstations, please change initial parameters!!!!") spot_found = False - break + if spot_found: train_stations[trainstation_node].append((station_x, station_y)) @@ -714,13 +714,12 @@ def sparse_rail_generator(num_cities=5, num_intersections=4, num_trainstations=2 (station_x, station_y)) # Check if connection was made if len(connection) == 0: - if len(train_stations) > 0: + if len(train_stations[trainstation_node]) > 0: train_stations[trainstation_node].pop(-1) else: built_num_trainstation += 1 # Adjust the number of agents if you could not build enough trainstations - if num_agents > built_num_trainstation: num_agents = built_num_trainstation warnings.warn("sparse_rail_generator: num_agents > nr_start_goal, changing num_agents")