schedule_generators.py 15.8 KB
Newer Older
u214892's avatar
u214892 committed
1
"""Schedule generators (railway undertaking, "EVU")."""
2
import warnings
u214892's avatar
u214892 committed
3 4 5
from typing import Tuple, List, Callable, Mapping, Optional, Any

import numpy as np
6
from numpy.random.mtrand import RandomState
u214892's avatar
u214892 committed
7 8 9

from flatland.core.grid.grid4_utils import get_new_position
from flatland.core.transition_map import GridTransitionMap
u229589's avatar
u229589 committed
10
from flatland.envs.agent_utils import EnvAgent
11
from flatland.envs.schedule_utils import Schedule
12
from flatland.envs import persistence
u214892's avatar
u214892 committed
13 14

AgentPosition = Tuple[int, int]
15
ScheduleGenerator = Callable[[GridTransitionMap, int, Optional[Any], Optional[int]], Schedule]
u214892's avatar
u214892 committed
16 17


18
def speed_initialization_helper(nb_agents: int, speed_ratio_map: Mapping[float, float] = None,
19
                                seed: int = None, np_random: RandomState = None) -> List[float]:
u214892's avatar
u214892 committed
20 21
    """
    Parameters
u214892's avatar
u214892 committed
22
    ----------
u214892's avatar
u214892 committed
23 24 25 26 27 28 29 30 31 32 33 34
    nb_agents : int
        The number of agents to generate a speed for
    speed_ratio_map : Mapping[float,float]
        A map of speeds mappint to their ratio of appearance. The ratios must sum up to 1.

    Returns
    -------
    List[float]
        A list of size nb_agents of speeds with the corresponding probabilistic ratios.
    """
    if speed_ratio_map is None:
        return [1.0] * nb_agents
Erik Nygren's avatar
Erik Nygren committed
35

u214892's avatar
u214892 committed
36 37 38 39
    nb_classes = len(speed_ratio_map.keys())
    speed_ratio_map_as_list: List[Tuple[float, float]] = list(speed_ratio_map.items())
    speed_ratios = list(map(lambda t: t[1], speed_ratio_map_as_list))
    speeds = list(map(lambda t: t[0], speed_ratio_map_as_list))
40
    return list(map(lambda index: speeds[index], np_random.choice(nb_classes, nb_agents, p=speed_ratios)))
u214892's avatar
u214892 committed
41 42


43 44 45 46 47 48 49 50 51 52 53 54 55 56
class BaseSchedGen(object):
    def __init__(self, speed_ratio_map: Mapping[float, float] = None, seed: int = 1):
        self.speed_ratio_map = speed_ratio_map
        self.seed = seed

    def generate(self, rail: GridTransitionMap, num_agents: int, hints: Any=None, num_resets: int = 0,
        np_random: RandomState = None) -> Schedule:
        pass

    def __call__(self, *args, **kwargs):
        return self.generate(*args, **kwargs)



Erik Nygren's avatar
Erik Nygren committed
57
def complex_schedule_generator(speed_ratio_map: Mapping[float, float] = None, seed: int = 1) -> ScheduleGenerator:
58
    """
59

60 61 62 63 64 65 66 67
    Generator used to generate the levels of Round 1 in the Flatland Challenge. It can only be used together
    with complex_rail_generator. It places agents at end and start points provided by the rail generator.
    It assigns speeds to the different agents according to the speed_ratio_map
    :param speed_ratio_map: Speed ratios of all agents. They are probabilities of all different speeds and have to
            add up to 1.
    :param seed: Initiate random seed generator
    :return:
    """
68

69 70
    def generator(rail: GridTransitionMap, num_agents: int, hints: Any = None, num_resets: int = 0,
                  np_random: RandomState = None) -> Schedule:
71 72 73 74 75 76 77 78 79
        """

        The generator that assigns tasks to all the agents
        :param rail: Rail infrastructure given by the rail_generator
        :param num_agents: Number of agents to include in the schedule
        :param hints: Hints provided by the rail_generator These include positions of start/target positions
        :param num_resets: How often the generator has been reset.
        :return: Returns the generator to the rail constructor
        """
80
        # Todo: Remove parameters and variables not used for next version, Issue: <https://gitlab.aicrowd.com/flatland/flatland/issues/305>
81 82
        _runtime_seed = seed + num_resets

u214892's avatar
u214892 committed
83 84 85 86 87 88 89
        start_goal = hints['start_goal']
        start_dir = hints['start_dir']
        agents_position = [sg[0] for sg in start_goal[:num_agents]]
        agents_target = [sg[1] for sg in start_goal[:num_agents]]
        agents_direction = start_dir[:num_agents]

        if speed_ratio_map:
