changed a lot of things but chrafting

ai_economist/foundation/base/world.py

@@ -91,7 +91,10 @@ class Maps:
 
             else:
                 raise NotImplementedError
-
+        self.reset_agent_maps(n_agents)
+       
+    def reset_agent_maps(self,n_agents):
+        self.n_agents=n_agents
         self._idx_map = np.stack(
             [i * np.ones(shape=self.size) for i in range(self.n_agents)]
         )
@@ -378,17 +381,8 @@ class World:
         self.multi_action_mode_planner = bool(multi_action_mode_planner)
         self._agent_class_idx_map={}
         #create agents
-        self.agent_composition=agent_composition
-        self.n_agents=0
-        self._agents = []
-        for k,v in agent_composition.items():
-            self._agent_class_idx_map[k]=[]
-            for offset in range(v):
-                agent_class=agent_registry.get(k)
-                agent=agent_class(self.n_agents,self.multi_action_mode_agents)
-                self._agents.append(agent)
-                self._agent_class_idx_map[k].append(str(self.n_agents))
-                self.n_agents+=1
+        self.create_agents(agent_composition)
+        
         self.maps = Maps(world_size, self.n_agents, world_resources, world_landmarks)
 
         planner_class = agent_registry.get("BasicPlanner")
@@ -402,6 +396,37 @@ class World:
         self.cuda_function_manager = None
         self.cuda_data_manager = None
 
+    def create_agents(self, agent_composition):
+        """create_agents creates the world agent db with the given compostition."""
+        self.agent_composition=agent_composition
+        self.n_agents=0
+        self._agents = []
+        for k,v in agent_composition.items():
+            self._agent_class_idx_map[k]=[]
+            for offset in range(v):
+                agent_class=agent_registry.get(k)
+                agent=agent_class(self.n_agents,self.multi_action_mode_agents)
+                self._agents.append(agent)
+                self._agent_class_idx_map[k].append(str(self.n_agents))
+                self.n_agents+=1
+
+    def apply_agent_db(self):
+        """Applys current agent db into lookup maps inside world and map itself. Enables insertion of new agents into existing env."""
+        self.n_agents=len(self._agents)
+        self._agent_class_idx_map={}
+        self.maps.reset_agent_maps(self.n_agents) # reset map lookups
+        #create mapping dict
+        for idx in range(self.n_agents):
+            cls=self.get_agent_class(idx)
+            agent=self._agents[idx]
+            if cls in self._agent_class_idx_map:
+                self._agent_class_idx_map[cls].append(idx)
+            else:
+                self._agent_class_idx_map[cls]=[idx]
+            # apply agent locs db to maps
+            if "loc" in agent.state:
+                self.maps.set_agent_loc(agent,*agent.loc)
+
     @property
     def agents(self):
         """Return a list of the agent objects in the world (sorted by index)."""