HumanOutputFormat: make length configurable, throw error if keys alias (#756)

* Make HumanOutputFormat length configurable and bump to 36 by default

* Add test case

* Updated changelog

* Blacken

* Blacken code

* Fix GitLab CI: switch to Docker container with new black version

* Incorporate suggestion

* Add class docstring

* Dummy commit to retrigger GitLab

Co-authored-by: Anssi <kaneran21@hotmail.com>

docs/misc/changelog.rst

@@ -14,6 +14,8 @@ Breaking Changes:
 
 New Features:
 ^^^^^^^^^^^^^
+- Makes the length of keys and values in `HumanOutputFormat` configurable,
+  depending on desired maximum width of output.
 
 SB3-Contrib
 ^^^^^^^^^^^
@@ -21,6 +23,9 @@ SB3-Contrib
 Bug Fixes:
 ^^^^^^^^^^
 - Fixed a bug in ``VecMonitor``. The monitor did not consider the ``info_keywords`` during stepping (@ScheiklP)
+- Fixed a bug in ``HumanOutputFormat``. Distinct keys truncated to the same prefix would overwrite each others value,
+  resulting in only one being output. This now raises an error (this should only affect a small fraction of use cases
+  with very long keys.)
 
 Deprecations:
 ^^^^^^^^^^^^^

stable_baselines3/common/logger.py

@@ -114,12 +114,24 @@ class SeqWriter(object):
 
 
 class HumanOutputFormat(KVWriter, SeqWriter):
-    def __init__(self, filename_or_file: Union[str, TextIO]):
+    """A human-readable output format producing ASCII tables of key-value pairs.
+
+    Set attribute `max_length` to change the maximum length of keys and values
+    to write to output (or specify it when calling `__init__`).
+    """
+
+    def __init__(self, filename_or_file: Union[str, TextIO], max_length: int = 36):
         """
         log to a file, in a human readable format
 
         :param filename_or_file: the file to write the log to
+        :param max_length: the maximum length of keys and values to write to output.
+            Outputs longer than this will be truncated. An error will be raised
+            if multiple keys are truncated to the same value. The maximum output
+            width will be ``2*max_length + 7``. The default of 36 produces output
+            no longer than 79 characters wide.
         """
+        self.max_length = max_length
         if isinstance(filename_or_file, str):
             self.file = open(filename_or_file, "wt")
             self.own_file = True
@@ -159,7 +171,12 @@ class HumanOutputFormat(KVWriter, SeqWriter):
             if tag is not None and tag in key:
                 key = str("   " + key[len(tag) :])
 
-            key2str[self._truncate(key)] = self._truncate(value_str)
+            truncated_key = self._truncate(key)
+            if truncated_key in key2str:
+                raise ValueError(
+                    f"Key '{key}' truncated to " f"'{truncated_key}' that already exists. Consider increasing `max_length`."
+                )
+            key2str[truncated_key] = self._truncate(value_str)
 
         # Find max widths
         if len(key2str) == 0:
@@ -182,9 +199,10 @@ class HumanOutputFormat(KVWriter, SeqWriter):
         # Flush the output to the file
         self.file.flush()
 
-    @classmethod
-    def _truncate(cls, string: str, max_length: int = 23) -> str:
-        return string[: max_length - 3] + "..." if len(string) > max_length else string
+    def _truncate(self, string: str) -> str:
+        if len(string) > self.max_length:
+            string = string[: self.max_length - 3] + "..."
+        return string
 
     def write_sequence(self, sequence: List) -> None:
         sequence = list(sequence)

tests/test_logger.py

@@ -295,6 +295,37 @@ def test_report_figure_to_unsupported_format_raises_error(tmp_path, unsupported_
     writer.close()
 
 
+def test_key_length(tmp_path):
+    writer = make_output_format("stdout", tmp_path)
+    assert writer.max_length == 36
+    long_prefix = "a" * writer.max_length
+
+    ok_dict = {
+        # keys truncated but not aliased -- OK
+        "a" + long_prefix: 42,
+        "b" + long_prefix: 42,
+        # values truncated and aliased -- also OK
+        "foobar": long_prefix + "a",
+        "fizzbuzz": long_prefix + "b",
+    }
+    ok_excluded = {k: None for k in ok_dict}
+    writer.write(ok_dict, ok_excluded)
+
+    long_key_dict = {
+        long_prefix + "a": 42,
+        "foobar": "sdf",
+        long_prefix + "b": 42,
+    }
+    long_key_excluded = {k: None for k in long_key_dict}
+    # keys truncated and aliased -- not OK
+    with pytest.raises(ValueError, match="Key.*truncated"):
+        writer.write(long_key_dict, long_key_excluded)
+
+    # Just long enough to not be truncated now
+    writer.max_length += 1
+    writer.write(long_key_dict, long_key_excluded)
+
+
 class TimeDelayEnv(gym.Env):
     """
     Gym env for testing FPS logging.