added scenario self stop, agent setup, action masking in PPO

ai_economist/foundation/base/base_env.py

@@ -234,7 +234,7 @@ class BaseEnvironment(ABC):
         self.num_agents = (
             n_agents + n_planners
         )  # used in the warp_drive env wrapper (+ 1 for the planner)
-
+        
         # Components must be a tuple/list where each element is either a...
         #   tuple: ('Component Name', {Component kwargs})
         #   dict : {'Component Name': {Component kwargs}}
@@ -345,11 +345,11 @@ class BaseEnvironment(ABC):
        
         self.world.planner.register_inventory(self.resources)
         self.world.planner.register_components(self._components)
-        self.apply_scenario_config_to_agents()
+        self.reapply_scenario_config_to_agents()
        
 
         self._completions = 0
-
+        self._finish_episode=False
         self._last_ep_metrics = None
 
         # For dense logging
@@ -366,7 +366,7 @@ class BaseEnvironment(ABC):
         # into a single agent with index 'a'
         self.collate_agent_step_and_reset_data = collate_agent_step_and_reset_data
     
-    def apply_scenario_config_to_agents(self):
+    def reapply_scenario_config_to_agents(self):
         # Register the components with the agents
         # to finish setting up their state/action spaces.
         for agent in self.world.agents:
@@ -506,6 +506,8 @@ class BaseEnvironment(ABC):
 
     # Getters & Setters
     # -----------------
+    def set_finish_episode(self,done):
+        self._finish_episode=done
 
     def get_component(self, component_name):
         """
@@ -909,6 +911,9 @@ class BaseEnvironment(ABC):
         # Reset the timestep counter
         self.world.timestep = 0
 
+        # Reset done flag
+        self._finish_episode=False
+
         # Perform the scenario reset,
         # which includes resetting the world and agent states
         self.reset_starting_layout()
@@ -1021,7 +1026,7 @@ class BaseEnvironment(ABC):
             flatten_masks=self._flatten_masks,
         )
         rew = self._generate_rewards()
-        done = {"__all__": self.world.timestep >= self._episode_length}
+        done = {"__all__": self.world.timestep >= self._episode_length | self._finish_episode}
         info = {k: {} for k in obs.keys()}
 
         if self._dense_log_this_episode: