collecting latest step as a stat

ml-agents/mlagents/plugins/stats_writer.py

@@ -31,6 +31,7 @@ def get_default_stats_writers(run_options: RunOptions) -> List[StatsWriter]:
         TensorboardWriter(
             checkpoint_settings.write_path,
             clear_past_data=not checkpoint_settings.resume,
+            hidden_keys=["Is Training", "Step"],
         ),
         GaugeWriter(),
         ConsoleWriter(),

ml-agents/mlagents/trainers/stats.py

@@ -1,6 +1,6 @@
 from collections import defaultdict
 from enum import Enum
-from typing import List, Dict, NamedTuple, Any
+from typing import List, Dict, NamedTuple, Any, Optional
 import numpy as np
 import abc
 import os
@@ -14,7 +14,6 @@
 from torch.utils.tensorboard import SummaryWriter
 from mlagents.torch_utils.globals import get_rank
 
-
 logger = get_logger(__name__)
 
 
@@ -212,24 +211,34 @@ 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: Optional[List[str]] = None,
+    ):
         """
         A StatsWriter that writes to a Tensorboard summary.
 
         :param base_dir: The directory within which to place all the summaries. Tensorboard files will be written to a
         {base_dir}/{category} directory.
         :param clear_past_data: Whether or not to clean up existing Tensorboard files associated with the base_dir and
         category.
+        :param hidden_keys: If provided, Tensorboard Writer won't write statistics identified with these Keys in
+        Tensorboard summary.
         """
         self.summary_writers: Dict[str, SummaryWriter] = {}
         self.base_dir: str = base_dir
         self._clear_past_data = clear_past_data
+        self.hidden_keys: List[str] = hidden_keys if hidden_keys is not None else []
 
     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/tests/test_stats.py

@@ -129,6 +129,31 @@ def test_tensorboard_writer_clear(tmp_path):
     assert len(os.listdir(os.path.join(tmp_path, "category1"))) == 1
 
 
+@mock.patch("mlagents.trainers.stats.SummaryWriter")
+def test_tensorboard_writer_hidden_keys(mock_summary):
+    # Test write_stats
+    category = "category1"
+    with tempfile.TemporaryDirectory(prefix="unittest-") as base_dir:
+        tb_writer = TensorboardWriter(
+            base_dir, clear_past_data=False, hidden_keys="hiddenKey"
+        )
+        statssummary1 = StatsSummary(
+            full_dist=[1.0], aggregation_method=StatsAggregationMethod.AVERAGE
+        )
+        tb_writer.write_stats("category1", {"hiddenKey": statssummary1}, 10)
+
+        # Test that the filewriter has been created and the directory has been created.
+        filewriter_dir = "{basedir}/{category}".format(
+            basedir=base_dir, category=category
+        )
+        assert os.path.exists(filewriter_dir)
+        mock_summary.assert_called_once_with(filewriter_dir)
+
+        # Test that the filewriter was not written to since we used the hidden key.
+        mock_summary.return_value.add_scalar.assert_not_called()
+        mock_summary.return_value.flush.assert_not_called()
+
+
 def test_gauge_stat_writer_sanitize():
     assert GaugeWriter.sanitize_string("Policy/Learning Rate") == "Policy.LearningRate"
     assert (

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

@@ -211,6 +211,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:
         """