Fix the reporting of histogram stats and adding a test (Unity-Technol…

…ogies#5410) * Fix the reporting of histogram stats and adding a test * Appending to the Changelog
surfnerd · Jun 1, 2021 · c7a1799 · c7a1799
1 parent 8b6dd3d
commit c7a1799
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 0 deletions.
diff --git a/com.unity.ml-agents/CHANGELOG.md b/com.unity.ml-agents/CHANGELOG.md
@@ -19,6 +19,7 @@ and this project adheres to
 - The calculation of the target entropy of SAC with continuous actions was incorrect and has been fixed. (#5372)
 - RigidBodySensorComponent now displays a warning if it's used in a way that won't generate useful observations. (#5387)
 - Update the documentation with a note saying that `GridSensor` does not work in 2D environments. (#5396)
+- Fixed an issue where the histogram stats would not be reported correctly in TensorBoard. (#5410)
 
 
 ## [2.0.0-exp.1] - 2021-04-22

diff --git a/ml-agents/mlagents/trainers/agent_processor.py b/ml-agents/mlagents/trainers/agent_processor.py
@@ -15,6 +15,7 @@
     StatsAggregationMethod,
     EnvironmentStats,
 )
+from mlagents.trainers.exception import UnityTrainerException
 from mlagents.trainers.trajectory import AgentStatus, Trajectory, AgentExperience
 from mlagents.trainers.policy import Policy
 from mlagents.trainers.action_info import ActionInfo, ActionInfoOutputs
@@ -438,8 +439,14 @@ def record_environment_stats(
                     self._stats_reporter.add_stat(stat_name, val, agg_type)
                 elif agg_type == StatsAggregationMethod.SUM:
                     self._stats_reporter.add_stat(stat_name, val, agg_type)
+                elif agg_type == StatsAggregationMethod.HISTOGRAM:
+                    self._stats_reporter.add_stat(stat_name, val, agg_type)
                 elif agg_type == StatsAggregationMethod.MOST_RECENT:
                     # In order to prevent conflicts between multiple environments,
                     # only stats from the first environment are recorded.
                     if worker_id == 0:
                         self._stats_reporter.set_stat(stat_name, val)
+                else:
+                    raise UnityTrainerException(
+                        f"Unknown StatsAggregationMethod encountered. {agg_type}"
+                    )
diff --git a/ml-agents/mlagents/trainers/tests/test_stats.py b/ml-agents/mlagents/trainers/tests/test_stats.py
@@ -5,6 +5,7 @@
 import unittest
 import time
 
+
 from mlagents.trainers.stats import (
     StatsReporter,
     TensorboardWriter,
@@ -15,6 +16,8 @@
     StatsAggregationMethod,
 )
 
+from mlagents.trainers.env_manager import AgentManager
+
 
 def test_stat_reporter_add_summary_write():
     # Test add_writer
@@ -107,6 +110,27 @@ def test_tensorboard_writer(mock_summary):
         assert mock_summary.return_value.add_text.call_count >= 1
 
 
+@pytest.mark.parametrize("aggregation_type", list(StatsAggregationMethod))
+def test_agent_manager_stats_report(aggregation_type):
+    stats_reporter = StatsReporter("recorder_name")
+    manager = AgentManager(None, "behaviorName", stats_reporter)
+
+    values = range(5)
+
+    env_stats = {"stat": [(i, aggregation_type) for i in values]}
+    manager.record_environment_stats(env_stats, 0)
+    summary = stats_reporter.get_stats_summaries("stat")
+    aggregation_result = {
+        StatsAggregationMethod.AVERAGE: sum(values) / len(values),
+        StatsAggregationMethod.MOST_RECENT: values[-1],
+        StatsAggregationMethod.SUM: sum(values),
+        StatsAggregationMethod.HISTOGRAM: sum(values) / len(values),
+    }
+
+    assert summary.aggregated_value == aggregation_result[aggregation_type]
+    stats_reporter.write_stats(0)
+
+
 def test_tensorboard_writer_clear(tmp_path):
     tb_writer = TensorboardWriter(tmp_path, clear_past_data=False)
     statssummary1 = StatsSummary(