diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py index 0853712e1a985fd3e39e88190eeffb14d141eff1..65cf0849d85d29cd785e26daccda62192bbe0ed8 100644 --- a/flatland/envs/rail_env.py +++ b/flatland/envs/rail_env.py @@ -406,18 +406,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: