collecting latest step as a stat

ml-agents/mlagents/trainers/stats.py

@@ -14,7 +14,6 @@
 from torch.utils.tensorboard import SummaryWriter
 from mlagents.torch_utils.globals import get_rank
 
-
 logger = get_logger(__name__)
 
 
@@ -212,7 +211,12 @@ def add_property(
 
 
 class TensorboardWriter(StatsWriter):
-    def __init__(self, base_dir: str, clear_past_data: bool = False):
+    def __init__(
+        self,
+        base_dir: str,
+        clear_past_data: bool = False,
+        hidden_keys: List[str] = None,
+    ):
         """
         A StatsWriter that writes to a Tensorboard summary.
 
@@ -224,12 +228,17 @@ def __init__(self, base_dir: str, clear_past_data: bool = False):
         self.summary_writers: Dict[str, SummaryWriter] = {}
         self.base_dir: str = base_dir
         self._clear_past_data = clear_past_data
+        if hidden_keys is None:
+            hidden_keys = ["Is Training", "Step"]
+        self.hidden_keys: List[str] = hidden_keys
 
     def write_stats(
         self, category: str, values: Dict[str, StatsSummary], step: int
     ) -> None:
         self._maybe_create_summary_writer(category)
         for key, value in values.items():
+            if key in self.hidden_keys:
+                continue
             self.summary_writers[category].add_scalar(
                 f"{key}", value.aggregated_value, step
             )

ml-agents/mlagents/trainers/trainer/rl_trainer.py

@@ -209,6 +209,7 @@ def _increment_step(self, n_steps: int, name_behavior_id: str) -> None:
         p = self.get_policy(name_behavior_id)
         if p:
             p.increment_step(n_steps)
+        self.stats_reporter.set_stat("Step", float(self.get_step))
 
     def _get_next_interval_step(self, interval: int) -> int:
         """