diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py index 8eb6aede1ed7bdf87aa557a213f87e4bf8b1115c..f8534640846eb182131e0f0d443f752721fd6672 100644 --- a/flatland/envs/rail_env.py +++ b/flatland/envs/rail_env.py @@ -399,18 +399,18 @@ class RailEnv(Environment): agent.malfunction_data['malfunction'] -= 1 return True - # Restart fixed agents - if agent.malfunction_data['malfunction'] < 1 and agent.malfunction_data['next_malfunction'] > 0 and not agent.malfunction_data['fixed']: - agent.malfunction_data['next_malfunction'] -= 1 - agent.malfunction_data['fixed'] = True - if 'moving_before_malfunction' in agent.malfunction_data: - self.agents[i_agent].moving = agent.malfunction_data['moving_before_malfunction'] - return False - - # Agent has been running smoothly - elif agent.malfunction_data['malfunction'] < 1 and agent.malfunction_data['next_malfunction'] > 0: - agent.malfunction_data['next_malfunction'] -= 1 - return False + if agent.malfunction_data['malfunction'] < 1 and agent.malfunction_data['next_malfunction'] > 0: + # Restart fixed agents + if not agent.malfunction_data['fixed']: + agent.malfunction_data['next_malfunction'] -= 1 + agent.malfunction_data['fixed'] = True + if 'moving_before_malfunction' in agent.malfunction_data: + self.agents[i_agent].moving = agent.malfunction_data['moving_before_malfunction'] + return False + else: + # Agent has been running smoothly + agent.malfunction_data['next_malfunction'] -= 1 + return False # Break agents that have next_malfunction if agent.malfunction_data['malfunction'] < 1 and agent.malfunction_data['next_malfunction'] < 1: