From e56c0e3475493f6c05686ceca7d8a0e69c9955aa Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 26 Sep 2022 11:46:34 +0100
Subject: [PATCH 01/11] Add support for prefix context manager in logger (from
 #529)

---
 src/imitation/util/logger.py | 148 +++++++++++++++++++++++------------
 tests/util/test_logger.py    |  80 ++++++++++---------
 2 files changed, 145 insertions(+), 83 deletions(-)

diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index b60b4d74e..f2bd33b9d 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -3,45 +3,14 @@
 import contextlib
 import datetime
 import os
-import sys
 import tempfile
-from typing import Any, Dict, Generator, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Generator, List, Optional, Sequence, Tuple, Union
 
 import stable_baselines3.common.logger as sb_logger
 
 from imitation.data import types
 
 
-def make_output_format(
-    _format: str,
-    log_dir: str,
-    log_suffix: str = "",
-    max_length: int = 40,
-) -> sb_logger.KVWriter:
-    """Returns a logger for the requested format.
-
-    Args:
-        _format: the requested format to log to
-            ('stdout', 'log', 'json' or 'csv' or 'tensorboard').
-        log_dir: the logging directory.
-        log_suffix: the suffix for the log file.
-        max_length: the maximum length beyond which the keys get truncated.
-
-    Returns:
-        the logger.
-    """
-    os.makedirs(log_dir, exist_ok=True)
-    if _format == "stdout":
-        return sb_logger.HumanOutputFormat(sys.stdout, max_length=max_length)
-    elif _format == "log":
-        return sb_logger.HumanOutputFormat(
-            os.path.join(log_dir, f"log{log_suffix}.txt"),
-            max_length=max_length,
-        )
-    else:
-        return sb_logger.make_output_format(_format, log_dir, log_suffix)
-
-
 def _build_output_formats(
     folder: str,
     format_strs: Sequence[str],
@@ -50,7 +19,7 @@ def _build_output_formats(
 
     Args:
         folder: Path to directory that logs are written to.
-        format_strs: A list of output format strings. For details on available
+        format_strs: An list of output format strings. For details on available
             output formats see `stable_baselines3.logger.make_output_format`.
 
     Returns:
@@ -62,7 +31,7 @@ def _build_output_formats(
         if f == "wandb":
             output_formats.append(WandbOutputFormat())
         else:
-            output_formats.append(make_output_format(f, folder))
+            output_formats.append(sb_logger.make_output_format(f, folder))
     return output_formats
 
 
@@ -72,8 +41,57 @@ class HierarchicalLogger(sb_logger.Logger):
     `self.accumulate_means` creates a context manager. While in this context,
     values are loggged to a sub-logger, with only mean values recorded in the
     top-level (root) logger.
+
+    >>> import tempfile
+    >>> with tempfile.TemporaryDirectory() as dir:
+    ...     logger: HierarchicalLogger = configure(dir, ('log',))
+    ...     # record the key value pair (loss, 1.0) to path `dir`
+    ...     # at step 1.
+    ...     logger.record("loss", 1.0)
+    ...     logger.dump(step=1)
+    ...     with logger.accumulate_means("dataset"):
+    ...         # record the key value pair `("raw/dataset/entropy", 5.0)` to path
+    ...         # `dir/raw/dataset` at step 100
+    ...         logger.record("entropy", 5.0)
+    ...         logger.dump(step=100)
+    ...         # record the key value pair `("raw/dataset/entropy", 6.0)` to path
+    ...         # `dir/raw/dataset` at step 200
+    ...         logger.record("entropy", 6.0)
+    ...         logger.dump(step=200)
+    ...     # record the key value pair `("mean/dataset/entropy", 5.5)` to path
+    ...     # `dir` at step 1.
+    ...     logger.dump(step=1)
+    ...     with logger.add_prefix("foo"), logger.accumulate_means("bar"):
+    ...         # record the key value pair ("raw/foo/bar/biz", 42.0) to path
+    ...         # `dir/raw/foo/bar` at step 2000
+    ...         logger.record("biz", 42.0)
+    ...         logger.dump(step=2000)
+    ...     # record the key value pair `("mean/foo/bar/biz", 42.0)` to path
+    ...     # `dir` at step 1.
+    ...     logger.dump(step=1)
+    ...     with open(os.path.join(dir, 'log.txt')) as f:
+    ...         print(f.read())
+    -------------------
+    | loss | 1        |
+    -------------------
+    ---------------------------------
+    | mean/              |          |
+    |    dataset/entropy | 5.5      |
+    ---------------------------------
+    -----------------------------
+    | mean/          |          |
+    |    foo/bar/biz | 42       |
+    -----------------------------
+    <BLANKLINE>
     """
 
+    default_logger: sb_logger.Logger
+    current_logger: Optional[sb_logger.Logger]
+    _cached_loggers: Dict[str, sb_logger.Logger]
+    _prefixes: List[str]
+    _subdir: Optional[str]
+    _name: Optional[str]
+
     def __init__(
         self,
         default_logger: sb_logger.Logger,
@@ -93,7 +111,9 @@ def __init__(
         self.default_logger = default_logger
         self.current_logger = None
         self._cached_loggers = {}
+        self._prefixes = []
         self._subdir = None
+        self._name = None
         self.format_strs = format_strs
         super().__init__(folder=self.default_logger.dir, output_formats=[])
 
@@ -103,27 +123,56 @@ def _update_name_to_maps(self) -> None:
         self.name_to_excluded = self._logger.name_to_excluded
 
     @contextlib.contextmanager
-    def accumulate_means(self, subdir: types.AnyPath) -> Generator[None, None, None]:
+    def add_prefix(self, prefix: str) -> Generator[None, None, None]:
+        """Add a prefix to the subdirectory used to accumulate means.
+
+        This prefix only applies when a `accumulate_means` context is active. If there
+        are multiple active prefixes, then they are concatenated.
+
+        Args:
+            prefix: The prefix to add to the named sub.
+
+        Yields:
+            None when the context manager is entered
+
+        Raises:
+            RuntimeError: if accumulate means context is already active.
+        """
+        if self.current_logger is not None:
+            raise RuntimeError(
+                "Cannot add prefix when accumulate_means context is already active.",
+            )
+
+        try:
+            self._prefixes.append(prefix)
+            yield
+        finally:
+            self._prefixes.pop()
+
+    @contextlib.contextmanager
+    def accumulate_means(self, name: str) -> Generator[None, None, None]:
         """Temporarily modifies this HierarchicalLogger to accumulate means values.
 
-        During this context, `self.record(key, value)` writes the "raw" values in
-        "{self.default_logger.log_dir}/{subdir}" under the key "raw/{subdir}/{key}".
-        At the same time, any call to `self.record` will also accumulate mean values
-        on the default logger by calling
-        `self.default_logger.record_mean(f"mean/{subdir}/{key}", value)`.
+        Within this context manager, `self.record(key, value)` writes the "raw" values
+        in `f"{self.default_logger.log_dir}/{prefix}/{name}"` under the key
+        `"raw/{prefix}/{name}/{key}"`. At the same time, any call to `self.record` will
+        also accumulate mean values on the default logger by calling
+        `self.default_logger.record_mean(f"mean/{prefix}/{name}/{key}", value)`.
 
-        During the context, `self.record(key, value)` will write the "raw" values in
-        `"{self.default_logger.log_dir}/subdir"` under the key "raw/{subdir}/key".
+        Multiple prefixes may be active at once. In this case the `prefix` is simply the
+        concatenation of each of the active prefixes in the order they
+        where created e.g. if the active `prefixes` are ['foo', 'bar'] then
+        the `prefix` is 'foo/bar'.
 
         After the context exits, calling `self.dump()` will write the means
         of all the "raw" values accumulated during this context to
-        `self.default_logger` under keys with the prefix `mean/{subdir}/`
+        `self.default_logger` under keys of the form  `mean/{prefix}/{name}/{key}`
 
         Note that the behavior of other logging methods, `log` and `record_mean`
         are unmodified and will go straight to the default logger.
 
         Args:
-            subdir: A string key which determines the `folder` where raw data is
+            name: A string key which determines the `folder` where raw data is
                 written and temporary logging prefixes for raw and mean data. Entering
                 an `accumulate_means` context in the future with the same `subdir`
                 will safely append to logs written in this folder rather than
@@ -139,10 +188,11 @@ def accumulate_means(self, subdir: types.AnyPath) -> Generator[None, None, None]
         if self.current_logger is not None:
             raise RuntimeError("Nested `accumulate_means` context")
 
+        subdir = os.path.join(*self._prefixes, name)
+
         if subdir in self._cached_loggers:
             logger = self._cached_loggers[subdir]
         else:
-            subdir = types.path_to_str(subdir)
             folder = os.path.join(self.default_logger.dir, "raw", subdir)
             os.makedirs(folder, exist_ok=True)
             output_formats = _build_output_formats(folder, self.format_strs)
@@ -152,20 +202,22 @@ def accumulate_means(self, subdir: types.AnyPath) -> Generator[None, None, None]
         try:
             self.current_logger = logger
             self._subdir = subdir
+            self._name = name
             self._update_name_to_maps()
             yield
         finally:
             self.current_logger = None
             self._subdir = None
+            self._name = None
             self._update_name_to_maps()
 
     def record(self, key, val, exclude=None):
         if self.current_logger is not None:  # In accumulate_means context.
             assert self._subdir is not None
-            raw_key = "/".join(["raw", self._subdir, key])
+            raw_key = "/".join(["raw", *self._prefixes, self._name, key])
             self.current_logger.record(raw_key, val, exclude)
 
-            mean_key = "/".join(["mean", self._subdir, key])
+            mean_key = "/".join(["mean", *self._prefixes, self._name, key])
             self.default_logger.record_mean(mean_key, val, exclude)
         else:  # Not in accumulate_means context.
             self.default_logger.record(key, val, exclude)
@@ -269,4 +321,4 @@ def configure(
     default_logger = sb_logger.Logger(folder, list(output_formats))
     hier_format_strs = [f for f in format_strs if f != "wandb"]
     hier_logger = HierarchicalLogger(default_logger, hier_format_strs)
-    return hier_logger
+    return hier_logger
\ No newline at end of file
diff --git a/tests/util/test_logger.py b/tests/util/test_logger.py
index 4048c44d3..ee0674141 100644
--- a/tests/util/test_logger.py
+++ b/tests/util/test_logger.py
@@ -1,7 +1,6 @@
 """Tests `imitation.util.logger`."""
 
 import csv
-import json
 import os.path as osp
 from collections import defaultdict
 
@@ -22,40 +21,13 @@ def _csv_to_dict(csv_path: str) -> dict:
     return result
 
 
-def _json_to_dict(json_path: str) -> dict:
-    r"""Loads the saved json logging file and convert it to expected dict format.
-
-    Args:
-        json_path: Path of the json log file.
-            Stored in the format - '{"A": 1, "B": 1}\n{"A": 2}\n{"B": 3}\n'
-
-    Returns:
-        dictionary in the format - `{"A": [1, 2, ""], "B": [1, "", 3]}`
-    """
-    result = defaultdict(list)
-    with open(json_path, "r") as f:
-        all_line_dicts = [json.loads(line) for line in f.readlines()]
-    # get all the keys in the dict so as to add "" if the key is not present in a line
-    all_keys = set().union(*[list(line_dict.keys()) for line_dict in all_line_dicts])
-
-    for line_dict in all_line_dicts:
-        for key in all_keys:
-            result[key].append(line_dict.get(key, ""))
-    return result
-
-
 def _compare_csv_lines(csv_path: str, expect: dict):
     observed = _csv_to_dict(csv_path)
     assert expect == observed
 
 
-def _compare_json_lines(json_path: str, expect: dict):
-    observed = _json_to_dict(json_path)
-    assert expect == observed
-
-
 def test_no_accum(tmpdir):
-    hier_logger = logger.configure(tmpdir, ["csv", "json"])
+    hier_logger = logger.configure(tmpdir, ["csv"])
     assert hier_logger.get_dir() == tmpdir
 
     # Check that the recorded "A": -1 is overwritten by "A": 1 in the next line.
@@ -71,12 +43,6 @@ def test_no_accum(tmpdir):
     hier_logger.dump()
     expect = {"A": [1, 2, ""], "B": [1, "", 3]}
     _compare_csv_lines(osp.join(tmpdir, "progress.csv"), expect)
-    _compare_json_lines(osp.join(tmpdir, "progress.json"), expect)
-
-
-def test_raise_unknown_format():
-    with pytest.raises(ValueError, match=r"Unknown format specified:.*"):
-        logger.make_output_format("txt", "log_dir")
 
 
 def test_free_form(tmpdir):
@@ -229,3 +195,47 @@ def test_hard(tmpdir):
     _compare_csv_lines(osp.join(tmpdir, "progress.csv"), expect_default)
     _compare_csv_lines(osp.join(tmpdir, "raw", "gen", "progress.csv"), expect_raw_gen)
     _compare_csv_lines(osp.join(tmpdir, "raw", "disc", "progress.csv"), expect_raw_disc)
+
+
+def test_prefix(tmpdir):
+    hier_logger = logger.configure(tmpdir)
+
+    with hier_logger.add_prefix("foo"), hier_logger.accumulate_means("bar"):
+        hier_logger.record("A", 1)
+        hier_logger.record("B", 2)
+        hier_logger.dump()
+
+    hier_logger.record("no_context", 1)
+
+    with hier_logger.accumulate_means("blat"):
+        hier_logger.record("C", 3)
+        hier_logger.dump()
+
+    hier_logger.dump()
+
+    expect_raw_foo_bar = {
+        "raw/foo/bar/A": [1],
+        "raw/foo/bar/B": [2],
+    }
+    expect_raw_blat = {
+        "raw/blat/C": [3],
+    }
+    expect_default = {
+        "mean/foo/bar/A": [1],
+        "mean/foo/bar/B": [2],
+        "mean/blat/C": [3],
+        "no_context": [1],
+    }
+
+    _compare_csv_lines(osp.join(tmpdir, "progress.csv"), expect_default)
+    _compare_csv_lines(
+        osp.join(tmpdir, "raw", "foo", "bar", "progress.csv"),
+        expect_raw_foo_bar,
+    )
+    _compare_csv_lines(osp.join(tmpdir, "raw", "blat", "progress.csv"), expect_raw_blat)
+
+
+def test_cant_add_prefix_within_accumulate_means(tmpdir):
+    h = logger.configure(tmpdir)
+    with pytest.raises(RuntimeError), h.accumulate_means("foo"), h.add_prefix("bar"):
+        pass  # pragma: no cover
\ No newline at end of file

From 939c52f9a7dca1ccf899b6b341cccaff783ed069 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 10:52:19 +0100
Subject: [PATCH 02/11] Added back accidentally removed code

---
 src/imitation/util/logger.py | 72 ++++++++++++++++++++++++++++++++++--
 tests/util/test_logger.py    | 35 +++++++++++++++++-
 2 files changed, 101 insertions(+), 6 deletions(-)

diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index f2bd33b9d..27b800f8a 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -3,6 +3,7 @@
 import contextlib
 import datetime
 import os
+import sys
 import tempfile
 from typing import Any, Dict, Generator, List, Optional, Sequence, Tuple, Union
 
@@ -11,6 +12,34 @@
 from imitation.data import types
 
 
+def make_output_format(
+    _format: str,
+    log_dir: str,
+    log_suffix: str = "",
+    max_length: int = 50,
+) -> sb_logger.KVWriter:
+    """Returns a logger for the requested format.
+    Args:
+        _format: the requested format to log to
+            ('stdout', 'log', 'json' or 'csv' or 'tensorboard').
+        log_dir: the logging directory.
+        log_suffix: the suffix for the log file.
+        max_length: the maximum length beyond which the keys get truncated.
+    Returns:
+        the logger.
+    """
+    os.makedirs(log_dir, exist_ok=True)
+    if _format == "stdout":
+        return sb_logger.HumanOutputFormat(sys.stdout, max_length=max_length)
+    elif _format == "log":
+        return sb_logger.HumanOutputFormat(
+            os.path.join(log_dir, f"log{log_suffix}.txt"),
+            max_length=max_length,
+        )
+    else:
+        return sb_logger.make_output_format(_format, log_dir, log_suffix)
+
+
 def _build_output_formats(
     folder: str,
     format_strs: Sequence[str],
@@ -31,7 +60,7 @@ def _build_output_formats(
         if f == "wandb":
             output_formats.append(WandbOutputFormat())
         else:
-            output_formats.append(sb_logger.make_output_format(f, folder))
+            output_formats.append(make_output_format(f, folder))
     return output_formats
 
 
@@ -89,6 +118,7 @@ class HierarchicalLogger(sb_logger.Logger):
     current_logger: Optional[sb_logger.Logger]
     _cached_loggers: Dict[str, sb_logger.Logger]
     _prefixes: List[str]
+    _key_prefixes: List[str]
     _subdir: Optional[str]
     _name: Optional[str]
 
@@ -112,6 +142,7 @@ def __init__(
         self.current_logger = None
         self._cached_loggers = {}
         self._prefixes = []
+        self._key_prefixes = []
         self._subdir = None
         self._name = None
         self.format_strs = format_strs
@@ -149,6 +180,35 @@ def add_prefix(self, prefix: str) -> Generator[None, None, None]:
         finally:
             self._prefixes.pop()
 
+    def get_prefixes(self) -> str:
+        return "/".join(self._prefixes)
+
+    @contextlib.contextmanager
+    def add_key_prefix(self, prefix: str) -> Generator[None, None, None]:
+        """Add a prefix to the keys logged during an accumulate_means context.
+
+        This prefix only applies when a `accumulate_means` context is active.
+
+        Args:
+            prefix: The prefix to add to the keys.
+
+        Yields:
+            None when the context manager is entered
+
+        Raises:
+            RuntimeError: if accumulate means context is already active.
+        """
+        if self.current_logger is None:
+            raise RuntimeError(
+                "Cannot add key prefix when accumulate_means context is not active.",
+            )
+
+        try:
+            self._key_prefixes.append(prefix)
+            yield
+        finally:
+            self._key_prefixes.pop()
+
     @contextlib.contextmanager
     def accumulate_means(self, name: str) -> Generator[None, None, None]:
         """Temporarily modifies this HierarchicalLogger to accumulate means values.
@@ -214,10 +274,14 @@ def accumulate_means(self, name: str) -> Generator[None, None, None]:
     def record(self, key, val, exclude=None):
         if self.current_logger is not None:  # In accumulate_means context.
             assert self._subdir is not None
-            raw_key = "/".join(["raw", *self._prefixes, self._name, key])
+            raw_key = "/".join(
+                ["raw", *self._prefixes, self._name, *self._key_prefixes, key]
+            )
             self.current_logger.record(raw_key, val, exclude)
 
-            mean_key = "/".join(["mean", *self._prefixes, self._name, key])
+            mean_key = "/".join(
+                ["mean", *self._prefixes, self._name, *self._key_prefixes, key]
+            )
             self.default_logger.record_mean(mean_key, val, exclude)
         else:  # Not in accumulate_means context.
             self.default_logger.record(key, val, exclude)
@@ -321,4 +385,4 @@ def configure(
     default_logger = sb_logger.Logger(folder, list(output_formats))
     hier_format_strs = [f for f in format_strs if f != "wandb"]
     hier_logger = HierarchicalLogger(default_logger, hier_format_strs)
-    return hier_logger
\ No newline at end of file
+    return hier_logger
diff --git a/tests/util/test_logger.py b/tests/util/test_logger.py
index ee0674141..b600fc915 100644
--- a/tests/util/test_logger.py
+++ b/tests/util/test_logger.py
@@ -21,13 +21,38 @@ def _csv_to_dict(csv_path: str) -> dict:
     return result
 
 
+def _json_to_dict(json_path: str) -> dict:
+    r"""Loads the saved json logging file and convert it to expected dict format.
+    Args:
+        json_path: Path of the json log file.
+            Stored in the format - '{"A": 1, "B": 1}\n{"A": 2}\n{"B": 3}\n'
+    Returns:
+        dictionary in the format - `{"A": [1, 2, ""], "B": [1, "", 3]}`
+    """
+    result = defaultdict(list)
+    with open(json_path, "r") as f:
+        all_line_dicts = [json.loads(line) for line in f.readlines()]
+    # get all the keys in the dict so as to add "" if the key is not present in a line
+    all_keys = set().union(*[list(line_dict.keys()) for line_dict in all_line_dicts])
+
+    for line_dict in all_line_dicts:
+        for key in all_keys:
+            result[key].append(line_dict.get(key, ""))
+    return result
+
+
 def _compare_csv_lines(csv_path: str, expect: dict):
     observed = _csv_to_dict(csv_path)
     assert expect == observed
 
 
+def _compare_json_lines(json_path: str, expect: dict):
+    observed = _json_to_dict(json_path)
+    assert expect == observed
+
+
 def test_no_accum(tmpdir):
-    hier_logger = logger.configure(tmpdir, ["csv"])
+    hier_logger = logger.configure(tmpdir, ["csv", "json"])
     assert hier_logger.get_dir() == tmpdir
 
     # Check that the recorded "A": -1 is overwritten by "A": 1 in the next line.
@@ -43,6 +68,12 @@ def test_no_accum(tmpdir):
     hier_logger.dump()
     expect = {"A": [1, 2, ""], "B": [1, "", 3]}
     _compare_csv_lines(osp.join(tmpdir, "progress.csv"), expect)
+    _compare_json_lines(osp.join(tmpdir, "progress.json"), expect)
+
+
+def test_raise_unknown_format():
+    with pytest.raises(ValueError, match=r"Unknown format specified:.*"):
+        logger.make_output_format("txt", "log_dir")
 
 
 def test_free_form(tmpdir):
@@ -238,4 +269,4 @@ def test_prefix(tmpdir):
 def test_cant_add_prefix_within_accumulate_means(tmpdir):
     h = logger.configure(tmpdir)
     with pytest.raises(RuntimeError), h.accumulate_means("foo"), h.add_prefix("bar"):
-        pass  # pragma: no cover
\ No newline at end of file
+        pass  # pragma: no cover

From da5c0384c37cc234386a700cae504c27358a381f Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 10:52:58 +0100
Subject: [PATCH 03/11] Replaced preference comparisons prefix with ctx manager

---
 .../algorithms/preference_comparisons.py      | 101 +++++++++---------
 1 file changed, 49 insertions(+), 52 deletions(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 092f3eeda..9f8b2f250 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -7,6 +7,7 @@
 import math
 import pickle
 import random
+import re
 from collections import defaultdict
 from typing import (
     Any,
@@ -1152,7 +1153,6 @@ def _train(
         self,
         dataset: PreferenceDataset,
         epoch_multiplier: float = 1.0,
-        prefix: Optional[str] = None,
     ) -> None:
         """Trains for `epoch_multiplier * self.epochs` epochs over `dataset`."""
         if self.regularizer is not None and self.regularizer.val_split is not None:
@@ -1180,71 +1180,65 @@ def _train(
         epochs = round(self.epochs * epoch_multiplier)
 
         assert epochs > 0, "Must train for at least one epoch."
-        epoch_num = 0
         with self.logger.accumulate_means("reward"):
             for epoch_num in tqdm(range(epochs), desc="Training reward model"):
-                logger_prefix = self._get_logger_key(prefix, f"epoch-{epoch_num}")
-                train_loss = 0.0
-                for fragment_pairs, preferences in dataloader:
-                    self.optim.zero_grad()
-                    loss = self._training_inner_loop(
-                        fragment_pairs,
-                        preferences,
-                        prefix=f"{logger_prefix}/train",
-                    )
-                    train_loss += loss.item()
-                    if self.regularizer:
-                        self.regularizer.regularize_and_backward(loss)
-                    else:
-                        loss.backward()
-                    self.optim.step()
-
-                if not self.requires_regularizer_update:
-                    continue
-                assert val_dataloader is not None
-                assert self.regularizer is not None
-
-                val_loss = 0.0
-                for fragment_pairs, preferences in val_dataloader:
-                    loss = self._training_inner_loop(
-                        fragment_pairs,
-                        preferences,
-                        prefix=f"{logger_prefix}/val",
-                    )
-                    val_loss += loss.item()
-                self.regularizer.update_params(train_loss, val_loss)
+                with self.logger.add_key_prefix(f"epoch-{epoch_num}"):
+                    train_loss = 0.0
+                    for fragment_pairs, preferences in dataloader:
+                        self.optim.zero_grad()
+                        with self.logger.add_key_prefix("train"):
+                            loss = self._training_inner_loop(
+                                fragment_pairs,
+                                preferences,
+                            )
+                        train_loss += loss.item()
+                        if self.regularizer:
+                            self.regularizer.regularize_and_backward(loss)
+                        else:
+                            loss.backward()
+                        self.optim.step()
+
+                    if not self.requires_regularizer_update:
+                        continue
+                    assert val_dataloader is not None
+                    assert self.regularizer is not None
+
+                    val_loss = 0.0
+                    for fragment_pairs, preferences in val_dataloader:
+                        with self.logger.add_key_prefix("val"):
+                            val_loss += self._training_inner_loop(
+                                fragment_pairs, preferences
+                            ).item()
+                    self.regularizer.update_params(train_loss, val_loss)
 
         # after training all the epochs,
         # record also the final value in a separate key for easy access.
         keys = list(self.logger.name_to_value.keys())
+        outer_prefix = self.logger.get_prefixes()
+        outer_prefix = outer_prefix + "/" if outer_prefix else ""
         for key in keys:
-            if key.startswith("mean/reward/" + logger_prefix):
+            base_path = f"{outer_prefix}reward/"  # existing prefix + accum_means ctx
+            epoch_path = f"mean/{base_path}epoch-{epoch_num}/"  # mean for last epoch
+            final_path = f"{base_path}final/"  # path to record last epoch
+            pattern = rf"{epoch_path}(.+)"
+            if regex_match := re.match(pattern, key):
+                (key_name,) = regex_match.groups()
                 val = self.logger.name_to_value[key]
-                new_key = key.replace(
-                    "mean/reward/" + logger_prefix,
-                    "reward/" + self._get_logger_key(prefix, "final"),
-                )
+                new_key = f"{final_path}{key_name}"
                 self.logger.record(new_key, val)
 
     def _training_inner_loop(
         self,
         fragment_pairs: Sequence[TrajectoryPair],
         preferences: np.ndarray,
-        prefix: Optional[str] = None,
     ) -> th.Tensor:
         output = self.loss.forward(fragment_pairs, preferences, self._preference_model)
         loss = output.loss
-        self.logger.record(self._get_logger_key(prefix, "loss"), loss.item())
+        self.logger.record("loss", loss.item())
         for name, value in output.metrics.items():
-            self.logger.record(self._get_logger_key(prefix, name), value.item())
+            self.logger.record(name, value.item())
         return loss
 
-    # TODO(juan) refactor & remove once #529 is merged.
-    def _get_logger_key(self, mode: Optional[str], key: str) -> str:
-        if mode is None:
-            return key
-        return f"{mode}/{key}"
-
 
 class EnsembleTrainer(BasicRewardTrainer):
     """Train a reward ensemble."""
@@ -1309,7 +1303,11 @@ def __init__(
         if seed:
             self.rng = self.rng.manual_seed(seed)
 
-    @RewardTrainer.logger.setter
+    @property
+    def logger(self):
+        return super().logger
+
+    @logger.setter
     def logger(self, custom_logger):
         self._logger = custom_logger
         for member_trainer in self.member_trainers:
@@ -1326,11 +1324,10 @@ def _train(self, dataset: PreferenceDataset, epoch_multiplier: float = 1.0) -> N
         for member_idx in range(len(self.member_trainers)):
             # sampler gives new indexes on every call
             bagging_dataset = data_th.Subset(dataset, list(sampler))
-            self.member_trainers[member_idx]._train(
-                bagging_dataset,
-                epoch_multiplier,
-                prefix=f"member-{member_idx}",
-            )
+            with self.logger.add_prefix(f"member-{member_idx}"):
+                self.member_trainers[member_idx].train(
+                    bagging_dataset, epoch_multiplier=epoch_multiplier
+                )
 
         # average the metrics across the member models
         metrics = defaultdict(list)

From 9652189d211bb0966d1f2c8af4ccaa80498f7046 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 11:41:27 +0100
Subject: [PATCH 04/11] Fixed errors

---
 src/imitation/algorithms/preference_comparisons.py | 12 +++++++-----
 src/imitation/util/logger.py                       | 14 +++++++++++---
 tests/util/test_logger.py                          |  1 +
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 9f8b2f250..fa8b1ac4b 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -1215,7 +1215,6 @@ def _train(
         # record also the final value in a separate key for easy access.
         keys = list(self.logger.name_to_value.keys())
         outer_prefix = self.logger.get_prefixes()
-        outer_prefix = outer_prefix + "/" if outer_prefix else ""
         for key in keys:
             base_path = f"{outer_prefix}reward/"  # existing prefix + accum_means ctx
             epoch_path = f"mean/{base_path}epoch-{epoch_num}/"  # mean for last epoch
@@ -1333,10 +1332,10 @@ def _train(self, dataset: PreferenceDataset, epoch_multiplier: float = 1.0) -> N
         metrics = defaultdict(list)
         keys = list(self.logger.name_to_value.keys())
         for key in keys:
-            if key.startswith("reward/member-") and "final" in key:
+            if re.match(r"member-(\d+)/reward/(.+)", key) and "final" in key:
                 val = self.logger.name_to_value[key]
                 key_list = key.split("/")
-                key_list.pop(1)
+                key_list.pop(0)
                 metrics["/".join(key_list)].append(val)
 
         for k, v in metrics.items():
@@ -1596,8 +1595,11 @@ def train(
                 epoch_multiplier = self.initial_epoch_multiplier
 
             self.reward_trainer.train(self.dataset, epoch_multiplier=epoch_multiplier)
-            reward_loss = self.logger.name_to_value["reward/final/train/loss"]
-            reward_accuracy = self.logger.name_to_value["reward/final/train/accuracy"]
+            base_key = self.logger.get_prefixes() + "reward/final/train"
+            assert f"{base_key}/loss" in self.logger.name_to_value
+            assert f"{base_key}/accuracy" in self.logger.name_to_value
+            reward_loss = self.logger.name_to_value[f"{base_key}/loss"]
+            reward_accuracy = self.logger.name_to_value[f"{base_key}/accuracy"]
 
             ###################
             # Train the agent #
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 27b800f8a..579da3d9a 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -12,6 +12,13 @@
 from imitation.data import types
 
 
+class HumanOutputFormat(sb_logger.HumanOutputFormat):
+    def write(self, key_values: Dict, key_excluded: Dict, step: int = 0) -> None:
+        # replace / in keys with : so that they are not truncated and give errors.
+        key_values = {k.replace("/", ":"): v for k, v in key_values.items()}
+        super().write(key_values, key_excluded, step)
+
+
 def make_output_format(
     _format: str,
     log_dir: str,
@@ -30,9 +37,9 @@ def make_output_format(
     """
     os.makedirs(log_dir, exist_ok=True)
     if _format == "stdout":
-        return sb_logger.HumanOutputFormat(sys.stdout, max_length=max_length)
+        return HumanOutputFormat(sys.stdout, max_length=max_length)
     elif _format == "log":
-        return sb_logger.HumanOutputFormat(
+        return HumanOutputFormat(
             os.path.join(log_dir, f"log{log_suffix}.txt"),
             max_length=max_length,
         )
@@ -181,7 +188,8 @@ def add_prefix(self, prefix: str) -> Generator[None, None, None]:
             self._prefixes.pop()
 
     def get_prefixes(self) -> str:
-        return "/".join(self._prefixes)
+        prefixes = "/".join(self._prefixes)
+        return prefixes + "/" if prefixes else ""
 
     @contextlib.contextmanager
     def add_key_prefix(self, prefix: str) -> Generator[None, None, None]:
diff --git a/tests/util/test_logger.py b/tests/util/test_logger.py
index b600fc915..8f75ed09b 100644
--- a/tests/util/test_logger.py
+++ b/tests/util/test_logger.py
@@ -1,6 +1,7 @@
 """Tests `imitation.util.logger`."""
 
 import csv
+import json
 import os.path as osp
 from collections import defaultdict
 

From b7e44cb1d5ddcca599a69765ab82bfe470313bd9 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 12:15:50 +0100
Subject: [PATCH 05/11] Docstring fixes

---
 src/imitation/algorithms/preference_comparisons.py | 6 ++++--
 src/imitation/util/logger.py                       | 8 ++++++--
 tests/util/test_logger.py                          | 2 ++
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index fa8b1ac4b..47f9dc34e 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -1207,7 +1207,8 @@ def _train(
                     for fragment_pairs, preferences in val_dataloader:
                         with self.logger.add_key_prefix("val"):
                             val_loss += self._training_inner_loop(
-                                fragment_pairs, preferences
+                                fragment_pairs,
+                                preferences,
                             ).item()
                     self.regularizer.update_params(train_loss, val_loss)
 
@@ -1325,7 +1326,8 @@ def _train(self, dataset: PreferenceDataset, epoch_multiplier: float = 1.0) -> N
             bagging_dataset = data_th.Subset(dataset, list(sampler))
             with self.logger.add_prefix(f"member-{member_idx}"):
                 self.member_trainers[member_idx].train(
-                    bagging_dataset, epoch_multiplier=epoch_multiplier
+                    bagging_dataset,
+                    epoch_multiplier=epoch_multiplier,
                 )
 
         # average the metrics across the member models
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 579da3d9a..17a51bbcd 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -13,6 +13,8 @@
 
 
 class HumanOutputFormat(sb_logger.HumanOutputFormat):
+    """Wrapper for Stable Baselines logger that prevents truncation of key prefixes."""
+
     def write(self, key_values: Dict, key_excluded: Dict, step: int = 0) -> None:
         # replace / in keys with : so that they are not truncated and give errors.
         key_values = {k.replace("/", ":"): v for k, v in key_values.items()}
@@ -26,12 +28,14 @@ def make_output_format(
     max_length: int = 50,
 ) -> sb_logger.KVWriter:
     """Returns a logger for the requested format.
+
     Args:
         _format: the requested format to log to
             ('stdout', 'log', 'json' or 'csv' or 'tensorboard').
         log_dir: the logging directory.
         log_suffix: the suffix for the log file.
         max_length: the maximum length beyond which the keys get truncated.
+
     Returns:
         the logger.
     """
@@ -283,12 +287,12 @@ def record(self, key, val, exclude=None):
         if self.current_logger is not None:  # In accumulate_means context.
             assert self._subdir is not None
             raw_key = "/".join(
-                ["raw", *self._prefixes, self._name, *self._key_prefixes, key]
+                ["raw", *self._prefixes, self._name, *self._key_prefixes, key],
             )
             self.current_logger.record(raw_key, val, exclude)
 
             mean_key = "/".join(
-                ["mean", *self._prefixes, self._name, *self._key_prefixes, key]
+                ["mean", *self._prefixes, self._name, *self._key_prefixes, key],
             )
             self.default_logger.record_mean(mean_key, val, exclude)
         else:  # Not in accumulate_means context.
diff --git a/tests/util/test_logger.py b/tests/util/test_logger.py
index 8f75ed09b..5639e64e0 100644
--- a/tests/util/test_logger.py
+++ b/tests/util/test_logger.py
@@ -24,9 +24,11 @@ def _csv_to_dict(csv_path: str) -> dict:
 
 def _json_to_dict(json_path: str) -> dict:
     r"""Loads the saved json logging file and convert it to expected dict format.
+
     Args:
         json_path: Path of the json log file.
             Stored in the format - '{"A": 1, "B": 1}\n{"A": 2}\n{"B": 3}\n'
+
     Returns:
         dictionary in the format - `{"A": [1, 2, ""], "B": [1, "", 3]}`
     """

From e69bd786d26fe1586b3a781fd3f489c687e7a566 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 28 Sep 2022 12:16:36 +0100
Subject: [PATCH 06/11] Address PR comments

---
 .../algorithms/preference_comparisons.py      |  6 +-
 src/imitation/util/logger.py                  | 68 +++++++++----------
 tests/util/test_logger.py                     | 54 +++++++++++++--
 3 files changed, 86 insertions(+), 42 deletions(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 47f9dc34e..7d37b337b 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -1215,7 +1215,7 @@ def _train(
         # after training all the epochs,
         # record also the final value in a separate key for easy access.
         keys = list(self.logger.name_to_value.keys())
-        outer_prefix = self.logger.get_prefixes()
+        outer_prefix = self.logger.get_accumulate_prefixes()
         for key in keys:
             base_path = f"{outer_prefix}reward/"  # existing prefix + accum_means ctx
             epoch_path = f"mean/{base_path}epoch-{epoch_num}/"  # mean for last epoch
@@ -1324,7 +1324,7 @@ def _train(self, dataset: PreferenceDataset, epoch_multiplier: float = 1.0) -> N
         for member_idx in range(len(self.member_trainers)):
             # sampler gives new indexes on every call
             bagging_dataset = data_th.Subset(dataset, list(sampler))
-            with self.logger.add_prefix(f"member-{member_idx}"):
+            with self.logger.add_accumulate_prefix(f"member-{member_idx}"):
                 self.member_trainers[member_idx].train(
                     bagging_dataset,
                     epoch_multiplier=epoch_multiplier,
@@ -1597,7 +1597,7 @@ def train(
                 epoch_multiplier = self.initial_epoch_multiplier
 
             self.reward_trainer.train(self.dataset, epoch_multiplier=epoch_multiplier)
-            base_key = self.logger.get_prefixes() + "reward/final/train"
+            base_key = self.logger.get_accumulate_prefixes() + "reward/final/train"
             assert f"{base_key}/loss" in self.logger.name_to_value
             assert f"{base_key}/accuracy" in self.logger.name_to_value
             reward_loss = self.logger.name_to_value[f"{base_key}/loss"]
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 17a51bbcd..2803a5707 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -12,15 +12,6 @@
 from imitation.data import types
 
 
-class HumanOutputFormat(sb_logger.HumanOutputFormat):
-    """Wrapper for Stable Baselines logger that prevents truncation of key prefixes."""
-
-    def write(self, key_values: Dict, key_excluded: Dict, step: int = 0) -> None:
-        # replace / in keys with : so that they are not truncated and give errors.
-        key_values = {k.replace("/", ":"): v for k, v in key_values.items()}
-        super().write(key_values, key_excluded, step)
-
-
 def make_output_format(
     _format: str,
     log_dir: str,
@@ -41,9 +32,9 @@ def make_output_format(
     """
     os.makedirs(log_dir, exist_ok=True)
     if _format == "stdout":
-        return HumanOutputFormat(sys.stdout, max_length=max_length)
+        return sb_logger.HumanOutputFormat(sys.stdout, max_length=max_length)
     elif _format == "log":
-        return HumanOutputFormat(
+        return sb_logger.HumanOutputFormat(
             os.path.join(log_dir, f"log{log_suffix}.txt"),
             max_length=max_length,
         )
@@ -59,7 +50,7 @@ def _build_output_formats(
 
     Args:
         folder: Path to directory that logs are written to.
-        format_strs: An list of output format strings. For details on available
+        format_strs: A list of output format strings. For details on available
             output formats see `stable_baselines3.logger.make_output_format`.
 
     Returns:
@@ -101,7 +92,7 @@ class HierarchicalLogger(sb_logger.Logger):
     ...     # record the key value pair `("mean/dataset/entropy", 5.5)` to path
     ...     # `dir` at step 1.
     ...     logger.dump(step=1)
-    ...     with logger.add_prefix("foo"), logger.accumulate_means("bar"):
+    ...     with logger.add_accumulate_prefix("foo"), logger.accumulate_means("bar"):
     ...         # record the key value pair ("raw/foo/bar/biz", 42.0) to path
     ...         # `dir/raw/foo/bar` at step 2000
     ...         logger.record("biz", 42.0)
@@ -128,7 +119,7 @@ class HierarchicalLogger(sb_logger.Logger):
     default_logger: sb_logger.Logger
     current_logger: Optional[sb_logger.Logger]
     _cached_loggers: Dict[str, sb_logger.Logger]
-    _prefixes: List[str]
+    _accumulate_prefixes: List[str]
     _key_prefixes: List[str]
     _subdir: Optional[str]
     _name: Optional[str]
@@ -152,7 +143,7 @@ def __init__(
         self.default_logger = default_logger
         self.current_logger = None
         self._cached_loggers = {}
-        self._prefixes = []
+        self._accumulate_prefixes = []
         self._key_prefixes = []
         self._subdir = None
         self._name = None
@@ -165,7 +156,7 @@ def _update_name_to_maps(self) -> None:
         self.name_to_excluded = self._logger.name_to_excluded
 
     @contextlib.contextmanager
-    def add_prefix(self, prefix: str) -> Generator[None, None, None]:
+    def add_accumulate_prefix(self, prefix: str) -> Generator[None, None, None]:
         """Add a prefix to the subdirectory used to accumulate means.
 
         This prefix only applies when a `accumulate_means` context is active. If there
@@ -186,13 +177,13 @@ def add_prefix(self, prefix: str) -> Generator[None, None, None]:
             )
 
         try:
-            self._prefixes.append(prefix)
+            self._accumulate_prefixes.append(prefix)
             yield
         finally:
-            self._prefixes.pop()
+            self._accumulate_prefixes.pop()
 
-    def get_prefixes(self) -> str:
-        prefixes = "/".join(self._prefixes)
+    def get_accumulate_prefixes(self) -> str:
+        prefixes = "/".join(self._accumulate_prefixes)
         return prefixes + "/" if prefixes else ""
 
     @contextlib.contextmanager
@@ -200,6 +191,7 @@ def add_key_prefix(self, prefix: str) -> Generator[None, None, None]:
         """Add a prefix to the keys logged during an accumulate_means context.
 
         This prefix only applies when a `accumulate_means` context is active.
+        If there are multiple active prefixes, then they are concatenated.
 
         Args:
             prefix: The prefix to add to the keys.
@@ -225,26 +217,34 @@ def add_key_prefix(self, prefix: str) -> Generator[None, None, None]:
     def accumulate_means(self, name: str) -> Generator[None, None, None]:
         """Temporarily modifies this HierarchicalLogger to accumulate means values.
 
-        Within this context manager, `self.record(key, value)` writes the "raw" values
-        in `f"{self.default_logger.log_dir}/{prefix}/{name}"` under the key
-        `"raw/{prefix}/{name}/{key}"`. At the same time, any call to `self.record` will
-        also accumulate mean values on the default logger by calling
-        `self.default_logger.record_mean(f"mean/{prefix}/{name}/{key}", value)`.
+        Within this context manager, ``self.record(key, value)`` writes the "raw" values
+        in ``f"{self.default_logger.log_dir}/[{accumulate_prefix}/]{name}"`` under the 
+        key ``"raw/[{accumulate_prefix}/]{name}/[{key_prefix}/]{key}"``, where
+        ``accumulate_prefix`` is the concatenation of all prefixes added by
+        ``add_accumulate_prefix`` and ``key_prefix`` is the concatenation of all
+        prefixes added by ``add_key_prefix``, if any. At the same time, any call to
+        ``self.record`` will also accumulate mean values on the default logger by
+        calling::
+
+            self.default_logger.record_mean(
+                f"mean/[{accumulate_prefix}/]{name}/[{key_prefix}/]{key}",
+                value,
+            )
 
         Multiple prefixes may be active at once. In this case the `prefix` is simply the
         concatenation of each of the active prefixes in the order they
-        where created e.g. if the active `prefixes` are ['foo', 'bar'] then
-        the `prefix` is 'foo/bar'.
+        were created e.g. if the active prefixes are ``['foo', 'bar']`` then
+        the prefix is ``'foo/bar'``.
 
-        After the context exits, calling `self.dump()` will write the means
+        After the context exits, calling ``self.dump()`` will write the means
         of all the "raw" values accumulated during this context to
-        `self.default_logger` under keys of the form  `mean/{prefix}/{name}/{key}`
+        ``self.default_logger`` under keys of the form  ``mean/{prefix}/{name}/{key}``
 
-        Note that the behavior of other logging methods, `log` and `record_mean`
+        Note that the behavior of other logging methods, ``log`` and ``record_mean``
         are unmodified and will go straight to the default logger.
 
         Args:
-            name: A string key which determines the `folder` where raw data is
+            name: A string key which determines the ``folder`` where raw data is
                 written and temporary logging prefixes for raw and mean data. Entering
                 an `accumulate_means` context in the future with the same `subdir`
                 will safely append to logs written in this folder rather than
@@ -260,7 +260,7 @@ def accumulate_means(self, name: str) -> Generator[None, None, None]:
         if self.current_logger is not None:
             raise RuntimeError("Nested `accumulate_means` context")
 
-        subdir = os.path.join(*self._prefixes, name)
+        subdir = os.path.join(*self._accumulate_prefixes, name)
 
         if subdir in self._cached_loggers:
             logger = self._cached_loggers[subdir]
@@ -287,12 +287,12 @@ def record(self, key, val, exclude=None):
         if self.current_logger is not None:  # In accumulate_means context.
             assert self._subdir is not None
             raw_key = "/".join(
-                ["raw", *self._prefixes, self._name, *self._key_prefixes, key],
+                ["raw", *self._accumulate_prefixes, self._name, *self._key_prefixes, key],
             )
             self.current_logger.record(raw_key, val, exclude)
 
             mean_key = "/".join(
-                ["mean", *self._prefixes, self._name, *self._key_prefixes, key],
+                ["mean", *self._accumulate_prefixes, self._name, *self._key_prefixes, key],
             )
             self.default_logger.record_mean(mean_key, val, exclude)
         else:  # Not in accumulate_means context.
diff --git a/tests/util/test_logger.py b/tests/util/test_logger.py
index 5639e64e0..ce8532998 100644
--- a/tests/util/test_logger.py
+++ b/tests/util/test_logger.py
@@ -163,8 +163,8 @@ def test_name_to_value(tmpdir):
 def test_hard(tmpdir):
     hier_logger = logger.configure(tmpdir)
 
-    # Part One: Test logging outside of the accumulating scope, and within scopes
-    # with two different different logging keys (including a repeat).
+    # Part One: Test logging outside the accumulating scope, and within scopes
+    # with two different logging keys (including a repeat).
 
     hier_logger.record("no_context", 1)
 
@@ -231,10 +231,10 @@ def test_hard(tmpdir):
     _compare_csv_lines(osp.join(tmpdir, "raw", "disc", "progress.csv"), expect_raw_disc)
 
 
-def test_prefix(tmpdir):
+def test_accumulate_prefix(tmpdir):
     hier_logger = logger.configure(tmpdir)
 
-    with hier_logger.add_prefix("foo"), hier_logger.accumulate_means("bar"):
+    with hier_logger.add_accumulate_prefix("foo"), hier_logger.accumulate_means("bar"):
         hier_logger.record("A", 1)
         hier_logger.record("B", 2)
         hier_logger.dump()
@@ -269,7 +269,51 @@ def test_prefix(tmpdir):
     _compare_csv_lines(osp.join(tmpdir, "raw", "blat", "progress.csv"), expect_raw_blat)
 
 
+def test_key_prefix(tmpdir):
+    hier_logger = logger.configure(tmpdir)
+
+    with hier_logger.accumulate_means("foo"), hier_logger.add_key_prefix("bar"):
+        hier_logger.record("A", 1)
+        hier_logger.record("B", 2)
+        hier_logger.dump()
+
+    hier_logger.record("no_context", 1)
+
+    with hier_logger.accumulate_means("blat"):
+        hier_logger.record("C", 3)
+        hier_logger.dump()
+
+    hier_logger.dump()
+
+    expect_raw_foo_bar = {
+        "raw/foo/bar/A": [1],
+        "raw/foo/bar/B": [2],
+    }
+    expect_raw_blat = {
+        "raw/blat/C": [3],
+    }
+    expect_default = {
+        "mean/foo/bar/A": [1],
+        "mean/foo/bar/B": [2],
+        "mean/blat/C": [3],
+        "no_context": [1],
+    }
+
+    _compare_csv_lines(osp.join(tmpdir, "progress.csv"), expect_default)
+    _compare_csv_lines(
+        osp.join(tmpdir, "raw", "foo", "progress.csv"),
+        expect_raw_foo_bar,
+    )
+    _compare_csv_lines(osp.join(tmpdir, "raw", "blat", "progress.csv"), expect_raw_blat)
+
+
 def test_cant_add_prefix_within_accumulate_means(tmpdir):
     h = logger.configure(tmpdir)
-    with pytest.raises(RuntimeError), h.accumulate_means("foo"), h.add_prefix("bar"):
+    with pytest.raises(RuntimeError), h.accumulate_means("foo"), h.add_accumulate_prefix("bar"):
+        pass  # pragma: no cover
+
+
+def test_cant_add_key_prefix_outside_accumulate_means(tmpdir):
+    h = logger.configure(tmpdir)
+    with pytest.raises(RuntimeError), h.add_key_prefix("bar"):
         pass  # pragma: no cover

From f4f14f0356046ecf24d7b5e3c5d5a5e6a2798ac9 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 28 Sep 2022 12:30:56 +0100
Subject: [PATCH 07/11] Point SB3 to master to include bug fix

---
 setup.py                     |  6 +++++-
 src/imitation/util/logger.py | 18 +++++++++++++++---
 tests/util/test_logger.py    |  4 +++-
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/setup.py b/setup.py
index 2f5ad4b6c..58dca83fc 100644
--- a/setup.py
+++ b/setup.py
@@ -23,7 +23,11 @@
 PYTYPE = ["pytype==2022.7.26"] if IS_NOT_WINDOWS else []
 if IS_NOT_WINDOWS:
     # TODO(adam): use this for Windows as well once PyPI is at >=1.6.1
-    STABLE_BASELINES3 = "stable-baselines3>=1.6.0"
+    STABLE_BASELINES3 = (
+        # "stable-baselines3>=1.6.0"
+        "stable-baselines3@git+"
+        "https://github.com/DLR-RM/stable-baselines3.git@master"
+    )
 else:
     STABLE_BASELINES3 = (
         "stable-baselines3@git+"
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 2803a5707..ea59fed6e 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -218,7 +218,7 @@ def accumulate_means(self, name: str) -> Generator[None, None, None]:
         """Temporarily modifies this HierarchicalLogger to accumulate means values.
 
         Within this context manager, ``self.record(key, value)`` writes the "raw" values
-        in ``f"{self.default_logger.log_dir}/[{accumulate_prefix}/]{name}"`` under the 
+        in ``f"{self.default_logger.log_dir}/[{accumulate_prefix}/]{name}"`` under the
         key ``"raw/[{accumulate_prefix}/]{name}/[{key_prefix}/]{key}"``, where
         ``accumulate_prefix`` is the concatenation of all prefixes added by
         ``add_accumulate_prefix`` and ``key_prefix`` is the concatenation of all
@@ -287,12 +287,24 @@ def record(self, key, val, exclude=None):
         if self.current_logger is not None:  # In accumulate_means context.
             assert self._subdir is not None
             raw_key = "/".join(
-                ["raw", *self._accumulate_prefixes, self._name, *self._key_prefixes, key],
+                [
+                    "raw",
+                    *self._accumulate_prefixes,
+                    self._name,
+                    *self._key_prefixes,
+                    key,
+                ],
             )
             self.current_logger.record(raw_key, val, exclude)
 
             mean_key = "/".join(
-                ["mean", *self._accumulate_prefixes, self._name, *self._key_prefixes, key],
+                [
+                    "mean",
+                    *self._accumulate_prefixes,
+                    self._name,
+                    *self._key_prefixes,
+                    key,
+                ],
             )
             self.default_logger.record_mean(mean_key, val, exclude)
         else:  # Not in accumulate_means context.
diff --git a/tests/util/test_logger.py b/tests/util/test_logger.py
index ce8532998..a4ea315af 100644
--- a/tests/util/test_logger.py
+++ b/tests/util/test_logger.py
@@ -309,7 +309,9 @@ def test_key_prefix(tmpdir):
 
 def test_cant_add_prefix_within_accumulate_means(tmpdir):
     h = logger.configure(tmpdir)
-    with pytest.raises(RuntimeError), h.accumulate_means("foo"), h.add_accumulate_prefix("bar"):
+    with pytest.raises(RuntimeError), h.accumulate_means(
+        "foo"
+    ), h.add_accumulate_prefix("bar"):
         pass  # pragma: no cover
 
 

From 146979be71c687488062ae88d0a9f166c8aa734f Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 29 Sep 2022 12:25:39 +0100
Subject: [PATCH 08/11] Format / fix tests for context manager

---
 tests/util/test_logger.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tests/util/test_logger.py b/tests/util/test_logger.py
index a4ea315af..7f076df9c 100644
--- a/tests/util/test_logger.py
+++ b/tests/util/test_logger.py
@@ -309,13 +309,13 @@ def test_key_prefix(tmpdir):
 
 def test_cant_add_prefix_within_accumulate_means(tmpdir):
     h = logger.configure(tmpdir)
-    with pytest.raises(RuntimeError), h.accumulate_means(
-        "foo"
-    ), h.add_accumulate_prefix("bar"):
-        pass  # pragma: no cover
+    with pytest.raises(RuntimeError):
+        with h.accumulate_means("foo"), h.add_accumulate_prefix("bar"):
+            pass  # pragma: no cover
 
 
 def test_cant_add_key_prefix_outside_accumulate_means(tmpdir):
     h = logger.configure(tmpdir)
-    with pytest.raises(RuntimeError), h.add_key_prefix("bar"):
-        pass  # pragma: no cover
+    with pytest.raises(RuntimeError):
+        with h.add_key_prefix("bar"):
+            pass  # pragma: no cover

From ebc3138fa4137468970ac3c476d954b3e6d6f999 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 29 Sep 2022 12:31:39 +0100
Subject: [PATCH 09/11] Switch to sb3 1.6.1

---
 setup.py | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/setup.py b/setup.py
index 58dca83fc..0f11e9716 100644
--- a/setup.py
+++ b/setup.py
@@ -21,18 +21,13 @@
     "autorom[accept-rom-license]~=0.4.2",
 ]
 PYTYPE = ["pytype==2022.7.26"] if IS_NOT_WINDOWS else []
-if IS_NOT_WINDOWS:
-    # TODO(adam): use this for Windows as well once PyPI is at >=1.6.1
-    STABLE_BASELINES3 = (
-        # "stable-baselines3>=1.6.0"
-        "stable-baselines3@git+"
-        "https://github.com/DLR-RM/stable-baselines3.git@master"
-    )
-else:
-    STABLE_BASELINES3 = (
-        "stable-baselines3@git+"
-        "https://github.com/DLR-RM/stable-baselines3.git@master"
-    )
+STABLE_BASELINES3 = (
+    "stable-baselines3>=1.6.1"
+)
+# STABLE_BASELINES3 = (
+#     "stable-baselines3@git+"
+#     "https://github.com/DLR-RM/stable-baselines3.git@master"
+# )
 
 # pinned to 0.21 until https://github.com/DLR-RM/stable-baselines3/pull/780 goes
 # upstream.

From 28991dafcfc6eb662b6d47d4a46d506327af8f96 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 29 Sep 2022 12:45:28 +0100
Subject: [PATCH 10/11] Formatting

---
 setup.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index 0f11e9716..e612d5f24 100644
--- a/setup.py
+++ b/setup.py
@@ -21,9 +21,7 @@
     "autorom[accept-rom-license]~=0.4.2",
 ]
 PYTYPE = ["pytype==2022.7.26"] if IS_NOT_WINDOWS else []
-STABLE_BASELINES3 = (
-    "stable-baselines3>=1.6.1"
-)
+STABLE_BASELINES3 = "stable-baselines3>=1.6.1"
 # STABLE_BASELINES3 = (
 #     "stable-baselines3@git+"
 #     "https://github.com/DLR-RM/stable-baselines3.git@master"

From 6fd29c34f90eec47a8dab9ea89cb4abeeff9914a Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 29 Sep 2022 19:03:15 +0100
Subject: [PATCH 11/11] Remove comment

---
 setup.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/setup.py b/setup.py
index e612d5f24..a2483ecc3 100644
--- a/setup.py
+++ b/setup.py
@@ -22,11 +22,6 @@
 ]
 PYTYPE = ["pytype==2022.7.26"] if IS_NOT_WINDOWS else []
 STABLE_BASELINES3 = "stable-baselines3>=1.6.1"
-# STABLE_BASELINES3 = (
-#     "stable-baselines3@git+"
-#     "https://github.com/DLR-RM/stable-baselines3.git@master"
-# )
-
 # pinned to 0.21 until https://github.com/DLR-RM/stable-baselines3/pull/780 goes
 # upstream.
 GYM_VERSION_SPECIFIER = "==0.21.0"