diff --git a/flatland/envs/observations.py b/flatland/envs/observations.py
index 80dc73adf417b43d2740b80e9b86fe5ba5ce257f..fecca331b7e8253d5196f9bfd881a40b4ab672d6 100644
--- a/flatland/envs/observations.py
+++ b/flatland/envs/observations.py
@@ -397,8 +397,9 @@ class TreeObsForRailEnv(ObservationBuilder):
             total_transitions = bin(self.env.rail.get_full_transitions(*position)).count("1")
             num_transitions = np.count_nonzero(cell_transitions)
             exploring = False
+
             # Detect Switches that can only be used by other agents.
-            if total_transitions > 2 > num_transitions:
+            if total_transitions > 2 > num_transitions and tot_dist < unusable_switch:
                 unusable_switch = tot_dist
 
             if num_transitions == 1:
diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py
index e098160d7a14f9b7cc07b7f365e6af573c3d7d56..4c44afb74129d8bbe0d301e2198188858693dd28 100644
--- a/flatland/envs/rail_env.py
+++ b/flatland/envs/rail_env.py
@@ -446,7 +446,7 @@ class RailEnv(Environment):
         # agents are always reset as not moving
         self.agents_static = [EnvAgentStatic(d[0], d[1], d[2], moving=False) for d in data[b"agents_static"]]
         self.agents = [EnvAgent(d[0], d[1], d[2], d[3], d[4]) for d in data[b"agents"]]
-        if hasattr(self.obs_builder, 'distance_map'):
+        if hasattr(self.obs_builder, 'distance_map') and b"distance_maps" in data.keys():
             self.obs_builder.distance_map = data[b"distance_maps"]
         # setup with loaded data
         self.height, self.width = self.rail.grid.shape
diff --git a/notebooks/complex_scene.pkl b/notebooks/complex_scene.pkl
index b5c272477f53794d78a896c33d7c91e5b8cb0ea3..9bad5f9674b7a7b7e792c4e4805bce51ced35f7c 100644
Binary files a/notebooks/complex_scene.pkl and b/notebooks/complex_scene.pkl differ