Merge pull request flatland-association#21 from flatland-association/feature/graph-env-first-steps

Update to breaking changes in flatland-rl #257 (feature/graph env first steps).

Author: chenkins

flatland_baselines/deadlock_avoidance_heuristic/observation/full_env_observation.py

@@ -2,10 +2,11 @@
 from flatland.envs.rail_env import RailEnv
 
 
-class FullEnvObservation(ObservationBuilder[RailEnv]):
+class FullEnvObservation(ObservationBuilder[RailEnv, RailEnv]):
     """
     Returns full env as observation.
     """
+
     def __init__(self):
         pass
 
@@ -15,5 +16,5 @@ def get(self, handle: AgentHandle = 0) -> ObservationType:
     def reset(self):
         pass
 
-    def set_env(self,env):
+    def set_env(self, env):
         self.env = env

flatland_baselines/deadlock_avoidance_heuristic/policy/deadlock_avoidance_policy.py

@@ -90,7 +90,7 @@ def callback(self, handle, agent, position, direction, action, possible_transiti
     @_enable_flatland_deadlock_avoidance_policy_lru_cache(maxsize=100000)
     def _is_no_switch_cell(self, position) -> bool:
         for new_dir in range(4):
-            possible_transitions = self.env.rail.get_transitions(*position, new_dir)
+            possible_transitions = self.env.rail.get_transitions((position, new_dir))
             num_transitions = fast_count_nonzero(possible_transitions)
             if num_transitions > 1:
                 return False

flatland_baselines/deadlock_avoidance_heuristic/utils/flatland/shortest_distance_walker.py

@@ -39,7 +39,7 @@ def walk(self, handle, position, direction):
         if self.distance_map is None:
             self.distance_map = self.env.distance_map.get()
 
-        possible_transitions = self.env.rail.get_transitions(*position, direction)
+        possible_transitions = self.env.rail.get_transitions((position, direction))
         num_transitions = fast_count_nonzero(possible_transitions)
         if num_transitions == 1:
             new_direction = fast_argmax(possible_transitions)