added possibility to specify cell types relative proportions when generating random rails

examples/temporary_example.py

@@ -10,13 +10,34 @@ random.seed(1)
 np.random.seed(1)
 
 """
+transition_probability = [1.0,  # empty cell - Case 0
+                          3.0,  # Case 1 - straight
+                          1.0,  # Case 2 - simple switch
+                          3.0,  # Case 3 - diamond drossing
+                          2.0,  # Case 4 - single slip
+                          1.0,  # Case 5 - double slip
+                          1.0,  # Case 6 - symmetrical
+                          1.0]  # Case 7 - dead end
+"""
+transition_probability = [1.0,  # empty cell - Case 0
+                          1.0,  # Case 1 - straight
+                          1.0,  # Case 2 - simple switch
+                          1.0,  # Case 3 - diamond drossing
+                          1.0,  # Case 4 - single slip
+                          1.0,  # Case 5 - double slip
+                          1.0,  # Case 6 - symmetrical
+                          1.0]  # Case 7 - dead end
+
 # Example generate a random rail
-env = RailEnv(width=20, height=20, rail_generator=random_rail_generator, number_of_agents=10)
+env = RailEnv(width=20,
+              height=20,
+              rail_generator=random_rail_generator(cell_type_relative_proportion=transition_probability),
+              number_of_agents=10)
 env.reset()
 
 env_renderer = RenderTool(env)
 env_renderer.renderEnv(show=True)
-"""
+
 
 # Example generate a rail given a manual specification,
 # a map of tuples (cell_type, rotation)
@@ -26,7 +47,7 @@ specs = [[(0, 0), (0, 0), (0, 0), (0, 0), (7, 0), (0, 0)],
 env = RailEnv(width=6,
               height=2,
               rail_generator=rail_from_manual_specifications_generator(specs),
-              number_of_agents=2,
+              number_of_agents=1,
               obs_builder_object=TreeObsForRailEnv(max_depth=2))
 
 handle = env.get_agent_handles()

flatland/core/env_observation_builder.py

@@ -72,7 +72,7 @@ class TreeObsForRailEnv(ObservationBuilder):
         # Update local lookup table for all agents' target locations
         self.location_has_target = {}
         for loc in self.env.agents_target:
-            self.location_has_target[(loc[0],loc[1])] = 1
+            self.location_has_target[(loc[0], loc[1])] = 1
 
     def _distance_map_walker(self, position, target_nr):
         """
@@ -292,8 +292,6 @@ class TreeObsForRailEnv(ObservationBuilder):
             if position in self.location_has_target:
                 other_target_encountered = True
 
-
-
             # #############################
             # #############################
 
@@ -354,10 +352,8 @@ class TreeObsForRailEnv(ObservationBuilder):
                            0,
                            self.distance_map[handle, position[0], position[1], direction]]
 
-
         # TODO:
 
-
         # #############################
         # #############################
 
@@ -368,9 +364,9 @@ class TreeObsForRailEnv(ObservationBuilder):
                                                                (branch_direction+2) % 4):
                 # Swap forward and back in case of dead-end, so that an agent can learn that going forward takes
                 # it back
-                new_cell = self._new_position(position, (branch_direction+2)%4)
+                new_cell = self._new_position(position, (branch_direction+2) % 4)
 
-                branch_observation = self._explore_branch(handle, new_cell, (branch_direction+2)%4, depth+1)
+                branch_observation = self._explore_branch(handle, new_cell, (branch_direction+2) % 4, depth+1)
                 observation = observation + branch_observation
 
             elif last_isSwitch and self.env.rail.get_transition((position[0], position[1], direction),