diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py
index b31be94f4a475db5175478ab222aefcae718374f..d2b47911ace505ee5ebf0f73e13749c9f810c57d 100644
--- a/flatland/envs/rail_env.py
+++ b/flatland/envs/rail_env.py
@@ -202,7 +202,7 @@ class RailEnv(Environment):
             self.rewards_dict = {i: r + global_reward for i, r in self.rewards_dict.items()}
             return self._get_observations(), self.rewards_dict, self.dones, {}
 
-        for i_agent, agent in enumerate(self.get_num_agents()):
+        for i_agent, agent in enumerate(self.agents):
             agent.old_direction = agent.direction
             agent.old_position = agent.position
             if self.dones[i_agent]:  # this agent has already completed...