reward logging fix

aLadNamedPat · Sep 1, 2024 · c41028e · c41028e
1 parent ffc41ed
commit c41028e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/Dreamer.py b/Dreamer.py
@@ -225,6 +225,7 @@ def agent_update(
     def rollout(
         self,
     ):
+        total_rewards = 0
         for t in range(self.batch_train_freq):
             self.num_timesteps += 1
             done = False
@@ -258,8 +259,10 @@ def rollout(
             self.replayBuffer.add(self.last_obs, action, timestep.reward, obs, done)
             self.last_obs = obs
 
-            # Log rewards incrementally
-            wandb.log({"reward": timestep.reward, "num_timesteps": self.num_timesteps})
+            total_rewards += timestep.reward
+
+        # Log total rewards after every rollout
+        wandb.log({"total_rewards": total_rewards, "num_timesteps": self.num_timesteps})
 
     def train(
         self,

diff --git a/__pycache__/Dreamer.cpython-310.pyc b/__pycache__/Dreamer.cpython-310.pyc