diff --git a/flatland/core/env.py b/flatland/core/env.py
index 32691f507f4cb5586f10b5645cc22ece718edc21..3618d965a39b5a71fd1cf24aa81f2f876d5c6365 100644
--- a/flatland/core/env.py
+++ b/flatland/core/env.py
@@ -99,12 +99,6 @@ class Environment:
"""
raise NotImplementedError()
- def render(self):
- """
- Perform rendering of the environment.
- """
- raise NotImplementedError()
-
def get_agent_handles(self):
"""
Returns a list of agents' handles to be used as keys in the step()
diff --git a/flatland/core/env_observation_builder.py b/flatland/core/env_observation_builder.py
index b30c2b1f5ddab079c9b6c41e35f03c69ed4162c3..53e7a068b73f9907217777251bce0fdd704603be 100644
--- a/flatland/core/env_observation_builder.py
+++ b/flatland/core/env_observation_builder.py
@@ -30,6 +30,27 @@ class ObservationBuilder:
"""
raise NotImplementedError()
+ def get_many(self, handles=[]):
+ """
+ Called whenever an observation has to be computed for the `env' environment, for each agent with handle
+ in the `handles' list.
+
+ Parameters
+ -------
+ handles : list of handles (optional)
+ List with the handles of the agents for which to compute the observation vector.
+
+ Returns
+ -------
+ function
+ A dictionary of observation structures, specific to the corresponding environment, with handles from
+ `handles' as keys.
+ """
+ observations = {}
+ for h in handles:
+ observations[h] = self.get(h)
+ return observations
+
def get(self, handle=0):
"""
Called whenever an observation has to be computed for the `env' environment, possibly
diff --git a/flatland/envs/observations.py b/flatland/envs/observations.py
index 676051d8338534c704ef98c90bee08a2836d4cfb..a3d88d773db9edaa0777e2aee94593a0392a956c 100644
--- a/flatland/envs/observations.py
+++ b/flatland/envs/observations.py
@@ -17,7 +17,7 @@ class TreeObsForRailEnv(ObservationBuilder):
network to simplify the representation of the state of the environment for each agent.
"""
- def __init__(self, max_depth):
+ def __init__(self, max_depth, predictor=None):
self.max_depth = max_depth
# Compute the size of the returned observation vector
@@ -30,6 +30,7 @@ class TreeObsForRailEnv(ObservationBuilder):
self.observation_space = [size * self.observation_dim]
self.location_has_agent = {}
self.location_has_agent_direction = {}
+ self.predictor = predictor
self.agents_previous_reset = None
@@ -167,6 +168,20 @@ class TreeObsForRailEnv(ObservationBuilder):
elif movement == 3: # WEST
return (position[0], position[1] - 1)
+ def get_many(self, handles=[]):
+ """
+ Called whenever an observation has to be computed for the `env' environment, for each agent with handle
+ in the `handles' list.
+ """
+
+ # TODO: @Erik this is where the predictions should be computed, storing any temporary data inside this object.
+ if self.predictor:
+ print(self.predictor.get(0))
+ observations = {}
+ for h in handles:
+ observations[h] = self.get(h)
+ return observations
+
def get(self, handle):
"""
Computes the current observation for agent `handle' in env
@@ -524,6 +539,11 @@ class TreeObsForRailEnv(ObservationBuilder):
agent_data.extend(tmp_agent_data)
return tree_data, distance_data, agent_data
+ def _set_env(self, env):
+ self.env = env
+ if self.predictor:
+ self.predictor._set_env(self.env)
+
class GlobalObsForRailEnv(ObservationBuilder):
"""
diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py
index 6cd645147a1a870c760c3bc60dd364903a2065d2..5d20a5d9f38230f353b0a9616c49ede333206c49 100644
--- a/flatland/envs/rail_env.py
+++ b/flatland/envs/rail_env.py
@@ -58,7 +58,6 @@ class RailEnv(Environment):
rail_generator=random_rail_generator(),
number_of_agents=1,
obs_builder_object=TreeObsForRailEnv(max_depth=2),
- prediction_builder_object=None
):
"""
Environment init.
@@ -99,10 +98,6 @@ class RailEnv(Environment):
self.obs_builder = obs_builder_object
self.obs_builder._set_env(self)
- self.prediction_builder = prediction_builder_object
- if self.prediction_builder:
- self.prediction_builder._set_env(self)
-
self.action_space = [1]
self.observation_space = self.obs_builder.observation_space # updated on resets?
@@ -297,10 +292,6 @@ class RailEnv(Environment):
np.equal(new_position, [agent2.position for agent2 in self.agents]).all(1))
return cell_isFree, new_cell_isValid, new_direction, new_position, transition_isValid
- def predict(self):
- if not self.prediction_builder:
- return {}
- return self.prediction_builder.get()
def check_action(self, agent, action):
transition_isValid = None
@@ -330,20 +321,9 @@ class RailEnv(Environment):
return new_direction, transition_isValid
def _get_observations(self):
- self.obs_dict = {}
- self.debug_obs_dict = {}
- for iAgent in range(self.get_num_agents()):
- self.obs_dict[iAgent] = self.obs_builder.get(iAgent)
+ self.obs_dict = self.obs_builder.get_many(list(range(self.get_num_agents())))
return self.obs_dict
- def _get_predictions(self):
- if not self.prediction_builder:
- return {}
- return {}
-
- def render(self):
- # TODO:
- pass
def get_full_state_msg(self):
grid_data = self.rail.grid.tolist()
diff --git a/flatland/utils/graphics_pil.py b/flatland/utils/graphics_pil.py
index 9041dc6391dd604566c1c48a40a01dcf9f56e987..bca964c960ee55d0d110e9c9e71d6b4481d93215 100644
--- a/flatland/utils/graphics_pil.py
+++ b/flatland/utils/graphics_pil.py
@@ -395,7 +395,9 @@ class PILSVG(PILGL):
(0, 3): "svg/Zug_2_Weiche_#0091ea.svg"
}
- # "paint" color of the train images we load
+ # "paint" color of the train images we load - this is the color we will change.
+ # a3BaseColor = self.rgb_s2i("0091ea")
+ # temporary workaround for trains / agents renamed with different colour:
a3BaseColor = self.rgb_s2i("d50000")
self.dPilZug = {}