diff --git a/flatland/envs/observations.py b/flatland/envs/observations.py
index 80dc73adf417b43d2740b80e9b86fe5ba5ce257f..3e8583d996635bf58552922707815984f8e76f85 100644
--- a/flatland/envs/observations.py
+++ b/flatland/envs/observations.py
@@ -396,18 +396,16 @@ class TreeObsForRailEnv(ObservationBuilder):
             cell_transitions = self.env.rail.get_transitions(*position, direction)
             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:
                 # Check if dead-end, or if we can go forward along direction
-                nbits = 0
-                tmp = self.env.rail.get_full_transitions(*position)
-                while tmp > 0:
-                    nbits += (tmp & 1)
-                    tmp = tmp >> 1
+                nbits = total_transitions
                 if nbits == 1:
                     # Dead-end!
                     last_is_dead_end = True
@@ -434,8 +432,6 @@ class TreeObsForRailEnv(ObservationBuilder):
 
         # `position' is either a terminal node or a switch
 
-        observation = []
-
         # #############################
         # #############################
         # Modify here to append new / different features for each visited cell!
@@ -463,6 +459,7 @@ class TreeObsForRailEnv(ObservationBuilder):
                            other_agent_same_direction,
                            other_agent_opposite_direction
                            ]
+
         else:
             observation = [own_target_encountered,
                            other_target_encountered,
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