diff --git a/flatland/envs/observations.py b/flatland/envs/observations.py
index 548c3e29a72aaa0390ab73f5550738f46f52c6c4..65708e5faa47f6c4db981fb3e4e93bd05fc211f9 100644
--- a/flatland/envs/observations.py
+++ b/flatland/envs/observations.py
@@ -161,6 +161,7 @@ class TreeObsForRailEnv(ObservationBuilder):
 
         # Update local lookup table for all agents' positions
         self.location_has_agent = {tuple(agent.position): 1 for agent in self.env.agents}
+        # TODO: Update this to handle number of agents at same location
         self.location_has_agent_direction = {tuple(agent.position): agent.direction for agent in self.env.agents}
         self.location_has_agent_speed = {tuple(agent.position): agent.speed_data['speed'] for agent in self.env.agents}
         self.location_has_agent_malfunction = {tuple(agent.position): agent.malfunction_data['malfunction'] for agent in