diff --git a/flatland/envs/rail_generators.py b/flatland/envs/rail_generators.py index de7804024dba7540786c90f9537f9babcede206b..147b1bd14a5943f87f4d78f37e9ea7e1722c9087 100644 --- a/flatland/envs/rail_generators.py +++ b/flatland/envs/rail_generators.py @@ -781,7 +781,7 @@ def sparse_rail_generator(num_cities=5, num_intersections=4, num_trainstations=2 node_available_target.append(len(train_stations[node_idx])) # Assign agents to slots - for agent_idx in range(nb_nodes): + for agent_idx in range(num_agents): avail_start_nodes = [idx for idx, val in enumerate(node_available_start) if val > 0] avail_target_nodes = [idx for idx, val in enumerate(node_available_target) if val > 0] start_node = np.random.choice(avail_start_nodes) diff --git a/flatland/envs/schedule_generators.py b/flatland/envs/schedule_generators.py index b3576a2bec77f75afc9331cc6c190649590a990c..8c1834901a1e132d6adfa54cdddb286aafbef2a7 100644 --- a/flatland/envs/schedule_generators.py +++ b/flatland/envs/schedule_generators.py @@ -68,6 +68,7 @@ def sparse_schedule_generator(speed_ratio_map: Mapping[float, float] = None) -> agents_position = [] agents_target = [] agents_direction = [] + for agent_idx in range(num_agents): # Set target for agent current_target_node = agent_start_targets_nodes[agent_idx][1]