minor updates to file loading and features in observation space

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:

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