Improve type hints.

Author: chenkins

flatland/envs/rail_env.py

@@ -99,7 +99,7 @@ class RailEnv(Environment):
     def __init__(self,
                  width,
                  height,
-                 rail_generator=None,
+                 rail_generator: "RailGenerator" = None,
                  line_generator: "LineGenerator" = None,  # : line_gen.LineGenerator = line_gen.random_line_generator(),
                  number_of_agents=2,
                  obs_builder_object: ObservationBuilder = GlobalObsForRailEnv(),

tests/test_over_under_passes.py

@@ -15,11 +15,9 @@
 def test_diamond_crossing_without_over_and_underpasses(rendering: bool = False):
     rail, rail_map, optionals = make_diamond_crossing_rail()
 
-    # TODO better way to init state?
     env = RailEnv(
         width=rail_map.shape[1],
         height=rail_map.shape[0],
-        # TODO typing
         rail_generator=rail_from_grid_transition_map(rail, optionals),
         line_generator=sparse_line_generator(),
         number_of_agents=2,
@@ -96,7 +94,6 @@ def test_diamond_crossing_with_over_and_underpasses(rendering: bool = False):
     env = RailEnv(
         width=rail_map.shape[1],
         height=rail_map.shape[0],
-        # TODO typing
         rail_generator=rail_from_grid_transition_map(rail, optionals),
         line_generator=sparse_line_generator(),
         number_of_agents=2,
@@ -172,18 +169,15 @@ def test_diamond_crossing_with_over_and_underpasses(rendering: bool = False):
 def test_diamond_crossing_with_over_and_underpasses_head_on(rendering: bool = False):
     rail, rail_map, optionals = make_diamond_crossing_rail()
 
-    # TODO better way to init state?
     env = RailEnv(
         width=rail_map.shape[1],
         height=rail_map.shape[0],
-        # TODO typing
         rail_generator=rail_from_grid_transition_map(rail, optionals),
         line_generator=sparse_line_generator(),
         number_of_agents=2,
         obs_builder_object=TreeObsForRailEnv(max_depth=2, predictor=ShortestPathPredictorForRailEnv()),
         record_steps=True
     )
-
     env.reset()
     env._max_episode_steps = 5