Malfunction data generator does not work

Hi @mlerik!

Please, check the behavior of the last agent in my new simulation. It seems that new_malfunction is not shown correctly on step number 5. I see that this problem occurs on the first steps only (but all agents have entered the environment).

import time

import numpy as np

from flatland.envs.observations import TreeObsForRailEnv, GlobalObsForRailEnv
from flatland.envs.predictions import ShortestPathPredictorForRailEnv
from flatland.envs.rail_env import RailEnv
from flatland.envs.rail_generators import sparse_rail_generator
from flatland.envs.schedule_generators import sparse_schedule_generator
from flatland.utils.rendertools import RenderTool, AgentRenderVariant

np.random.seed(1)

stochastic_data = {'prop_malfunction': 0.5,  # Percentage of defective agents
                   'malfunction_rate': 30,  # Rate of malfunction occurence
                   'min_duration': 3,  # Minimal duration of malfunction
                   'max_duration': 10  # Max duration of malfunction
                   }

TreeObservation = TreeObsForRailEnv(max_depth=2, predictor=ShortestPathPredictorForRailEnv())

speed_ration_map = {1.: 0.25,  # Fast passenger train
                    1. / 2.: 0.25,  # Fast freight train
                    1. / 3.: 0.25,  # Slow commuter train
                    1. / 4.: 0.25}  # Slow freight train

env = RailEnv(width=60,
              height=60,
              rail_generator=sparse_rail_generator(max_num_cities=12,
                                                   # Number of cities in map (where train stations are)
                                                   seed=14,  # Random seed
                                                   grid_mode=False,
                                                   max_rails_between_cities=2,
                                                   max_rails_in_city=6,
                                                   ),
              schedule_generator=sparse_schedule_generator(speed_ration_map),
              number_of_agents=5,
              stochastic_data=stochastic_data,  # Malfunction data generator
              obs_builder_object=GlobalObsForRailEnv(),
              remove_agents_at_target=True
              )

obs = env.reset()

MY_ACTION = []
MY_ACTION.append({0: 4, 1: 2, 2: 4, 3: 4, 4: 2})
MY_ACTION.append({0: 2, 1: 2, 2: 2, 3: 2, 4: 2})
MY_ACTION.append({0: 2, 1: 2, 2: 2, 3: 2, 4: 2})
MY_ACTION.append({0: 2, 1: 2, 2: 2, 3: 2, 4: 2})
MY_ACTION.append({0: 2, 1: 2, 2: 3, 3: 2, 4: 4})
MY_ACTION.append({0: 2, 1: 2, 2: 1, 3: 2, 4: 2})
MY_ACTION.append({1: 4, 4: 4})

for step in range(7):
    print("========== step number ", step, " ==========", sep = "")
    action_dict = MY_ACTION[step]
    print("my action ", action_dict)
    for ind in range(env.get_num_agents()):
        print(env.agents[ind].malfunction_data)
    next_obs, all_rewards, done, _ = env.step(action_dict)