From 09d4e0f28532a3a46066d29bfb2f444aa3a367b5 Mon Sep 17 00:00:00 2001
From: Erik Nygren <erik.nygren@sbb.ch>
Date: Sat, 5 Oct 2019 10:35:28 -0400
Subject: [PATCH] added info to reset function
---
flatland/envs/rail_env.py | 13 ++++++++++++-
1 file changed, 12 insertions(+), 1 deletion(-)
diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py
index 23f0dd57..1c2df4ea 100644
--- a/flatland/envs/rail_env.py
+++ b/flatland/envs/rail_env.py
@@ -307,8 +307,19 @@ class RailEnv(Environment):
self.obs_builder.reset()
self.distance_map.reset(self.agents, self.rail)
+ info_dict = {
+ 'action_required': {
+ i: (agent.status == RailAgentStatus.READY_TO_DEPART or (
+ agent.status == RailAgentStatus.ACTIVE and agent.speed_data['position_fraction'] == 0.0))
+ for i, agent in enumerate(self.agents)},
+ 'malfunction': {
+ i: self.agents[i].malfunction_data['malfunction'] for i in range(self.get_num_agents())
+ },
+ 'speed': {i: self.agents[i].speed_data['speed'] for i in range(self.get_num_agents())},
+ 'status': {i: agent.status for i, agent in enumerate(self.agents)}
+ }
# Return the new observation vectors for each agent
- return self._get_observations()
+ return self._get_observations(), info_dict
def _agent_malfunction(self, i_agent) -> bool:
"""
--
GitLab