diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py
index 8cbab637d2b14c7f47bec05b743ad3dcc87aeb86..b5debcd409373f3c839c7faadf8931151a8baa7a 100644
--- a/flatland/envs/rail_env.py
+++ b/flatland/envs/rail_env.py
@@ -488,6 +488,8 @@ class RailEnv(Environment):
             self.rewards_dict = {i: self.global_reward for i in range(self.get_num_agents())}
         if (self._max_episode_steps is not None) and (self._elapsed_steps >= self._max_episode_steps):
             self.dones["__all__"] = True
+            for i_agent in range(self.get_num_agents()):
+                self.dones[i] = True
 
         return self._get_observations(), self.rewards_dict, self.dones, info_dict