need to update tree observation as well

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