diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py
index b5debcd409373f3c839c7faadf8931151a8baa7a..cb4da95a81758625f1f45ecb535e47d834fd5005 100644
--- a/flatland/envs/rail_env.py
+++ b/flatland/envs/rail_env.py
@@ -489,7 +489,7 @@ class RailEnv(Environment):
         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
+                self.dones[i_agent] = True
 
         return self._get_observations(), self.rewards_dict, self.dones, info_dict