Add regression test diamond crossing without over and underpasses.

Author: chenkins

flatland/utils/simple_rail.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from flatland.core.grid.rail_env_grid import RailEnvTransitions
+from flatland.core.grid.rail_env_grid import RailEnvTransitions, RailEnvTransitionsEnum
 from flatland.core.transition_map import GridTransitionMap
 
 
@@ -42,16 +42,16 @@ def make_simple_rail() -> Tuple[GridTransitionMap, np.array, Dict]:
     rail = GridTransitionMap(width=rail_map.shape[1],
                              height=rail_map.shape[0], transitions=transitions)
     rail.grid = rail_map
-    city_positions = [(0,3), (6, 6)]
+    city_positions = [(3, 9), (6, 2)]
     train_stations = [
-                      [( (0, 3), 0 ) ],
-                      [( (6, 6), 0 ) ],
-                     ]
+        [((3, 9), 1)],
+        [((6, 2), 2)],
+    ]
     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
 
@@ -93,16 +93,16 @@ def make_disconnected_simple_rail() -> Tuple[GridTransitionMap, np.array, Dict]:
     rail = GridTransitionMap(width=rail_map.shape[1],
                              height=rail_map.shape[0], transitions=transitions)
     rail.grid = rail_map
-    city_positions = [(0,3), (6, 6)]
+    city_positions = [(0, 3), (6, 6)]
     train_stations = [
-                      [( (0, 3), 0 ) ],
-                      [( (6, 6), 0 ) ],
-                     ]
+        [((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
 
@@ -141,16 +141,16 @@ def make_simple_rail2() -> Tuple[GridTransitionMap, np.array, Dict]:
     rail = GridTransitionMap(width=rail_map.shape[1],
                              height=rail_map.shape[0], transitions=transitions)
     rail.grid = rail_map
-    city_positions = [(0,3), (6, 6)]
+    city_positions = [(0, 3), (6, 6)]
     train_stations = [
-                      [( (0, 3), 0 ) ],
-                      [( (6, 6), 0 ) ],
-                     ]
+        [((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
 
@@ -190,16 +190,16 @@ def make_simple_rail_unconnected() -> Tuple[GridTransitionMap, np.array, Dict]:
     rail = GridTransitionMap(width=rail_map.shape[1],
                              height=rail_map.shape[0], transitions=transitions)
     rail.grid = rail_map
-    city_positions = [(0,3), (6, 6)]
+    city_positions = [(0, 3), (6, 6)]
     train_stations = [
-                      [( (0, 3), 0 ) ],
-                      [( (6, 6), 0 ) ],
-                     ]
+        [((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
 
@@ -245,16 +245,16 @@ def make_simple_rail_with_alternatives() -> Tuple[GridTransitionMap, np.array, D
     rail = GridTransitionMap(width=rail_map.shape[1],
                              height=rail_map.shape[0], transitions=transitions)
     rail.grid = rail_map
-    city_positions = [(0,3), (6, 6)]
+    city_positions = [(0, 3), (6, 6)]
     train_stations = [
-                      [( (0, 3), 0 ) ],
-                      [( (6, 6), 0 ) ],
-                     ]
+        [((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
 
@@ -294,20 +294,21 @@ def make_invalid_simple_rail() -> Tuple[GridTransitionMap, np.array, Dict[str, A
     rail = GridTransitionMap(width=rail_map.shape[1],
                              height=rail_map.shape[0], transitions=transitions)
     rail.grid = rail_map
-    city_positions = [(0,3), (6, 6)]
+    city_positions = [(0, 3), (6, 6)]
     train_stations = [
-                      [( (0, 3), 0 ) ],
-                      [( (6, 6), 0 ) ],
-                     ]
+        [((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
 
-def make_oval_rail() -> Tuple[GridTransitionMap, np.array]:
+
+def make_oval_rail() -> Tuple[GridTransitionMap, np.array, Any]:
     transitions = RailEnvTransitions()
     cells = transitions.transition_list
 
@@ -322,7 +323,7 @@ def make_oval_rail() -> Tuple[GridTransitionMap, np.array]:
     rail_map = np.array(
         [[empty] * 9] +
         [[empty] + [right_turn_from_south] + [horizontal_straight] * 5 + [right_turn_from_west] + [empty]] +
-        [[empty] + [vertical_straight] + [empty] * 5 + [vertical_straight] + [empty]]+
+        [[empty] + [vertical_straight] + [empty] * 5 + [vertical_straight] + [empty]] +
         [[empty] + [vertical_straight] + [empty] * 5 + [vertical_straight] + [empty]] +
         [[empty] + [right_turn_from_east] + [horizontal_straight] * 5 + [right_turn_from_north] + [empty]] +
         [[empty] * 9], dtype=np.uint16)
@@ -341,4 +342,39 @@ def make_oval_rail() -> Tuple[GridTransitionMap, np.array]:
                     'city_orientations': city_orientations
                     }
     optionals = {'agents_hints': agents_hints}
-    return  rail, rail_map, optionals
+    return rail, rail_map, optionals
+
+
+def make_diamond_crossing_rail() -> Tuple[GridTransitionMap, np.array, Dict]:
+    # We instantiate a very simple rail network on a 6x10 grid:
+    # Note that  some cells have invalid RailEnvTransitions!
+    #        |
+    #        |
+    #    _ _ | _ _ _ _ _ _ _
+    #        |
+    #        |
+    #        |
+    transitions = RailEnvTransitions()
+    rail_map = np.array(
+        [[RailEnvTransitionsEnum.empty] * 2 + [RailEnvTransitionsEnum.dead_end_from_south] + [RailEnvTransitionsEnum.empty] * 7] +
+        [[RailEnvTransitionsEnum.empty] * 2 + [RailEnvTransitionsEnum.vertical_straight] + [RailEnvTransitionsEnum.empty] * 7] * 2 +
+        [[RailEnvTransitionsEnum.dead_end_from_east] + [RailEnvTransitionsEnum.horizontal_straight] * 1 + [RailEnvTransitionsEnum.diamond_crossing] * 1 + [
+            RailEnvTransitionsEnum.horizontal_straight] * 6 + [RailEnvTransitionsEnum.dead_end_from_west]] +
+        [[RailEnvTransitionsEnum.empty] * 2 + [RailEnvTransitionsEnum.vertical_straight] + [RailEnvTransitionsEnum.empty] * 7] * 2 +
+        [[RailEnvTransitionsEnum.empty] * 2 + [RailEnvTransitionsEnum.dead_end_from_north] + [RailEnvTransitionsEnum.empty] * 7]
+        , dtype=np.uint16)
+    rail = GridTransitionMap(width=rail_map.shape[1],
+                             height=rail_map.shape[0], transitions=transitions)
+    rail.grid = rail_map
+    city_positions = [(1, 4), (4, 4)]
+    train_stations = [
+        [((1, 4), 0)],
+        [((4, 4), 0)],
+    ]
+    city_orientations = [1, 3]
+    agents_hints = {'city_positions': city_positions,
+                    'train_stations': train_stations,
+                    'city_orientations': city_orientations
+                    }
+    optionals = {'agents_hints': agents_hints}
+    return rail, rail_map, optionals

tests/test_over_under_passes.py

@@ -0,0 +1,86 @@
+import time
+
+from flatland.envs.line_generators import sparse_line_generator
+from flatland.envs.observations import TreeObsForRailEnv
+from flatland.envs.predictions import ShortestPathPredictorForRailEnv
+from flatland.envs.rail_env import RailEnv
+from flatland.envs.rail_env_action import RailEnvActions
+from flatland.envs.rail_generators import rail_from_grid_transition_map
+from flatland.envs.rail_trainrun_data_structures import Waypoint
+from flatland.envs.step_utils.states import TrainState
+from flatland.utils.rendertools import RenderTool
+from flatland.utils.simple_rail import make_diamond_crossing_rail
+
+
+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,
+        obs_builder_object=TreeObsForRailEnv(max_depth=2, predictor=ShortestPathPredictorForRailEnv()),
+        record_steps=True
+    )
+
+    env.reset()
+    env._max_episode_steps = 555
+
+    # set the initial position
+    agent_0 = env.agents[0]
+    agent_0.initial_position = (3, 1)  # just ahead of diamond crossing facing east
+    agent_0.position = (3, 1)  # just ahead of diamond crossing facing east
+    agent_0.direction = 1  # east
+    agent_0.initial_direction = 1  # east
+    agent_0.target = (3, 9)  # east dead-end
+    agent_0.moving = True
+    agent_0.latest_arrival = 999
+    agent_0._set_state(TrainState.MOVING)
+
+    agent_1 = env.agents[1]
+    agent_1.initial_position = (2, 2)  # just ahead of diamond crossing facing south
+    agent_1.position = (2, 2)  # just ahead of diamond crossing facing south
+    agent_1.direction = 2  # south
+    agent_1.initial_direction = 2  # south
+    agent_1.target = (6, 2)  # south dead-end
+    agent_1.moving = True
+    agent_1.latest_arrival = 999
+    agent_1._set_state(TrainState.MOVING)
+
+    env.distance_map._compute(env.agents, env.rail)
+    done = False
+    env_renderer = None
+    if rendering:
+        env_renderer = RenderTool(env)
+    while not done:
+        _, _, dones, _ = env.step({
+            0: RailEnvActions.MOVE_FORWARD,
+            1: RailEnvActions.MOVE_FORWARD,
+        })
+        done = dones["__all__"]
+        if env_renderer is not None:
+            env_renderer.render_env(show=True, show_observations=False, show_predictions=False)
+            time.sleep(0.2)
+
+    waypoints = []
+    for agent_states in env.cur_episode:
+        cur = []
+        for agent_state in agent_states:
+            r, c, d, _, _, _ = agent_state
+            cur.append(Waypoint((r, c), d))
+        waypoints.append(cur)
+    expected = [
+        [Waypoint(position=(3, 2), direction=1), Waypoint(position=(2, 2), direction=2)],
+        [Waypoint(position=(3, 3), direction=1), Waypoint(position=(3, 2), direction=2)],
+        [Waypoint(position=(3, 4), direction=1), Waypoint(position=(4, 2), direction=2)],
+        [Waypoint(position=(3, 5), direction=1), Waypoint(position=(5, 2), direction=2)],
+        [Waypoint(position=(3, 6), direction=1), Waypoint(position=(0, 0), direction=2)],
+        [Waypoint(position=(3, 7), direction=1), Waypoint(position=(0, 0), direction=2)],
+        [Waypoint(position=(3, 8), direction=1), Waypoint(position=(0, 0), direction=2)],
+        [Waypoint(position=(0, 0), direction=1), Waypoint(position=(0, 0), direction=2)]
+    ]
+    assert expected == waypoints, waypoints