updated test_environments.py to import correct library

flatland/envs/observations.py

@@ -396,6 +396,9 @@ class TreeObsForRailEnv(ObservationBuilder):
                 observation = observation + branch_observation
                 if len(branch_visited) != 0:
                     visited.union(branch_visited)
+                    print(visited)
+                    print(branch_visited)
+                    print("-------")
             else:
                 num_cells_to_fill_in = 0
                 pow4 = 1

tests/test_environments.py

@@ -7,7 +7,7 @@ from flatland.envs.generators import rail_from_GridTransitionMap_generator
 from flatland.envs.generators import complex_rail_generator
 from flatland.core.transitions import Grid4Transitions
 from flatland.core.transition_map import GridTransitionMap
-from flatland.core.env_observation_builder import GlobalObsForRailEnv
+from flatland.envs.observations import GlobalObsForRailEnv
 from flatland.envs.agent_utils import EnvAgent
 
 """Tests for `flatland` package."""