From e22e372306aa47a9b4b9c0f4055d5d15d26401a1 Mon Sep 17 00:00:00 2001 From: Juan Rocamonde Date: Wed, 28 Sep 2022 12:06:07 +0200 Subject: [PATCH] Fix duplicate key error in HumanOutputFormat (#1079) * Fix duplicate key error in HumanOutputFormat * Update changelog * Add test * Update changelog.rst Co-authored-by: Adam Gleave Co-authored-by: Adam Gleave --- docs/misc/changelog.rst | 2 +- stable_baselines3/common/logger.py | 11 ++++++----- tests/test_logger.py | 9 +++++++++ 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 8fd853b..43f9a08 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -32,7 +32,7 @@ Bug Fixes: - Fixed loading saved model with different number of envrionments - Removed ``forward()`` abstract method declaration from ``common.policies.BaseModel`` (already defined in ``torch.nn.Module``) to fix type errors in subclasses (@Rocamonde) - Fixed the return type of ``.load()`` and ``.learn()`` methods in ``BaseAlgorithm`` so that they now use ``TypeVar`` (@Rocamonde) - +- Fixed an issue where keys with different tags but the same key raised an error in ``common.logger.HumanOutputFormat`` (@Rocamonde and @AdamGleave) Deprecations: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/logger.py b/stable_baselines3/common/logger.py index c1e8433..31e5655 100644 --- a/stable_baselines3/common/logger.py +++ b/stable_baselines3/common/logger.py @@ -193,30 +193,31 @@ class HumanOutputFormat(KVWriter, SeqWriter): if key.find("/") > 0: # Find tag and add it to the dict tag = key[: key.find("/") + 1] - key2str[self._truncate(tag)] = "" + key2str[(tag, self._truncate(tag))] = "" # Remove tag from key if tag is not None and tag in key: key = str(" " + key[len(tag) :]) truncated_key = self._truncate(key) - if truncated_key in key2str: + if (tag, truncated_key) in key2str: raise ValueError( f"Key '{key}' truncated to '{truncated_key}' that already exists. Consider increasing `max_length`." ) - key2str[truncated_key] = self._truncate(value_str) + key2str[(tag, truncated_key)] = self._truncate(value_str) # Find max widths if len(key2str) == 0: warnings.warn("Tried to write empty key-value dict") return else: - key_width = max(map(len, key2str.keys())) + tagless_keys = map(lambda x: x[1], key2str.keys()) + key_width = max(map(len, tagless_keys)) val_width = max(map(len, key2str.values())) # Write out the data dashes = "-" * (key_width + val_width + 7) lines = [dashes] - for key, value in key2str.items(): + for (_, key), value in key2str.items(): key_space = " " * (key_width - len(key)) val_space = " " * (val_width - len(value)) lines.append(f"| {key}{key_space} | {value}{val_space} |") diff --git a/tests/test_logger.py b/tests/test_logger.py index 516a622..92b65e8 100644 --- a/tests/test_logger.py +++ b/tests/test_logger.py @@ -1,4 +1,5 @@ import os +import sys import time from typing import Sequence from unittest import mock @@ -409,3 +410,11 @@ def test_fps_no_div_zero(algo): with mock.patch("time.time_ns", lambda: 42.0): model = algo("MlpPolicy", "CartPole-v1") model.learn(total_timesteps=100) + + +def test_human_output_format_no_crash_on_same_keys_different_tags(): + o = HumanOutputFormat(sys.stdout, max_length=60) + o.write( + {"key1/foo": "value1", "key1/bar": "value2", "key2/bizz": "value3", "key2/foo": "value4"}, + {"key1/foo": None, "key2/bizz": None, "key1/bar": None, "key2/foo": None}, + )