90
            speeds = speed_initialization_helper(num_agents, speed_ratio_map, seed=_runtime_seed, np_random=np_random)
u214892's avatar
u214892 committed
91 92
        else:
            speeds = [1.0] * len(agents_position)
93
        # Compute max number of steps with given schedule
Erik Nygren's avatar
Erik Nygren committed
94 95
        extra_time_factor = 1.5  # Factor to allow for more then minimal time
        max_episode_steps = int(extra_time_factor * rail.height * rail.width)
u214892's avatar
u214892 committed
96

97
        return Schedule(agent_positions=agents_position, agent_directions=agents_direction,
98 99
                        agent_targets=agents_target, agent_speeds=speeds, agent_malfunction_rates=None,
                        max_episode_steps=max_episode_steps)
u214892's avatar
u214892 committed
100 101 102 103

    return generator


Erik Nygren's avatar
Erik Nygren committed
104
def sparse_schedule_generator(speed_ratio_map: Mapping[float, float] = None, seed: int = 1) -> ScheduleGenerator:
105 106 107 108
    return SparseSchedGen(speed_ratio_map, seed)


class SparseSchedGen(BaseSchedGen):
109 110 111 112 113 114 115 116
    """

    This is the schedule generator which is used for Round 2 of the Flatland challenge. It produces schedules
    to railway networks provided by sparse_rail_generator.
    :param speed_ratio_map: Speed ratios of all agents. They are probabilities of all different speeds and have to
            add up to 1.
    :param seed: Initiate random seed generator
    """
117

118
    def generate(self, rail: GridTransitionMap, num_agents: int, hints: Any = None, num_resets: int = 0,
119
                  np_random: RandomState = None) -> Schedule:
120 121 122 123 124 125 126 127 128
        """

        The generator that assigns tasks to all the agents
        :param rail: Rail infrastructure given by the rail_generator
        :param num_agents: Number of agents to include in the schedule
        :param hints: Hints provided by the rail_generator These include positions of start/target positions
        :param num_resets: How often the generator has been reset.
        :return: Returns the generator to the rail constructor
        """
129

130
        _runtime_seed = self.seed + num_resets
131

132
        train_stations = hints['train_stations']
Erik Nygren's avatar
Erik Nygren committed
133 134
        city_positions = hints['city_positions']
        city_orientation = hints['city_orientations']
135
        max_num_agents = hints['num_agents']
136
        city_orientations = hints['city_orientations']
137 138 139
        if num_agents > max_num_agents:
            num_agents = max_num_agents
            warnings.warn("Too many agents! Changes number of agents.")
140 141 142 143
        # Place agents and targets within available train stations
        agents_position = []
        agents_target = []
        agents_direction = []
144

145 146 147 148 149 150 151
        for agent_idx in range(num_agents):
            infeasible_agent = True
            tries = 0
            while infeasible_agent:
                tries += 1
                infeasible_agent = False
                # Set target for agent
152
                city_idx = np_random.choice(len(city_positions), 2, replace=False)
Erik Nygren's avatar
Erik Nygren committed
153 154
                start_city = city_idx[0]
                target_city = city_idx[1]
155

156 157
                start_idx = np_random.choice(np.arange(len(train_stations[start_city])))
                target_idx = np_random.choice(np.arange(len(train_stations[target_city])))
158
                start = train_stations[start_city][start_idx]
159
                target = train_stations[target_city][target_idx]
Erik Nygren's avatar
Erik Nygren committed
160

161
                while start[1] % 2 != 0:
162
                    start_idx = np_random.choice(np.arange(len(train_stations[start_city])))
163 164
                    start = train_stations[start_city][start_idx]
                while target[1] % 2 != 1:
165
                    target_idx = np_random.choice(np.arange(len(train_stations[target_city])))
166
                    target = train_stations[target_city][target_idx]
Erik Nygren's avatar
Erik Nygren committed
167 168
                possible_orientations = [city_orientation[start_city],
                                         (city_orientation[start_city] + 2) % 4]
169
                agent_orientation = np_random.choice(possible_orientations)
170 171 172 173 174 175 176
                if not rail.check_path_exists(start[0], agent_orientation, target[0]):
                    agent_orientation = (agent_orientation + 2) % 4
                if not (rail.check_path_exists(start[0], agent_orientation, target[0])):
                    infeasible_agent = True
                if tries >= 100:
                    warnings.warn("Did not find any possible path, check your parameters!!!")
                    break
177
            agents_position.append((start[0][0], start[0][1]))
