fix score

markub3327 · Oct 12, 2023 · 37fa500 · 37fa500
1 parent bf36280
commit 37fa500
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/rl_toolkit/agents/dueling_dqn/agent.py b/rl_toolkit/agents/dueling_dqn/agent.py
@@ -134,10 +134,14 @@ def collect(self, writer, policy):
         action = np.array(action, copy=False, dtype=self._env.action_space.dtype)
 
         # Perform action
-        new_obs, ext_reward, terminated, truncated, _ = self._env.step(action)
+        new_obs, ext_reward, terminated, truncated, info = self._env.step(action)
 
         # Update variables
-        self._episode_reward += ext_reward
+        try:
+            self._episode_reward = info["score"]
+        except KeyError:
+            self._episode_reward += ext_reward
+
         self._episode_steps += 1
         self._total_steps += 1