Fix: Masked DQN Grayscale

BluemlJ · Jan 16, 2025 · 79a80c6 · 79a80c6
1 parent 06545d5
commit 79a80c6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cleanrl/dqn_atari_oc.py b/cleanrl/dqn_atari_oc.py
@@ -342,7 +342,7 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
                 if global_step % 100 == 0:
                     writer.add_scalar("losses/td_loss", loss, global_step)
                     writer.add_scalar("losses/q_values", old_val.mean().item(), global_step)
-                    print("SPS:", int(global_step / (time.time() - start_time)))
+                    # print("SPS:", int(global_step / (time.time() - start_time)))
                     writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)
 
                 # optimize the model
@@ -363,7 +363,7 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
                 writer.add_scalar("charts/Episodic_New_Reward", enewr / count, global_step)
             writer.add_scalar("charts/Episodic_Original_Reward", eorgr / count, global_step)
             writer.add_scalar("charts/Episodic_Length", elength / count, global_step)
-            #pbar.set_description(f"Reward: {eorgr / count:.1f}")
+            pbar.set_description(f"Reward: {eorgr / count:.1f}")
             elength = 0
             eorgr = 0
             enewr = 0