Erik Nygren's avatar
Erik Nygren committed
178
            agents_target.append((target[0][0], target[0][1]))
179

Erik Nygren's avatar
Erik Nygren committed
180
            agents_direction.append(agent_orientation)
181
            # Orient the agent correctly
182

183 184
        if self.speed_ratio_map:
            speeds = speed_initialization_helper(num_agents, self.speed_ratio_map, seed=_runtime_seed, np_random=np_random)
185 186
        else:
            speeds = [1.0] * len(agents_position)
Erik Nygren's avatar
Erik Nygren committed
187 188 189

        # We add multiply factors to the max number of time steps to simplify task in Flatland challenge.
        # These factors might change in the future.
190 191 192 193
        timedelay_factor = 4
        alpha = 2
        max_episode_steps = int(
            timedelay_factor * alpha * (rail.width + rail.height + num_agents / len(city_positions)))
194

195
        return Schedule(agent_positions=agents_position, agent_directions=agents_direction,
196 197
                        agent_targets=agents_target, agent_speeds=speeds, agent_malfunction_rates=None,
                        max_episode_steps=max_episode_steps)
198 199


200 201
def random_schedule_generator(speed_ratio_map: Mapping[float, float] = None, seed: int = 1) -> ScheduleGenerator:
    return RandomSchedGen(speed_ratio_map, seed)
202

203 204 205

class RandomSchedGen(BaseSchedGen):
    
u214892's avatar
u214892 committed
206
    """
u214892's avatar
u214892 committed
207
    Given a `rail` GridTransitionMap, return a random placement of agents (initial position, direction and target).
u214892's avatar
u214892 committed
208 209

    Parameters
u214892's avatar
u214892 committed
210
    ----------
u214892's avatar
u214892 committed
211 212 213
        speed_ratio_map : Optional[Mapping[float, float]]
            A map of speeds mapping to their ratio of appearance. The ratios must sum up to 1.

u214892's avatar
u214892 committed
214 215 216
    Returns
    -------
        Tuple[List[Tuple[int,int]], List[Tuple[int,int]], List[Tuple[int,int]], List[float]]
u214892's avatar
u214892 committed
217
            initial positions, directions, targets speeds
u214892's avatar
u214892 committed
218 219
    """

220
    def generate(self, rail: GridTransitionMap, num_agents: int, hints: Any = None, num_resets: int = 0,
221
                  np_random: RandomState = None) -> Schedule:
222
        _runtime_seed = self.seed + num_resets
Erik Nygren's avatar
Erik Nygren committed
223

u214892's avatar
u214892 committed
224 225 226 227 228 229
        valid_positions = []
        for r in range(rail.height):
            for c in range(rail.width):
                if rail.get_full_transitions(r, c) > 0:
                    valid_positions.append((r, c))
        if len(valid_positions) == 0:
230
            return Schedule(agent_positions=[], agent_directions=[],
231
                            agent_targets=[], agent_speeds=[], agent_malfunction_rates=None, max_episode_steps=0)
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
232 233

        if len(valid_positions) < num_agents:
234
            warnings.warn("schedule_generators: len(valid_positions) < num_agents")
235
            return Schedule(agent_positions=[], agent_directions=[],
236
                            agent_targets=[], agent_speeds=[], agent_malfunction_rates=None, max_episode_steps=0)
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
237

238
        agents_position_idx = [i for i in np_random.choice(len(valid_positions), num_agents, replace=False)]
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
239
        agents_position = [valid_positions[agents_position_idx[i]] for i in range(num_agents)]
240
        agents_target_idx = [i for i in np_random.choice(len(valid_positions), num_agents, replace=False)]
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
241
        agents_target = [valid_positions[agents_target_idx[i]] for i in range(num_agents)]
242
        update_agents = np.zeros(num_agents)
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
243

u214892's avatar
u214892 committed
244
        re_generate = True
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
245
        cnt = 0
246
        while re_generate:
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
247
            cnt += 1
248
            if cnt > 1:
249 250 251
                print("re_generate cnt={}".format(cnt))
            if cnt > 1000:
                raise Exception("After 1000 re_generates still not success, giving up.")
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
252 253 254
            # update position
            for i in range(num_agents):
                if update_agents[i] == 1:
255
                    x = np.setdiff1d(np.arange(len(valid_positions)), agents_position_idx)
256
                    agents_position_idx[i] = np_random.choice(x)
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
257
                    agents_position[i] = valid_positions[agents_position_idx[i]]
258
                    x = np.setdiff1d(np.arange(len(valid_positions)), agents_target_idx)
