diff --git a/metadrive/envs/base_env.py b/metadrive/envs/base_env.py index 360cb8614..963937efa 100644 --- a/metadrive/envs/base_env.py +++ b/metadrive/envs/base_env.py @@ -736,6 +736,16 @@ def agent(self): ) return self.agents[DEFAULT_AGENT] + @property + def agents_including_just_terminated(self): + """ + Return all agents that occupy some space in current environments + :return: Dict[agent_id:vehicle] + """ + ret = self.agent_manager.active_agents + ret.update(self.agent_manager.just_terminated_agents) + return ret + def setup_engine(self): """ Engine setting after launching diff --git a/metadrive/manager/agent_manager.py b/metadrive/manager/agent_manager.py index f29884ffd..cae3a3223 100644 --- a/metadrive/manager/agent_manager.py +++ b/metadrive/manager/agent_manager.py @@ -245,7 +245,13 @@ def just_terminated_agents(self): return ret def get_agent(self, agent_name, raise_error=True): - object_name = self.agent_to_object(agent_name) + try: + object_name = self.agent_to_object(agent_name) + except KeyError: + if raise_error: + raise ValueError("Object {} not found!".format(object_name)) + else: + return None if object_name in self._active_objects: return self._active_objects[object_name] elif object_name in self._dying_objects: