adding current lr to stats for monitoring

Signed-off-by: Shakti Kumar <[email protected]>
jerrodparker20 · Mar 30, 2020 · 9b293c9 · 9b293c9
1 parent 4f1da87
commit 9b293c9
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/train.py b/train.py
@@ -382,7 +382,8 @@ def learn(
             "baseline_loss": 0,
             "entropy_loss": 0,
             "num_unpadded_steps": 0,
-            "len_max_traj": 0
+            "len_max_traj": 0,
+            "learning_rate":optimizer.param_groups[0]['lr']
         }
 
         logging.debug('AT LEARN')
@@ -755,6 +756,7 @@ def lr_lambda(epoch):
         "pg_loss",
         "baseline_loss",
         "entropy_loss",
+        "learning_rate",
     ]
     logger.info("# Step\t%s", "\t".join(stat_keys))
 
@@ -905,13 +907,12 @@ def checkpoint():
             #     best_val_loss = val_loss
 
             logging.info(
-                "Steps %i @ %.1f SPS. Loss %f. %sStats:\n%s\n%s",
+                "Steps %i @ %.1f SPS. Loss %f. %sStats:\n%s",
                 step,
                 sps,
                 total_loss,
                 mean_return,
                 pprint.pformat(stats),
-                optimizer.params_group[0]['lr']
             )
     except KeyboardInterrupt:
         return  # Try joining actors then quit.