259
                    agents_target_idx[i] = np_random.choice(x)
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
260 261
                    agents_target[i] = valid_positions[agents_target_idx[i]]
            update_agents = np.zeros(num_agents)
u214892's avatar
u214892 committed
262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278

            # agents_direction must be a direction for which a solution is
            # guaranteed.
            agents_direction = [0] * num_agents
            re_generate = False
            for i in range(num_agents):
                valid_movements = []
                for direction in range(4):
                    position = agents_position[i]
                    moves = rail.get_transitions(position[0], position[1], direction)
                    for move_index in range(4):
                        if moves[move_index]:
                            valid_movements.append((direction, move_index))

                valid_starting_directions = []
                for m in valid_movements:
                    new_position = get_new_position(agents_position[i], m[1])
279
                    if m[0] not in valid_starting_directions and rail.check_path_exists(new_position, m[1],
280
                                                                                        agents_target[i]):
u214892's avatar
u214892 committed
281 282 283
                        valid_starting_directions.append(m[0])

                if len(valid_starting_directions) == 0:
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
284
                    update_agents[i] = 1
Erik Nygren's avatar
Erik Nygren committed
285 286
                    warnings.warn(
                        "reset position for agent[{}]: {} -> {}".format(i, agents_position[i], agents_target[i]))
287
                    re_generate = True
Egli Adrian (IT-SCI-API-PFI)'s avatar
Egli Adrian (IT-SCI-API-PFI) committed
288
                    break
u214892's avatar
u214892 committed
289 290
                else:
                    agents_direction[i] = valid_starting_directions[
291
                        np_random.choice(len(valid_starting_directions), 1)[0]]
u214892's avatar
u214892 committed
292

293 294
        agents_speed = speed_initialization_helper(num_agents, self.speed_ratio_map, seed=_runtime_seed, 
            np_random=np_random)
295 296

        # Compute max number of steps with given schedule
Erik Nygren's avatar
Erik Nygren committed
297 298
        extra_time_factor = 1.5  # Factor to allow for more then minimal time
        max_episode_steps = int(extra_time_factor * rail.height * rail.width)
299

300
        return Schedule(agent_positions=agents_position, agent_directions=agents_direction,
301 302
                        agent_targets=agents_target, agent_speeds=agents_speed, agent_malfunction_rates=None,
                        max_episode_steps=max_episode_steps)
u214892's avatar
u214892 committed
303 304 305



u214892's avatar
u214892 committed
306
def schedule_from_file(filename, load_from_package=None) -> ScheduleGenerator:
u214892's avatar
u214892 committed
307 308 309 310
    """
    Utility to load pickle file

    Parameters
u214892's avatar
u214892 committed
311
    ----------
u214892's avatar
u214892 committed
312 313 314 315 316 317 318 319
    input_file : Pickle file generated by env.save() or editor

    Returns
    -------
    Tuple[List[Tuple[int,int]], List[Tuple[int,int]], List[Tuple[int,int]], List[float]]
        initial positions, directions, targets speeds
    """

320 321
    def generator(rail: GridTransitionMap, num_agents: int, hints: Any = None, num_resets: int = 0,
                  np_random: RandomState = None) -> Schedule:
322 323 324 325

        env_dict = persistence.RailEnvPersister.load_env_dict(filename, load_from_package=load_from_package)

        max_episode_steps = env_dict.get("max_episode_steps", 0)
326 327 328 329
        if (max_episode_steps==0):
            print("This env file has no max_episode_steps (deprecated) - setting to 100")
            max_episode_steps = 100
            
330 331
        agents = env_dict["agents"]

332 333
        #print("schedule generator from_file - agents: ", agents)

u214892's avatar
u214892 committed
334
        # setup with loaded data
u229589's avatar
u229589 committed
335
        agents_position = [a.initial_position for a in agents]
336 337 338 339

        # this logic is wrong - we should really load the initial_direction as the direction.
        #agents_direction = [a.direction for a in agents]
        agents_direction = [a.initial_direction for a in agents]
u229589's avatar
u229589 committed
340 341
        agents_target = [a.target for a in agents]
        agents_speed = [a.speed_data['speed'] for a in agents]
342 343

        # Malfunctions from here are not used.  They have their own generator.
344
        #agents_malfunction = [a.malfunction_data['malfunction_rate'] for a in agents]
u229589's avatar
u229589 committed
345

346
        return Schedule(agent_positions=agents_position, agent_directions=agents_direction,
347 348
                        agent_targets=agents_target, agent_speeds=agents_speed, 
                        agent_malfunction_rates=None,
349
                        max_episode_steps=max_episode_steps)
u214892's avatar
u214892 committed
350 351

    return generator