fix: Fixed reward class

UFRN-URNAI · Aug 13, 2024 · 63d565a · 63d565a
1 parent e8ae2bd
commit 63d565a
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/urnai/agents/agent_base.py b/urnai/agents/agent_base.py
@@ -45,7 +45,7 @@ def reset(self, episode=0) -> None:
         self.previous_state = None
         self.action_space.reset()
         self.model.ep_reset(episode)
-        # self.reward.reset()
+        self.reward.reset()
         self.state_space.reset()
 
     def learn(self, obs, reward, done) -> None:

diff --git a/urnai/rewards/reward_base.py b/urnai/rewards/reward_base.py
@@ -10,3 +10,6 @@ class RewardBase(ABC):
 
     @abstractmethod
     def get_reward(self, obs, reward, done) -> int: ...
+
+    @abstractmethod
+    def reset(self) -> None: ...
diff --git a/urnai/sc2/rewards/collectables.py b/urnai/sc2/rewards/collectables.py
@@ -27,6 +27,10 @@ def get_reward(self, obs) -> int:
 
         self.previous_state = obs
         return reward
+
+    def reset(self) -> None:
+        self.previous_state = None
+        self.old_collectable_counter = STATE_MAXIMUM_NUMBER_OF_MINERAL_SHARDS
 
     def filter_non_mineral_shard_units(self, obs):
         filtered_map = np.zeros((len(obs.feature_minimap[0]),