Skip to content
Snippets Groups Projects
Commit 5df0f5de authored by u214892's avatar u214892
Browse files

#237 added missing doc string

parent f8ef8acc
No related branches found
No related tags found
No related merge requests found
from flatland.core.env_observation_builder import ObservationBuilder
from flatland.envs.observations import TreeObsForRailEnv from flatland.envs.observations import TreeObsForRailEnv
from flatland.envs.predictions import ShortestPathPredictorForRailEnv from flatland.envs.predictions import ShortestPathPredictorForRailEnv
from flatland.envs.rail_env import RailEnv from flatland.envs.rail_env import RailEnv
...@@ -5,15 +6,34 @@ from flatland.envs.rail_generators import rail_from_file ...@@ -5,15 +6,34 @@ from flatland.envs.rail_generators import rail_from_file
from flatland.envs.schedule_generators import schedule_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: if obs_builder_object is None:
obs_builder_object = TreeObsForRailEnv( obs_builder_object = TreeObsForRailEnv(
max_depth=2, max_depth=2,
predictor=ShortestPathPredictorForRailEnv(max_depth=10)) predictor=ShortestPathPredictorForRailEnv(max_depth=10))
environment = RailEnv(width=1, environment = RailEnv(width=1, # will be overridden when loading from file
height=1, height=1, # will be overridden when loading from file
rail_generator=rail_from_file(file_name, load_from_package), 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), schedule_generator=schedule_from_file(file_name, load_from_package),
obs_builder_object=obs_builder_object) obs_builder_object=obs_builder_object)
return environment return environment
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment