Currently, _remove_agent is incredibly slow for large models. When an agent is removed, the code uses a for loop to shift every single subsequent agent down by one index in the dictionaries, and then does an $O(N)$ array slice to shift the spatial data up. Because agents in a continuous space do not need to be stored in any specific order, we can make this removal $O(1)$ using the classic "Swap and Pop" algorithm. We simply take the last agent in the array, move it into the slot of the removed agent, and pop the end of the array.
|
def _remove_agent(self, agent: Agent) -> None: |
|
"""Remove an agent from the space. |
|
|
|
This method is automatically called by ContinuousSpaceAgent.remove. |
|
|
|
""" |
|
index = self._agent_to_index[agent] |
|
self._agent_to_index.pop(agent, None) |
|
self._index_to_agent.pop(index, None) |
|
del self.active_agents[index] |
|
|
|
# Shift all subsequent agents up by 1 |
|
for agent in self.active_agents[index::]: |
|
old_index = self._agent_to_index[agent] |
|
self._agent_to_index[agent] = old_index - 1 |
|
self._index_to_agent[old_index - 1] = agent |
|
|
|
# Clean up the stale entry from the last shifted agent |
|
if len(self.active_agents) > index: |
|
self._index_to_agent.pop(len(self.active_agents), None) |
|
|
|
# we move all data below the removed agent one row up |
|
self._agent_positions[index : self._n_agents - 1] = self._agent_positions[ |
|
index + 1 : self._n_agents |
|
] |
|
self._n_agents -= 1 |
|
self.agent_positions = self._agent_positions[0 : self._n_agents] |
Or Am I missing something? Please let me know.
