move logging

pytorch · Dec 5, 2023 · 2346870 · 2346870
1 parent fb6f61a
commit 2346870
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/examples/dqn/dqn_atari.py b/examples/dqn/dqn_atari.py
@@ -140,6 +140,9 @@ def main(cfg: "DictConfig"):  # noqa: F821
             )
 
         if collected_frames < init_random_frames:
+            if logger:
+                for key, value in log_info.items():
+                    logger.log_scalar(key, value, step=collected_frames)
             continue
 
         # optimization steps

diff --git a/examples/dqn/dqn_cartpole.py b/examples/dqn/dqn_cartpole.py
@@ -126,7 +126,11 @@ def main(cfg: "DictConfig"):  # noqa: F821
             )
 
         if collected_frames < init_random_frames:
-            continue
+            if collected_frames < init_random_frames:
+                if logger:
+                    for key, value in log_info.items():
+                        logger.log_scalar(key, value, step=collected_frames)
+                continue
 
         # optimization steps
         training_start = time.time()