examples/custom_observation_example_02_SingleAgentNavigationObs.py

+import getopt
+import random
+import sys
+import time
+from typing import List
+
+import numpy as np
+
+from flatland.core.env_observation_builder import ObservationBuilder
+from flatland.core.grid.grid4_utils import get_new_position
+from flatland.envs.line_generators import sparse_line_generator
+from flatland.envs.rail_env import RailEnv
+from flatland.envs.rail_generators import sparse_rail_generator
+from flatland.utils.misc import str2bool
+from flatland.utils.rendertools import RenderTool
+
+random.seed(100)
+np.random.seed(100)
+
+
+class SingleAgentNavigationObs(ObservationBuilder):
+    """
+    We build a representation vector with 3 binary components, indicating which of the 3 available directions
+    for each agent (Left, Forward, Right) lead to the shortest path to its target.
+    E.g., if taking the Left branch (if available) is the shortest route to the agent's target, the observation vector
+    will be [1, 0, 0].
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def reset(self):
+        pass
+
+    def get(self, handle: int = 0) -> List[int]:
+        agent = self.env.agents[handle]
+
+        if agent.position:
+            possible_transitions = self.env.rail.get_transitions(*agent.position, agent.direction)
+        else:
+            possible_transitions = self.env.rail.get_transitions(*agent.initial_position, agent.direction)
+
+        num_transitions = np.count_nonzero(possible_transitions)
+
+        # Start from the current orientation, and see which transitions are available;
+        # organize them as [left, forward, right], relative to the current orientation
+        # If only one transition is possible, the forward branch is aligned with it.
+        if num_transitions == 1:
+            observation = [0, 1, 0]
+        else:
+            min_distances = []
+            for direction in [(agent.direction + i) % 4 for i in range(-1, 2)]:
+                if possible_transitions[direction]:
+                    new_position = get_new_position(agent.position, direction)
+                    min_distances.append(
+                        self.env.distance_map.get()[handle, new_position[0], new_position[1], direction])
+                else:
+                    min_distances.append(np.inf)
+
+            observation = [0, 0, 0]
+            observation[np.argmin(min_distances)] = 1
+
+        return observation
+
+
+def create_env():
+    nAgents = 1
+    n_cities = 2
+    max_rails_between_cities = 2
+    max_rails_in_city = 4
+    seed = 0
+    env = RailEnv(
+        width=30,
+        height=40,
+        rail_generator=sparse_rail_generator(
+            max_num_cities=n_cities,
+            seed=seed,
+            grid_mode=True,
+            max_rails_between_cities=max_rails_between_cities,
+            max_rail_pairs_in_city=max_rails_in_city
+        ),
+        line_generator=sparse_line_generator(),
+        number_of_agents=nAgents,
+        obs_builder_object=SingleAgentNavigationObs()
+    )
+    return env
+
+
+def custom_observation_example_02_SingleAgentNavigationObs(sleep_for_animation, do_rendering):
+    env = create_env()
+    obs, info = env.reset()
+
+    env_renderer = None
+    if do_rendering:
+        env_renderer = RenderTool(env)
+        env_renderer.render_env(show=True, frames=True, show_observations=False)
+
+    for step in range(100):
+        action = np.argmax(obs[0]) + 1
+        obs, all_rewards, done, _ = env.step({0: action})
+        print("Rewards: ", all_rewards, "  [done=", done, "]")
+
+        if env_renderer is not None:
+            env_renderer.render_env(show=True, frames=True, show_observations=True)
+        if sleep_for_animation:
+            time.sleep(0.1)
+        if done["__all__"]:
+            break
+    if env_renderer is not None:
+        env_renderer.close_window()
+
+
+def main(args):
+    try:
+        opts, args = getopt.getopt(args, "", ["sleep-for-animation=", "do_rendering=", ""])
+    except getopt.GetoptError as err:
+        print(str(err))  # will print something like "option -a not recognized"
+        sys.exit(2)
+    sleep_for_animation = True
+    do_rendering = True
+    for o, a in opts:
+        if o in ("--sleep-for-animation"):
+            sleep_for_animation = str2bool(a)
+        elif o in ("--do_rendering"):
+            do_rendering = str2bool(a)
+        else:
+            assert False, "unhandled option"
+
+    # execute example
+    custom_observation_example_02_SingleAgentNavigationObs(sleep_for_animation, do_rendering)
+
+
+if __name__ == '__main__':
+    if 'argv' in globals():
+        main(argv)
+    else:
+        main(sys.argv[1:])

examples/custom_observation_example_03_ObservePredictions.py

+import getopt
+import random
+import sys
+import time
+from typing import Optional, List, Dict
+
+import numpy as np
+
+from flatland.core.env import Environment
+from flatland.core.env_observation_builder import ObservationBuilder
+from flatland.core.grid.grid_utils import coordinate_to_position
+from flatland.envs.line_generators import sparse_line_generator
+from flatland.envs.predictions import ShortestPathPredictorForRailEnv
+from flatland.envs.rail_env import RailEnv
+from flatland.envs.rail_generators import sparse_rail_generator
+from flatland.utils.misc import str2bool
+from flatland.utils.ordered_set import OrderedSet
+from flatland.utils.rendertools import RenderTool
+
+random.seed(100)
+np.random.seed(100)
+
+
+class ObservePredictions(ObservationBuilder):
+    """
+    We use the provided ShortestPathPredictor to illustrate the usage of predictors in your custom observation.
+    """
+
+    def __init__(self, predictor):
+        super().__init__()
+        self.predictor = predictor
+
+    def reset(self):
+        pass
+
+    def get_many(self, handles: Optional[List[int]] = None) -> Dict[int, np.ndarray]:
+        '''
+        Because we do not want to call the predictor seperately for every agent we implement the get_many function
+        Here we can call the predictor just ones for all the agents and use the predictions to generate our observations
+        :param handles:
+        :return:
+        '''
+
+        self.predictions = self.predictor.get()
+
+        self.predicted_pos = {}
+
+        if handles is None:
+            handles = []
+
+        for t in range(len(self.predictions[0])):
+            pos_list = []
+            for a in handles:
+                pos_list.append(self.predictions[a][t][1:3])
+            # We transform (x,y) coodrinates to a single integer number for simpler comparison
+            self.predicted_pos.update({t: coordinate_to_position(self.env.width, pos_list)})
+
+        observations = super().get_many(handles)
+
+        return observations
+
+    def get(self, handle: int = 0) -> np.ndarray:
+        '''
+        Lets write a simple observation which just indicates whether or not the own predicted path
+        overlaps with other predicted paths at any time. This is useless for the task of navigation but might
+        help when looking for conflicts. A more complex implementation can be found in the TreeObsForRailEnv class
+
+        Each agent recieves an observation of length 10, where each element represents a prediction step and its value
+        is:
+         - 0 if no overlap is happening
+         - 1 where n i the number of other paths crossing the predicted cell
+
+        :param handle: handeled as an index of an agent
+        :return: Observation of handle
+        '''
+
+        observation = np.zeros(10)
+
+        # We are going to track what cells where considered while building the obervation and make them accesible
+        # For rendering
+
+        visited = OrderedSet()
+        for _idx in range(10):
+            # Check if any of the other prediction overlap with agents own predictions
+            x_coord = self.predictions[handle][_idx][1]
+            y_coord = self.predictions[handle][_idx][2]
+
+            # We add every observed cell to the observation rendering
+            visited.add((x_coord, y_coord))
+            if self.predicted_pos[_idx][handle] in np.delete(self.predicted_pos[_idx], handle, 0):
+                # We detect if another agent is predicting to pass through the same cell at the same predicted time
+                observation[handle] = 1
+
+        # This variable will be access by the renderer to visualize the observation
+        self.env.dev_obs_dict[handle] = visited
+
+        return observation
+
+    def set_env(self, env: Environment):
+        super().set_env(env)
+        if self.predictor:
+            self.predictor.set_env(self.env)
+
+
+def create_env(custom_obs_builder):
+    nAgents = 3
+    n_cities = 2
+    max_rails_between_cities = 4
+    max_rails_in_city = 2
+    seed = 0
+    env = RailEnv(
+        width=30,
+        height=30,
+        rail_generator=sparse_rail_generator(
+            max_num_cities=n_cities,
+            seed=seed,
+            grid_mode=True,
+            max_rails_between_cities=max_rails_between_cities,
+            max_rail_pairs_in_city=max_rails_in_city
+        ),
+        line_generator=sparse_line_generator(),
+        number_of_agents=nAgents,
+        obs_builder_object=custom_obs_builder
+    )
+    return env
+
+
+def custom_observation_example_03_ObservePredictions(sleep_for_animation, do_rendering):
+    # Initiate the Predictor
+    custom_predictor = ShortestPathPredictorForRailEnv(10)
+
+    # Pass the Predictor to the observation builder
+    custom_obs_builder = ObservePredictions(custom_predictor)
+
+    # Initiate Environment
+    env = create_env(custom_obs_builder)
+    obs, info = env.reset()
+
+    env_renderer = None
+    if do_rendering:
+        env_renderer = RenderTool(env)
+        # We render the initial step and show the obsered cells as colored boxes
+        env_renderer.render_env(show=True, frames=True, show_observations=True, show_predictions=False)
+
+    action_dict = {}
+    for step in range(100):
+        for a in range(env.get_num_agents()):
+            action = np.random.randint(0, 5)
+            action_dict[a] = action
+        obs, all_rewards, done, _ = env.step(action_dict)
+        print("Rewards: ", all_rewards, "  [done=", done, "]")
+        if env_renderer is not None:
+            env_renderer.render_env(show=True, frames=True, show_observations=True, show_predictions=False)
+        if sleep_for_animation:
+            time.sleep(0.5)
+
+        if done["__all__"]:
+            print("All done!")
+            break
+
+    if env_renderer is not None:
+        env_renderer.close_window()
+
+
+def main(args):
+    try:
+        opts, args = getopt.getopt(args, "", ["sleep-for-animation=", "do_rendering=", ""])
+    except getopt.GetoptError as err:
+        print(str(err))  # will print something like "option -a not recognized"
+        sys.exit(2)
+    sleep_for_animation = True
+    do_rendering = True
+    for o, a in opts:
+        if o in ("--sleep-for-animation"):
+            sleep_for_animation = str2bool(a)
+        elif o in ("--do_rendering"):
+            do_rendering = str2bool(a)
+        else:
+            assert False, "unhandled option"
+
+    # execute example
+    custom_observation_example_03_ObservePredictions(sleep_for_animation, do_rendering)
+
+if __name__ == '__main__':
+    if 'argv' in globals():
+        main(argv)
+    else:
+        main(sys.argv[1:])

examples/custom_railmap_example.py

+import getopt
 import random
+import sys
+import time
+from typing import Tuple
 
 import numpy as np
 
+from flatland.core.env_observation_builder import DummyObservationBuilder
 from flatland.core.grid.rail_env_grid import RailEnvTransitions
 from flatland.core.transition_map import GridTransitionMap
+from flatland.envs.line_generators import sparse_line_generator
 from flatland.envs.rail_env import RailEnv
+from flatland.envs.rail_generators import rail_from_grid_transition_map
+from flatland.utils.misc import str2bool
 from flatland.utils.rendertools import RenderTool
 
-random.seed(100)
-np.random.seed(100)
 
+def custom_rail_map() -> Tuple[GridTransitionMap, np.array]:
+    # We instantiate a very simple rail network on a 7x10 grid:
+    #  0 1 2 3 4 5 6 7 8 9  10
+    # 0        /-------------\
+    # 1        |             |
+    # 2        |             |
+    # 3 _ _ _ /_  _ _        |
+    # 4              \   ___ /
+    # 5               |/
+    # 6               |
+    # 7               |
+    transitions = RailEnvTransitions()
+    cells = transitions.transition_list
 
-def custom_rail_generator():
-    def generator(width, height, num_agents=0, num_resets=0):
-        rail_trans = RailEnvTransitions()
-        grid_map = GridTransitionMap(width=width, height=height, transitions=rail_trans)
-        rail_array = grid_map.grid
-        rail_array.fill(0)
-        new_tran = rail_trans.set_transition(1, 1, 1, 1)
-        print(new_tran)
-        agents_positions = []
-        agents_direction = []
-        agents_target = []
-        rail_array[0, 0] = new_tran
-        rail_array[0, 1] = new_tran
-        return grid_map, agents_positions, agents_direction, agents_target
+    empty = cells[0]
+    dead_end_from_south = cells[7]
+    right_turn_from_south = cells[8]
+    right_turn_from_west = transitions.rotate_transition(right_turn_from_south, 90)
+    right_turn_from_north = transitions.rotate_transition(right_turn_from_south, 180)
+    dead_end_from_west = transitions.rotate_transition(dead_end_from_south, 90)
+    dead_end_from_north = transitions.rotate_transition(dead_end_from_south, 180)
+    dead_end_from_east = transitions.rotate_transition(dead_end_from_south, 270)
+    vertical_straight = cells[1]
+    simple_switch_north_left = cells[2]
+    simple_switch_north_right = cells[10]
+    simple_switch_left_east = transitions.rotate_transition(simple_switch_north_left, 90)
+    horizontal_straight = transitions.rotate_transition(vertical_straight, 90)
+    double_switch_south_horizontal_straight = horizontal_straight + cells[6]
+    double_switch_north_horizontal_straight = transitions.rotate_transition(
+        double_switch_south_horizontal_straight, 180)
+    rail_map = np.array(
+        [[empty] * 3 + [right_turn_from_south] + [horizontal_straight] * 5 + [right_turn_from_west]] +
+        [[empty] * 3 + [vertical_straight] + [empty] * 5 + [vertical_straight]] * 2 +
+        [[dead_end_from_east] + [horizontal_straight] * 2 + [simple_switch_left_east] + [horizontal_straight] * 2 + [
+            right_turn_from_west] + [empty] * 2 + [vertical_straight]] +
+        [[empty] * 6 + [simple_switch_north_right] + [horizontal_straight] * 2 + [right_turn_from_north]] +
+        [[empty] * 6 + [vertical_straight] + [empty] * 3] +
+        [[empty] * 6 + [dead_end_from_north] + [empty] * 3], dtype=np.uint16)
+    rail = GridTransitionMap(width=rail_map.shape[1],
+                             height=rail_map.shape[0], transitions=transitions)
+    rail.grid = rail_map
+    city_positions = [(0, 3), (6, 6)]
+    train_stations = [
+        [((0, 3), 0)],
+        [((6, 6), 0)],
+    ]
+    city_orientations = [0, 2]
+    agents_hints = {'city_positions': city_positions,
+                    'train_stations': train_stations,
+                    'city_orientations': city_orientations
+                    }
+    optionals = {'agents_hints': agents_hints}
+    return rail, rail_map, optionals
 
-    return generator
 
+def create_env():
+    rail, rail_map, optionals = custom_rail_map()
+    env = RailEnv(width=rail_map.shape[1],
+                  height=rail_map.shape[0],
+                  rail_generator=rail_from_grid_transition_map(rail, optionals),
+                  line_generator=sparse_line_generator(),
+                  number_of_agents=2,
+                  obs_builder_object=DummyObservationBuilder(),
+                  )
+    return env
 
-env = RailEnv(width=6,
-              height=4,
-              rail_generator=custom_rail_generator(),
-              number_of_agents=1)
 
-env.reset()
+def custom_railmap_example(sleep_for_animation, do_rendering):
+    random.seed(100)
+    np.random.seed(100)
 
-env_renderer = RenderTool(env)
-env_renderer.renderEnv(show=True)
+    env = create_env()
+    env.reset()
 
-input("Press Enter to continue...")
+    if do_rendering:
+        env_renderer = RenderTool(env)
+        env_renderer.render_env(show=True, show_observations=False)
+        env_renderer.close_window()
+
+    if sleep_for_animation:
+        time.sleep(1)
+
+    # uncomment to keep the renderer open
+    # input("Press Enter to continue...")
+
+
+def main(args):
+    try:
+        opts, args = getopt.getopt(args, "", ["sleep-for-animation=", "do_rendering=", ""])
+    except getopt.GetoptError as err:
+        print(str(err))  # will print something like "option -a not recognized"
+        sys.exit(2)
+    sleep_for_animation = True
+    do_rendering = True
+    for o, a in opts:
+        if o in ("--sleep-for-animation"):
+            sleep_for_animation = str2bool(a)
+        elif o in ("--do_rendering"):
+            do_rendering = str2bool(a)
+        else:
+            assert False, "unhandled option"
+
+    # execute example
+    custom_railmap_example(sleep_for_animation, do_rendering)
+
+
+if __name__ == '__main__':
+    if 'argv' in globals():
+        main(argv)
+    else:
+        main(sys.argv[1:])

examples/example_basic_elements_test.py

-import random
-
-import numpy as np
-
-from examples.demo import Demo
-
-random.seed(1)
-np.random.seed(1)
-
-if __name__ == "__main__":
-    Demo.run_basic_elements_test()

examples/example_flatland_000.py

-import random
-
-import numpy as np
-
-from examples.demo import Demo
-
-random.seed(1)
-np.random.seed(1)
-
-if __name__ == "__main__":
-    Demo.run_example_flatland_000()

examples/example_flatland_001.py

-import random
-
-import numpy as np
-
-from examples.demo import Demo
-
-random.seed(1)
-np.random.seed(1)
-
-if __name__ == "__main__":
-    Demo.run_example_flatland_001()

examples/example_network_000.py

-import random
-
-import numpy as np
-
-from examples.demo import Demo
-
-random.seed(1)
-np.random.seed(1)
-
-if __name__ == "__main__":
-    Demo.run_example_network_000()

examples/example_network_001.py

-import random
-
-import numpy as np
-
-from examples.demo import Demo
-
-random.seed(1)
-np.random.seed(1)
-
-if __name__ == "__main__":
-    Demo.run_example_network_001()

examples/example_network_002.py

-import random
-
-import numpy as np
-
-from examples.demo import Demo
-
-random.seed(1)
-np.random.seed(1)
-
-if __name__ == "__main__":
-    Demo.run_example_network_002()

examples/example_network_003.py

-import random
-
-import numpy as np
-
-from examples.demo import Demo
-
-random.seed(1)
-np.random.seed(1)
-
-if __name__ == "__main__":
-    Demo.run_example_network_003()