diff --git a/flatland/envs/distance_map.py b/flatland/envs/distance_map.py
index 228d3992429cb0db4d5bd977a7c41214c3857b28..22721407f059ff2e02d907110cb9f982aa9d599e 100644
--- a/flatland/envs/distance_map.py
+++ b/flatland/envs/distance_map.py
@@ -24,8 +24,6 @@ class DistanceMap:
         """
         self.distance_map = distance_map
 
-
-
     def get(self) -> np.ndarray:
         """
         Get the distance map
@@ -35,11 +33,6 @@ class DistanceMap:
 
             nb_agents = len(self.agents)
             compute_distance_map = True
-            if self.agents_previous_computation is not None and nb_agents == len(self.agents_previous_computation):
-                compute_distance_map = False
-                for i in range(nb_agents):
-                    if self.agents[i].target != self.agents_previous_computation[i].target:
-                        compute_distance_map = True
             # Don't compute the distance map if it was loaded
             if self.agents_previous_computation is None and self.distance_map is not None:
                 compute_distance_map = False
@@ -52,8 +45,6 @@ class DistanceMap:
 
         return self.distance_map
 
-
-
     def reset(self, agents: List[EnvAgent], rail: GridTransitionMap):
         """
         Reset the distance map
diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py
index f8c54f33088a97f989cd10a120ed848bbf9bf4ae..b518d2f889d335e64495dafadcc26dc8f2f90082 100644
--- a/flatland/envs/rail_env.py
+++ b/flatland/envs/rail_env.py
@@ -240,9 +240,6 @@ class RailEnv(Environment):
         #  can we not put 'self.rail_generator(..)' into 'if regen_rail or self.rail is None' condition?
         rail, optionals = self.rail_generator(self.width, self.height, self.get_num_agents(), self.num_resets)
 
-        if optionals and 'distance_map' in optionals:
-            self.distance_map.set(optionals['distance_map'])
-
         if regen_rail or self.rail is None:
             self.rail = rail
             self.height, self.width = self.rail.grid.shape
@@ -252,6 +249,11 @@ class RailEnv(Environment):
                     check = self.rail.cell_neighbours_valid(rc_pos, True)
                     if not check:
                         warnings.warn("Invalid grid at {} -> {}".format(rc_pos, check))
+        # TODO https://gitlab.aicrowd.com/flatland/flatland/issues/172
+        #  hacky: we must re-compute the distance map and not use the initial distance_map loaded from file by
+        #  rail_from_file!!!
+        elif optionals and 'distance_map' in optionals:
+            self.distance_map.set(optionals['distance_map'])
 
         if replace_agents:
             agents_hints = None