feat: Add default reward based on win/lose

strakam · Dec 6, 2024 · 7b2b5b8 · 7b2b5b8
1 parent 0e540b6
commit 7b2b5b8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/generals/envs/gymnasium_generals.py b/generals/envs/gymnasium_generals.py
@@ -105,7 +105,6 @@ def step(self, action: Action) -> tuple[Observation, SupportsFloat, bool, bool,
         actions = {self.agent_id: action, self.npc.id: npc_action}
 
         observations, infos = self.game.step(actions)
-        infos = {agent_id: {} for agent_id in self.agent_ids}
 
         # From observations of all agents, pick only those relevant for the main agent
         obs = observations[self.agent_id].as_dict()
@@ -131,6 +130,8 @@ def _default_reward(
         done: bool,
         info: Info,
     ) -> Reward:
+        if done:
+            return 1 if info["is_winner"] else -1
         return 0
 
     def close(self) -> None: