env steps (#307)

alex-petrenko · Nov 19, 2024 · e9589e4 · e9589e4
1 parent 46402c3
commit e9589e4
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/sample_factory/algo/learning/learner.py b/sample_factory/algo/learning/learner.py
@@ -843,6 +843,7 @@ def _record_summaries(self, train_loop_vars) -> AttrDict:
         self.last_summary_time = time.time()
         stats = AttrDict()
 
+        stats.env_steps = self.env_steps
         stats.lr = self.curr_lr
         stats.actual_lr = train_loop_vars.actual_lr  # potentially scaled because of masked data
 

diff --git a/sample_factory/utils/wandb_utils.py b/sample_factory/utils/wandb_utils.py
@@ -56,6 +56,14 @@ def init_wandb_func():
 
     wandb.config.update(cfg, allow_val_change=True)
 
+    wandb.define_metric("train/env_steps")
+    wandb.define_metric("train/*", step_metric="train/env_steps")
+    wandb.define_metric("perf/*", step_metric="train/env_steps")
+    wandb.define_metric("len/*", step_metric="train/env_steps")
+    wandb.define_metric("policy_stats/*", step_metric="train/env_steps")
+    wandb.define_metric("reward/*", step_metric="train/env_steps")
+    wandb.define_metric("stats/*", step_metric="train/env_steps")
+
 
 def finish_wandb(cfg):
     if cfg.with_wandb: