Flag for saving distance map in save option

flatland/envs/rail_env.py

@@ -705,16 +705,20 @@ class RailEnv(Environment):
 
         return msgpack.packb(msg_data, use_bin_type=True)
 
-    def save(self, filename):
-        if self.distance_map.get() is not None:
-            if len(self.distance_map.get()) > 0:
-                with open(filename, "wb") as file_out:
-                    file_out.write(self.get_full_state_dist_msg())
+    def save(self, filename,save_distance_maps=False):
+        if save_distance_maps == True:
+            if self.distance_map.get() is not None:
+                if len(self.distance_map.get()) > 0:
+                    with open(filename, "wb") as file_out:
+                        file_out.write(self.get_full_state_dist_msg())
+                else:
+                    with open(filename, "wb") as file_out:
+                        file_out.write(self.get_full_state_msg())
             else:
                 with open(filename, "wb") as file_out:
                     file_out.write(self.get_full_state_msg())
         else:
-            with open(filename, "wb") as file_out:
+            with open(filename,"wb") as file_out:
                 file_out.write(self.get_full_state_msg())
 
     def load(self, filename):