diff --git a/flatland/envs/rail_env_utils.py b/flatland/envs/rail_env_utils.py
index dc1cff12c0c8b1860859208a13d6403734a2d2ad..7a814891635e86330c5168d944ae5dea421ea5bc 100644
--- a/flatland/envs/rail_env_utils.py
+++ b/flatland/envs/rail_env_utils.py
@@ -1,3 +1,4 @@
+from flatland.core.env_observation_builder import ObservationBuilder
 from flatland.envs.observations import TreeObsForRailEnv
 from flatland.envs.predictions import ShortestPathPredictorForRailEnv
 from flatland.envs.rail_env import RailEnv
@@ -5,15 +6,34 @@ from flatland.envs.rail_generators import rail_from_file
 from flatland.envs.schedule_generators import schedule_from_file
 
 
-def load_flatland_environment_from_file(file_name, load_from_package=None, obs_builder_object=None):
+def load_flatland_environment_from_file(file_name: str,
+                                        load_from_package: str = None,
+                                        obs_builder_object: ObservationBuilder = None) -> RailEnv:
+    """
+    Parameters
+    ----------
+    file_name : str
+        The pickle file.
+    load_from_package : str
+        The python module to import from. Example: 'env_data.tests'
+        This requires that there are `__init__.py` files in the folder structure we load the file from.
+    obs_builder_object: ObservationBuilder
+        The obs builder for the `RailEnv` that is created.
+
+
+    Returns
+    -------
+    RailEnv
+        The environment loaded from the pickle file.
+    """
     if obs_builder_object is None:
         obs_builder_object = TreeObsForRailEnv(
             max_depth=2,
             predictor=ShortestPathPredictorForRailEnv(max_depth=10))
-    environment = RailEnv(width=1,
-                          height=1,
+    environment = RailEnv(width=1,  # will be overridden when loading from file
+                          height=1,  # will be overridden when loading from file
                           rail_generator=rail_from_file(file_name, load_from_package),
-                          number_of_agents=1,
+                          number_of_agents=1,  # will be overridden when loading from file
                           schedule_generator=schedule_from_file(file_name, load_from_package),
                           obs_builder_object=obs_builder_object)
     return environment