From fe6b219646bcaa796f974f663f1a44683774e393 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 23 Aug 2022 20:20:24 +0100
Subject: [PATCH 001/187] Initial mypy configuration

---
 Makefile | 20 ++++++++++++++++++++
 mypy.ini |  2 ++
 setup.py |  1 +
 3 files changed, 23 insertions(+)
 create mode 100644 Makefile
 create mode 100644 mypy.ini

diff --git a/Makefile b/Makefile
new file mode 100644
index 000000000..551d0dac6
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,20 @@
+SRC_FILES=src/ tests/ experiments/ examples/ docs/conf.py setup.py
+
+typecheck:
+	mypy ${SRC_FILES}
+	pytype -j auto "${SRC_FILES}"
+
+
+shellcheck:
+	find . -path ./venv -prune -o -name '*.sh' -print0 | xargs -0 shellcheck
+
+formatcheck:
+	flake8 --darglint-ignore-regex '.*' "${SRC_FILES[@]}"
+	black --check --diff "${SRC_FILES[@]}"
+	codespell -I .codespell.skip --skip='*.pyc,tests/testdata/*,*.ipynb,*.csv' "${SRC_FILES[@]}"
+
+docscheck:
+	pushd docs/
+	make clean
+	make html
+	popd
\ No newline at end of file
diff --git a/mypy.ini b/mypy.ini
new file mode 100644
index 000000000..e886c0858
--- /dev/null
+++ b/mypy.ini
@@ -0,0 +1,2 @@
+[mypy]
+ignore_missing_imports = true
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 81a095452..bcd3cc00b 100644
--- a/setup.py
+++ b/setup.py
@@ -48,6 +48,7 @@
         # TODO: upgrade jupyter-client once
         #  https://github.com/jupyter/jupyter_client/issues/637 is fixed
         "jupyter-client~=6.1.12",
+        "mypy~=0.791",
         "pandas~=1.4.3",
         "pytest~=7.1.2",
         "pytest-cov~=3.0.0",

From 5b684a80d49780b4365828385667e32f577222f3 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 23 Aug 2022 21:23:02 +0100
Subject: [PATCH 002/187] Initial change to get the PR up

---
 src/imitation/algorithms/dagger.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index 8700e0b12..dacdcedef 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -372,11 +372,7 @@ def _load_all_demos(self):
         return demo_transitions, num_demos_by_round
 
     def _get_demo_paths(self, round_dir):
-        return [
-            os.path.join(round_dir, p)
-            for p in os.listdir(round_dir)
-            if p.endswith(".npz")
-        ]
+        return [round_dir / p for p in os.listdir(round_dir) if p.endswith(".npz")]
 
     def _demo_dir_path_for_round(self, round_num: Optional[int] = None) -> pathlib.Path:
         if round_num is None:

From 81d159483f26b764463ebeac4509e0466a612d7a Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Sat, 27 Aug 2022 16:39:49 +0100
Subject: [PATCH 003/187] Initial review at replacing os.path

---
 .../algorithms/adversarial/common.py          |  9 ++--
 src/imitation/algorithms/dagger.py            |  2 +-
 src/imitation/data/types.py                   |  5 +++
 src/imitation/policies/serialize.py           | 14 ++++---
 src/imitation/scripts/analyze.py              | 25 ++++++-----
 src/imitation/scripts/common/common.py        | 39 +++++++++--------
 .../scripts/common/demonstrations.py          | 16 +++----
 src/imitation/scripts/convert_trajs.py        |  7 ++--
 src/imitation/scripts/eval_policy.py          | 13 +++---
 src/imitation/scripts/parallel.py             |  8 ++--
 src/imitation/scripts/train_adversarial.py    | 25 ++++++-----
 src/imitation/scripts/train_imitation.py      |  4 +-
 .../scripts/train_preference_comparisons.py   | 29 +++++++------
 src/imitation/scripts/train_rl.py             | 19 +++++----
 src/imitation/util/logger.py                  | 42 ++++++++++++-------
 src/imitation/util/video_wrapper.py           | 10 ++---
 tests/data/test_types.py                      |  2 +-
 17 files changed, 151 insertions(+), 118 deletions(-)

diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
index eee6937e4..877a0cba2 100644
--- a/src/imitation/algorithms/adversarial/common.py
+++ b/src/imitation/algorithms/adversarial/common.py
@@ -4,6 +4,7 @@
 import dataclasses
 import logging
 import os
+import pathlib
 from typing import Callable, Mapping, Optional, Sequence, Tuple, Type
 
 import numpy as np
@@ -187,7 +188,7 @@ def __init__(
         self.venv = venv
         self.gen_algo = gen_algo
         self._reward_net = reward_net.to(gen_algo.device)
-        self._log_dir = log_dir
+        self._log_dir = pathlib.Path(log_dir).resolve()
 
         # Create graph for optimising/recording stats on discriminator
         self._disc_opt_cls = disc_opt_cls
@@ -201,9 +202,9 @@ def __init__(
 
         if self._init_tensorboard:
             logging.info("building summary directory at " + self._log_dir)
-            summary_dir = os.path.join(self._log_dir, "summary")
-            os.makedirs(summary_dir, exist_ok=True)
-            self._summary_writer = thboard.SummaryWriter(summary_dir)
+            summary_dir = self._log_dir / "summary"
+            summary_dir.mkdir(parents=True, exist_ok=True)
+            self._summary_writer = thboard.SummaryWriter(str(summary_dir))
 
         venv = self.venv_buffering = wrappers.BufferingWrapper(self.venv)
 
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index dacdcedef..f440ba349 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -382,7 +382,7 @@ def _demo_dir_path_for_round(self, round_num: Optional[int] = None) -> pathlib.P
     def _try_load_demos(self) -> None:
         """Load the dataset for this round into self.bc_trainer as a DataLoader."""
         demo_dir = self._demo_dir_path_for_round()
-        demo_paths = self._get_demo_paths(demo_dir) if os.path.isdir(demo_dir) else []
+        demo_paths = self._get_demo_paths(demo_dir) if demo_dir.is_dir() else []
         if len(demo_paths) == 0:
             raise NeedsDemosException(
                 f"No demos found for round {self.round_num} in dir '{demo_dir}'. "
diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index 363d86b0e..3ce236bb4 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -40,6 +40,11 @@ def path_to_str(path: AnyPath) -> str:
         return str(path)
 
 
+def path_to_pathlib(path: AnyPath) -> pathlib.Path:
+    """Converts a path to a pathlib.Path."""
+    return pathlib.Path(path_to_str(path))
+
+
 @dataclasses.dataclass(frozen=True)
 class Trajectory:
     """A trajectory, e.g. a one episode rollout from an expert policy."""
diff --git a/src/imitation/policies/serialize.py b/src/imitation/policies/serialize.py
index e43651a0e..ca925213a 100644
--- a/src/imitation/policies/serialize.py
+++ b/src/imitation/policies/serialize.py
@@ -124,7 +124,7 @@ def load_policy(
 
 
 def save_stable_model(
-    output_dir: str,
+    output_dir: pathlib.Path,
     model: base_class.BaseAlgorithm,
 ) -> None:
     """Serialize Stable Baselines model.
@@ -138,9 +138,10 @@ def save_stable_model(
     # Save each model in new directory in case we want to add metadata or other
     # information in future. (E.g. we used to save `VecNormalize` statistics here,
     # although that is no longer necessary.)
-    os.makedirs(output_dir, exist_ok=True)
-    model.save(os.path.join(output_dir, "model.zip"))
-    logging.info("Saved policy to %s", output_dir)
+    output_dir = output_dir.resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    model.save(str(output_dir / "model.zip"))
+    logging.info(f"Saved policy to {output_dir}")
 
 
 class SavePolicyCallback(callbacks.EventCallback):
@@ -152,7 +153,7 @@ class SavePolicyCallback(callbacks.EventCallback):
 
     def __init__(
         self,
-        policy_dir: str,
+        policy_dir: pathlib.Path,
         *args,
         **kwargs,
     ):
@@ -167,6 +168,7 @@ def __init__(
         self.policy_dir = policy_dir
 
     def _on_step(self) -> bool:
-        output_dir = os.path.join(self.policy_dir, f"{self.num_timesteps:012d}")
+        assert self.model is not None
+        output_dir = self.policy_dir / f"{self.num_timesteps:012d}"
         save_stable_model(output_dir, self.model)
         return True
diff --git a/src/imitation/scripts/analyze.py b/src/imitation/scripts/analyze.py
index 74617891c..1dd298499 100644
--- a/src/imitation/scripts/analyze.py
+++ b/src/imitation/scripts/analyze.py
@@ -6,6 +6,7 @@
 import logging
 import os
 import os.path as osp
+import pathlib
 import tempfile
 import warnings
 from collections import OrderedDict
@@ -97,24 +98,25 @@ def gather_tb_directories() -> dict:
     Raises:
         OSError: If the symlink cannot be created.
     """
-    os.makedirs("/tmp/analysis_tb", exist_ok=True)
-    tmp_dir = tempfile.mkdtemp(dir="/tmp/analysis_tb/")
+    tb_analysis_dir = pathlib.Path("/tmp/analysis_tb")
+    tb_analysis_dir.mkdir(exist_ok=True)
+    tmp_dir = pathlib.Path(tempfile.mkdtemp(dir=tb_analysis_dir))
 
     tb_dirs_count = 0
     for sd in _gather_sacred_dicts():
         # Expecting a path like "~/ray_results/{run_name}/sacred/1".
         # Want to search for all Tensorboard dirs inside
         # "~/ray_results/{run_name}".
-        sacred_dir = sd.sacred_dir.rstrip("/")
-        run_dir = osp.dirname(osp.dirname(sacred_dir))
-        run_name = osp.basename(run_dir)
+        sacred_dir = pathlib.Path(sd.sacred_dir)
+        run_dir = sacred_dir.parent.parent
+        run_name = run_dir.name
 
         # log is what we use as subdirectory in new code.
         # rl, tb, sb_tb all appear in old versions.
         for basename in ["log", "rl", "tb", "sb_tb"]:
             tb_src_dirs = tuple(
                 sacred_util.filter_subdirs(
-                    run_dir,
+                    str(run_dir),
                     lambda path: osp.basename(path) == basename,
                 ),
             )
@@ -122,12 +124,12 @@ def gather_tb_directories() -> dict:
                 assert len(tb_src_dirs) == 1, "expect at most one TB dir of each type"
                 tb_src_dir = tb_src_dirs[0]
 
-                symlinks_dir = osp.join(tmp_dir, basename)
-                os.makedirs(symlinks_dir, exist_ok=True)
+                symlinks_dir = tmp_dir / basename
+                symlinks_dir.mkdir(exist_ok=True)
 
-                tb_symlink = osp.join(symlinks_dir, run_name)
+                tb_symlink = symlinks_dir / run_name
                 try:
-                    os.symlink(tb_src_dir, tb_symlink)
+                    tb_symlink.symlink_to(tb_src_dir)
                 except OSError as e:
                     if os.name == "nt":  # Windows
                         msg = (
@@ -318,7 +320,8 @@ def _make_return_summary(stats: dict, prefix="") -> str:
 
 
 def main_console():
-    observer = FileStorageObserver(osp.join("output", "sacred", "analyze"))
+    observer_path = pathlib.Path.cwd() / "output" / "sacred" / "analyze"
+    observer = FileStorageObserver(observer_path)
     analysis_ex.observers.append(observer)
     analysis_ex.run_commandline()
 
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index c08eda766..ffaccb453 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -3,7 +3,8 @@
 import contextlib
 import logging
 import os
-from typing import Any, Mapping, Sequence, Tuple, Union
+import pathlib
+from typing import Any, Generator, Mapping, Sequence, Tuple, Union
 
 import sacred
 from stable_baselines3.common import vec_env
@@ -44,19 +45,20 @@ def update_log_format_strs(log_format_strs, log_format_strs_additional):
 
 
 @common_ingredient.config_hook
-def hook(config, command_name, logger):
+def hook(config, command_name: str, logger):
     del logger
     updates = {}
     if config["common"]["log_dir"] is None:
         env_sanitized = config["common"]["env_name"].replace("/", "_")
-        log_root = config["common"]["log_root"] or "output"
-        log_dir = os.path.join(
-            log_root,
-            command_name,
-            env_sanitized,
-            util.make_unique_timestamp(),
+        assert isinstance(env_sanitized, str)
+        config_log_root = config["common"]["log_root"]
+        log_root = (
+            pathlib.Path(config_log_root)
+            if config_log_root
+            else pathlib.Path.cwd() / "output"
         )
-        updates["log_dir"] = log_dir
+        log_dir = log_root / command_name / env_sanitized / util.make_unique_timestamp()
+        updates["log_dir"] = str(log_dir)
     return updates
 
 
@@ -79,7 +81,7 @@ def make_log_dir(
     _run,
     log_dir: str,
     log_level: Union[int, str],
-) -> str:
+) -> pathlib.Path:
     """Creates log directory and sets up symlink to Sacred logs.
 
     Args:
@@ -91,23 +93,24 @@ def make_log_dir(
     Returns:
         The `log_dir`. This avoids the caller needing to capture this argument.
     """
-    os.makedirs(log_dir, exist_ok=True)
+    _log_dir = pathlib.Path(log_dir)
+    _log_dir.mkdir(parents=True, exist_ok=True)
     # convert strings of digits to numbers; but leave levels like 'INFO' unmodified
     try:
         log_level = int(log_level)
     except ValueError:
         pass
     logging.basicConfig(level=log_level)
-    logger.info("Logging to %s", log_dir)
-    sacred_util.build_sacred_symlink(log_dir, _run)
-    return log_dir
+    logger.info("Logging to %s", _log_dir)
+    sacred_util.build_sacred_symlink(_log_dir, _run)
+    return _log_dir
 
 
 @common_ingredient.capture
 def setup_logging(
     _run,
     log_format_strs: Sequence[str],
-) -> Tuple[imit_logger.HierarchicalLogger, str]:
+) -> Tuple[imit_logger.HierarchicalLogger, pathlib.Path]:
     """Builds the imitation logger.
 
     Args:
@@ -119,9 +122,9 @@ def setup_logging(
     """
     log_dir = make_log_dir()
     if "wandb" in log_format_strs:
-        wb.wandb_init(log_dir=log_dir)
+        wb.wandb_init(log_dir=str(log_dir))
     custom_logger = imit_logger.configure(
-        folder=os.path.join(log_dir, "log"),
+        folder=log_dir / "log",
         format_strs=log_format_strs,
     )
     return custom_logger, log_dir
@@ -138,7 +141,7 @@ def make_venv(
     max_episode_steps: int,
     env_make_kwargs: Mapping[str, Any],
     **kwargs,
-) -> vec_env.VecEnv:
+) -> Generator[vec_env.VecEnv, None, None]:
     """Builds the vector environment.
 
     Args:
diff --git a/src/imitation/scripts/common/demonstrations.py b/src/imitation/scripts/common/demonstrations.py
index 16a0541a6..eb5c8b275 100644
--- a/src/imitation/scripts/common/demonstrations.py
+++ b/src/imitation/scripts/common/demonstrations.py
@@ -2,7 +2,8 @@
 
 import logging
 import os
-from typing import Optional, Sequence
+import pathlib
+from typing import Dict, Optional, Sequence
 
 import sacred
 
@@ -27,23 +28,22 @@ def fast():
     n_expert_demos = 1  # noqa: F841
 
 
-def guess_expert_dir(data_dir: str, env_name: str) -> str:
+def guess_expert_dir(data_dir: pathlib.Path, env_name: str) -> pathlib.Path:
+    assert data_dir.is_absolute()
     rollout_hint = env_name.rsplit("-", 1)[0].replace("/", "_").lower()
-    return os.path.join(data_dir, "expert_models", f"{rollout_hint}_0")
+    return data_dir / "expert_models" / f"{rollout_hint}_0"
 
 
 @demonstrations_ingredient.config_hook
 def hook(config, command_name, logger):
     """If rollout_path not set explicitly, then guess it based on environment name."""
     del command_name, logger
-    updates = {}
+    updates: Dict[str, str] = {}
     if config["demonstrations"]["rollout_path"] is None:
         data_dir = config["demonstrations"]["data_dir"]
         env_name = config["common"]["env_name"].replace("/", "_")
-        updates["rollout_path"] = os.path.join(
-            guess_expert_dir(data_dir, env_name),
-            "rollouts",
-            "final.pkl",
+        updates["rollout_path"] = str(
+            guess_expert_dir(data_dir, env_name) / "rollouts" / "final.pkl"
         )
     return updates
 
diff --git a/src/imitation/scripts/convert_trajs.py b/src/imitation/scripts/convert_trajs.py
index 5ae4eb0a8..7c1e9549a 100644
--- a/src/imitation/scripts/convert_trajs.py
+++ b/src/imitation/scripts/convert_trajs.py
@@ -10,12 +10,13 @@
 """
 
 import os
+import pathlib
 import warnings
 
 from imitation.data import types
 
 
-def update_traj_file_in_place(path: str) -> None:
+def update_traj_file_in_place(path: pathlib.Path) -> None:
     """Modifies trajectories pickle file in-place to update data to new format.
 
     The new data is saved as `Sequence[imitation.types.TrajectoryWithRew]`.
@@ -33,9 +34,9 @@ def update_traj_file_in_place(path: str) -> None:
         )
         trajs = types.load(path)
 
-    path, ext = os.path.splitext(path)
+    ext = path.suffix
     new_ext = ".npz" if ext in (".pkl", ".npz") else ext + ".npz"
-    types.save(path + new_ext, trajs)
+    types.save(path.with_suffix(new_ext), trajs)
 
 
 def main():
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 8a667df36..7569acd8a 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -1,8 +1,8 @@
 """Evaluate policies: render policy interactively, save videos, log episode return."""
 
 import logging
-import os
-import os.path as osp
+import pathlib
+import re
 import time
 from typing import Any, Mapping, Optional
 
@@ -40,12 +40,12 @@ def step_wait(self):
         return ob
 
 
-def video_wrapper_factory(log_dir: str, **kwargs):
+def video_wrapper_factory(log_dir: pathlib.Path, **kwargs):
     """Returns a function that wraps the environment in a video recorder."""
 
-    def f(env: gym.Env, i: int) -> gym.Env:
+    def f(env: gym.Env, i: int) -> video_wrapper.VideoWrapper:
         """Wraps `env` in a recorder saving videos to `{log_dir}/videos/{i}`."""
-        directory = os.path.join(log_dir, "videos", str(i))
+        directory = log_dir / "videos" / str(i)
         return video_wrapper.VideoWrapper(env, directory=directory, **kwargs)
 
     return f
@@ -116,7 +116,8 @@ def eval_policy(
 
 
 def main_console():
-    observer = FileStorageObserver(osp.join("output", "sacred", "eval_policy"))
+    observer_path = pathlib.Path.cwd() / "output" / "sacred" / "eval_policy"
+    observer = FileStorageObserver(observer_path)
     eval_policy_ex.observers.append(observer)
     eval_policy_ex.run_commandline()
 
diff --git a/src/imitation/scripts/parallel.py b/src/imitation/scripts/parallel.py
index a85ee121f..2533c5397 100644
--- a/src/imitation/scripts/parallel.py
+++ b/src/imitation/scripts/parallel.py
@@ -3,6 +3,7 @@
 import collections.abc
 import copy
 import os
+import pathlib
 from typing import Any, Callable, Mapping, Optional, Sequence
 
 import ray
@@ -96,9 +97,9 @@ def parallel(
             and "data_dir" not in base_config_updates.get("demonstrations", {})
         )
         if no_data_dir:
-            data_dir = os.path.join(os.getcwd(), "data/")
+            data_dir = pathlib.Path.cwd() / "data"
             base_config_updates = dict(base_config_updates)
-            base_config_updates["demonstrations.data_dir"] = data_dir
+            base_config_updates["demonstrations.data_dir"] = str(data_dir)
 
     trainable = _ray_tune_sacred_wrapper(
         sacred_ex_name,
@@ -229,7 +230,8 @@ def inner(config: Mapping[str, Any], reporter) -> Mapping[str, Any]:
 
 
 def main_console():
-    observer = FileStorageObserver(os.path.join("output", "sacred", "parallel"))
+    observer_path = pathlib.Path.cwd() / "output" / "sacred" / "parallel"
+    observer = FileStorageObserver(observer_path)
     parallel_ex.observers.append(observer)
     parallel_ex.run_commandline()
 
diff --git a/src/imitation/scripts/train_adversarial.py b/src/imitation/scripts/train_adversarial.py
index 1d24955ad..186f6da87 100644
--- a/src/imitation/scripts/train_adversarial.py
+++ b/src/imitation/scripts/train_adversarial.py
@@ -4,6 +4,7 @@
 import logging
 import os
 import os.path as osp
+import pathlib
 from typing import Any, Mapping, Optional, Type
 
 import sacred.commands
@@ -22,15 +23,15 @@
 logger = logging.getLogger("imitation.scripts.train_adversarial")
 
 
-def save(trainer, save_path):
+def save(trainer: common.AdversarialTrainer, save_path: pathlib.Path):
     """Save discriminator and generator."""
     # We implement this here and not in Trainer since we do not want to actually
     # serialize the whole Trainer (including e.g. expert demonstrations).
-    os.makedirs(save_path, exist_ok=True)
-    th.save(trainer.reward_train, os.path.join(save_path, "reward_train.pt"))
-    th.save(trainer.reward_test, os.path.join(save_path, "reward_test.pt"))
+    save_path.mkdir(parents=True, exist_ok=True)
+    th.save(trainer.reward_train, save_path / "reward_train.pt")
+    th.save(trainer.reward_test, save_path / "reward_test.pt")
     serialize.save_stable_model(
-        os.path.join(save_path, "gen_policy"),
+        save_path / "gen_policy",
         trainer.gen_algo,
     )
 
@@ -113,7 +114,8 @@ def train_adversarial(
         # So, support showing merged config from `train_adversarial {airl,gail}`.
         sacred.commands.print_config(_run)
 
-    custom_logger, log_dir = common_config.setup_logging()
+    custom_logger, _log_dir = common_config.setup_logging()
+    log_dir = pathlib.Path(_log_dir)
     expert_trajs = demonstrations.load_expert_trajs()
 
     with common_config.make_venv() as venv:
@@ -144,22 +146,22 @@ def train_adversarial(
             venv=venv,
             demonstrations=expert_trajs,
             gen_algo=gen_algo,
-            log_dir=log_dir,
+            log_dir=str(log_dir),
             reward_net=reward_net,
             custom_logger=custom_logger,
             **algorithm_kwargs,
         )
 
-        def callback(round_num):
+        def callback(round_num: int, /) -> None:
             if checkpoint_interval > 0 and round_num % checkpoint_interval == 0:
-                save(trainer, os.path.join(log_dir, "checkpoints", f"{round_num:05d}"))
+                save(trainer, log_dir / "checkpoints" / f"{round_num:05d}")
 
         trainer.train(total_timesteps, callback)
         imit_stats = train.eval_policy(trainer.policy, trainer.venv_train)
 
     # Save final artifacts.
     if checkpoint_interval >= 0:
-        save(trainer, os.path.join(log_dir, "checkpoints", "final"))
+        save(trainer, log_dir / "checkpoints" / "final")
 
     return {
         "imit_stats": imit_stats,
@@ -178,7 +180,8 @@ def airl():
 
 
 def main_console():
-    observer = FileStorageObserver(osp.join("output", "sacred", "train_adversarial"))
+    observer_path = pathlib.Path.cwd() / "output" / "sacred" / "train_adversarial"
+    observer = FileStorageObserver(observer_path)
     train_adversarial_ex.observers.append(observer)
     train_adversarial_ex.run_commandline()
 
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index 043268550..584150cd5 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -2,6 +2,7 @@
 
 import logging
 import os.path as osp
+import pathlib
 import warnings
 from typing import Any, Mapping, Optional, Type
 
@@ -195,7 +196,8 @@ def dagger() -> Mapping[str, Mapping[str, float]]:
 
 
 def main_console():
-    observer = FileStorageObserver(osp.join("output", "sacred", "train_dagger"))
+    observer_path = pathlib.Path.cwd() / "output" / "sacred" / "train_dagger"
+    observer = FileStorageObserver(observer_path)
     train_imitation_ex.observers.append(observer)
     train_imitation_ex.run_commandline()
 
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 2ace92bc9..2794f4d37 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -6,6 +6,7 @@
 
 import functools
 import os
+import pathlib
 from typing import Any, Mapping, Optional, Type, Union
 
 import torch as th
@@ -25,23 +26,23 @@
 
 def save_model(
     agent_trainer: preference_comparisons.AgentTrainer,
-    save_path: str,
+    save_path: pathlib.Path,
 ):
     """Save the model as model.pkl."""
     serialize.save_stable_model(
-        output_dir=os.path.join(save_path, "policy"),
+        output_dir=save_path / "policy",
         model=agent_trainer.algorithm,
     )
 
 
 def save_checkpoint(
     trainer: preference_comparisons.PreferenceComparisons,
-    save_path: str,
+    save_path: pathlib.Path,
     allow_save_policy: Optional[bool],
 ):
     """Save reward model and optionally policy."""
-    os.makedirs(save_path, exist_ok=True)
-    th.save(trainer.model, os.path.join(save_path, "reward_net.pt"))
+    save_path.mkdir(parents=True, exist_ok=True)
+    th.save(trainer.model, save_path / "reward_net.pt")
     if allow_save_policy:
         # Note: We should only save the model as model.pkl if `trajectory_generator`
         # contains one. Specifically we check if the `trajectory_generator` contains an
@@ -138,7 +139,8 @@ def train_preference_comparisons(
     Raises:
         ValueError: Inconsistency between config and deserialized policy normalization.
     """
-    custom_logger, log_dir = common.setup_logging()
+    custom_logger, _log_dir = common.setup_logging()
+    log_dir = pathlib.Path(_log_dir)
 
     with common.make_venv() as venv:
         reward_net = reward.make_reward_net(venv)
@@ -224,11 +226,7 @@ def save_callback(iteration_num):
             if checkpoint_interval > 0 and iteration_num % checkpoint_interval == 0:
                 save_checkpoint(
                     trainer=main_trainer,
-                    save_path=os.path.join(
-                        log_dir,
-                        "checkpoints",
-                        f"{iteration_num:04d}",
-                    ),
+                    save_path=log_dir / "checkpoints" / f"{iteration_num:04d}",
                     allow_save_policy=bool(trajectory_path is None),
                 )
 
@@ -244,13 +242,13 @@ def save_callback(iteration_num):
             results["rollout"] = train.eval_policy(agent, venv)
 
     if save_preferences:
-        main_trainer.dataset.save(os.path.join(log_dir, "preferences.pkl"))
+        main_trainer.dataset.save(log_dir / "preferences.pkl")
 
     # Save final artifacts.
     if checkpoint_interval >= 0:
         save_checkpoint(
             trainer=main_trainer,
-            save_path=os.path.join(log_dir, "checkpoints", "final"),
+            save_path=log_dir / "checkpoints" / "final",
             allow_save_policy=bool(trajectory_path is None),
         )
 
@@ -258,9 +256,10 @@ def save_callback(iteration_num):
 
 
 def main_console():
-    observer = FileStorageObserver(
-        os.path.join("output", "sacred", "train_preference_comparisons"),
+    observer_path = (
+        pathlib.Path.cwd() / "output" / "sacred" / "train_preference_comparisons"
     )
+    observer = FileStorageObserver(observer_path)
     train_preference_comparisons_ex.observers.append(observer)
     train_preference_comparisons_ex.run_commandline()
 
diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
index 2540ac79c..f458d2872 100644
--- a/src/imitation/scripts/train_rl.py
+++ b/src/imitation/scripts/train_rl.py
@@ -11,6 +11,7 @@
 import logging
 import os
 import os.path as osp
+import pathlib
 import warnings
 from typing import Any, Mapping, Optional
 
@@ -87,11 +88,12 @@ def train_rl(
     Returns:
         The return value of `rollout_stats()` using the final policy.
     """
-    custom_logger, log_dir = common.setup_logging()
-    rollout_dir = osp.join(log_dir, "rollouts")
-    policy_dir = osp.join(log_dir, "policies")
-    os.makedirs(rollout_dir, exist_ok=True)
-    os.makedirs(policy_dir, exist_ok=True)
+    custom_logger, _log_dir = common.setup_logging()
+    log_dir = pathlib.Path(_log_dir)
+    rollout_dir = log_dir / "rollouts"
+    policy_dir = log_dir / "policies"
+    rollout_dir.mkdir(parents=True, exist_ok=True)
+    policy_dir.mkdir(parents=True, exist_ok=True)
 
     post_wrappers = [lambda env, idx: wrappers.RolloutInfoWrapper(env)]
     with common.make_venv(post_wrappers=post_wrappers) as venv:
@@ -137,14 +139,14 @@ def train_rl(
 
         # Save final artifacts after training is complete.
         if rollout_save_final:
-            save_path = osp.join(rollout_dir, "final.pkl")
+            save_path = rollout_dir / "final.pkl"
             sample_until = rollout.make_sample_until(
                 rollout_save_n_timesteps,
                 rollout_save_n_episodes,
             )
             types.save(save_path, rollout.rollout(rl_algo, venv, sample_until))
         if policy_save_final:
-            output_dir = os.path.join(policy_dir, "final")
+            output_dir = policy_dir / "final"
             serialize.save_stable_model(output_dir, rl_algo)
 
         # Final evaluation of expert policy.
@@ -152,7 +154,8 @@ def train_rl(
 
 
 def main_console():
-    observer = FileStorageObserver(osp.join("output", "sacred", "train_rl"))
+    observer_path = pathlib.Path.cwd() / "output" / "sacred" / "train_rl"
+    observer = FileStorageObserver(observer_path)
     train_rl_ex.observers.append(observer)
     train_rl_ex.run_commandline()
 
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 8875cb211..440c39bea 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -3,8 +3,9 @@
 import contextlib
 import datetime
 import os
+import pathlib
 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
 
@@ -12,7 +13,7 @@
 
 
 def _build_output_formats(
-    folder: str,
+    folder: pathlib.Path,
     format_strs: Sequence[str],
 ) -> Sequence[sb_logger.KVWriter]:
     """Build output formats for initializing a Stable Baselines Logger.
@@ -25,13 +26,13 @@ def _build_output_formats(
     Returns:
         A list of output formats, one corresponding to each `format_strs`.
     """
-    os.makedirs(folder, exist_ok=True)
-    output_formats = []
+    folder.mkdir(parents=True, exist_ok=True)
+    output_formats: List[sb_logger.KVWriter] = []
     for f in format_strs:
         if f == "wandb":
             output_formats.append(WandbOutputFormat())
         else:
-            output_formats.append(sb_logger.make_output_format(f, folder))
+            output_formats.append(sb_logger.make_output_format(f, str(folder)))
     return output_formats
 
 
@@ -43,6 +44,12 @@ class HierarchicalLogger(sb_logger.Logger):
     top-level (root) logger.
     """
 
+    default_logger: sb_logger.Logger
+    current_logger: Optional[sb_logger.Logger]
+    _cached_loggers: Dict[str, sb_logger.Logger]
+    _subdir: Optional[str]
+    _format_strs: Sequence[str]
+
     def __init__(
         self,
         default_logger: sb_logger.Logger,
@@ -108,19 +115,21 @@ def accumulate_means(self, subdir: types.AnyPath) -> Generator[None, None, None]
         if self.current_logger is not None:
             raise RuntimeError("Nested `accumulate_means` context")
 
-        if subdir in self._cached_loggers:
-            logger = self._cached_loggers[subdir]
+        subdir_str = types.path_to_str(subdir)
+        if subdir_str in self._cached_loggers:
+            logger = self._cached_loggers[subdir_str]
         else:
-            subdir = types.path_to_str(subdir)
-            folder = os.path.join(self.default_logger.dir, "raw", subdir)
-            os.makedirs(folder, exist_ok=True)
+            default_logger_dir = self.default_logger.dir
+            assert default_logger_dir is not None
+            folder = pathlib.Path(default_logger_dir) / "raw" / subdir_str
+            folder.mkdir(exist_ok=True, parents=True)
             output_formats = _build_output_formats(folder, self.format_strs)
-            logger = sb_logger.Logger(folder, list(output_formats))
-            self._cached_loggers[subdir] = logger
+            logger = sb_logger.Logger(str(folder), list(output_formats))
+            self._cached_loggers[subdir_str] = logger
 
         try:
             self.current_logger = logger
-            self._subdir = subdir
+            self._subdir = subdir_str
             self._update_name_to_maps()
             yield
         finally:
@@ -227,15 +236,16 @@ def configure(
         The configured HierarchicalLogger instance.
     """
     if folder is None:
+        tempdir = pathlib.Path(tempfile.gettempdir())
         now = datetime.datetime.now()
         timestamp = now.strftime("imitation-%Y-%m-%d-%H-%M-%S-%f")
-        folder = os.path.join(tempfile.gettempdir(), timestamp)
+        folder = tempdir / timestamp
     else:
-        folder = types.path_to_str(folder)
+        folder = types.path_to_pathlib(folder)
     if format_strs is None:
         format_strs = ["stdout", "log", "csv"]
     output_formats = _build_output_formats(folder, format_strs)
-    default_logger = sb_logger.Logger(folder, list(output_formats))
+    default_logger = sb_logger.Logger(str(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
diff --git a/src/imitation/util/video_wrapper.py b/src/imitation/util/video_wrapper.py
index 5fc2ae4f5..4a297ce73 100644
--- a/src/imitation/util/video_wrapper.py
+++ b/src/imitation/util/video_wrapper.py
@@ -1,6 +1,7 @@
 """Wrapper to record rendered video frames from an environment."""
 
 import os
+import pathlib
 
 import gym
 from gym.wrappers.monitoring import video_recorder
@@ -33,8 +34,8 @@ def __init__(
         self.video_recorder = None
         self.single_video = single_video
 
-        self.directory = os.path.abspath(directory)
-        os.makedirs(self.directory)
+        self.directory = pathlib.Path(types.path_to_str(directory)).resolve()
+        self.directory.mkdir(parents=True, exist_ok=True)
 
     def _reset_video_recorder(self) -> None:
         """Creates a video recorder if one does not already exist.
@@ -53,10 +54,7 @@ def _reset_video_recorder(self) -> None:
             # No video recorder -- start a new one.
             self.video_recorder = video_recorder.VideoRecorder(
                 env=self.env,
-                base_path=os.path.join(
-                    self.directory,
-                    "video.{:06}".format(self.episode_id),
-                ),
+                base_path=str(self.directory / f"video.{self.episode_id:06}"),
                 metadata={"episode_id": self.episode_id},
             )
 
diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index af8398ba3..b5672869a 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -111,7 +111,7 @@ def _check_transitions_get_item(trans, key):
 @contextlib.contextmanager
 def pushd(dir_path):
     """Change directory temporarily inside context."""
-    orig_dir = os.getcwd()
+    orig_dir = pathlib.Path.cwd()
     try:
         os.chdir(dir_path)
         yield

From a818941ae269cfd3b583e99d75a2db06a444cd88 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Sun, 28 Aug 2022 13:35:31 +0100
Subject: [PATCH 004/187] Bug fixes from tests

---
 src/imitation/algorithms/adversarial/common.py | 2 +-
 src/imitation/scripts/common/demonstrations.py | 3 ++-
 src/imitation/scripts/convert_trajs.py         | 3 ++-
 src/imitation/scripts/eval_policy.py           | 3 ++-
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
index 877a0cba2..ba244812d 100644
--- a/src/imitation/algorithms/adversarial/common.py
+++ b/src/imitation/algorithms/adversarial/common.py
@@ -201,7 +201,7 @@ def __init__(
         )
 
         if self._init_tensorboard:
-            logging.info("building summary directory at " + self._log_dir)
+            logging.info(f"building summary directory at {self._log_dir}")
             summary_dir = self._log_dir / "summary"
             summary_dir.mkdir(parents=True, exist_ok=True)
             self._summary_writer = thboard.SummaryWriter(str(summary_dir))
diff --git a/src/imitation/scripts/common/demonstrations.py b/src/imitation/scripts/common/demonstrations.py
index eb5c8b275..c97eec58c 100644
--- a/src/imitation/scripts/common/demonstrations.py
+++ b/src/imitation/scripts/common/demonstrations.py
@@ -28,7 +28,8 @@ def fast():
     n_expert_demos = 1  # noqa: F841
 
 
-def guess_expert_dir(data_dir: pathlib.Path, env_name: str) -> pathlib.Path:
+def guess_expert_dir(data_dir: str, env_name: str) -> pathlib.Path:
+    data_dir = pathlib.Path(data_dir)
     assert data_dir.is_absolute()
     rollout_hint = env_name.rsplit("-", 1)[0].replace("/", "_").lower()
     return data_dir / "expert_models" / f"{rollout_hint}_0"
diff --git a/src/imitation/scripts/convert_trajs.py b/src/imitation/scripts/convert_trajs.py
index 7c1e9549a..8eabc7edd 100644
--- a/src/imitation/scripts/convert_trajs.py
+++ b/src/imitation/scripts/convert_trajs.py
@@ -16,7 +16,7 @@
 from imitation.data import types
 
 
-def update_traj_file_in_place(path: pathlib.Path) -> None:
+def update_traj_file_in_place(path: str) -> None:
     """Modifies trajectories pickle file in-place to update data to new format.
 
     The new data is saved as `Sequence[imitation.types.TrajectoryWithRew]`.
@@ -34,6 +34,7 @@ def update_traj_file_in_place(path: pathlib.Path) -> None:
         )
         trajs = types.load(path)
 
+    path = pathlib.Path(path)
     ext = path.suffix
     new_ext = ".npz" if ext in (".pkl", ".npz") else ext + ".npz"
     types.save(path.with_suffix(new_ext), trajs)
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 7569acd8a..28bc76da1 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -110,7 +110,8 @@ def eval_policy(
         trajs = rollout.generate_trajectories(policy, venv, sample_until)
 
     if rollout_save_path:
-        types.save(rollout_save_path.replace("{log_dir}", log_dir), trajs)
+
+        types.save(log_dir / rollout_save_path.lstrip("{log_dir}/"), trajs)
 
     return rollout.rollout_stats(trajs)
 

From ad658776f22005e0e3066d250d4a9f0daa37ee13 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 16:21:42 +0100
Subject: [PATCH 005/187] Fix types: test_envs.py

---
 tests/test_envs.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_envs.py b/tests/test_envs.py
index fa53931d1..eb55db036 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -1,4 +1,5 @@
 """Tests for `imitation.envs.*`."""
+from typing import List
 
 import gym
 import numpy as np
@@ -16,8 +17,7 @@
     if env_spec.id.startswith("imitation/")
 ]
 
-DETERMINISTIC_ENVS = []
-
+DETERMINISTIC_ENVS: List[str] = []
 
 env = pytest.fixture(seals_test.make_env_fixture(skip_fn=pytest.skip))
 

From 240b6009245283b73c3047ac1784def07aae678f Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 16:21:54 +0100
Subject: [PATCH 006/187] Fix types: conftest.py

---
 tests/conftest.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 206d04039..b27244678 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -33,7 +33,12 @@ def load_or_train_ppo(
     os.makedirs(os.path.dirname(cache_path), exist_ok=True)
     with FileLock(cache_path + ".lock"):
         try:
-            return PPO.load(cache_path, venv)
+            ppo = PPO.load(cache_path, venv)
+            # TODO(juan): sb3 return type for .load() is too general.
+            #  See https://github.com/DLR-RM/stable-baselines3/issues/1040
+            #  remove the line blow once fixed.
+            assert isinstance(ppo, PPO)
+            return ppo
         except (OSError, AssertionError, pickle.PickleError):  # pragma: no cover
             # Note, when loading models from older stable-baselines versions, we can get
             # AssertionErrors.
@@ -110,6 +115,7 @@ def train_cartpole_expert(cartpole_env) -> Optional[PPO]:  # pragma: no cover
         )
         policy.learn(100000)
         mean_reward, _ = evaluate_policy(policy, cartpole_env, 10)
+        assert isinstance(mean_reward, float)
         if mean_reward >= 500:
             return policy
     return None
@@ -169,6 +175,7 @@ def train_pendulum_expert(pendulum_env) -> Optional[PPO]:  # pragma: no cover
         )
         policy.learn(int(1e5))
         mean_reward, _ = evaluate_policy(policy, pendulum_env, 10)
+        assert isinstance(mean_reward, float)
         if mean_reward >= -185:
             return policy
     return None

From 45eac477e9dcc040ce0451535b2fc08307e3eb35 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 16:22:27 +0100
Subject: [PATCH 007/187] Fix types: tests/util

---
 tests/util/test_wb_logger.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/util/test_wb_logger.py b/tests/util/test_wb_logger.py
index fe755ad31..392f21fdc 100644
--- a/tests/util/test_wb_logger.py
+++ b/tests/util/test_wb_logger.py
@@ -87,7 +87,7 @@ def finish(self):
 mock_wandb = MockWandb()
 
 
-@mock.patch.object(wandb, "__init__", mock_wandb.__init__)
+@mock.patch.object(wandb, "__init__", mock_wandb.__init__)  # type: ignore
 @mock.patch.object(wandb, "init", mock_wandb.init)
 @mock.patch.object(wandb, "log", mock_wandb.log)
 @mock.patch.object(wandb, "finish", mock_wandb.finish)

From 7e4cf7b5ef70a448399962a3b9eb8dda58e43fd3 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 16:22:41 +0100
Subject: [PATCH 008/187] Fix types: tests/scripts

---
 tests/scripts/test_scripts.py | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/tests/scripts/test_scripts.py b/tests/scripts/test_scripts.py
index 6a959b1e9..6b0694f92 100644
--- a/tests/scripts/test_scripts.py
+++ b/tests/scripts/test_scripts.py
@@ -15,7 +15,7 @@
 import sys
 import tempfile
 from collections import Counter
-from typing import Dict, List, Mapping, Optional
+from typing import Any, Dict, List, Mapping, Optional
 from unittest import mock
 
 import numpy as np
@@ -464,6 +464,7 @@ def _check_train_ex_result(result: dict):
     assert "monitor_return_mean" not in expert_stats
 
     imit_stats = result.get("imit_stats")
+    assert isinstance(imit_stats, dict)
     _check_rollout_stats(imit_stats)
 
 
@@ -592,8 +593,8 @@ def test_transfer_learning(tmpdir: str) -> None:
     Args:
         tmpdir: Temporary directory to save results to.
     """
-    tmpdir = pathlib.Path(tmpdir)
-    log_dir_train = tmpdir / "train"
+    tmpdir_path = pathlib.Path(tmpdir)
+    log_dir_train = tmpdir_path / "train"
     run = train_adversarial.train_adversarial_ex.run(
         command_name="airl",
         named_configs=["cartpole"] + ALGO_FAST_CONFIGS["adversarial"],
@@ -607,7 +608,7 @@ def test_transfer_learning(tmpdir: str) -> None:
 
     _check_rollout_stats(run.result["imit_stats"])
 
-    log_dir_data = tmpdir / "train_rl"
+    log_dir_data = tmpdir_path / "train_rl"
     reward_path = log_dir_train / "checkpoints" / "final" / "reward_test.pt"
     run = train_rl.train_rl_ex.run(
         named_configs=["cartpole"] + ALGO_FAST_CONFIGS["rl"],
@@ -641,9 +642,9 @@ def test_preference_comparisons_transfer_learning(
         tmpdir: Temporary directory to save results to.
         named_configs_dict: Named configs for preference_comparisons and rl.
     """
-    tmpdir = pathlib.Path(tmpdir)
+    tmpdir_path = pathlib.Path(tmpdir)
 
-    log_dir_train = tmpdir / "train"
+    log_dir_train = tmpdir_path / "train"
     run = train_preference_comparisons.train_preference_comparisons_ex.run(
         named_configs=["pendulum"]
         + ALGO_FAST_CONFIGS["preference_comparison"]
@@ -661,7 +662,7 @@ def test_preference_comparisons_transfer_learning(
         reward_type = "RewardNet_unnormalized"
         load_reward_kwargs = {}
 
-    log_dir_data = tmpdir / "train_rl"
+    log_dir_data = tmpdir_path / "train_rl"
     reward_path = log_dir_train / "checkpoints" / "final" / "reward_net.pt"
     agent_path = log_dir_train / "checkpoints" / "final" / "policy"
     run = train_rl.train_rl_ex.run(
@@ -680,8 +681,10 @@ def test_preference_comparisons_transfer_learning(
 
 def test_train_rl_double_normalization(tmpdir: str):
     venv = util.make_vec_env("CartPole-v1", n_envs=1, parallel=False)
-    net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
-    net = reward_nets.NormalizedRewardNet(net, networks.RunningNorm)
+    basic_reward_net = reward_nets.BasicRewardNet(
+        venv.observation_space, venv.action_space
+    )
+    net = reward_nets.NormalizedRewardNet(basic_reward_net, networks.RunningNorm)
     tmppath = os.path.join(tmpdir, "reward.pt")
     th.save(net, tmppath)
 
@@ -774,14 +777,14 @@ def test_parallel_arg_errors(tmpdir):
 
 
 def _generate_test_rollouts(tmpdir: str, env_named_config: str) -> pathlib.Path:
-    tmpdir = pathlib.Path(tmpdir)
+    tmpdir_path = pathlib.Path(tmpdir)
     train_rl.train_rl_ex.run(
         named_configs=[env_named_config] + ALGO_FAST_CONFIGS["rl"],
         config_updates=dict(
             common=dict(log_dir=tmpdir),
         ),
     )
-    rollout_path = tmpdir / "rollouts/final.pkl"
+    rollout_path = tmpdir_path / "rollouts/final.pkl"
     return rollout_path.absolute()
 
 
@@ -846,7 +849,7 @@ def _run_train_bc_for_test_analyze_imit(run_name, sacred_logs_dir, log_dir):
     ),
 )
 def test_analyze_imitation(tmpdir: str, run_names: List[str], run_sacred_fn):
-    sacred_logs_dir = tmpdir = pathlib.Path(tmpdir)
+    sacred_logs_dir = tmpdir_path = pathlib.Path(tmpdir)
 
     # Generate sacred logs (other logs are put in separate tmpdir for deletion).
     for i, run_name in enumerate(run_names):
@@ -862,8 +865,8 @@ def check(run_name: Optional[str], count: int) -> None:
                 source_dirs=[sacred_logs_dir],
                 env_name="CartPole-v1",
                 run_name=run_name,
-                csv_output_path=tmpdir / "analysis.csv",
-                tex_output_path=tmpdir / "analysis.tex",
+                csv_output_path=tmpdir_path / "analysis.csv",
+                tex_output_path=tmpdir_path / "analysis.tex",
                 print_table=True,
             ),
         )
@@ -881,7 +884,7 @@ def test_analyze_gather_tb(tmpdir: str):
     if os.name == "nt":  # pragma: no cover
         pytest.skip("gather_tb uses symlinks: not supported by Windows")
 
-    config_updates = dict(local_dir=tmpdir, run_name="test")
+    config_updates: Dict[str, Any] = dict(local_dir=tmpdir, run_name="test")
     config_updates.update(PARALLEL_CONFIG_LOW_RESOURCE)
     parallel_run = parallel.parallel_ex.run(
         named_configs=["generate_test_data"],

From 587be34430bc580c325c76215855fc1b11b348df Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 16:23:01 +0100
Subject: [PATCH 009/187] Fix types: tests/rewards

---
 tests/rewards/test_reward_fn.py   |  2 +-
 tests/rewards/test_reward_nets.py | 11 +++++------
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/tests/rewards/test_reward_fn.py b/tests/rewards/test_reward_fn.py
index 5e14aa3d6..98a3f130c 100644
--- a/tests/rewards/test_reward_fn.py
+++ b/tests/rewards/test_reward_fn.py
@@ -7,7 +7,7 @@
 
 OBS = np.random.randint(0, 10, (64, 100))
 ACTS = NEXT_OBS = OBS
-DONES = np.zeros(64, dtype=np.bool)
+DONES = np.zeros(64, dtype=np.bool)  # type: ignore
 
 
 def _funky_reward_fn(obs, act, next_obs, done):
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index f6084d3ce..9c51e96be 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -35,7 +35,6 @@ def _potential(x):
     "RewardNet_unshaped",
 ]
 
-
 # Reward net classes, allowed kwargs
 MAKE_REWARD_NET = [
     reward_nets.BasicRewardNet,
@@ -43,7 +42,6 @@ def _potential(x):
     testing_reward_nets.make_ensemble,
 ]
 
-
 MAKE_BASIC_REWARD_NET_WRAPPERS = [
     lambda base: reward_nets.ShapedRewardNet(base, _potential, 0.99),
     lambda base: reward_nets.NormalizedRewardNet(base, networks.RunningNorm),
@@ -59,7 +57,6 @@ def _potential(x):
     networks.RunningNorm,
 ]
 
-
 NumpyTransitions = Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
 
 
@@ -82,9 +79,9 @@ def torch_transitions() -> TorchTransitions:
     """A batch of states, actions, next_states, and dones as th.Tensors for Env2D."""
     return (
         th.zeros((10, 5, 5)),
-        th.zeros((10, 1), dtype=int),
+        th.zeros((10, 1), dtype=th.int),
         th.zeros((10, 5, 5)),
-        th.zeros((10,), dtype=bool),
+        th.zeros((10,), dtype=th.bool),
     )
 
 
@@ -541,7 +538,9 @@ def test_wrappers_pass_on_kwargs(
         env_2d.observation_space,
         env_2d.action_space,
     )
-    basic_reward_net.predict_processed = mock.Mock(return_value=np.zeros((10,)))
+    # TODO(juan) reassigning a method is very bad practice. I added
+    #  type ignore for now but is there not a better way to do this?
+    basic_reward_net.predict_processed = mock.Mock(return_value=np.zeros((10,)))  # type: ignore
     wrapped_reward_net = make_wrapper(
         basic_reward_net,
     )

From e79305299b7f38d6c503e02080b2315f6a340b91 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 16:23:22 +0100
Subject: [PATCH 010/187] Fix types: tests/policies

---
 tests/policies/test_policies.py              |  6 +++++-
 tests/policies/test_replay_buffer_wrapper.py | 21 ++++++++++----------
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index f095edb7c..d48940042 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -2,6 +2,7 @@
 
 import functools
 import pathlib
+from typing import cast
 
 import gym
 import numpy as np
@@ -161,7 +162,10 @@ def test_normalize_features_extractor(obs_space: gym.Space) -> None:
 
     for i in range(10):
         obs = th.as_tensor([obs_space.sample()])
-        obs = preprocessing.preprocess_obs(obs, obs_space)
+        # TODO(juan) the cast below is because preprocess_obs has too general a type.
+        #  this should be replaced with an overload or a generic.
+        obs = cast(th.Tensor, preprocessing.preprocess_obs(obs, obs_space))
+        assert isinstance(obs, th.Tensor)
         flattened_obs = obs.flatten(1, -1)
         extracted = {k: extractor(obs) for k, extractor in extractors.items()}
         for k, v in extracted.items():
diff --git a/tests/policies/test_replay_buffer_wrapper.py b/tests/policies/test_replay_buffer_wrapper.py
index 4ae73d39e..1d6ec001a 100644
--- a/tests/policies/test_replay_buffer_wrapper.py
+++ b/tests/policies/test_replay_buffer_wrapper.py
@@ -32,21 +32,20 @@ def make_algo_with_wrapped_buffer(
     buffer_size: int = 100,
 ) -> off_policy_algorithm.OffPolicyAlgorithm:
     venv = util.make_vec_env("Pendulum-v1", n_envs=1)
-    rl_kwargs = dict(
-        replay_buffer_class=ReplayBufferRewardWrapper,
-        replay_buffer_kwargs=dict(
-            replay_buffer_class=replay_buffer_class,
-            reward_fn=zero_reward_fn,
-        ),
-        buffer_size=buffer_size,
-    )
     rl_algo = rl_cls(
         policy=policy_cls,
         policy_kwargs=dict(),
         env=venv,
         seed=42,
-        **rl_kwargs,
-    )
+        # we ignore the type below because sb3 has a bug (forgot to put Type[...])
+        # https://github.com/DLR-RM/stable-baselines3/issues/1039
+        replay_buffer_class=ReplayBufferRewardWrapper,  # type: ignore
+        replay_buffer_kwargs=dict(
+            replay_buffer_class=replay_buffer_class,
+            reward_fn=zero_reward_fn,
+        ),
+        buffer_size=buffer_size,
+    )  # type: ignore
     return rl_algo
 
 
@@ -56,7 +55,7 @@ def test_invalid_args():
         match=r".*unexpected keyword argument 'replay_buffer_class'.*",
     ):
         make_algo_with_wrapped_buffer(
-            rl_cls=sb3.PPO,
+            rl_cls=sb3.PPO,  # type: ignore
             policy_cls=policies.ActorCriticPolicy,
             replay_buffer_class=buffers.ReplayBuffer,
         )

From 9f367c0d8e7e5ee0e27d8ea22003d50a252c3117 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 16:24:31 +0100
Subject: [PATCH 011/187] Incorrect decorator in update_stats method form
 networks.py::BaseNorm

---
 src/imitation/util/networks.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index a59a6be2a..e6b0e221f 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -1,8 +1,8 @@
 """Helper methods to build and run neural networks."""
+import abc
 import collections
 import contextlib
 import functools
-from abc import ABC, abstractclassmethod
 from typing import Iterable, Optional, Type
 
 import torch as th
@@ -44,7 +44,7 @@ def forward(self, x):
         return new_value
 
 
-class BaseNorm(nn.Module, ABC):
+class BaseNorm(nn.Module, abc.ABC):
     """Base class for layers that try to normalize the input to mean 0 and variance 1.
 
     Similar to BatchNorm, LayerNorm, etc. but whereas they only use statistics from
@@ -88,7 +88,7 @@ def forward(self, x: th.Tensor) -> th.Tensor:
 
         return (x - self.running_mean) / th.sqrt(self.running_var + self.eps)
 
-    @abstractclassmethod
+    @abc.abstractmethod
     def update_stats(self, batch: th.Tensor) -> None:
         """Update `self.running_mean`, `self.running_var` and `self.count`."""
 

From b8c4cb18d6742e91f480731dd9ce9d1d729cbbcd Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 16:24:48 +0100
Subject: [PATCH 012/187] Fix types: tests/algorithms (adersarial and bc)

---
 tests/algorithms/test_adversarial.py | 92 ++++++++++++++--------------
 tests/algorithms/test_bc.py          | 20 +++---
 2 files changed, 58 insertions(+), 54 deletions(-)

diff --git a/tests/algorithms/test_adversarial.py b/tests/algorithms/test_adversarial.py
index fc06235fb..806e773c8 100644
--- a/tests/algorithms/test_adversarial.py
+++ b/tests/algorithms/test_adversarial.py
@@ -2,7 +2,7 @@
 
 import contextlib
 import os
-from typing import Any, Mapping
+from typing import Any, Mapping, Union, Type
 
 import numpy as np
 import pytest
@@ -42,7 +42,7 @@
 EXPERT_BATCH_SIZES = [1, 128]
 
 
-@pytest.fixture(params=ALGORITHM_KWARGS.values(), ids=ALGORITHM_KWARGS.keys())
+@pytest.fixture(params=ALGORITHM_KWARGS.values(), ids=list(ALGORITHM_KWARGS.keys()))
 def _algorithm_kwargs(request):
     """Auto-parametrizes `_rl_algorithm_cls` for the `trainer` fixture."""
     return dict(request.param)
@@ -55,15 +55,16 @@ def expert_transitions(cartpole_expert_trajectories):
 
 @contextlib.contextmanager
 def make_trainer(
-    algorithm_kwargs: Mapping[str, Any],
-    tmpdir: str,
-    expert_transitions: types.Transitions,
-    expert_batch_size: int = 1,
-    env_name: str = "seals/CartPole-v0",
-    num_envs: int = 1,
-    parallel: bool = False,
-    convert_dataset: bool = False,
+        algorithm_kwargs: Mapping[str, Any],
+        tmpdir: str,
+        expert_transitions: types.Transitions,
+        expert_batch_size: int = 1,
+        env_name: str = "seals/CartPole-v0",
+        num_envs: int = 1,
+        parallel: bool = False,
+        convert_dataset: bool = False,
 ):
+    expert_data: Union[th_data.DataLoader, th_data.Dataset]
     if convert_dataset:
         expert_data = th_data.DataLoader(
             expert_transitions,
@@ -78,7 +79,7 @@ def make_trainer(
     venv = util.make_vec_env(env_name, n_envs=num_envs, parallel=parallel)
     model_cls = algorithm_kwargs["model_class"]
     gen_algo = model_cls(algorithm_kwargs["policy_class"], venv)
-    reward_net_cls = reward_nets.BasicRewardNet
+    reward_net_cls: Type[reward_nets.RewardNet] = reward_nets.BasicRewardNet
     if algorithm_kwargs["algorithm_cls"] == airl.AIRL:
         reward_net_cls = reward_nets.BasicShapedRewardNet
     reward_net = reward_net_cls(venv.observation_space, venv.action_space)
@@ -126,7 +127,7 @@ def test_airl_fail_fast(custom_logger, tmpdir):
         )
 
 
-@pytest.fixture(params=ALGORITHM_KWARGS.values(), ids=ALGORITHM_KWARGS.keys())
+@pytest.fixture(params=ALGORITHM_KWARGS.values(), ids=list(ALGORITHM_KWARGS.keys()))
 def trainer(request, tmpdir, expert_transitions):
     with make_trainer(request.param, tmpdir, expert_transitions) as trainer:
         yield trainer
@@ -138,8 +139,8 @@ def test_train_disc_no_samples_error(trainer: common.AdversarialTrainer):
 
 
 def test_train_disc_unequal_expert_gen_samples_error(
-    trainer: common.AdversarialTrainer,
-    expert_transitions: types.Transitions,
+        trainer: common.AdversarialTrainer,
+        expert_transitions: types.Transitions,
 ):
     """Test that train_disc raises error when n_gen != n_expert samples."""
     if len(expert_transitions) < 2:  # pragma: no cover
@@ -170,20 +171,20 @@ def _expert_batch_size(request):
 
 @pytest.fixture
 def trainer_parametrized(
-    _algorithm_kwargs,
-    _parallel,
-    _convert_dataset,
-    _expert_batch_size,
-    tmpdir,
-    expert_transitions,
-):
-    with make_trainer(
         _algorithm_kwargs,
+        _parallel,
+        _convert_dataset,
+        _expert_batch_size,
         tmpdir,
         expert_transitions,
-        parallel=_parallel,
-        convert_dataset=_convert_dataset,
-        expert_batch_size=_expert_batch_size,
+):
+    with make_trainer(
+            _algorithm_kwargs,
+            tmpdir,
+            expert_transitions,
+            parallel=_parallel,
+            convert_dataset=_convert_dataset,
+            expert_batch_size=_expert_batch_size,
     ) as trainer:
         yield trainer
 
@@ -201,8 +202,8 @@ def test_train_disc_step_no_crash(trainer_parametrized, _expert_batch_size):
 
 
 def test_train_gen_train_disc_no_crash(
-    trainer_parametrized: common.AdversarialTrainer,
-    n_updates: int = 2,
+        trainer_parametrized: common.AdversarialTrainer,
+        n_updates: int = 2,
 ) -> None:
     trainer_parametrized.train_gen(n_updates * trainer_parametrized.gen_train_timesteps)
     trainer_parametrized.train_disc()
@@ -210,26 +211,26 @@ def test_train_gen_train_disc_no_crash(
 
 @pytest.fixture
 def trainer_batch_sizes(
-    _algorithm_kwargs,
-    _expert_batch_size,
-    tmpdir,
-    expert_transitions,
-):
-    with make_trainer(
         _algorithm_kwargs,
+        _expert_batch_size,
         tmpdir,
         expert_transitions,
-        expert_batch_size=_expert_batch_size,
+):
+    with make_trainer(
+            _algorithm_kwargs,
+            tmpdir,
+            expert_transitions,
+            expert_batch_size=_expert_batch_size,
     ) as trainer:
         yield trainer
 
 
 def test_train_disc_improve_D(
-    trainer_batch_sizes,
-    tmpdir,
-    expert_transitions,
-    _expert_batch_size,
-    n_steps=3,
+        trainer_batch_sizes,
+        tmpdir,
+        expert_transitions,
+        _expert_batch_size,
+        n_steps=3,
 ):
     expert_samples = expert_transitions[:_expert_batch_size]
     expert_samples = types.dataclass_quick_asdict(expert_samples)
@@ -262,18 +263,18 @@ def trainer_diverse_env(_algorithm_kwargs, _env_name, tmpdir, expert_transitions
     if _algorithm_kwargs["model_class"] == stable_baselines3.DQN:
         pytest.skip("DQN does not support all environments.")
     with make_trainer(
-        _algorithm_kwargs,
-        tmpdir,
-        expert_transitions,
-        env_name=_env_name,
+            _algorithm_kwargs,
+            tmpdir,
+            expert_transitions,
+            env_name=_env_name,
     ) as trainer:
         yield trainer
 
 
 @pytest.mark.parametrize("n_timesteps", [2, 4, 10])
 def test_logits_expert_is_high_log_policy_act_prob(
-    trainer_diverse_env: common.AdversarialTrainer,
-    n_timesteps: int,
+        trainer_diverse_env: common.AdversarialTrainer,
+        n_timesteps: int,
 ):
     """Smoke test calling `logits_expert_is_high` on `AdversarialTrainer`.
 
@@ -300,6 +301,7 @@ def test_logits_expert_is_high_log_policy_act_prob(
     log_act_prob_non_none = th.as_tensor(log_act_prob_non_none).to(obs.device)
 
     for log_act_prob in [None, log_act_prob_non_none]:
+        maybe_error_ctx: contextlib.AbstractContextManager
         if isinstance(trainer_diverse_env, airl.AIRL) and log_act_prob is None:
             maybe_error_ctx = pytest.raises(TypeError, match="Non-None.*required.*")
         else:
diff --git a/tests/algorithms/test_bc.py b/tests/algorithms/test_bc.py
index 3e01a12e2..3e1405ff6 100644
--- a/tests/algorithms/test_bc.py
+++ b/tests/algorithms/test_bc.py
@@ -42,11 +42,11 @@ def __iter__(self):
 
 @pytest.fixture
 def trainer(
-    batch_size,
-    cartpole_venv,
-    expert_data_type,
-    custom_logger,
-    cartpole_expert_trajectories,
+        batch_size,
+        cartpole_venv,
+        expert_data_type,
+        custom_logger,
+        cartpole_expert_trajectories,
 ):
     trans = rollout.flatten_trajectories(cartpole_expert_trajectories)
     if expert_data_type == "data_loader":
@@ -101,6 +101,7 @@ def test_bc(trainer: bc.BC, cartpole_venv):
         15,
         return_episode_rewards=True,
     )
+    assert isinstance(novice_rewards, (list, tuple))
 
     trainer.train(
         n_epochs=1,
@@ -114,6 +115,7 @@ def test_bc(trainer: bc.BC, cartpole_venv):
         15,
         return_episode_rewards=True,
     )
+    assert isinstance(rewards_after_training, (list, tuple))
     assert reward_improvement.is_significant_reward_improvement(
         novice_rewards,
         rewards_after_training,
@@ -159,10 +161,10 @@ def __iter__(self):
 
 @pytest.mark.parametrize("no_yield_after_iter", [0, 1, 5])
 def test_bc_data_loader_empty_iter_error(
-    cartpole_venv: vec_env.VecEnv,
-    no_yield_after_iter: bool,
-    custom_logger: logger.HierarchicalLogger,
-    cartpole_expert_trajectories,
+        cartpole_venv: vec_env.VecEnv,
+        no_yield_after_iter: bool,
+        custom_logger: logger.HierarchicalLogger,
+        cartpole_expert_trajectories,
 ) -> None:
     """Check that we error out if the DataLoader suddenly stops yielding any batches.
 

From f4d62d204380de35eb490bcd1fb0f0faa4649aa4 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 19:05:50 +0100
Subject: [PATCH 013/187] Fix types: tests/algorithms (dagger and pc)

---
 tests/algorithms/test_dagger.py               |  83 ++++++-------
 .../algorithms/test_preference_comparisons.py | 116 +++++++++---------
 2 files changed, 101 insertions(+), 98 deletions(-)

diff --git a/tests/algorithms/test_dagger.py b/tests/algorithms/test_dagger.py
index 1bb204e2c..dd255565c 100644
--- a/tests/algorithms/test_dagger.py
+++ b/tests/algorithms/test_dagger.py
@@ -23,8 +23,8 @@
 
 @pytest.fixture(params=[True, False])
 def maybe_pendulum_expert_trajectories(
-    pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
-    request,
+        pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
+        request,
 ) -> Optional[Sequence[TrajectoryWithRew]]:
     keep_trajs = request.param
     if keep_trajs:
@@ -104,12 +104,12 @@ def get_random_acts(obs):
 
 
 def _build_dagger_trainer(
-    tmpdir,
-    venv,
-    beta_schedule,
-    expert_policy,
-    pendulum_expert_rollouts: List[TrajectoryWithRew],
-    custom_logger,
+        tmpdir,
+        venv,
+        beta_schedule,
+        expert_policy,
+        pendulum_expert_rollouts: List[TrajectoryWithRew],
+        custom_logger,
 ):
     del expert_policy
     if pendulum_expert_rollouts is not None:
@@ -133,12 +133,12 @@ def _build_dagger_trainer(
 
 
 def _build_simple_dagger_trainer(
-    tmpdir,
-    venv,
-    beta_schedule,
-    expert_policy,
-    pendulum_expert_rollouts: List[TrajectoryWithRew],
-    custom_logger,
+        tmpdir,
+        venv,
+        beta_schedule,
+        expert_policy,
+        pendulum_expert_rollouts: Optional[List[TrajectoryWithRew]],
+        custom_logger,
 ):
     bc_trainer = bc.BC(
         observation_space=venv.observation_space,
@@ -164,13 +164,13 @@ def beta_schedule(request):
 
 @pytest.fixture(params=[_build_dagger_trainer, _build_simple_dagger_trainer])
 def init_trainer_fn(
-    request,
-    tmpdir,
-    pendulum_venv,
-    beta_schedule,
-    pendulum_expert_policy,
-    maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
-    custom_logger,
+        request,
+        tmpdir,
+        pendulum_venv,
+        beta_schedule,
+        pendulum_expert_policy,
+        maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
+        custom_logger,
 ):
     # Provide a trainer initialization fixture in addition `trainer` fixture below
     # for tests that want to initialize multiple DAggerTrainer.
@@ -192,12 +192,12 @@ def trainer(init_trainer_fn):
 
 @pytest.fixture
 def simple_dagger_trainer(
-    tmpdir,
-    pendulum_venv,
-    beta_schedule,
-    pendulum_expert_policy,
-    maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
-    custom_logger,
+        tmpdir,
+        pendulum_venv,
+        beta_schedule,
+        pendulum_expert_policy,
+        maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
+        custom_logger,
 ):
     return _build_simple_dagger_trainer(
         tmpdir,
@@ -210,14 +210,15 @@ def simple_dagger_trainer(
 
 
 def test_trainer_needs_demos_exception_error(
-    trainer,
-    maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
+        trainer,
+        maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
 ):
     assert trainer.round_num == 0
     error_ctx = pytest.raises(dagger.NeedsDemosException)
+    ctx: contextlib.AbstractContextManager
     if maybe_pendulum_expert_trajectories is not None and isinstance(
-        trainer,
-        dagger.SimpleDAggerTrainer,
+            trainer,
+            dagger.SimpleDAggerTrainer,
     ):
         # In this case, demos should be preloaded and we shouldn't experience
         # the NeedsDemoException error.
@@ -332,10 +333,10 @@ def test_trainer_save_reload(tmpdir, init_trainer_fn, pendulum_venv):
 
 @pytest.mark.parametrize("num_episodes", [1, 4])
 def test_simple_dagger_trainer_train(
-    simple_dagger_trainer: dagger.SimpleDAggerTrainer,
-    pendulum_venv,
-    num_episodes: int,
-    tmpdir: str,
+        simple_dagger_trainer: dagger.SimpleDAggerTrainer,
+        pendulum_venv,
+        num_episodes: int,
+        tmpdir: str,
 ):
     episode_length = 200  # for Pendulum-v1
     rollout_min_episodes = 2
@@ -365,12 +366,12 @@ def test_policy_save_reload(tmpdir, trainer):
 
 
 def test_simple_dagger_space_mismatch_error(
-    tmpdir,
-    pendulum_venv,
-    beta_schedule,
-    pendulum_expert_policy,
-    maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
-    custom_logger,
+        tmpdir,
+        pendulum_venv,
+        beta_schedule,
+        pendulum_expert_policy,
+        maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
+        custom_logger,
 ):
     class MismatchedSpace(gym.spaces.Space):
         """Dummy space that is not equal to any other space."""
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 24bdf263a..3c7bb99cc 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -64,14 +64,16 @@ def agent_trainer(agent, reward_net, venv):
 
 
 def _check_trajs_equal(
-    trajs1: Sequence[types.TrajectoryWithRew],
-    trajs2: Sequence[types.TrajectoryWithRew],
+        trajs1: Sequence[types.TrajectoryWithRew],
+        trajs2: Sequence[types.TrajectoryWithRew],
 ):
     assert len(trajs1) == len(trajs2)
     for traj1, traj2 in zip(trajs1, trajs2):
         assert np.array_equal(traj1.obs, traj2.obs)
         assert np.array_equal(traj1.acts, traj2.acts)
         assert np.array_equal(traj1.rews, traj2.rews)
+        assert traj1.infos is not None
+        assert traj2.infos is not None
         assert np.array_equal(traj1.infos, traj2.infos)
         assert traj1.terminal == traj2.terminal
 
@@ -86,15 +88,15 @@ def test_mismatched_spaces(venv, agent):
         other_venv.action_space,
     )
     with pytest.raises(
-        ValueError,
-        match="spaces do not match",
+            ValueError,
+            match="spaces do not match",
     ):
         preference_comparisons.AgentTrainer(agent, bad_reward_net, venv)
 
 
 def test_trajectory_dataset_seeding(
-    cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
-    num_samples: int = 400,
+        cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+        num_samples: int = 400,
 ):
     dataset1 = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
@@ -121,8 +123,8 @@ def test_trajectory_dataset_seeding(
 # CartPole max episode length is 200
 @pytest.mark.parametrize("num_steps", [0, 199, 200, 201, 400])
 def test_trajectory_dataset_len(
-    cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
-    num_steps: int,
+        cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+        num_steps: int,
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
@@ -136,7 +138,7 @@ def test_trajectory_dataset_len(
 
 
 def test_trajectory_dataset_too_long(
-    cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+        cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
@@ -147,8 +149,8 @@ def test_trajectory_dataset_too_long(
 
 
 def test_trajectory_dataset_shuffle(
-    cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
-    num_steps: int = 400,
+        cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+        num_steps: int = 400,
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
@@ -165,25 +167,25 @@ def test_transitions_left_in_buffer(agent_trainer):
     # with transitions.
     agent_trainer.buffering_wrapper.n_transitions = 2
     with pytest.raises(
-        RuntimeError,
-        match=re.escape(
-            "There are 2 transitions left in the buffer. "
-            "Call AgentTrainer.sample() first to clear them.",
-        ),
+            RuntimeError,
+            match=re.escape(
+                "There are 2 transitions left in the buffer. "
+                "Call AgentTrainer.sample() first to clear them.",
+            ),
     ):
         agent_trainer.train(steps=1)
 
 
 @pytest.mark.parametrize(
     "schedule",
-    ["constant", "hyperbolic", "inverse_quadratic", lambda t: 1 / (1 + t**3)],
+    ["constant", "hyperbolic", "inverse_quadratic", lambda t: 1 / (1 + t ** 3)],
 )
 def test_trainer_no_crash(
-    agent_trainer,
-    reward_net,
-    random_fragmenter,
-    custom_logger,
-    schedule,
+        agent_trainer,
+        reward_net,
+        random_fragmenter,
+        custom_logger,
+        schedule,
 ):
     main_trainer = preference_comparisons.PreferenceComparisons(
         agent_trainer,
@@ -213,8 +215,8 @@ def test_reward_ensemble_trainer_raises_type_error(venv):
     loss = preference_comparisons.CrossEntropyRewardLoss(preference_model)
 
     with pytest.raises(
-        TypeError,
-        match=r"RewardEnsemble expected by EnsembleTrainer not .*",
+            TypeError,
+            match=r"RewardEnsemble expected by EnsembleTrainer not .*",
     ):
         preference_comparisons.EnsembleTrainer(
             reward_net,
@@ -223,10 +225,10 @@ def test_reward_ensemble_trainer_raises_type_error(venv):
 
 
 def test_correct_reward_trainer_used_by_default(
-    agent_trainer,
-    reward_net,
-    random_fragmenter,
-    custom_logger,
+        agent_trainer,
+        reward_net,
+        random_fragmenter,
+        custom_logger,
 ):
     main_trainer = preference_comparisons.PreferenceComparisons(
         agent_trainer,
@@ -252,10 +254,10 @@ def test_correct_reward_trainer_used_by_default(
 
 
 def test_init_raises_error_when_trying_use_improperly_wrapped_ensemble(
-    agent_trainer,
-    venv,
-    random_fragmenter,
-    custom_logger,
+        agent_trainer,
+        venv,
+        random_fragmenter,
+        custom_logger,
 ):
     reward_net = testing_reward_nets.make_ensemble(
         venv.observation_space,
@@ -267,8 +269,8 @@ def test_init_raises_error_when_trying_use_improperly_wrapped_ensemble(
         r"AddSTDRewardWrapper but found NormalizedRewardNet."
     )
     with pytest.raises(
-        ValueError,
-        match=rgx,
+            ValueError,
+            match=rgx,
     ):
         preference_comparisons.PreferenceComparisons(
             agent_trainer,
@@ -348,8 +350,8 @@ def test_fragments_too_short_error(agent_trainer):
         warning_threshold=0,
     )
     with pytest.raises(
-        ValueError,
-        match="No trajectories are long enough for the desired fragment length.",
+            ValueError,
+            match="No trajectories are long enough for the desired fragment length.",
     ):
         # the only important bit is that fragment_length is higher than
         # we'll ever reach
@@ -408,11 +410,11 @@ def test_store_and_load_preference_dataset(agent_trainer, random_fragmenter, tmp
 
 
 def test_exploration_no_crash(
-    agent,
-    reward_net,
-    venv,
-    random_fragmenter,
-    custom_logger,
+        agent,
+        reward_net,
+        venv,
+        random_fragmenter,
+        custom_logger,
 ):
     agent_trainer = preference_comparisons.AgentTrainer(
         agent,
@@ -434,11 +436,11 @@ def test_exploration_no_crash(
 
 @pytest.mark.parametrize("uncertainty_on", UNCERTAINTY_ON)
 def test_active_fragmenter_discount_rate_no_crash(
-    agent_trainer,
-    venv,
-    random_fragmenter,
-    uncertainty_on,
-    custom_logger,
+        agent_trainer,
+        venv,
+        random_fragmenter,
+        uncertainty_on,
+        custom_logger,
 ):
     # also use a non-zero noise probability to check that doesn't cause errors
     reward_net = reward_nets.RewardEnsemble(
@@ -491,7 +493,7 @@ def test_active_fragmenter_discount_rate_no_crash(
 
 
 @pytest.fixture
-def ensemble_preference_model(venv) -> preference_comparisons.PreferenceComparisons:
+def ensemble_preference_model(venv) -> preference_comparisons.PreferenceModel:
     reward_net = reward_nets.RewardEnsemble(
         venv.observation_space,
         venv.action_space,
@@ -509,7 +511,7 @@ def ensemble_preference_model(venv) -> preference_comparisons.PreferenceComparis
 
 
 @pytest.fixture
-def preference_model(venv) -> preference_comparisons.PreferenceComparisons:
+def preference_model(venv) -> preference_comparisons.PreferenceModel:
     reward_net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     return preference_comparisons.PreferenceModel(
         model=reward_net,
@@ -520,19 +522,19 @@ def preference_model(venv) -> preference_comparisons.PreferenceComparisons:
 
 
 def test_probability_model_raises_error_when_ensemble_member_index_not_provided(
-    ensemble_preference_model,
+        ensemble_preference_model,
 ):
     assert ensemble_preference_model.is_ensemble
     with pytest.raises(
-        ValueError,
-        match="`ensemble_member_index` required for ensemble models",
+            ValueError,
+            match="`ensemble_member_index` required for ensemble models",
     ):
         ensemble_preference_model([])
 
 
 def test_active_fragmenter_uncertainty_on_not_supported_error(
-    ensemble_preference_model,
-    random_fragmenter,
+        ensemble_preference_model,
+        random_fragmenter,
 ):
     re_match = r".* not supported\.\n\s+`uncertainty_on` should be from .*"
     with pytest.raises(ValueError, match=re_match):
@@ -556,12 +558,12 @@ def test_active_fragmenter_uncertainty_on_not_supported_error(
 
 
 def test_active_selection_raises_error_when_initialized_without_an_ensemble(
-    preference_model,
-    random_fragmenter,
+        preference_model,
+        random_fragmenter,
 ):
     with pytest.raises(
-        ValueError,
-        match=r"Preference model not wrapped over an ensemble.*",
+            ValueError,
+            match=r"Preference model not wrapped over an ensemble.*",
     ):
         preference_comparisons.ActiveSelectionFragmenter(
             preference_model=preference_model,

From 50fde43a79bf034757ddeb1c6bb60e9e6d26b232 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 19:06:04 +0100
Subject: [PATCH 014/187] Fix types: tests/data

---
 tests/data/test_rollout.py  |  15 ++--
 tests/data/test_types.py    | 136 ++++++++++++++++++------------------
 tests/data/test_wrappers.py |  28 ++++----
 3 files changed, 91 insertions(+), 88 deletions(-)

diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index 38fce6ba5..17a603398 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -36,14 +36,15 @@ def step(self, action):
 
 
 def _sample_fixed_length_trajectories(
-    episode_lengths: Sequence[int],
-    min_episodes: int,
-    policy_type: str = "policy",
-    **kwargs,
+        episode_lengths: Sequence[int],
+        min_episodes: int,
+        policy_type: str = "policy",
+        **kwargs,
 ) -> Sequence[types.Trajectory]:
     venv = vec_env.DummyVecEnv(
         [functools.partial(TerminalSentinelEnv, length) for length in episode_lengths],
     )
+    policy: rollout.AnyPolicy
     if policy_type == "policy":
         policy = RandomPolicy(venv.observation_space, venv.action_space)
     elif policy_type == "callable":
@@ -114,9 +115,9 @@ def test_complete_trajectories(policy_type) -> None:
     ],
 )
 def test_unbiased_trajectories(
-    episode_lengths: Sequence[int],
-    min_episodes: int,
-    expected_counts: Mapping[int, int],
+        episode_lengths: Sequence[int],
+        min_episodes: int,
+        expected_counts: Mapping[int, int],
 ) -> None:
     """Checks trajectories are sampled without bias towards shorter episodes.
 
diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index af8398ba3..b9f8fbe2c 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -1,13 +1,12 @@
 """Tests of `imitation.data.types`."""
 
-
 import contextlib
 import copy
 import dataclasses
 import os
 import pathlib
 import pickle
-from typing import Any, Callable
+from typing import Any, Callable, Sequence
 
 import gym
 import numpy as np
@@ -27,7 +26,8 @@
 LENGTHS = [0, 1, 2, 10]
 
 
-def _check_1d_shape(fn: Callable[[np.ndarray], Any], length: float, expected_msg: str):
+def _check_1d_shape(fn: Callable[[np.ndarray], Any], length: int, expected_msg: str):
+    assert isinstance(length, int)
     for shape in [(), (length, 1), (length, 2), (length - 1,), (length + 1,)]:
         with pytest.raises(ValueError, match=expected_msg):
             fn(np.zeros(shape))
@@ -35,9 +35,9 @@ def _check_1d_shape(fn: Callable[[np.ndarray], Any], length: float, expected_msg
 
 @pytest.fixture
 def trajectory(
-    obs_space: gym.Space,
-    act_space: gym.Space,
-    length: int,
+        obs_space: gym.Space,
+        act_space: gym.Space,
+        length: int,
 ) -> types.Trajectory:
     """Fixture to generate trajectory of length `length` iid sampled from spaces."""
     if length == 0:
@@ -57,9 +57,9 @@ def trajectory_rew(trajectory: types.Trajectory) -> types.TrajectoryWithRew:
 
 @pytest.fixture
 def transitions_min(
-    obs_space: gym.Space,
-    act_space: gym.Space,
-    length: int,
+        obs_space: gym.Space,
+        act_space: gym.Space,
+        length: int,
 ) -> types.TransitionsMinimal:
     obs = np.array([obs_space.sample() for _ in range(length)])
     acts = np.array([act_space.sample() for _ in range(length)])
@@ -69,9 +69,9 @@ def transitions_min(
 
 @pytest.fixture
 def transitions(
-    transitions_min: types.TransitionsMinimal,
-    obs_space: gym.Space,
-    length: int,
+        transitions_min: types.TransitionsMinimal,
+        obs_space: gym.Space,
+        length: int,
 ) -> types.Transitions:
     """Fixture to generate transitions of length `length` iid sampled from spaces."""
     next_obs = np.array([obs_space.sample() for _ in range(length)])
@@ -85,8 +85,8 @@ def transitions(
 
 @pytest.fixture
 def transitions_rew(
-    transitions: types.Transitions,
-    length: int,
+        transitions: types.Transitions,
+        length: int,
 ) -> types.TransitionsWithRew:
     """Like `transitions` but with reward randomly sampled from a Gaussian."""
     rews = np.random.randn(length)
@@ -129,10 +129,10 @@ class TestData:
     """
 
     def test_valid_trajectories(
-        self,
-        trajectory: types.Trajectory,
-        trajectory_rew: types.TrajectoryWithRew,
-        length: int,
+            self,
+            trajectory: types.Trajectory,
+            trajectory_rew: types.TrajectoryWithRew,
+            length: int,
     ) -> None:
         """Checks trajectories can be created for a variety of lengths and spaces."""
         trajs = [trajectory, trajectory_rew]
@@ -141,9 +141,9 @@ def test_valid_trajectories(
             assert len(traj) == length
 
     def test_traj_unequal_to_other_types(
-        self,
-        trajectory: types.Trajectory,
-        trajectory_rew: types.TrajectoryWithRew,
+            self,
+            trajectory: types.Trajectory,
+            trajectory_rew: types.TrajectoryWithRew,
     ) -> None:
         """Test trajectories unequal to objects of different types."""
         for t in [trajectory, trajectory_rew]:
@@ -155,9 +155,9 @@ def test_traj_unequal_to_other_types(
         assert trajectory != trajectory_rew
 
     def test_traj_equal_to_self_and_copies(
-        self,
-        trajectory: types.Trajectory,
-        trajectory_rew: types.TrajectoryWithRew,
+            self,
+            trajectory: types.Trajectory,
+            trajectory_rew: types.TrajectoryWithRew,
     ) -> None:
         """Test that trajectories are equal to themselves and copies."""
         for t in [trajectory, trajectory_rew]:
@@ -167,10 +167,10 @@ def test_traj_equal_to_self_and_copies(
             assert t == copy.copy(t)
 
     def test_traj_unequal_to_perturbations(
-        self,
-        trajectory: types.Trajectory,
-        trajectory_rew: types.TrajectoryWithRew,
-        length: int,
+            self,
+            trajectory: types.Trajectory,
+            trajectory_rew: types.TrajectoryWithRew,
+            length: int,
     ) -> None:
         """Test that trajectories unequal to perturbed versions."""
         # Unequal to a copy of itself truncated
@@ -199,15 +199,16 @@ def test_traj_unequal_to_perturbations(
     @pytest.mark.parametrize("use_rewards", [False, True])
     @pytest.mark.parametrize("use_chdir", [False, True])
     def test_save_trajectories(
-        self,
-        trajectory: types.Trajectory,
-        trajectory_rew: types.TrajectoryWithRew,
-        use_chdir,
-        tmpdir,
-        use_pickle,
-        use_rewards,
-        type_safe,
+            self,
+            trajectory: types.Trajectory,
+            trajectory_rew: types.TrajectoryWithRew,
+            use_chdir,
+            tmpdir,
+            use_pickle,
+            use_rewards,
+            type_safe,
     ):
+        chdir_context: contextlib.AbstractContextManager
         """Check that trajectories are properly saved."""
         if use_chdir:
             # Test no relative path without directory edge-case.
@@ -234,12 +235,13 @@ def test_save_trajectories(
                     with pytest.raises(ValueError):
                         types.save(save_path, [trajectory, trajectory_rew])
 
+            loaded_trajs: Sequence[types.Trajectory]
             if type_safe:
                 if use_rewards:
                     loaded_trajs = types.load_with_rewards(save_path)
                 else:
                     with pytest.raises(ValueError):
-                        loaded_trajs = types.load_with_rewards(save_path)
+                        types.load_with_rewards(save_path)
                     loaded_trajs = types.load(save_path)
             else:
                 loaded_trajs = types.load(save_path)
@@ -249,32 +251,32 @@ def test_save_trajectories(
                 assert t1 == t2
 
     def test_invalid_trajectories(
-        self,
-        trajectory: types.Trajectory,
-        trajectory_rew: types.TrajectoryWithRew,
+            self,
+            trajectory: types.Trajectory,
+            trajectory_rew: types.TrajectoryWithRew,
     ) -> None:
         """Checks input validation catches space and dtype related errors."""
         trajs = [trajectory, trajectory_rew]
         for traj in trajs:
             with pytest.raises(
-                ValueError,
-                match=r"expected one more observations than actions.*",
+                    ValueError,
+                    match=r"expected one more observations than actions.*",
             ):
                 dataclasses.replace(traj, obs=traj.obs[:-1])
             with pytest.raises(
-                ValueError,
-                match=r"expected one more observations than actions.*",
+                    ValueError,
+                    match=r"expected one more observations than actions.*",
             ):
                 dataclasses.replace(traj, acts=traj.acts[:-1])
 
             with pytest.raises(
-                ValueError,
-                match=r"infos when present must be present for each action.*",
+                    ValueError,
+                    match=r"infos when present must be present for each action.*",
             ):
                 dataclasses.replace(traj, infos=traj.infos[:-1])
             with pytest.raises(
-                ValueError,
-                match=r"infos when present must be present for each action.*",
+                    ValueError,
+                    match=r"infos when present must be present for each action.*",
             ):
                 dataclasses.replace(traj, obs=traj.obs[:-1], acts=traj.acts[:-1])
 
@@ -291,12 +293,12 @@ def test_invalid_trajectories(
             )
 
     def test_valid_transitions(
-        self,
-        transitions_min: types.TransitionsMinimal,
-        transitions: types.Transitions,
-        transitions_rew: types.TransitionsWithRew,
-        length: int,
-        n_checks: int = 20,
+            self,
+            transitions_min: types.TransitionsMinimal,
+            transitions: types.Transitions,
+            transitions_rew: types.TransitionsWithRew,
+            length: int,
+            n_checks: int = 20,
     ) -> None:
         """Checks initialization, indexing, and slicing sanity."""
         for trans in [transitions_min, transitions, transitions_rew]:
@@ -320,11 +322,11 @@ def test_valid_transitions(
                 _check_transitions_get_item(trans, s)
 
     def test_invalid_transitions(
-        self,
-        transitions_min: types.Transitions,
-        transitions: types.Transitions,
-        transitions_rew: types.TransitionsWithRew,
-        length: int,
+            self,
+            transitions_min: types.Transitions,
+            transitions: types.Transitions,
+            transitions_rew: types.TransitionsWithRew,
+            length: int,
     ) -> None:
         """Checks input validation catches space and dtype related errors."""
         if length == 0:
@@ -332,26 +334,26 @@ def test_invalid_transitions(
 
         for trans in [transitions_min, transitions, transitions_rew]:
             with pytest.raises(
-                ValueError,
-                match=r"obs and acts must have same number of timesteps:.*",
+                    ValueError,
+                    match=r"obs and acts must have same number of timesteps:.*",
             ):
                 dataclasses.replace(trans, acts=trans.acts[:-1])
             with pytest.raises(
-                ValueError,
-                match=r"obs and infos must have same number of timesteps:.*",
+                    ValueError,
+                    match=r"obs and infos must have same number of timesteps:.*",
             ):
                 dataclasses.replace(trans, infos=[{}] * (length - 1))
 
         for trans in [transitions, transitions_rew]:
             with pytest.raises(
-                ValueError,
-                match=r"obs and next_obs must have same shape:.*",
+                    ValueError,
+                    match=r"obs and next_obs must have same shape:.*",
             ):
                 dataclasses.replace(trans, next_obs=np.zeros((len(trans), 4, 2)))
 
             with pytest.raises(
-                ValueError,
-                match=r"obs and next_obs must have the same dtype:.*",
+                    ValueError,
+                    match=r"obs and next_obs must have the same dtype:.*",
             ):
                 dataclasses.replace(
                     trans,
diff --git a/tests/data/test_wrappers.py b/tests/data/test_wrappers.py
index 364bebb33..e1dca8ed6 100644
--- a/tests/data/test_wrappers.py
+++ b/tests/data/test_wrappers.py
@@ -49,23 +49,23 @@ def step(self, action):
 
 
 def _make_buffering_venv(
-    error_on_premature_reset: bool,
+        error_on_premature_reset: bool,
 ) -> BufferingWrapper:
     venv = DummyVecEnv([_CountingEnv] * 2)
-    venv = BufferingWrapper(venv, error_on_premature_reset)
-    venv.reset()
-    return venv
+    wrapped_venv = BufferingWrapper(venv, error_on_premature_reset)
+    wrapped_venv.reset()
+    return wrapped_venv
 
 
-def _assert_equal_scrambled_vectors(a: np.ndarray, b: np.ndarray) -> bool:
-    """Returns True if `a` and `b` are identical up to sorting."""
+def _assert_equal_scrambled_vectors(a: np.ndarray, b: np.ndarray) -> None:
+    """Raises AssertionError if `a` and `b` are not identical up to sorting."""
     assert a.shape == b.shape
     assert a.ndim == 1
     np.testing.assert_allclose(np.sort(a), np.sort(b))
 
 
 def _join_transitions(
-    trans_list: Sequence[types.TransitionsWithRew],
+        trans_list: Sequence[types.TransitionsWithRew],
 ) -> types.TransitionsWithRew:
     def concat(x):
         return np.concatenate(list(x))
@@ -90,9 +90,9 @@ def concat(x):
 @pytest.mark.parametrize("n_steps", [1, 2, 20, 21])
 @pytest.mark.parametrize("extra_pop_timesteps", [(), (1,), (4, 8)])
 def test_pop(
-    episode_lengths: Sequence[int],
-    n_steps: int,
-    extra_pop_timesteps: Sequence[int],
+        episode_lengths: Sequence[int],
+        n_steps: int,
+        extra_pop_timesteps: Sequence[int],
 ) -> None:
     """Check pop_transitions() results for BufferWrapper.
 
@@ -166,13 +166,13 @@ def make_env(ep_len):
     transitions_list.append(venv_buffer.pop_transitions())
 
     # Build expected transitions
-    expect_obs = []
+    expect_obs_list = []
     for ep_len in episode_lengths:
         n_complete, remainder = divmod(n_steps, ep_len)
-        expect_obs.extend([np.arange(ep_len)] * n_complete)
-        expect_obs.append(np.arange(remainder))
+        expect_obs_list.extend([np.arange(ep_len)] * n_complete)
+        expect_obs_list.append(np.arange(remainder))
 
-    expect_obs = np.concatenate(expect_obs)
+    expect_obs = np.concatenate(expect_obs_list)
     expect_next_obs = expect_obs + 1
     expect_acts = expect_obs * 2.1
     expect_rews = expect_next_obs * 10

From 97bc6aa0a69abc794e9b12abed1a5974b908b15e Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 31 Aug 2022 19:17:06 +0100
Subject: [PATCH 015/187] Linting

---
 tests/algorithms/test_adversarial.py          |  84 ++++++------
 tests/algorithms/test_bc.py                   |  18 +--
 tests/algorithms/test_dagger.py               |  82 ++++++------
 .../algorithms/test_preference_comparisons.py | 110 +++++++--------
 tests/data/test_rollout.py                    |  14 +-
 tests/data/test_types.py                      | 126 +++++++++---------
 tests/data/test_wrappers.py                   |  10 +-
 7 files changed, 222 insertions(+), 222 deletions(-)

diff --git a/tests/algorithms/test_adversarial.py b/tests/algorithms/test_adversarial.py
index 806e773c8..d8a5712f2 100644
--- a/tests/algorithms/test_adversarial.py
+++ b/tests/algorithms/test_adversarial.py
@@ -2,7 +2,7 @@
 
 import contextlib
 import os
-from typing import Any, Mapping, Union, Type
+from typing import Any, Mapping, Type, Union
 
 import numpy as np
 import pytest
@@ -55,14 +55,14 @@ def expert_transitions(cartpole_expert_trajectories):
 
 @contextlib.contextmanager
 def make_trainer(
-        algorithm_kwargs: Mapping[str, Any],
-        tmpdir: str,
-        expert_transitions: types.Transitions,
-        expert_batch_size: int = 1,
-        env_name: str = "seals/CartPole-v0",
-        num_envs: int = 1,
-        parallel: bool = False,
-        convert_dataset: bool = False,
+    algorithm_kwargs: Mapping[str, Any],
+    tmpdir: str,
+    expert_transitions: types.Transitions,
+    expert_batch_size: int = 1,
+    env_name: str = "seals/CartPole-v0",
+    num_envs: int = 1,
+    parallel: bool = False,
+    convert_dataset: bool = False,
 ):
     expert_data: Union[th_data.DataLoader, th_data.Dataset]
     if convert_dataset:
@@ -139,8 +139,8 @@ def test_train_disc_no_samples_error(trainer: common.AdversarialTrainer):
 
 
 def test_train_disc_unequal_expert_gen_samples_error(
-        trainer: common.AdversarialTrainer,
-        expert_transitions: types.Transitions,
+    trainer: common.AdversarialTrainer,
+    expert_transitions: types.Transitions,
 ):
     """Test that train_disc raises error when n_gen != n_expert samples."""
     if len(expert_transitions) < 2:  # pragma: no cover
@@ -171,20 +171,20 @@ def _expert_batch_size(request):
 
 @pytest.fixture
 def trainer_parametrized(
+    _algorithm_kwargs,
+    _parallel,
+    _convert_dataset,
+    _expert_batch_size,
+    tmpdir,
+    expert_transitions,
+):
+    with make_trainer(
         _algorithm_kwargs,
-        _parallel,
-        _convert_dataset,
-        _expert_batch_size,
         tmpdir,
         expert_transitions,
-):
-    with make_trainer(
-            _algorithm_kwargs,
-            tmpdir,
-            expert_transitions,
-            parallel=_parallel,
-            convert_dataset=_convert_dataset,
-            expert_batch_size=_expert_batch_size,
+        parallel=_parallel,
+        convert_dataset=_convert_dataset,
+        expert_batch_size=_expert_batch_size,
     ) as trainer:
         yield trainer
 
@@ -202,8 +202,8 @@ def test_train_disc_step_no_crash(trainer_parametrized, _expert_batch_size):
 
 
 def test_train_gen_train_disc_no_crash(
-        trainer_parametrized: common.AdversarialTrainer,
-        n_updates: int = 2,
+    trainer_parametrized: common.AdversarialTrainer,
+    n_updates: int = 2,
 ) -> None:
     trainer_parametrized.train_gen(n_updates * trainer_parametrized.gen_train_timesteps)
     trainer_parametrized.train_disc()
@@ -211,26 +211,26 @@ def test_train_gen_train_disc_no_crash(
 
 @pytest.fixture
 def trainer_batch_sizes(
+    _algorithm_kwargs,
+    _expert_batch_size,
+    tmpdir,
+    expert_transitions,
+):
+    with make_trainer(
         _algorithm_kwargs,
-        _expert_batch_size,
         tmpdir,
         expert_transitions,
-):
-    with make_trainer(
-            _algorithm_kwargs,
-            tmpdir,
-            expert_transitions,
-            expert_batch_size=_expert_batch_size,
+        expert_batch_size=_expert_batch_size,
     ) as trainer:
         yield trainer
 
 
 def test_train_disc_improve_D(
-        trainer_batch_sizes,
-        tmpdir,
-        expert_transitions,
-        _expert_batch_size,
-        n_steps=3,
+    trainer_batch_sizes,
+    tmpdir,
+    expert_transitions,
+    _expert_batch_size,
+    n_steps=3,
 ):
     expert_samples = expert_transitions[:_expert_batch_size]
     expert_samples = types.dataclass_quick_asdict(expert_samples)
@@ -263,18 +263,18 @@ def trainer_diverse_env(_algorithm_kwargs, _env_name, tmpdir, expert_transitions
     if _algorithm_kwargs["model_class"] == stable_baselines3.DQN:
         pytest.skip("DQN does not support all environments.")
     with make_trainer(
-            _algorithm_kwargs,
-            tmpdir,
-            expert_transitions,
-            env_name=_env_name,
+        _algorithm_kwargs,
+        tmpdir,
+        expert_transitions,
+        env_name=_env_name,
     ) as trainer:
         yield trainer
 
 
 @pytest.mark.parametrize("n_timesteps", [2, 4, 10])
 def test_logits_expert_is_high_log_policy_act_prob(
-        trainer_diverse_env: common.AdversarialTrainer,
-        n_timesteps: int,
+    trainer_diverse_env: common.AdversarialTrainer,
+    n_timesteps: int,
 ):
     """Smoke test calling `logits_expert_is_high` on `AdversarialTrainer`.
 
diff --git a/tests/algorithms/test_bc.py b/tests/algorithms/test_bc.py
index 3e1405ff6..200f2101b 100644
--- a/tests/algorithms/test_bc.py
+++ b/tests/algorithms/test_bc.py
@@ -42,11 +42,11 @@ def __iter__(self):
 
 @pytest.fixture
 def trainer(
-        batch_size,
-        cartpole_venv,
-        expert_data_type,
-        custom_logger,
-        cartpole_expert_trajectories,
+    batch_size,
+    cartpole_venv,
+    expert_data_type,
+    custom_logger,
+    cartpole_expert_trajectories,
 ):
     trans = rollout.flatten_trajectories(cartpole_expert_trajectories)
     if expert_data_type == "data_loader":
@@ -161,10 +161,10 @@ def __iter__(self):
 
 @pytest.mark.parametrize("no_yield_after_iter", [0, 1, 5])
 def test_bc_data_loader_empty_iter_error(
-        cartpole_venv: vec_env.VecEnv,
-        no_yield_after_iter: bool,
-        custom_logger: logger.HierarchicalLogger,
-        cartpole_expert_trajectories,
+    cartpole_venv: vec_env.VecEnv,
+    no_yield_after_iter: bool,
+    custom_logger: logger.HierarchicalLogger,
+    cartpole_expert_trajectories,
 ) -> None:
     """Check that we error out if the DataLoader suddenly stops yielding any batches.
 
diff --git a/tests/algorithms/test_dagger.py b/tests/algorithms/test_dagger.py
index dd255565c..f31d2edd3 100644
--- a/tests/algorithms/test_dagger.py
+++ b/tests/algorithms/test_dagger.py
@@ -23,8 +23,8 @@
 
 @pytest.fixture(params=[True, False])
 def maybe_pendulum_expert_trajectories(
-        pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
-        request,
+    pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
+    request,
 ) -> Optional[Sequence[TrajectoryWithRew]]:
     keep_trajs = request.param
     if keep_trajs:
@@ -104,12 +104,12 @@ def get_random_acts(obs):
 
 
 def _build_dagger_trainer(
-        tmpdir,
-        venv,
-        beta_schedule,
-        expert_policy,
-        pendulum_expert_rollouts: List[TrajectoryWithRew],
-        custom_logger,
+    tmpdir,
+    venv,
+    beta_schedule,
+    expert_policy,
+    pendulum_expert_rollouts: List[TrajectoryWithRew],
+    custom_logger,
 ):
     del expert_policy
     if pendulum_expert_rollouts is not None:
@@ -133,12 +133,12 @@ def _build_dagger_trainer(
 
 
 def _build_simple_dagger_trainer(
-        tmpdir,
-        venv,
-        beta_schedule,
-        expert_policy,
-        pendulum_expert_rollouts: Optional[List[TrajectoryWithRew]],
-        custom_logger,
+    tmpdir,
+    venv,
+    beta_schedule,
+    expert_policy,
+    pendulum_expert_rollouts: Optional[List[TrajectoryWithRew]],
+    custom_logger,
 ):
     bc_trainer = bc.BC(
         observation_space=venv.observation_space,
@@ -164,13 +164,13 @@ def beta_schedule(request):
 
 @pytest.fixture(params=[_build_dagger_trainer, _build_simple_dagger_trainer])
 def init_trainer_fn(
-        request,
-        tmpdir,
-        pendulum_venv,
-        beta_schedule,
-        pendulum_expert_policy,
-        maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
-        custom_logger,
+    request,
+    tmpdir,
+    pendulum_venv,
+    beta_schedule,
+    pendulum_expert_policy,
+    maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
+    custom_logger,
 ):
     # Provide a trainer initialization fixture in addition `trainer` fixture below
     # for tests that want to initialize multiple DAggerTrainer.
@@ -192,12 +192,12 @@ def trainer(init_trainer_fn):
 
 @pytest.fixture
 def simple_dagger_trainer(
-        tmpdir,
-        pendulum_venv,
-        beta_schedule,
-        pendulum_expert_policy,
-        maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
-        custom_logger,
+    tmpdir,
+    pendulum_venv,
+    beta_schedule,
+    pendulum_expert_policy,
+    maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
+    custom_logger,
 ):
     return _build_simple_dagger_trainer(
         tmpdir,
@@ -210,15 +210,15 @@ def simple_dagger_trainer(
 
 
 def test_trainer_needs_demos_exception_error(
-        trainer,
-        maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
+    trainer,
+    maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
 ):
     assert trainer.round_num == 0
     error_ctx = pytest.raises(dagger.NeedsDemosException)
     ctx: contextlib.AbstractContextManager
     if maybe_pendulum_expert_trajectories is not None and isinstance(
-            trainer,
-            dagger.SimpleDAggerTrainer,
+        trainer,
+        dagger.SimpleDAggerTrainer,
     ):
         # In this case, demos should be preloaded and we shouldn't experience
         # the NeedsDemoException error.
@@ -333,10 +333,10 @@ def test_trainer_save_reload(tmpdir, init_trainer_fn, pendulum_venv):
 
 @pytest.mark.parametrize("num_episodes", [1, 4])
 def test_simple_dagger_trainer_train(
-        simple_dagger_trainer: dagger.SimpleDAggerTrainer,
-        pendulum_venv,
-        num_episodes: int,
-        tmpdir: str,
+    simple_dagger_trainer: dagger.SimpleDAggerTrainer,
+    pendulum_venv,
+    num_episodes: int,
+    tmpdir: str,
 ):
     episode_length = 200  # for Pendulum-v1
     rollout_min_episodes = 2
@@ -366,12 +366,12 @@ def test_policy_save_reload(tmpdir, trainer):
 
 
 def test_simple_dagger_space_mismatch_error(
-        tmpdir,
-        pendulum_venv,
-        beta_schedule,
-        pendulum_expert_policy,
-        maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
-        custom_logger,
+    tmpdir,
+    pendulum_venv,
+    beta_schedule,
+    pendulum_expert_policy,
+    maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
+    custom_logger,
 ):
     class MismatchedSpace(gym.spaces.Space):
         """Dummy space that is not equal to any other space."""
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 3c7bb99cc..ca85ce114 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -64,8 +64,8 @@ def agent_trainer(agent, reward_net, venv):
 
 
 def _check_trajs_equal(
-        trajs1: Sequence[types.TrajectoryWithRew],
-        trajs2: Sequence[types.TrajectoryWithRew],
+    trajs1: Sequence[types.TrajectoryWithRew],
+    trajs2: Sequence[types.TrajectoryWithRew],
 ):
     assert len(trajs1) == len(trajs2)
     for traj1, traj2 in zip(trajs1, trajs2):
@@ -88,15 +88,15 @@ def test_mismatched_spaces(venv, agent):
         other_venv.action_space,
     )
     with pytest.raises(
-            ValueError,
-            match="spaces do not match",
+        ValueError,
+        match="spaces do not match",
     ):
         preference_comparisons.AgentTrainer(agent, bad_reward_net, venv)
 
 
 def test_trajectory_dataset_seeding(
-        cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
-        num_samples: int = 400,
+    cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+    num_samples: int = 400,
 ):
     dataset1 = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
@@ -123,8 +123,8 @@ def test_trajectory_dataset_seeding(
 # CartPole max episode length is 200
 @pytest.mark.parametrize("num_steps", [0, 199, 200, 201, 400])
 def test_trajectory_dataset_len(
-        cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
-        num_steps: int,
+    cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+    num_steps: int,
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
@@ -138,7 +138,7 @@ def test_trajectory_dataset_len(
 
 
 def test_trajectory_dataset_too_long(
-        cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+    cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
@@ -149,8 +149,8 @@ def test_trajectory_dataset_too_long(
 
 
 def test_trajectory_dataset_shuffle(
-        cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
-        num_steps: int = 400,
+    cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+    num_steps: int = 400,
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
@@ -167,25 +167,25 @@ def test_transitions_left_in_buffer(agent_trainer):
     # with transitions.
     agent_trainer.buffering_wrapper.n_transitions = 2
     with pytest.raises(
-            RuntimeError,
-            match=re.escape(
-                "There are 2 transitions left in the buffer. "
-                "Call AgentTrainer.sample() first to clear them.",
-            ),
+        RuntimeError,
+        match=re.escape(
+            "There are 2 transitions left in the buffer. "
+            "Call AgentTrainer.sample() first to clear them.",
+        ),
     ):
         agent_trainer.train(steps=1)
 
 
 @pytest.mark.parametrize(
     "schedule",
-    ["constant", "hyperbolic", "inverse_quadratic", lambda t: 1 / (1 + t ** 3)],
+    ["constant", "hyperbolic", "inverse_quadratic", lambda t: 1 / (1 + t**3)],
 )
 def test_trainer_no_crash(
-        agent_trainer,
-        reward_net,
-        random_fragmenter,
-        custom_logger,
-        schedule,
+    agent_trainer,
+    reward_net,
+    random_fragmenter,
+    custom_logger,
+    schedule,
 ):
     main_trainer = preference_comparisons.PreferenceComparisons(
         agent_trainer,
@@ -215,8 +215,8 @@ def test_reward_ensemble_trainer_raises_type_error(venv):
     loss = preference_comparisons.CrossEntropyRewardLoss(preference_model)
 
     with pytest.raises(
-            TypeError,
-            match=r"RewardEnsemble expected by EnsembleTrainer not .*",
+        TypeError,
+        match=r"RewardEnsemble expected by EnsembleTrainer not .*",
     ):
         preference_comparisons.EnsembleTrainer(
             reward_net,
@@ -225,10 +225,10 @@ def test_reward_ensemble_trainer_raises_type_error(venv):
 
 
 def test_correct_reward_trainer_used_by_default(
-        agent_trainer,
-        reward_net,
-        random_fragmenter,
-        custom_logger,
+    agent_trainer,
+    reward_net,
+    random_fragmenter,
+    custom_logger,
 ):
     main_trainer = preference_comparisons.PreferenceComparisons(
         agent_trainer,
@@ -254,10 +254,10 @@ def test_correct_reward_trainer_used_by_default(
 
 
 def test_init_raises_error_when_trying_use_improperly_wrapped_ensemble(
-        agent_trainer,
-        venv,
-        random_fragmenter,
-        custom_logger,
+    agent_trainer,
+    venv,
+    random_fragmenter,
+    custom_logger,
 ):
     reward_net = testing_reward_nets.make_ensemble(
         venv.observation_space,
@@ -269,8 +269,8 @@ def test_init_raises_error_when_trying_use_improperly_wrapped_ensemble(
         r"AddSTDRewardWrapper but found NormalizedRewardNet."
     )
     with pytest.raises(
-            ValueError,
-            match=rgx,
+        ValueError,
+        match=rgx,
     ):
         preference_comparisons.PreferenceComparisons(
             agent_trainer,
@@ -350,8 +350,8 @@ def test_fragments_too_short_error(agent_trainer):
         warning_threshold=0,
     )
     with pytest.raises(
-            ValueError,
-            match="No trajectories are long enough for the desired fragment length.",
+        ValueError,
+        match="No trajectories are long enough for the desired fragment length.",
     ):
         # the only important bit is that fragment_length is higher than
         # we'll ever reach
@@ -410,11 +410,11 @@ def test_store_and_load_preference_dataset(agent_trainer, random_fragmenter, tmp
 
 
 def test_exploration_no_crash(
-        agent,
-        reward_net,
-        venv,
-        random_fragmenter,
-        custom_logger,
+    agent,
+    reward_net,
+    venv,
+    random_fragmenter,
+    custom_logger,
 ):
     agent_trainer = preference_comparisons.AgentTrainer(
         agent,
@@ -436,11 +436,11 @@ def test_exploration_no_crash(
 
 @pytest.mark.parametrize("uncertainty_on", UNCERTAINTY_ON)
 def test_active_fragmenter_discount_rate_no_crash(
-        agent_trainer,
-        venv,
-        random_fragmenter,
-        uncertainty_on,
-        custom_logger,
+    agent_trainer,
+    venv,
+    random_fragmenter,
+    uncertainty_on,
+    custom_logger,
 ):
     # also use a non-zero noise probability to check that doesn't cause errors
     reward_net = reward_nets.RewardEnsemble(
@@ -522,19 +522,19 @@ def preference_model(venv) -> preference_comparisons.PreferenceModel:
 
 
 def test_probability_model_raises_error_when_ensemble_member_index_not_provided(
-        ensemble_preference_model,
+    ensemble_preference_model,
 ):
     assert ensemble_preference_model.is_ensemble
     with pytest.raises(
-            ValueError,
-            match="`ensemble_member_index` required for ensemble models",
+        ValueError,
+        match="`ensemble_member_index` required for ensemble models",
     ):
         ensemble_preference_model([])
 
 
 def test_active_fragmenter_uncertainty_on_not_supported_error(
-        ensemble_preference_model,
-        random_fragmenter,
+    ensemble_preference_model,
+    random_fragmenter,
 ):
     re_match = r".* not supported\.\n\s+`uncertainty_on` should be from .*"
     with pytest.raises(ValueError, match=re_match):
@@ -558,12 +558,12 @@ def test_active_fragmenter_uncertainty_on_not_supported_error(
 
 
 def test_active_selection_raises_error_when_initialized_without_an_ensemble(
-        preference_model,
-        random_fragmenter,
+    preference_model,
+    random_fragmenter,
 ):
     with pytest.raises(
-            ValueError,
-            match=r"Preference model not wrapped over an ensemble.*",
+        ValueError,
+        match=r"Preference model not wrapped over an ensemble.*",
     ):
         preference_comparisons.ActiveSelectionFragmenter(
             preference_model=preference_model,
diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index 17a603398..4b4f7ee4f 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -36,10 +36,10 @@ def step(self, action):
 
 
 def _sample_fixed_length_trajectories(
-        episode_lengths: Sequence[int],
-        min_episodes: int,
-        policy_type: str = "policy",
-        **kwargs,
+    episode_lengths: Sequence[int],
+    min_episodes: int,
+    policy_type: str = "policy",
+    **kwargs,
 ) -> Sequence[types.Trajectory]:
     venv = vec_env.DummyVecEnv(
         [functools.partial(TerminalSentinelEnv, length) for length in episode_lengths],
@@ -115,9 +115,9 @@ def test_complete_trajectories(policy_type) -> None:
     ],
 )
 def test_unbiased_trajectories(
-        episode_lengths: Sequence[int],
-        min_episodes: int,
-        expected_counts: Mapping[int, int],
+    episode_lengths: Sequence[int],
+    min_episodes: int,
+    expected_counts: Mapping[int, int],
 ) -> None:
     """Checks trajectories are sampled without bias towards shorter episodes.
 
diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index b9f8fbe2c..ba3c4db10 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -35,9 +35,9 @@ def _check_1d_shape(fn: Callable[[np.ndarray], Any], length: int, expected_msg:
 
 @pytest.fixture
 def trajectory(
-        obs_space: gym.Space,
-        act_space: gym.Space,
-        length: int,
+    obs_space: gym.Space,
+    act_space: gym.Space,
+    length: int,
 ) -> types.Trajectory:
     """Fixture to generate trajectory of length `length` iid sampled from spaces."""
     if length == 0:
@@ -57,9 +57,9 @@ def trajectory_rew(trajectory: types.Trajectory) -> types.TrajectoryWithRew:
 
 @pytest.fixture
 def transitions_min(
-        obs_space: gym.Space,
-        act_space: gym.Space,
-        length: int,
+    obs_space: gym.Space,
+    act_space: gym.Space,
+    length: int,
 ) -> types.TransitionsMinimal:
     obs = np.array([obs_space.sample() for _ in range(length)])
     acts = np.array([act_space.sample() for _ in range(length)])
@@ -69,9 +69,9 @@ def transitions_min(
 
 @pytest.fixture
 def transitions(
-        transitions_min: types.TransitionsMinimal,
-        obs_space: gym.Space,
-        length: int,
+    transitions_min: types.TransitionsMinimal,
+    obs_space: gym.Space,
+    length: int,
 ) -> types.Transitions:
     """Fixture to generate transitions of length `length` iid sampled from spaces."""
     next_obs = np.array([obs_space.sample() for _ in range(length)])
@@ -85,8 +85,8 @@ def transitions(
 
 @pytest.fixture
 def transitions_rew(
-        transitions: types.Transitions,
-        length: int,
+    transitions: types.Transitions,
+    length: int,
 ) -> types.TransitionsWithRew:
     """Like `transitions` but with reward randomly sampled from a Gaussian."""
     rews = np.random.randn(length)
@@ -129,10 +129,10 @@ class TestData:
     """
 
     def test_valid_trajectories(
-            self,
-            trajectory: types.Trajectory,
-            trajectory_rew: types.TrajectoryWithRew,
-            length: int,
+        self,
+        trajectory: types.Trajectory,
+        trajectory_rew: types.TrajectoryWithRew,
+        length: int,
     ) -> None:
         """Checks trajectories can be created for a variety of lengths and spaces."""
         trajs = [trajectory, trajectory_rew]
@@ -141,9 +141,9 @@ def test_valid_trajectories(
             assert len(traj) == length
 
     def test_traj_unequal_to_other_types(
-            self,
-            trajectory: types.Trajectory,
-            trajectory_rew: types.TrajectoryWithRew,
+        self,
+        trajectory: types.Trajectory,
+        trajectory_rew: types.TrajectoryWithRew,
     ) -> None:
         """Test trajectories unequal to objects of different types."""
         for t in [trajectory, trajectory_rew]:
@@ -155,9 +155,9 @@ def test_traj_unequal_to_other_types(
         assert trajectory != trajectory_rew
 
     def test_traj_equal_to_self_and_copies(
-            self,
-            trajectory: types.Trajectory,
-            trajectory_rew: types.TrajectoryWithRew,
+        self,
+        trajectory: types.Trajectory,
+        trajectory_rew: types.TrajectoryWithRew,
     ) -> None:
         """Test that trajectories are equal to themselves and copies."""
         for t in [trajectory, trajectory_rew]:
@@ -167,10 +167,10 @@ def test_traj_equal_to_self_and_copies(
             assert t == copy.copy(t)
 
     def test_traj_unequal_to_perturbations(
-            self,
-            trajectory: types.Trajectory,
-            trajectory_rew: types.TrajectoryWithRew,
-            length: int,
+        self,
+        trajectory: types.Trajectory,
+        trajectory_rew: types.TrajectoryWithRew,
+        length: int,
     ) -> None:
         """Test that trajectories unequal to perturbed versions."""
         # Unequal to a copy of itself truncated
@@ -199,14 +199,14 @@ def test_traj_unequal_to_perturbations(
     @pytest.mark.parametrize("use_rewards", [False, True])
     @pytest.mark.parametrize("use_chdir", [False, True])
     def test_save_trajectories(
-            self,
-            trajectory: types.Trajectory,
-            trajectory_rew: types.TrajectoryWithRew,
-            use_chdir,
-            tmpdir,
-            use_pickle,
-            use_rewards,
-            type_safe,
+        self,
+        trajectory: types.Trajectory,
+        trajectory_rew: types.TrajectoryWithRew,
+        use_chdir,
+        tmpdir,
+        use_pickle,
+        use_rewards,
+        type_safe,
     ):
         chdir_context: contextlib.AbstractContextManager
         """Check that trajectories are properly saved."""
@@ -251,32 +251,32 @@ def test_save_trajectories(
                 assert t1 == t2
 
     def test_invalid_trajectories(
-            self,
-            trajectory: types.Trajectory,
-            trajectory_rew: types.TrajectoryWithRew,
+        self,
+        trajectory: types.Trajectory,
+        trajectory_rew: types.TrajectoryWithRew,
     ) -> None:
         """Checks input validation catches space and dtype related errors."""
         trajs = [trajectory, trajectory_rew]
         for traj in trajs:
             with pytest.raises(
-                    ValueError,
-                    match=r"expected one more observations than actions.*",
+                ValueError,
+                match=r"expected one more observations than actions.*",
             ):
                 dataclasses.replace(traj, obs=traj.obs[:-1])
             with pytest.raises(
-                    ValueError,
-                    match=r"expected one more observations than actions.*",
+                ValueError,
+                match=r"expected one more observations than actions.*",
             ):
                 dataclasses.replace(traj, acts=traj.acts[:-1])
 
             with pytest.raises(
-                    ValueError,
-                    match=r"infos when present must be present for each action.*",
+                ValueError,
+                match=r"infos when present must be present for each action.*",
             ):
                 dataclasses.replace(traj, infos=traj.infos[:-1])
             with pytest.raises(
-                    ValueError,
-                    match=r"infos when present must be present for each action.*",
+                ValueError,
+                match=r"infos when present must be present for each action.*",
             ):
                 dataclasses.replace(traj, obs=traj.obs[:-1], acts=traj.acts[:-1])
 
@@ -293,12 +293,12 @@ def test_invalid_trajectories(
             )
 
     def test_valid_transitions(
-            self,
-            transitions_min: types.TransitionsMinimal,
-            transitions: types.Transitions,
-            transitions_rew: types.TransitionsWithRew,
-            length: int,
-            n_checks: int = 20,
+        self,
+        transitions_min: types.TransitionsMinimal,
+        transitions: types.Transitions,
+        transitions_rew: types.TransitionsWithRew,
+        length: int,
+        n_checks: int = 20,
     ) -> None:
         """Checks initialization, indexing, and slicing sanity."""
         for trans in [transitions_min, transitions, transitions_rew]:
@@ -322,11 +322,11 @@ def test_valid_transitions(
                 _check_transitions_get_item(trans, s)
 
     def test_invalid_transitions(
-            self,
-            transitions_min: types.Transitions,
-            transitions: types.Transitions,
-            transitions_rew: types.TransitionsWithRew,
-            length: int,
+        self,
+        transitions_min: types.Transitions,
+        transitions: types.Transitions,
+        transitions_rew: types.TransitionsWithRew,
+        length: int,
     ) -> None:
         """Checks input validation catches space and dtype related errors."""
         if length == 0:
@@ -334,26 +334,26 @@ def test_invalid_transitions(
 
         for trans in [transitions_min, transitions, transitions_rew]:
             with pytest.raises(
-                    ValueError,
-                    match=r"obs and acts must have same number of timesteps:.*",
+                ValueError,
+                match=r"obs and acts must have same number of timesteps:.*",
             ):
                 dataclasses.replace(trans, acts=trans.acts[:-1])
             with pytest.raises(
-                    ValueError,
-                    match=r"obs and infos must have same number of timesteps:.*",
+                ValueError,
+                match=r"obs and infos must have same number of timesteps:.*",
             ):
                 dataclasses.replace(trans, infos=[{}] * (length - 1))
 
         for trans in [transitions, transitions_rew]:
             with pytest.raises(
-                    ValueError,
-                    match=r"obs and next_obs must have same shape:.*",
+                ValueError,
+                match=r"obs and next_obs must have same shape:.*",
             ):
                 dataclasses.replace(trans, next_obs=np.zeros((len(trans), 4, 2)))
 
             with pytest.raises(
-                    ValueError,
-                    match=r"obs and next_obs must have the same dtype:.*",
+                ValueError,
+                match=r"obs and next_obs must have the same dtype:.*",
             ):
                 dataclasses.replace(
                     trans,
diff --git a/tests/data/test_wrappers.py b/tests/data/test_wrappers.py
index e1dca8ed6..14a7626c8 100644
--- a/tests/data/test_wrappers.py
+++ b/tests/data/test_wrappers.py
@@ -49,7 +49,7 @@ def step(self, action):
 
 
 def _make_buffering_venv(
-        error_on_premature_reset: bool,
+    error_on_premature_reset: bool,
 ) -> BufferingWrapper:
     venv = DummyVecEnv([_CountingEnv] * 2)
     wrapped_venv = BufferingWrapper(venv, error_on_premature_reset)
@@ -65,7 +65,7 @@ def _assert_equal_scrambled_vectors(a: np.ndarray, b: np.ndarray) -> None:
 
 
 def _join_transitions(
-        trans_list: Sequence[types.TransitionsWithRew],
+    trans_list: Sequence[types.TransitionsWithRew],
 ) -> types.TransitionsWithRew:
     def concat(x):
         return np.concatenate(list(x))
@@ -90,9 +90,9 @@ def concat(x):
 @pytest.mark.parametrize("n_steps", [1, 2, 20, 21])
 @pytest.mark.parametrize("extra_pop_timesteps", [(), (1,), (4, 8)])
 def test_pop(
-        episode_lengths: Sequence[int],
-        n_steps: int,
-        extra_pop_timesteps: Sequence[int],
+    episode_lengths: Sequence[int],
+    n_steps: int,
+    extra_pop_timesteps: Sequence[int],
 ) -> None:
     """Check pop_transitions() results for BufferWrapper.
 

From b992ab398eb0b9423fba5c76d34333f2a5a63773 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 1 Sep 2022 14:48:58 +0100
Subject: [PATCH 016/187] Linting

---
 tests/scripts/test_scripts.py | 62 +++++++++++++++++------------------
 1 file changed, 31 insertions(+), 31 deletions(-)

diff --git a/tests/scripts/test_scripts.py b/tests/scripts/test_scripts.py
index 6b0694f92..3e9fe014b 100644
--- a/tests/scripts/test_scripts.py
+++ b/tests/scripts/test_scripts.py
@@ -167,8 +167,8 @@ def test_train_preference_comparisons_sac(tmpdir):
         # make sure rl.sac named_config is called after rl.fast to overwrite
         # rl_kwargs.batch_size to None
         named_configs=["pendulum"]
-        + ALGO_FAST_CONFIGS["preference_comparison"]
-        + RL_SAC_NAMED_CONFIGS,
+                      + ALGO_FAST_CONFIGS["preference_comparison"]
+                      + RL_SAC_NAMED_CONFIGS,
         config_updates=config_updates,
     )
     assert run.config["rl"]["rl_cls"] is stable_baselines3.SAC
@@ -180,8 +180,8 @@ def test_train_preference_comparisons_sac(tmpdir):
     with pytest.raises(Exception, match=".*set 'batch_size' at top-level.*"):
         train_preference_comparisons.train_preference_comparisons_ex.run(
             named_configs=["pendulum"]
-            + RL_SAC_NAMED_CONFIGS
-            + ALGO_FAST_CONFIGS["preference_comparison"],
+                          + RL_SAC_NAMED_CONFIGS
+                          + ALGO_FAST_CONFIGS["preference_comparison"],
             config_updates=config_updates,
         )
 
@@ -201,8 +201,8 @@ def _run_reward_relabel_sac_preference_comparisons(buffer_cls):
             # make sure rl.sac named_config is called after rl.fast to overwrite
             # rl_kwargs.batch_size to None
             named_configs=["pendulum"]
-            + ALGO_FAST_CONFIGS["preference_comparison"]
-            + RL_SAC_NAMED_CONFIGS,
+                          + ALGO_FAST_CONFIGS["preference_comparison"]
+                          + RL_SAC_NAMED_CONFIGS,
             config_updates=config_updates,
         )
         return run
@@ -220,17 +220,17 @@ def _run_reward_relabel_sac_preference_comparisons(buffer_cls):
 @pytest.mark.parametrize(
     "named_configs",
     (
-        [],
-        ["reward.normalize_output_running"],
-        ["reward.normalize_output_disable"],
+            [],
+            ["reward.normalize_output_running"],
+            ["reward.normalize_output_disable"],
     ),
 )
 def test_train_preference_comparisons_reward_named_config(tmpdir, named_configs):
     config_updates = dict(common=dict(log_root=tmpdir))
     run = train_preference_comparisons.train_preference_comparisons_ex.run(
         named_configs=["cartpole"]
-        + ALGO_FAST_CONFIGS["preference_comparison"]
-        + named_configs,
+                      + ALGO_FAST_CONFIGS["preference_comparison"]
+                      + named_configs,
         config_updates=config_updates,
     )
     if "reward.normalize_output_running" in named_configs:
@@ -249,8 +249,8 @@ def test_train_preference_comparisons_active_learning(tmpdir, config):
     sacred.utils.recursive_update(config_updates, config)
     run = train_preference_comparisons.train_preference_comparisons_ex.run(
         named_configs=["cartpole"]
-        + ALGO_FAST_CONFIGS["preference_comparison"]
-        + ["reward.reward_ensemble"],
+                      + ALGO_FAST_CONFIGS["preference_comparison"]
+                      + ["reward.reward_ensemble"],
         config_updates=config_updates,
     )
     assert run.status == "COMPLETED"
@@ -275,8 +275,8 @@ def test_train_dagger_main(tmpdir):
         # PyTorch wants writeable arrays.
         # See https://github.com/HumanCompatibleAI/imitation/issues/219
         assert not (
-            warning.category == UserWarning
-            and "NumPy array is not writeable" in warning.message.args[0]
+                warning.category == UserWarning
+                and "NumPy array is not writeable" in warning.message.args[0]
         )
     assert run.status == "COMPLETED"
     assert isinstance(run.result, dict)
@@ -392,8 +392,8 @@ def test_train_rl_wb_logging(tmpdir):
     with pytest.raises(Exception, match=".*api_key not configured.*"):
         train_rl.train_rl_ex.run(
             named_configs=["cartpole"]
-            + ALGO_FAST_CONFIGS["rl"]
-            + ["common.wandb_logging"],
+                          + ALGO_FAST_CONFIGS["rl"]
+                          + ["common.wandb_logging"],
             config_updates=dict(
                 common=dict(log_root=tmpdir),
             ),
@@ -471,9 +471,9 @@ def _check_train_ex_result(result: dict):
 @pytest.mark.parametrize(
     "named_configs",
     (
-        [],
-        ["train.normalize_running", "reward.normalize_input_running"],
-        ["train.normalize_disable", "reward.normalize_input_disable"],
+            [],
+            ["train.normalize_running", "reward.normalize_input_running"],
+            ["train.normalize_disable", "reward.normalize_input_disable"],
     ),
 )
 @pytest.mark.parametrize("command", ("airl", "gail"))
@@ -530,7 +530,7 @@ def test_train_adversarial_sac(tmpdir, command):
     # Make sure rl.sac named_config is called after rl.fast to overwrite
     # rl_kwargs.batch_size to None
     named_configs = (
-        ["pendulum"] + ALGO_FAST_CONFIGS["adversarial"] + RL_SAC_NAMED_CONFIGS
+            ["pendulum"] + ALGO_FAST_CONFIGS["adversarial"] + RL_SAC_NAMED_CONFIGS
     )
     config_updates = {
         "common": dict(log_root=tmpdir),
@@ -625,14 +625,14 @@ def test_transfer_learning(tmpdir: str) -> None:
 @pytest.mark.parametrize(
     "named_configs_dict",
     (
-        dict(pc=[], rl=[]),
-        dict(pc=["rl.sac", "train.sac"], rl=["rl.sac", "train.sac"]),
-        dict(pc=["reward.reward_ensemble"], rl=[]),
+            dict(pc=[], rl=[]),
+            dict(pc=["rl.sac", "train.sac"], rl=["rl.sac", "train.sac"]),
+            dict(pc=["reward.reward_ensemble"], rl=[]),
     ),
 )
 def test_preference_comparisons_transfer_learning(
-    tmpdir: str,
-    named_configs_dict: Mapping[str, List[str]],
+        tmpdir: str,
+        named_configs_dict: Mapping[str, List[str]],
 ) -> None:
     """Transfer learning smoke test.
 
@@ -647,8 +647,8 @@ def test_preference_comparisons_transfer_learning(
     log_dir_train = tmpdir_path / "train"
     run = train_preference_comparisons.train_preference_comparisons_ex.run(
         named_configs=["pendulum"]
-        + ALGO_FAST_CONFIGS["preference_comparison"]
-        + named_configs_dict["pc"],
+                      + ALGO_FAST_CONFIGS["preference_comparison"]
+                      + named_configs_dict["pc"],
         config_updates=dict(common=dict(log_dir=log_dir_train)),
     )
     assert run.status == "COMPLETED"
@@ -682,7 +682,7 @@ def test_preference_comparisons_transfer_learning(
 def test_train_rl_double_normalization(tmpdir: str):
     venv = util.make_vec_env("CartPole-v1", n_envs=1, parallel=False)
     basic_reward_net = reward_nets.BasicRewardNet(
-        venv.observation_space, venv.action_space
+        venv.observation_space, venv.action_space,
     )
     net = reward_nets.NormalizedRewardNet(basic_reward_net, networks.RunningNorm)
     tmppath = os.path.join(tmpdir, "reward.pt")
@@ -844,8 +844,8 @@ def _run_train_bc_for_test_analyze_imit(run_name, sacred_logs_dir, log_dir):
 @pytest.mark.parametrize(
     "run_sacred_fn",
     (
-        _run_train_adv_for_test_analyze_imit,
-        _run_train_bc_for_test_analyze_imit,
+            _run_train_adv_for_test_analyze_imit,
+            _run_train_bc_for_test_analyze_imit,
     ),
 )
 def test_analyze_imitation(tmpdir: str, run_names: List[str], run_sacred_fn):

From 7a788cd88452a7f63c534310753552015e814714 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 1 Sep 2022 14:50:20 +0100
Subject: [PATCH 017/187] Fix types: algorithms/preference_comparisons.py

---
 .../algorithms/preference_comparisons.py      | 90 ++++++++++++-------
 1 file changed, 57 insertions(+), 33 deletions(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 24bc29895..abc3ad0ba 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -18,6 +18,7 @@
     Sequence,
     Tuple,
     Union,
+    cast,
 )
 
 import numpy as np
@@ -182,9 +183,11 @@ def __init__(
         self.algorithm.set_env(self.venv)
         # Unlike with BufferingWrapper, we should use `algorithm.get_env()` instead
         # of `venv` when interacting with `algorithm`.
+        algo_venv = self.algorithm.get_env()
+        assert algo_venv is not None
         policy_callable = rollout._policy_to_callable(
             self.algorithm,
-            self.algorithm.get_env(),
+            algo_venv,
             # By setting deterministic_policy to False, we ensure that the rollouts
             # are collected from a deterministic policy only if self.algorithm is
             # deterministic. If self.algorithm is stochastic, then policy_callable
@@ -193,7 +196,7 @@ def __init__(
         )
         self.exploration_wrapper = exploration_wrapper.ExplorationWrapper(
             policy_callable=policy_callable,
-            venv=self.algorithm.get_env(),
+            venv=algo_venv,
             random_prob=random_prob,
             switch_prob=switch_prob,
             seed=seed,
@@ -255,9 +258,11 @@ def sample(self, steps: int) -> Sequence[types.TrajectoryWithRew]:
             # here because 1) they might miss initial timesteps taken by the RL agent
             # and 2) their rewards are the ones provided by the reward model!
             # Instead, we collect the trajectories using the BufferingWrapper.
+            algo_venv = self.algorithm.get_env()
+            assert algo_venv is not None
             rollout.generate_trajectories(
                 self.algorithm,
-                self.algorithm.get_env(),
+                algo_venv,
                 sample_until=sample_until,
                 # By setting deterministic_policy to False, we ensure that the rollouts
                 # are collected from a deterministic policy only if self.algorithm is
@@ -270,30 +275,35 @@ def sample(self, steps: int) -> Sequence[types.TrajectoryWithRew]:
 
         agent_trajs = _get_trajectories(agent_trajs, agent_steps)
 
-        exploration_trajs = []
+        trajectories = list(agent_trajs)
+
         if exploration_steps > 0:
             self.logger.log(f"Sampling {exploration_steps} exploratory transitions.")
             sample_until = rollout.make_sample_until(
                 min_timesteps=exploration_steps,
                 min_episodes=None,
             )
+            algo_venv = self.algorithm.get_env()
+            assert algo_venv is not None
             rollout.generate_trajectories(
                 policy=self.exploration_wrapper,
-                venv=self.algorithm.get_env(),
+                venv=algo_venv,
                 sample_until=sample_until,
-                # buffering_wrapper collects rollouts from a non-deterministic policy
+                # buffering_wrapper collects rollouts from a non-deterministic policy,
                 # so we do that here as well for consistency.
                 deterministic_policy=False,
             )
             exploration_trajs, _ = self.buffering_wrapper.pop_finished_trajectories()
             exploration_trajs = _get_trajectories(exploration_trajs, exploration_steps)
-        # We call _get_trajectories separately on agent_trajs and exploration_trajs
-        # and then just concatenate. This could mean we return slightly too many
-        # transitions, but it gets the proportion of exploratory and agent transitions
-        # roughly right.
-        return list(agent_trajs) + list(exploration_trajs)
-
-    @TrajectoryGenerator.logger.setter
+            # We call _get_trajectories separately on agent_trajs and exploration_trajs
+            # and then just concatenate. This could mean we return slightly too many
+            # transitions, but it gets the proportion of exploratory and agent transitions
+            # roughly right.
+            trajectories.extend(list(exploration_trajs))
+        return trajectories
+
+    # Type ignore due to https://github.com/python/mypy/issues/5936
+    @TrajectoryGenerator.logger.setter  # type: ignore[attr-defined]
     def logger(self, value: imit_logger.HierarchicalLogger):
         self._logger = value
         self.algorithm.set_logger(self.logger)
@@ -318,6 +328,7 @@ def _get_trajectories(
     # Now we find the first index that gives us enough
     # total steps:
     idx = (steps_cumsum >= steps).argmax()
+    assert isinstance(idx, int)
     # we need to include the element at position idx
     trajectories = trajectories[: idx + 1]
     # sanity check
@@ -358,15 +369,19 @@ def __init__(
         self.noise_prob = noise_prob
         self.discount_factor = discount_factor
         self.threshold = threshold
-        self.is_ensemble, base_model = is_base_model_ensemble(self.model)
+        self.is_ensemble: bool
+        base_model = get_base_model(model)
+        self.is_ensemble = isinstance(base_model, reward_nets.RewardEnsemble)
         # if the base model is an ensemble model, then keep the base model as
         # model to get rewards from all networks
         if self.is_ensemble:
-            self.model = base_model
+            # For some reason, mypy is not able to infer the type of base_model
+            # when self.is_ensemble is True.
+            self.model = cast(reward_nets.RewardEnsemble, base_model)
             self.member_pref_models = []
             for member in self.model.members:
                 member_pref_model = PreferenceModel(
-                    member,
+                    cast(reward_nets.RewardNet, member),  # nn.ModuleList is not generic
                     self.noise_prob,
                     self.discount_factor,
                     self.threshold,
@@ -426,6 +441,10 @@ def forward(
             rews2 = self.rewards(trans2)
             probs[i] = self.probability(rews1, rews2)
             if gt_reward_available:
+                frag1, frag2 = cast(TrajectoryWithRew, frag1), cast(
+                    TrajectoryWithRew,
+                    frag2,
+                )
                 gt_rews_1 = th.from_numpy(frag1.rews)
                 gt_rews_2 = th.from_numpy(frag2.rews)
                 gt_probs[i] = self.probability(gt_rews_1, gt_rews_2)
@@ -451,7 +470,7 @@ def rewards(
         next_state = transitions.next_obs
         done = transitions.dones
         if self.is_ensemble:
-            rews = self.model.predict_processed_all(state, action, next_state, done)
+            rews = self.model.predict_processed_all(state, action, next_state, done)  # type: ignore[operator]
             assert rews.shape == (len(state), self.model.num_members)
             return util.safe_to_tensor(rews).to(self.model.device)
         else:
@@ -487,7 +506,7 @@ def probability(
         # factor of 1 to avoid unnecessary computation (especially
         # since this is the default setting).
         if self.discount_factor == 1:
-            returns_diff = (rews2 - rews1).sum(axis=0)
+            returns_diff = (rews2 - rews1).sum(axis=0)  # type: ignore[call-overload]
         else:
             discounts = self.discount_factor ** th.arange(len(rews1))
             if self.is_ensemble:
@@ -737,12 +756,12 @@ def variance_estimate(self, rews1, rews2) -> float:
             var_estimate = (returns1 - returns2).var().item()
         else:  # uncertainty_on is probability or label
             probs = self.preference_model.probability(rews1, rews2)
-            probs = probs.cpu().numpy()
-            assert probs.shape == (self.preference_model.model.num_members,)
+            probs_np = probs.cpu().numpy()
+            assert probs_np.shape == (self.preference_model.model.num_members,)
             if self.uncertainty_on == "probability":
-                var_estimate = probs.var()
+                var_estimate = probs_np.var()
             elif self.uncertainty_on == "label":  # uncertainty_on is label
-                preds = (probs > 0.5).astype(np.float32)
+                preds = (probs_np > 0.5).astype(np.float32)
                 # probability estimate of Bernoulli random variable
                 prob_estimate = preds.mean()
                 # variance estimate of Bernoulli random variable
@@ -917,7 +936,7 @@ def push(
         if preferences.shape != (len(fragments),):
             raise ValueError(
                 f"Unexpected preferences shape {preferences.shape}, "
-                f"expected {(len(fragments), )}",
+                f"expected {(len(fragments),)}",
             )
         if preferences.dtype != np.float32:
             raise ValueError("preferences should have dtype float32")
@@ -1220,8 +1239,8 @@ def _training_inner_loop(
         preferences: np.ndarray,
     ) -> th.Tensor:
         assert len(fragment_pairs) == preferences.shape[0]
-        losses = []
-        metrics = []
+        losses_list = []
+        metrics_list = []
         for member_idx in range(self._model.num_members):
             # sample fragments for training via bagging
             sample_idx = self.rng.choice(
@@ -1232,9 +1251,9 @@ def _training_inner_loop(
             sample_fragments = [fragment_pairs[i] for i in sample_idx]
             sample_preferences = preferences[sample_idx]
             output = self.loss.forward(sample_fragments, sample_preferences, member_idx)
-            losses.append(output.loss)
-            metrics.append(output.metrics)
-        losses = th.stack(losses)
+            losses_list.append(output.loss)
+            metrics_list.append(output.metrics)
+        losses = th.stack(losses_list)
         loss = losses.sum()
 
         self.logger.record("loss", loss.item())
@@ -1242,19 +1261,21 @@ def _training_inner_loop(
 
         # Turn metrics from a list of dictionaries into a dictionary of
         # tensors.
-        metrics = {k: th.stack([di[k] for di in metrics]) for k in metrics[0]}
+        metrics = {k: th.stack([di[k] for di in metrics_list]) for k in metrics_list[0]}
         for name, value in metrics.items():
             self.logger.record(name, value.mean().item())
 
         return loss
 
 
-def is_base_model_ensemble(reward_model):
+def get_base_model(
+    reward_model: reward_nets.RewardNet,
+) -> Union[reward_nets.RewardNet, reward_nets.RewardEnsemble]:
     base_model = reward_model
     while hasattr(base_model, "base"):
-        base_model = base_model.base
+        base_model = cast(reward_nets.RewardNet, base_model.base)
 
-    return isinstance(base_model, reward_nets.RewardEnsemble), base_model
+    return base_model
 
 
 def _make_reward_trainer(
@@ -1266,9 +1287,12 @@ def _make_reward_trainer(
     """Construct the correct type of reward trainer for this reward function."""
     if reward_trainer_kwargs is None:
         reward_trainer_kwargs = {}
-    is_ensemble, base_model = is_base_model_ensemble(reward_model)
+
+    base_model = get_base_model(reward_model)
+    is_ensemble = isinstance(base_model, reward_nets.RewardEnsemble)
 
     if is_ensemble:
+        base_model = cast(reward_nets.RewardEnsemble, base_model)
         # reward_model may include an AddSTDRewardWrapper for RL training; but we
         # must train directly on the base model for reward model training.
         is_base = reward_model is base_model

From 26c32c54f18d275a1a6aaf28d5dc9e63f8deb08b Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 2 Sep 2022 12:56:17 +0100
Subject: [PATCH 018/187] Fix types: algorithms/mce_irl.py

---
 src/imitation/algorithms/mce_irl.py | 105 +++++++++++++++-------------
 1 file changed, 58 insertions(+), 47 deletions(-)

diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index 2c7e29378..2644382e6 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -7,7 +7,7 @@
 """
 import collections
 import warnings
-from typing import Any, Iterable, Mapping, Optional, Tuple, Type, Union
+from typing import Any, Iterable, Mapping, Optional, Tuple, Type, Union, List
 
 import gym
 import numpy as np
@@ -24,10 +24,10 @@
 
 
 def mce_partition_fh(
-    env: resettable_env.TabularModelEnv,
-    *,
-    reward: Optional[np.ndarray] = None,
-    discount: float = 1.0,
+        env: resettable_env.TabularModelEnv,
+        *,
+        reward: Optional[np.ndarray] = None,
+        discount: float = 1.0,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     r"""Performs the soft Bellman backup for a finite-horizon MDP.
 
@@ -77,11 +77,11 @@ def mce_partition_fh(
 
 
 def mce_occupancy_measures(
-    env: resettable_env.TabularModelEnv,
-    *,
-    reward: Optional[np.ndarray] = None,
-    pi: Optional[np.ndarray] = None,
-    discount: float = 1.0,
+        env: resettable_env.TabularModelEnv,
+        *,
+        reward: Optional[np.ndarray] = None,
+        pi: Optional[np.ndarray] = None,
+        discount: float = 1.0,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """Calculate state visitation frequency Ds for each state s under a given policy pi.
 
@@ -140,13 +140,15 @@ def squeeze_r(r_output: th.Tensor) -> th.Tensor:
 
 class TabularPolicy(policies.BasePolicy):
     """A tabular policy. Cannot be trained -- prediction only."""
+    pi: np.ndarray
+    rng: np.random.RandomState
 
     def __init__(
-        self,
-        state_space: gym.Space,
-        action_space: gym.Space,
-        pi: np.ndarray,
-        rng: Optional[np.random.RandomState],
+            self,
+            state_space: gym.Space,
+            action_space: gym.Space,
+            pi: np.ndarray,
+            rng: Optional[np.random.RandomState],
     ):
         """Builds TabularPolicy.
 
@@ -162,8 +164,7 @@ def __init__(
         assert isinstance(action_space, gym.spaces.Discrete), "action not tabular"
         # What we call state space here is observation space in SB3 nomenclature.
         super().__init__(observation_space=state_space, action_space=action_space)
-        self.rng = rng or np.random
-        self.pi = None
+        self.rng = rng or np.random.RandomState()
         self.set_pi(pi)
 
     def set_pi(self, pi: np.ndarray) -> None:
@@ -176,15 +177,15 @@ def set_pi(self, pi: np.ndarray) -> None:
     def _predict(self, observation: th.Tensor, deterministic: bool = False):
         raise NotImplementedError("Should never be called as predict overridden.")
 
-    def forward(self, observation: th.Tensor, deterministic: bool = False):
+    def forward(self, observation: th.Tensor, deterministic: bool = False):  # type: ignore[override]
         raise NotImplementedError("Should never be called.")
 
-    def predict(
-        self,
-        observation: np.ndarray,
-        state: Optional[np.ndarray] = None,
-        mask: Optional[np.ndarray] = None,
-        deterministic: bool = False,
+    def predict(  # type: ignore[override]
+            self,
+            observation: np.ndarray,
+            state: Optional[np.ndarray] = None,
+            mask: Optional[np.ndarray] = None,
+            deterministic: bool = False,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Predict action to take in given state.
 
@@ -218,12 +219,12 @@ def predict(
         if mask is not None:
             timesteps[mask] = 0
 
-        actions = []
+        actions: List[int] = []
         for obs, t in zip(observation, timesteps):
             assert self.observation_space.contains(obs), "illegal state"
             dist = self.pi[t, obs, :]
             if deterministic:
-                actions.append(dist.argmax())
+                actions.append(int(dist.argmax()))
             else:
                 actions.append(self.rng.choice(len(dist), p=dist))
 
@@ -250,20 +251,20 @@ class MCEIRL(base.DemonstrationAlgorithm[types.TransitionsMinimal]):
     demo_state_om: Optional[np.ndarray]
 
     def __init__(
-        self,
-        demonstrations: Optional[MCEDemonstrations],
-        env: resettable_env.TabularModelEnv,
-        reward_net: reward_nets.RewardNet,
-        optimizer_cls: Type[th.optim.Optimizer] = th.optim.Adam,
-        optimizer_kwargs: Optional[Mapping[str, Any]] = None,
-        discount: float = 1.0,
-        linf_eps: float = 1e-3,
-        grad_l2_eps: float = 1e-4,
-        # TODO(adam): do we need log_interval or can just use record_mean...?
-        log_interval: Optional[int] = 100,
-        *,
-        rng: Optional[np.random.RandomState] = None,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            demonstrations: Optional[MCEDemonstrations],
+            env: resettable_env.TabularModelEnv,
+            reward_net: reward_nets.RewardNet,
+            optimizer_cls: Type[th.optim.Optimizer] = th.optim.Adam,
+            optimizer_kwargs: Optional[Mapping[str, Any]] = None,
+            discount: float = 1.0,
+            linf_eps: float = 1e-3,
+            grad_l2_eps: float = 1e-4,
+            # TODO(adam): do we need log_interval or can just use record_mean...?
+            log_interval: Optional[int] = 100,
+            *,
+            rng: Optional[np.random.RandomState] = None,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         r"""Creates MCE IRL.
 
@@ -334,10 +335,10 @@ def _set_demo_from_trajectories(self, trajs: Iterable[types.Trajectory]) -> None
         self.demo_state_om /= num_demos
 
     def _set_demo_from_obs(
-        self,
-        obses: np.ndarray,
-        dones: Optional[np.ndarray],
-        next_obses: Optional[np.ndarray],
+            self,
+            obses: np.ndarray,
+            dones: Optional[np.ndarray],
+            next_obses: Optional[np.ndarray],
     ) -> None:
         self.demo_state_om = np.zeros((self.env.n_states,))
 
@@ -351,6 +352,10 @@ def _set_demo_from_obs(
         # as they will not appear anywhere else; but ignore next observations
         # for all other states as they occur elsewhere in dataset.
         if next_obses is not None:
+            # TODO(juan) this is wrong; dones is an optional
+            #  and zip() cannot handle this when it's None.
+            #  either require an np.ndarray or do something else
+            #  if the value is None.
             for done, obs in zip(dones, next_obses):
                 if isinstance(done, th.Tensor):
                     done = done.item()  # must be scalar
@@ -376,6 +381,10 @@ def set_demonstrations(self, demonstrations: MCEDemonstrations) -> None:
         # Demonstrations are either trajectories or transitions;
         # we must compute occupancy measure from this.
         if isinstance(demonstrations, Iterable):
+            # TODO(juan) this is wrong; if this is a container (list-like)
+            #  object then a new fresh iterator will be generated every time,
+            #  but for a general iterator calling next() exhaust the
+            #  iterator.
             first_item = next(iter(demonstrations))
             if isinstance(first_item, types.Trajectory):
                 self._set_demo_from_trajectories(demonstrations)
@@ -401,13 +410,14 @@ def set_demonstrations(self, demonstrations: MCEDemonstrations) -> None:
             # Collect them together into one big NumPy array. This is inefficient,
             # we could compute the running statistics instead, but in practice do
             # not expect large dataset sizes together with MCE IRL.
-            collated = collections.defaultdict(list)
+            collated_list = collections.defaultdict(list)
             for batch in demonstrations:
                 assert isinstance(batch, Mapping)
                 for k in ("obs", "dones", "next_obs"):
                     if k in batch:
-                        collated[k].append(batch[k])
-            for k, v in collated.items():
+                        collated_list[k].append(batch[k])
+            collated = {}
+            for k, v in collated_list.items():
                 collated[k] = np.concatenate(v)
 
             assert "obs" in collated
@@ -474,6 +484,7 @@ def train(self, max_iter: int = 1000) -> np.ndarray:
             dtype=self.reward_net.dtype,
             device=self.reward_net.device,
         )
+        assert self.demo_state_om is not None
         assert self.demo_state_om.shape == (len(obs_mat),)
 
         with networks.training(self.reward_net):

From e6f91247c52e7558aef5c3003846245c65d4e999 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 2 Sep 2022 13:07:34 +0100
Subject: [PATCH 019/187] Formatting, fixed minor bug

---
 src/imitation/algorithms/mce_irl.py           | 77 ++++++++++---------
 .../algorithms/preference_comparisons.py      |  3 +-
 tests/scripts/test_scripts.py                 | 63 +++++++--------
 3 files changed, 72 insertions(+), 71 deletions(-)

diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index 2644382e6..3534bcef2 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -7,7 +7,7 @@
 """
 import collections
 import warnings
-from typing import Any, Iterable, Mapping, Optional, Tuple, Type, Union, List
+from typing import Any, Iterable, List, Mapping, Optional, Tuple, Type, Union
 
 import gym
 import numpy as np
@@ -24,10 +24,10 @@
 
 
 def mce_partition_fh(
-        env: resettable_env.TabularModelEnv,
-        *,
-        reward: Optional[np.ndarray] = None,
-        discount: float = 1.0,
+    env: resettable_env.TabularModelEnv,
+    *,
+    reward: Optional[np.ndarray] = None,
+    discount: float = 1.0,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     r"""Performs the soft Bellman backup for a finite-horizon MDP.
 
@@ -77,11 +77,11 @@ def mce_partition_fh(
 
 
 def mce_occupancy_measures(
-        env: resettable_env.TabularModelEnv,
-        *,
-        reward: Optional[np.ndarray] = None,
-        pi: Optional[np.ndarray] = None,
-        discount: float = 1.0,
+    env: resettable_env.TabularModelEnv,
+    *,
+    reward: Optional[np.ndarray] = None,
+    pi: Optional[np.ndarray] = None,
+    discount: float = 1.0,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """Calculate state visitation frequency Ds for each state s under a given policy pi.
 
@@ -140,15 +140,16 @@ def squeeze_r(r_output: th.Tensor) -> th.Tensor:
 
 class TabularPolicy(policies.BasePolicy):
     """A tabular policy. Cannot be trained -- prediction only."""
+
     pi: np.ndarray
     rng: np.random.RandomState
 
     def __init__(
-            self,
-            state_space: gym.Space,
-            action_space: gym.Space,
-            pi: np.ndarray,
-            rng: Optional[np.random.RandomState],
+        self,
+        state_space: gym.Space,
+        action_space: gym.Space,
+        pi: np.ndarray,
+        rng: Optional[np.random.RandomState],
     ):
         """Builds TabularPolicy.
 
@@ -181,11 +182,11 @@ def forward(self, observation: th.Tensor, deterministic: bool = False):  # type:
         raise NotImplementedError("Should never be called.")
 
     def predict(  # type: ignore[override]
-            self,
-            observation: np.ndarray,
-            state: Optional[np.ndarray] = None,
-            mask: Optional[np.ndarray] = None,
-            deterministic: bool = False,
+        self,
+        observation: np.ndarray,
+        state: Optional[np.ndarray] = None,
+        mask: Optional[np.ndarray] = None,
+        deterministic: bool = False,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Predict action to take in given state.
 
@@ -251,20 +252,20 @@ class MCEIRL(base.DemonstrationAlgorithm[types.TransitionsMinimal]):
     demo_state_om: Optional[np.ndarray]
 
     def __init__(
-            self,
-            demonstrations: Optional[MCEDemonstrations],
-            env: resettable_env.TabularModelEnv,
-            reward_net: reward_nets.RewardNet,
-            optimizer_cls: Type[th.optim.Optimizer] = th.optim.Adam,
-            optimizer_kwargs: Optional[Mapping[str, Any]] = None,
-            discount: float = 1.0,
-            linf_eps: float = 1e-3,
-            grad_l2_eps: float = 1e-4,
-            # TODO(adam): do we need log_interval or can just use record_mean...?
-            log_interval: Optional[int] = 100,
-            *,
-            rng: Optional[np.random.RandomState] = None,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        demonstrations: Optional[MCEDemonstrations],
+        env: resettable_env.TabularModelEnv,
+        reward_net: reward_nets.RewardNet,
+        optimizer_cls: Type[th.optim.Optimizer] = th.optim.Adam,
+        optimizer_kwargs: Optional[Mapping[str, Any]] = None,
+        discount: float = 1.0,
+        linf_eps: float = 1e-3,
+        grad_l2_eps: float = 1e-4,
+        # TODO(adam): do we need log_interval or can just use record_mean...?
+        log_interval: Optional[int] = 100,
+        *,
+        rng: Optional[np.random.RandomState] = None,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         r"""Creates MCE IRL.
 
@@ -335,10 +336,10 @@ def _set_demo_from_trajectories(self, trajs: Iterable[types.Trajectory]) -> None
         self.demo_state_om /= num_demos
 
     def _set_demo_from_obs(
-            self,
-            obses: np.ndarray,
-            dones: Optional[np.ndarray],
-            next_obses: Optional[np.ndarray],
+        self,
+        obses: np.ndarray,
+        dones: Optional[np.ndarray],
+        next_obses: Optional[np.ndarray],
     ) -> None:
         self.demo_state_om = np.zeros((self.env.n_states,))
 
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index abc3ad0ba..bd1061244 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -327,8 +327,7 @@ def _get_trajectories(
     steps_cumsum = np.cumsum([len(traj) for traj in trajectories])
     # Now we find the first index that gives us enough
     # total steps:
-    idx = (steps_cumsum >= steps).argmax()
-    assert isinstance(idx, int)
+    idx = int((steps_cumsum >= steps).argmax())
     # we need to include the element at position idx
     trajectories = trajectories[: idx + 1]
     # sanity check
diff --git a/tests/scripts/test_scripts.py b/tests/scripts/test_scripts.py
index 3e9fe014b..7b592f0c4 100644
--- a/tests/scripts/test_scripts.py
+++ b/tests/scripts/test_scripts.py
@@ -167,8 +167,8 @@ def test_train_preference_comparisons_sac(tmpdir):
         # make sure rl.sac named_config is called after rl.fast to overwrite
         # rl_kwargs.batch_size to None
         named_configs=["pendulum"]
-                      + ALGO_FAST_CONFIGS["preference_comparison"]
-                      + RL_SAC_NAMED_CONFIGS,
+        + ALGO_FAST_CONFIGS["preference_comparison"]
+        + RL_SAC_NAMED_CONFIGS,
         config_updates=config_updates,
     )
     assert run.config["rl"]["rl_cls"] is stable_baselines3.SAC
@@ -180,8 +180,8 @@ def test_train_preference_comparisons_sac(tmpdir):
     with pytest.raises(Exception, match=".*set 'batch_size' at top-level.*"):
         train_preference_comparisons.train_preference_comparisons_ex.run(
             named_configs=["pendulum"]
-                          + RL_SAC_NAMED_CONFIGS
-                          + ALGO_FAST_CONFIGS["preference_comparison"],
+            + RL_SAC_NAMED_CONFIGS
+            + ALGO_FAST_CONFIGS["preference_comparison"],
             config_updates=config_updates,
         )
 
@@ -201,8 +201,8 @@ def _run_reward_relabel_sac_preference_comparisons(buffer_cls):
             # make sure rl.sac named_config is called after rl.fast to overwrite
             # rl_kwargs.batch_size to None
             named_configs=["pendulum"]
-                          + ALGO_FAST_CONFIGS["preference_comparison"]
-                          + RL_SAC_NAMED_CONFIGS,
+            + ALGO_FAST_CONFIGS["preference_comparison"]
+            + RL_SAC_NAMED_CONFIGS,
             config_updates=config_updates,
         )
         return run
@@ -220,17 +220,17 @@ def _run_reward_relabel_sac_preference_comparisons(buffer_cls):
 @pytest.mark.parametrize(
     "named_configs",
     (
-            [],
-            ["reward.normalize_output_running"],
-            ["reward.normalize_output_disable"],
+        [],
+        ["reward.normalize_output_running"],
+        ["reward.normalize_output_disable"],
     ),
 )
 def test_train_preference_comparisons_reward_named_config(tmpdir, named_configs):
     config_updates = dict(common=dict(log_root=tmpdir))
     run = train_preference_comparisons.train_preference_comparisons_ex.run(
         named_configs=["cartpole"]
-                      + ALGO_FAST_CONFIGS["preference_comparison"]
-                      + named_configs,
+        + ALGO_FAST_CONFIGS["preference_comparison"]
+        + named_configs,
         config_updates=config_updates,
     )
     if "reward.normalize_output_running" in named_configs:
@@ -249,8 +249,8 @@ def test_train_preference_comparisons_active_learning(tmpdir, config):
     sacred.utils.recursive_update(config_updates, config)
     run = train_preference_comparisons.train_preference_comparisons_ex.run(
         named_configs=["cartpole"]
-                      + ALGO_FAST_CONFIGS["preference_comparison"]
-                      + ["reward.reward_ensemble"],
+        + ALGO_FAST_CONFIGS["preference_comparison"]
+        + ["reward.reward_ensemble"],
         config_updates=config_updates,
     )
     assert run.status == "COMPLETED"
@@ -275,8 +275,8 @@ def test_train_dagger_main(tmpdir):
         # PyTorch wants writeable arrays.
         # See https://github.com/HumanCompatibleAI/imitation/issues/219
         assert not (
-                warning.category == UserWarning
-                and "NumPy array is not writeable" in warning.message.args[0]
+            warning.category == UserWarning
+            and "NumPy array is not writeable" in warning.message.args[0]
         )
     assert run.status == "COMPLETED"
     assert isinstance(run.result, dict)
@@ -392,8 +392,8 @@ def test_train_rl_wb_logging(tmpdir):
     with pytest.raises(Exception, match=".*api_key not configured.*"):
         train_rl.train_rl_ex.run(
             named_configs=["cartpole"]
-                          + ALGO_FAST_CONFIGS["rl"]
-                          + ["common.wandb_logging"],
+            + ALGO_FAST_CONFIGS["rl"]
+            + ["common.wandb_logging"],
             config_updates=dict(
                 common=dict(log_root=tmpdir),
             ),
@@ -471,9 +471,9 @@ def _check_train_ex_result(result: dict):
 @pytest.mark.parametrize(
     "named_configs",
     (
-            [],
-            ["train.normalize_running", "reward.normalize_input_running"],
-            ["train.normalize_disable", "reward.normalize_input_disable"],
+        [],
+        ["train.normalize_running", "reward.normalize_input_running"],
+        ["train.normalize_disable", "reward.normalize_input_disable"],
     ),
 )
 @pytest.mark.parametrize("command", ("airl", "gail"))
@@ -530,7 +530,7 @@ def test_train_adversarial_sac(tmpdir, command):
     # Make sure rl.sac named_config is called after rl.fast to overwrite
     # rl_kwargs.batch_size to None
     named_configs = (
-            ["pendulum"] + ALGO_FAST_CONFIGS["adversarial"] + RL_SAC_NAMED_CONFIGS
+        ["pendulum"] + ALGO_FAST_CONFIGS["adversarial"] + RL_SAC_NAMED_CONFIGS
     )
     config_updates = {
         "common": dict(log_root=tmpdir),
@@ -625,14 +625,14 @@ def test_transfer_learning(tmpdir: str) -> None:
 @pytest.mark.parametrize(
     "named_configs_dict",
     (
-            dict(pc=[], rl=[]),
-            dict(pc=["rl.sac", "train.sac"], rl=["rl.sac", "train.sac"]),
-            dict(pc=["reward.reward_ensemble"], rl=[]),
+        dict(pc=[], rl=[]),
+        dict(pc=["rl.sac", "train.sac"], rl=["rl.sac", "train.sac"]),
+        dict(pc=["reward.reward_ensemble"], rl=[]),
     ),
 )
 def test_preference_comparisons_transfer_learning(
-        tmpdir: str,
-        named_configs_dict: Mapping[str, List[str]],
+    tmpdir: str,
+    named_configs_dict: Mapping[str, List[str]],
 ) -> None:
     """Transfer learning smoke test.
 
@@ -647,8 +647,8 @@ def test_preference_comparisons_transfer_learning(
     log_dir_train = tmpdir_path / "train"
     run = train_preference_comparisons.train_preference_comparisons_ex.run(
         named_configs=["pendulum"]
-                      + ALGO_FAST_CONFIGS["preference_comparison"]
-                      + named_configs_dict["pc"],
+        + ALGO_FAST_CONFIGS["preference_comparison"]
+        + named_configs_dict["pc"],
         config_updates=dict(common=dict(log_dir=log_dir_train)),
     )
     assert run.status == "COMPLETED"
@@ -682,7 +682,8 @@ def test_preference_comparisons_transfer_learning(
 def test_train_rl_double_normalization(tmpdir: str):
     venv = util.make_vec_env("CartPole-v1", n_envs=1, parallel=False)
     basic_reward_net = reward_nets.BasicRewardNet(
-        venv.observation_space, venv.action_space,
+        venv.observation_space,
+        venv.action_space,
     )
     net = reward_nets.NormalizedRewardNet(basic_reward_net, networks.RunningNorm)
     tmppath = os.path.join(tmpdir, "reward.pt")
@@ -844,8 +845,8 @@ def _run_train_bc_for_test_analyze_imit(run_name, sacred_logs_dir, log_dir):
 @pytest.mark.parametrize(
     "run_sacred_fn",
     (
-            _run_train_adv_for_test_analyze_imit,
-            _run_train_bc_for_test_analyze_imit,
+        _run_train_adv_for_test_analyze_imit,
+        _run_train_bc_for_test_analyze_imit,
     ),
 )
 def test_analyze_imitation(tmpdir: str, run_names: List[str], run_sacred_fn):

From e7806727f98c4f2d23e016eb89008ee0f1eae941 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 2 Sep 2022 13:38:53 +0100
Subject: [PATCH 020/187] Clarify why types are ignored

---
 tests/policies/test_replay_buffer_wrapper.py | 2 ++
 tests/util/test_wb_logger.py                 | 4 ++++
 2 files changed, 6 insertions(+)

diff --git a/tests/policies/test_replay_buffer_wrapper.py b/tests/policies/test_replay_buffer_wrapper.py
index 1d6ec001a..c08f5e46e 100644
--- a/tests/policies/test_replay_buffer_wrapper.py
+++ b/tests/policies/test_replay_buffer_wrapper.py
@@ -54,6 +54,8 @@ def test_invalid_args():
         TypeError,
         match=r".*unexpected keyword argument 'replay_buffer_class'.*",
     ):
+        # we ignore the type because we are intentionally
+        # passing the wrong type for the test
         make_algo_with_wrapped_buffer(
             rl_cls=sb3.PPO,  # type: ignore
             policy_cls=policies.ActorCriticPolicy,
diff --git a/tests/util/test_wb_logger.py b/tests/util/test_wb_logger.py
index 392f21fdc..9bd4a65a5 100644
--- a/tests/util/test_wb_logger.py
+++ b/tests/util/test_wb_logger.py
@@ -87,6 +87,10 @@ def finish(self):
 mock_wandb = MockWandb()
 
 
+# we ignore the type below as one should technically not access the
+# __init__ method directly but only by creating an instance.
+
+
 @mock.patch.object(wandb, "__init__", mock_wandb.__init__)  # type: ignore
 @mock.patch.object(wandb, "init", mock_wandb.init)
 @mock.patch.object(wandb, "log", mock_wandb.log)

From 38aef218533d41c6c07fc502483300d4f5eee4b9 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Sat, 3 Sep 2022 11:31:42 +0200
Subject: [PATCH 021/187] Started fixing types on algorithms/density.py

---
 src/imitation/algorithms/density.py | 54 +++++++++++++++++++++--------
 1 file changed, 39 insertions(+), 15 deletions(-)

diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index ad52e7ad3..e7e9464ad 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -6,7 +6,7 @@
 
 import enum
 import itertools
-from typing import Iterable, Mapping, Optional
+from typing import Dict, Iterable, Optional
 
 import numpy as np
 from gym.spaces.utils import flatten
@@ -39,7 +39,20 @@ class DensityAlgorithm(base.DemonstrationAlgorithm):
     and then computes a reward using the log of these probabilities.
     """
 
-    transitions: Mapping[Optional[int], np.ndarray]
+    is_stationary: bool
+    density_type: DensityType
+    venv: vec_env.VecEnv
+    transitions: Dict[Optional[int], np.ndarray]
+    kernel: str
+    kernel_bandwidth: float
+    standardise: bool
+
+    _scaler: Optional[preprocessing.StandardScaler]
+    _density_models: Dict[Optional[int], neighbors.KernelDensity]
+    rl_algo: Optional[base_class.BaseAlgorithm]
+    buffering_wrapper: wrappers.BufferingWrapper
+    venv_wrapped: reward_wrapper.RewardVecEnvWrapper
+    wrapper_callback: reward_wrapper.WrappedRewardCallback
 
     def __init__(
         self,
@@ -93,7 +106,7 @@ def __init__(
         self.is_stationary = is_stationary
         self.density_type = density_type
         self.venv = venv
-        self.transitions = {}
+        self.transitions = dict()
         super().__init__(
             demonstrations=demonstrations,
             custom_logger=custom_logger,
@@ -104,7 +117,7 @@ def __init__(
         self.kernel_bandwidth = kernel_bandwidth
         self.standardise = standardise_inputs
         self._scaler = None
-        self._density_models = {}
+        self._density_models = dict()
 
         self.rl_algo = rl_algo
         self.buffering_wrapper = wrappers.BufferingWrapper(self.venv)
@@ -121,6 +134,13 @@ def _set_demo_from_batch(
         next_obs_b: Optional[np.ndarray],
     ) -> None:
         next_obs_b = next_obs_b or itertools.repeat(None)
+        # TODO(juan) I think this assumes that if next_obs_b is an array,
+        #  it has at least two axes, and zip maps along the first one.
+        #  This is not checked for, and if the array has only one dimension
+        #  it would raise an obscure error within _preprocess_transition.
+        #  _preprocess_transition also requires next_obs to be an ndarray
+        #  that is part of the observation space, so not sure how
+        #  this isn't failing when next_obs_b is None.
         for obs, act, next_obs in zip(obs_b, act_b, next_obs_b):
             flat_trans = self._preprocess_transition(obs, act, next_obs)
             self.transitions.setdefault(None, []).append(flat_trans)
@@ -222,12 +242,16 @@ def _preprocess_transition(
         else:
             raise ValueError(f"Unknown density type {self.density_type}")
 
+    # TODO(juan) I opted for renaming the function signature to match the
+    #  rewards.reward_function.RewardFn protocol, but changing the protocol
+    #  or making the arguments positional-only are two other valid approaches
+    #  if the previous names are better suited.
     def __call__(
         self,
-        obs_b: np.ndarray,
-        act_b: np.ndarray,
-        next_obs_b: np.ndarray,
-        dones: np.ndarray,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
         steps: Optional[np.ndarray] = None,
     ) -> np.ndarray:
         r"""Compute reward from given (s,a,s') transition batch.
@@ -236,12 +260,12 @@ def __call__(
         VecEnvs.
 
         Args:
-            obs_b: current batch of observations.
-            act_b: batch of actions that agent took in response to those
+            state: current batch of observations.
+            action: batch of actions that agent took in response to those
                 observations.
-            next_obs_b: batch of observations encountered after the
+            next_state: batch of observations encountered after the
                 agent took those actions.
-            dones: is it terminal state?
+            done: is it terminal state?
             steps: What timestep is this from? Used if `self.is_stationary` is false,
                 otherwise ignored.
 
@@ -258,11 +282,11 @@ def __call__(
         if not self.is_stationary and steps is None:
             raise ValueError("steps must be provided with non-stationary models")
 
-        del dones  # TODO(adam): should we handle terminal state specially in any way?
+        del done  # TODO(adam): should we handle terminal state specially in any way?
 
         rew_list = []
-        assert len(obs_b) == len(act_b) and len(obs_b) == len(next_obs_b)
-        for idx, (obs, act, next_obs) in enumerate(zip(obs_b, act_b, next_obs_b)):
+        assert len(state) == len(action) and len(state) == len(next_state)
+        for idx, (obs, act, next_obs) in enumerate(zip(state, action, next_state)):
             flat_trans = self._preprocess_transition(obs, act, next_obs)
             scaled_padded_trans = self._scaler.transform(flat_trans[np.newaxis])
             if self.is_stationary:

From 2bdc902010f27e6cad8b619086a51a86ec5f272a Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Sun, 4 Sep 2022 12:22:52 +0200
Subject: [PATCH 022/187] Linting

---
 Makefile                                           | 13 ++++++++-----
 src/imitation/algorithms/mce_irl.py                |  6 +++++-
 src/imitation/algorithms/preference_comparisons.py |  6 +++---
 tests/rewards/test_reward_nets.py                  |  2 +-
 4 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/Makefile b/Makefile
index 551d0dac6..75887fad7 100644
--- a/Makefile
+++ b/Makefile
@@ -2,16 +2,19 @@ SRC_FILES=src/ tests/ experiments/ examples/ docs/conf.py setup.py
 
 typecheck:
 	mypy ${SRC_FILES}
-	pytype -j auto "${SRC_FILES}"
-
+	pytype -j auto ${SRC_FILES}
 
 shellcheck:
 	find . -path ./venv -prune -o -name '*.sh' -print0 | xargs -0 shellcheck
 
+format:
+	black ${SRC_FILES}
+	isort ${SRC_FILES}
+
 formatcheck:
-	flake8 --darglint-ignore-regex '.*' "${SRC_FILES[@]}"
-	black --check --diff "${SRC_FILES[@]}"
-	codespell -I .codespell.skip --skip='*.pyc,tests/testdata/*,*.ipynb,*.csv' "${SRC_FILES[@]}"
+	flake8 --darglint-ignore-regex '.*' ${SRC_FILES}
+	black --check --diff ${SRC_FILES}
+	codespell -I .codespell.skip --skip='*.pyc,tests/testdata/*,*.ipynb,*.csv' ${SRC_FILES}
 
 docscheck:
 	pushd docs/
diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index 3534bcef2..d8cf4e108 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -178,7 +178,11 @@ def set_pi(self, pi: np.ndarray) -> None:
     def _predict(self, observation: th.Tensor, deterministic: bool = False):
         raise NotImplementedError("Should never be called as predict overridden.")
 
-    def forward(self, observation: th.Tensor, deterministic: bool = False):  # type: ignore[override]
+    def forward(  # type: ignore[override]
+        self,
+        observation: th.Tensor,
+        deterministic: bool = False,
+    ):
         raise NotImplementedError("Should never be called.")
 
     def predict(  # type: ignore[override]
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index bd1061244..532a83237 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -297,8 +297,8 @@ def sample(self, steps: int) -> Sequence[types.TrajectoryWithRew]:
             exploration_trajs = _get_trajectories(exploration_trajs, exploration_steps)
             # We call _get_trajectories separately on agent_trajs and exploration_trajs
             # and then just concatenate. This could mean we return slightly too many
-            # transitions, but it gets the proportion of exploratory and agent transitions
-            # roughly right.
+            # transitions, but it gets the proportion of exploratory and agent
+            # transitions roughly right.
             trajectories.extend(list(exploration_trajs))
         return trajectories
 
@@ -469,7 +469,7 @@ def rewards(
         next_state = transitions.next_obs
         done = transitions.dones
         if self.is_ensemble:
-            rews = self.model.predict_processed_all(state, action, next_state, done)  # type: ignore[operator]
+            rews = self.model.predict_processed_all(state, action, next_state, done)
             assert rews.shape == (len(state), self.model.num_members)
             return util.safe_to_tensor(rews).to(self.model.device)
         else:
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index 9c51e96be..2323b7459 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -540,7 +540,7 @@ def test_wrappers_pass_on_kwargs(
     )
     # TODO(juan) reassigning a method is very bad practice. I added
     #  type ignore for now but is there not a better way to do this?
-    basic_reward_net.predict_processed = mock.Mock(return_value=np.zeros((10,)))  # type: ignore
+    basic_reward_net.predict_processed = mock.Mock(return_value=np.zeros((10,)))
     wrapped_reward_net = make_wrapper(
         basic_reward_net,
     )

From ccad60bee3ba1fa8ea13c1cd50e7779e04f2fdb6 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Sun, 4 Sep 2022 12:30:08 +0200
Subject: [PATCH 023/187] Linting (add back type ignore after reformatting)

---
 tests/rewards/test_reward_nets.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index 2323b7459..11f58e149 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -540,7 +540,9 @@ def test_wrappers_pass_on_kwargs(
     )
     # TODO(juan) reassigning a method is very bad practice. I added
     #  type ignore for now but is there not a better way to do this?
-    basic_reward_net.predict_processed = mock.Mock(return_value=np.zeros((10,)))
+    basic_reward_net.predict_processed = mock.Mock(  # type: ignore[assignment]
+        return_value=np.zeros((10,))
+    )
     wrapped_reward_net = make_wrapper(
         basic_reward_net,
     )

From a5f23c097ebef0f111fbc772a4ef130395fb4c12 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Sun, 4 Sep 2022 12:52:49 +0200
Subject: [PATCH 024/187] Fixed types: imitation/data/types.py

---
 src/imitation/data/types.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index 363d86b0e..a64c10264 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -345,22 +345,25 @@ def load(path: AnyPath) -> Sequence[Trajectory]:
     # .npz format and the old pickle based format. To tell the difference we need to
     # look at the type of the resulting object. If it's the new compressed format,
     # it should be a Mapping that we need to decode, whereas if it's the old format
-    # it's just the sequence of trajectories and we can return it directly.
+    # it's just the sequence of trajectories, and we can return it directly.
     data = np.load(path, allow_pickle=True)
     if isinstance(data, Sequence):  # old format
         warnings.warn("Loading old version of Trajectory's", DeprecationWarning)
         return data
     elif isinstance(data, Mapping):  # new format
         num_trajs = len(data["indices"])
-        fields = (
+        fields = [
             # Account for the extra obs in each trajectory
             np.split(data["obs"], data["indices"] + np.arange(num_trajs) + 1),
             np.split(data["acts"], data["indices"]),
             np.split(data["infos"], data["indices"]),
             data["terminal"],
-        )
+        ]
         if "rews" in data:
-            fields += (np.split(data["rews"], data["indices"]),)
+            fields = [
+                *fields,
+                np.split(data["rews"], data["indices"]),
+            ]
             return [TrajectoryWithRew(*args) for args in zip(*fields)]
         else:
             return [Trajectory(*args) for args in zip(*fields)]
@@ -395,6 +398,8 @@ def save(path: AnyPath, trajectories: Sequence[Trajectory]):
         ValueError: If the trajectories are not all of the same type, i.e. some are
             `Trajectory` and others are `TrajectoryWithRew`.
     """
+    if isinstance(path, bytes):
+        path = path.decode("utf-8")
     p = pathlib.Path(path)
     p.parent.mkdir(parents=True, exist_ok=True)
     tmp_path = f"{path}.tmp"

From 8ae929c0bdd66a9151d412b9d7fdb5bdd595eed7 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Sun, 4 Sep 2022 12:53:07 +0200
Subject: [PATCH 025/187] Fixed types (started): imitation/data/

---
 src/imitation/data/rollout.py  | 15 +++++++++------
 src/imitation/data/wrappers.py |  6 +++++-
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index 02dc74aed..c3af7fb6f 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -29,6 +29,8 @@ def unwrap_traj(traj: types.TrajectoryWithRew) -> types.TrajectoryWithRew:
     Returns:
         A copy of `traj` with replaced `obs` and `rews` fields.
     """
+    if traj.infos is None:
+        raise ValueError("Trajectory must have infos to unwrap")
     ep_info = traj.infos[-1]["rollout"]
     res = dataclasses.replace(traj, obs=ep_info["obs"], rews=ep_info["rews"])
     assert len(res.obs) == len(res.acts) + 1
@@ -88,11 +90,11 @@ def finish_trajectory(
         del self.partial_trajectories[key]
         out_dict_unstacked = collections.defaultdict(list)
         for part_dict in part_dicts:
-            for key, array in part_dict.items():
-                out_dict_unstacked[key].append(array)
+            for key_, array in part_dict.items():
+                out_dict_unstacked[key_].append(array)
         out_dict_stacked = {
-            key: np.stack(arr_list, axis=0)
-            for key, arr_list in out_dict_unstacked.items()
+            key_: np.stack(arr_list, axis=0)
+            for key_, arr_list in out_dict_unstacked.items()
         }
         traj = types.TrajectoryWithRew(**out_dict_stacked, terminal=terminal)
         assert traj.rews.shape[0] == traj.acts.shape[0] == traj.obs.shape[0] - 1
@@ -281,7 +283,7 @@ def get_actions(states):
             )
             return acts
 
-    elif isinstance(policy, Callable):
+    elif callable(policy):
         # When a policy callable is passed, by default we will use it directly.
         # We are not able to change the determinism of the policy when it is a
         # callable that only takes in the states.
@@ -311,7 +313,7 @@ def generate_trajectories(
     sample_until: GenTrajTerminationFn,
     *,
     deterministic_policy: bool = False,
-    rng: np.random.RandomState = np.random,
+    rng: Optional[np.random.RandomState] = None,
 ) -> Sequence[types.TrajectoryWithRew]:
     """Generate trajectory dictionaries from a policy and an environment.
 
@@ -335,6 +337,7 @@ def generate_trajectories(
         may be collected to avoid biasing process towards short episodes; the user
         should truncate if required.
     """
+    rng = rng or np.random.RandomState()
     get_actions = _policy_to_callable(policy, venv, deterministic_policy)
 
     # Collect rollout tuples.
diff --git a/src/imitation/data/wrappers.py b/src/imitation/data/wrappers.py
index cd5b4ac26..f811586bc 100644
--- a/src/imitation/data/wrappers.py
+++ b/src/imitation/data/wrappers.py
@@ -1,6 +1,6 @@
 """Environment wrappers for collecting rollouts."""
 
-from typing import Sequence, Tuple
+from typing import List, Optional, Sequence, Tuple
 
 import gym
 import numpy as np
@@ -15,6 +15,10 @@ class BufferingWrapper(VecEnvWrapper):
     Retrieve saved transitions using `pop_transitions()`.
     """
 
+    error_on_premature_event: bool
+    _trajectories: List[types.Trajectory]
+    _traj_accum: Optional[rollout.TrajectoryAccumulator]
+
     def __init__(self, venv: VecEnv, error_on_premature_reset: bool = True):
         """Builds BufferingWrapper.
 

From 19f023c4da53abea3d61331a17ef0b8a416b3fea Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 11:08:11 +0200
Subject: [PATCH 026/187] Fixed types: imitation/data/buffer.py

---
 src/imitation/data/buffer.py | 39 +++++++++++++++++++++---------------
 1 file changed, 23 insertions(+), 16 deletions(-)

diff --git a/src/imitation/data/buffer.py b/src/imitation/data/buffer.py
index 3f7f283d8..33fbbbaa7 100644
--- a/src/imitation/data/buffer.py
+++ b/src/imitation/data/buffer.py
@@ -111,8 +111,8 @@ def from_data(
             ValueError: `data` has items mapping to arrays differing in the
                 length of their first axis.
         """
-        data_capacities = [arr.shape[0] for arr in data.values()]
-        data_capacities = np.unique(data_capacities)
+        data_capacities_list = [arr.shape[0] for arr in data.values()]
+        data_capacities = np.unique(data_capacities_list)
         if len(data) == 0:
             raise ValueError("No keys in data.")
         if len(data_capacities) > 1:
@@ -150,11 +150,11 @@ def store(self, data: Mapping[str, np.ndarray], truncate_ok: bool = False) -> No
         if len(unexpected_keys) > 0:
             raise ValueError(f"Unexpected keys {unexpected_keys}")
 
-        n_samples = [arr.shape[0] for arr in data.values()]
-        n_samples = np.unique(n_samples)
-        if len(n_samples) > 1:
+        n_samples_list = [arr.shape[0] for arr in data.values()]
+        n_samples_np = np.unique(n_samples_list)
+        if len(n_samples_np) > 1:
             raise ValueError("Keys map to different length values.")
-        n_samples = n_samples[0]
+        n_samples = int(n_samples_np[0])
 
         if n_samples == 0:
             raise ValueError("Trying to store empty data.")
@@ -192,10 +192,10 @@ def _store_easy(self, data: Mapping[str, np.ndarray]) -> None:
             data: Same as in `self.store`'s docstring, except with the additional
                 constraint `size(data) <= self.capacity - self._idx`.
         """
-        n_samples = [arr.shape[0] for arr in data.values()]
-        n_samples = np.unique(n_samples)
-        assert len(n_samples) == 1
-        n_samples = n_samples[0]
+        n_samples_list = [arr.shape[0] for arr in data.values()]
+        n_samples_np = np.unique(n_samples_list)
+        assert len(n_samples_np) == 1
+        n_samples = int(n_samples_np[0])
 
         assert n_samples <= self.capacity - self._idx
         idx_hi = self._idx + n_samples
@@ -222,7 +222,7 @@ def sample(self, n_samples: int) -> Mapping[str, np.ndarray]:
         ind = np.random.randint(self.size(), size=n_samples)
         return {k: buffer[ind] for k, buffer in self._arrays.items()}
 
-    def size(self) -> Optional[int]:
+    def size(self) -> int:
         """Returns the number of samples stored in the buffer."""
         assert 0 <= self._n_data <= self.capacity
         return self._n_data
@@ -262,16 +262,23 @@ def __init__(
         """
         params = [obs_shape, act_shape, obs_dtype, act_dtype]
         if venv is not None:
-            if np.any([x is not None for x in params]):
-                raise ValueError("Specified shape or dtype and environment.")
+            if not all([x is None for x in params]):
+                raise ValueError(
+                    "Cannot specify both shape/dtype and also environment."
+                )
             obs_shape = tuple(venv.observation_space.shape)
             act_shape = tuple(venv.action_space.shape)
             obs_dtype = venv.observation_space.dtype
             act_dtype = venv.action_space.dtype
         else:
-            if np.any([x is None for x in params]):
+            if None in params:
                 raise ValueError("Shape or dtype missing and no environment specified.")
 
+        assert obs_shape is not None
+        assert act_shape is not None
+        assert obs_dtype is not None
+        assert act_dtype is not None
+
         self.capacity = capacity
         sample_shapes = {
             "obs": obs_shape,
@@ -284,8 +291,8 @@ def __init__(
             "obs": obs_dtype,
             "acts": act_dtype,
             "next_obs": obs_dtype,
-            "dones": bool,
-            "infos": np.object,
+            "dones": np.dtype(bool),
+            "infos": np.dtype(object),
         }
         self._buffer = Buffer(capacity, sample_shapes=sample_shapes, dtypes=dtypes)
 

From 184694597869c12d7efd007443af14040e886c95 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 11:33:57 +0200
Subject: [PATCH 027/187] Fixed bug in buffer.py

---
 src/imitation/data/buffer.py | 2 +-
 tests/data/test_buffer.py    | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/imitation/data/buffer.py b/src/imitation/data/buffer.py
index 33fbbbaa7..3b6148093 100644
--- a/src/imitation/data/buffer.py
+++ b/src/imitation/data/buffer.py
@@ -271,7 +271,7 @@ def __init__(
             obs_dtype = venv.observation_space.dtype
             act_dtype = venv.action_space.dtype
         else:
-            if None in params:
+            if any([x is None for x in params]):
                 raise ValueError("Shape or dtype missing and no environment specified.")
 
         assert obs_shape is not None
diff --git a/tests/data/test_buffer.py b/tests/data/test_buffer.py
index b64cf5fc7..632666f39 100644
--- a/tests/data/test_buffer.py
+++ b/tests/data/test_buffer.py
@@ -213,7 +213,9 @@ def test_buffer_init_errors():
 
 
 def test_replay_buffer_init_errors():
-    with pytest.raises(ValueError, match=r"Specified.* and environment"):
+    with pytest.raises(
+        ValueError, match=r"Cannot specify both shape/dtype and also environment"
+    ):
         ReplayBuffer(15, venv="MockEnv", obs_shape=(10, 10))
     with pytest.raises(ValueError, match=r"Shape or dtype missing.*"):
         ReplayBuffer(15, obs_shape=(10, 10), act_shape=(15,), obs_dtype=bool)

From ca4e844739dea911f26662e3e90a2f4d9e042003 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 11:34:35 +0200
Subject: [PATCH 028/187] Fixed types: imitation/data/rollout.py

---
 src/imitation/data/rollout.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index c3af7fb6f..a9ded8bcf 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -361,9 +361,11 @@ def generate_trajectories(
     #
     # To start with, all environments are active.
     active = np.ones(venv.num_envs, dtype=bool)
+    assert isinstance(obs, np.ndarray)
     while np.any(active):
         acts = get_actions(obs)
         obs, rews, dones, infos = venv.step(acts)
+        assert isinstance(obs, np.ndarray)
 
         # If an environment is inactive, i.e. the episode completed for that
         # environment after `sample_until(trajectories)` was true, then we do
@@ -392,7 +394,7 @@ def generate_trajectories(
     # `trajectories` sooner. Shuffle to avoid bias in order. This is important
     # when callees end up truncating the number of trajectories or transitions.
     # It is also cheap, since we're just shuffling pointers.
-    rng.shuffle(trajectories)
+    rng.shuffle(trajectories)  # type: ignore[arg-type]
 
     # Sanity checks.
     for trajectory in trajectories:
@@ -477,7 +479,7 @@ def flatten_trajectories(
         The trajectories flattened into a single batch of Transitions.
     """
     keys = ["obs", "next_obs", "acts", "dones", "infos"]
-    parts = {key: [] for key in keys}
+    parts: Mapping[str, List[np.ndarray]] = {key: [] for key in keys}
     for traj in trajectories:
         parts["acts"].append(traj.acts)
 

From 6204d4c9ff06c7e911a96323784348a479a31572 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 11:46:44 +0200
Subject: [PATCH 029/187] Fixed types: imitation/data/wrappers.py

---
 src/imitation/data/wrappers.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/imitation/data/wrappers.py b/src/imitation/data/wrappers.py
index f811586bc..380115438 100644
--- a/src/imitation/data/wrappers.py
+++ b/src/imitation/data/wrappers.py
@@ -4,6 +4,7 @@
 
 import gym
 import numpy as np
+import numpy.typing as npt
 from stable_baselines3.common.vec_env import VecEnv, VecEnvWrapper
 
 from imitation.data import rollout, types
@@ -16,8 +17,12 @@ class BufferingWrapper(VecEnvWrapper):
     """
 
     error_on_premature_event: bool
-    _trajectories: List[types.Trajectory]
+    _trajectories: List[types.TrajectoryWithRew]
+    _ep_lens: List[int]
+    _init_reset: bool
     _traj_accum: Optional[rollout.TrajectoryAccumulator]
+    _timesteps: Optional[npt.NDArray[np.int_]]
+    n_transitions: Optional[int]
 
     def __init__(self, venv: VecEnv, error_on_premature_reset: bool = True):
         """Builds BufferingWrapper.
@@ -39,9 +44,9 @@ def __init__(self, venv: VecEnv, error_on_premature_reset: bool = True):
 
     def reset(self, **kwargs):
         if (
-            self._init_reset
-            and self.error_on_premature_reset
-            and self.n_transitions > 0
+                self._init_reset
+                and self.error_on_premature_reset
+                and self.n_transitions > 0
         ):  # noqa: E127
             raise RuntimeError("BufferingWrapper reset() before samples were accessed")
         self._init_reset = True
@@ -85,6 +90,7 @@ def step_wait(self):
 
     def _finish_partial_trajectories(self) -> Sequence[types.TrajectoryWithRew]:
         """Finishes and returns partial trajectories in `self._traj_accum`."""
+        assert self._traj_accum is not None
         trajs = []
         for i in range(self.num_envs):
             # Check that we have any transitions at all.
@@ -103,7 +109,7 @@ def _finish_partial_trajectories(self) -> Sequence[types.TrajectoryWithRew]:
         return trajs
 
     def pop_finished_trajectories(
-        self,
+            self,
     ) -> Tuple[Sequence[types.TrajectoryWithRew], Sequence[int]]:
         """Pops recorded complete trajectories `trajs` and episode lengths `ep_lens`.
 
@@ -122,7 +128,7 @@ def pop_finished_trajectories(
         return trajectories, ep_lens
 
     def pop_trajectories(
-        self,
+            self,
     ) -> Tuple[Sequence[types.TrajectoryWithRew], Sequence[int]]:
         """Pops recorded trajectories `trajs` and episode lengths `ep_lens`.
 

From 44aa867666381dceac14375fe10c54628bb84af3 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 12:16:39 +0200
Subject: [PATCH 030/187] Improve makefile to support automatic cache cleaning

---
 Makefile | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/Makefile b/Makefile
index 75887fad7..b6d0e79c1 100644
--- a/Makefile
+++ b/Makefile
@@ -17,7 +17,24 @@ formatcheck:
 	codespell -I .codespell.skip --skip='*.pyc,tests/testdata/*,*.ipynb,*.csv' ${SRC_FILES}
 
 docscheck:
-	pushd docs/
-	make clean
-	make html
-	popd
\ No newline at end of file
+	pushd docs/ \
+	&& make clean \
+	&& make html \
+	&& popd
+
+cleandocs:
+	pushd docs/ \
+	&& make clean \
+	&& popd
+
+cleantests:
+	rm -rf output/
+	rm -rf quickstart/
+
+cleancache:
+	rm -rf .mypy_cache/
+	rm -rf .pytest_cache/
+	rm -rf .pytype/
+	rm -rf .hypothesis/
+
+clean: cleandocs cleancache cleantests
\ No newline at end of file

From 955c0182ec6ca12830ba1d8e9077c19ca57bbf4a Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 12:17:07 +0200
Subject: [PATCH 031/187] Fixed types: imitation/testing/

---
 src/imitation/data/wrappers.py              | 10 +++++-----
 src/imitation/testing/reward_improvement.py |  3 ++-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/imitation/data/wrappers.py b/src/imitation/data/wrappers.py
index 380115438..09ad42247 100644
--- a/src/imitation/data/wrappers.py
+++ b/src/imitation/data/wrappers.py
@@ -44,9 +44,9 @@ def __init__(self, venv: VecEnv, error_on_premature_reset: bool = True):
 
     def reset(self, **kwargs):
         if (
-                self._init_reset
-                and self.error_on_premature_reset
-                and self.n_transitions > 0
+            self._init_reset
+            and self.error_on_premature_reset
+            and self.n_transitions > 0
         ):  # noqa: E127
             raise RuntimeError("BufferingWrapper reset() before samples were accessed")
         self._init_reset = True
@@ -109,7 +109,7 @@ def _finish_partial_trajectories(self) -> Sequence[types.TrajectoryWithRew]:
         return trajs
 
     def pop_finished_trajectories(
-            self,
+        self,
     ) -> Tuple[Sequence[types.TrajectoryWithRew], Sequence[int]]:
         """Pops recorded complete trajectories `trajs` and episode lengths `ep_lens`.
 
@@ -128,7 +128,7 @@ def pop_finished_trajectories(
         return trajectories, ep_lens
 
     def pop_trajectories(
-            self,
+        self,
     ) -> Tuple[Sequence[types.TrajectoryWithRew], Sequence[int]]:
         """Pops recorded trajectories `trajs` and episode lengths `ep_lens`.
 
diff --git a/src/imitation/testing/reward_improvement.py b/src/imitation/testing/reward_improvement.py
index c8ad3780b..1b7028700 100644
--- a/src/imitation/testing/reward_improvement.py
+++ b/src/imitation/testing/reward_improvement.py
@@ -66,4 +66,5 @@ def mean_reward_improved_by(
     >>> mean_reward_improved_by([5, 8, 7], [8, 9, 10], 5)
     False
     """
-    return np.mean(new_rewards) - np.mean(old_rewards) >= min_improvement
+    improvement = np.mean(new_rewards) - np.mean(old_rewards)  # type: ignore
+    return improvement >= min_improvement

From 58535b85563ea380b54decbc88863243bab63600 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 12:37:59 +0200
Subject: [PATCH 032/187] Linting, fixed wrong return type in
 rewards.predict_processed_all

---
 src/imitation/data/buffer.py         | 2 +-
 src/imitation/rewards/reward_nets.py | 4 ++--
 tests/data/test_buffer.py            | 3 ++-
 tests/rewards/test_reward_nets.py    | 2 +-
 4 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/imitation/data/buffer.py b/src/imitation/data/buffer.py
index 3b6148093..33bef6afd 100644
--- a/src/imitation/data/buffer.py
+++ b/src/imitation/data/buffer.py
@@ -264,7 +264,7 @@ def __init__(
         if venv is not None:
             if not all([x is None for x in params]):
                 raise ValueError(
-                    "Cannot specify both shape/dtype and also environment."
+                    "Cannot specify both shape/dtype and also environment.",
                 )
             obs_shape = tuple(venv.observation_space.shape)
             act_shape = tuple(venv.action_space.shape)
diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index 2141122ba..85e04cce9 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -686,7 +686,7 @@ def predict_processed_all(
         next_state: np.ndarray,
         done: np.ndarray,
         **kwargs,
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    ) -> np.ndarray:
         """Get the results of predict processed on all of the members.
 
         Args:
@@ -705,7 +705,7 @@ def predict_processed_all(
             member.predict_processed(state, action, next_state, done, **kwargs)
             for member in self.members
         ]
-        rewards = np.stack(rewards, axis=-1)
+        rewards: np.ndarray = np.stack(rewards, axis=-1)
         assert rewards.shape == (batch_size, self.num_members)
         return rewards
 
diff --git a/tests/data/test_buffer.py b/tests/data/test_buffer.py
index 632666f39..db6c52c3f 100644
--- a/tests/data/test_buffer.py
+++ b/tests/data/test_buffer.py
@@ -214,7 +214,8 @@ def test_buffer_init_errors():
 
 def test_replay_buffer_init_errors():
     with pytest.raises(
-        ValueError, match=r"Cannot specify both shape/dtype and also environment"
+        ValueError,
+        match=r"Cannot specify both shape/dtype and also environment",
     ):
         ReplayBuffer(15, venv="MockEnv", obs_shape=(10, 10))
     with pytest.raises(ValueError, match=r"Shape or dtype missing.*"):
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index 11f58e149..93bc6c79a 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -541,7 +541,7 @@ def test_wrappers_pass_on_kwargs(
     # TODO(juan) reassigning a method is very bad practice. I added
     #  type ignore for now but is there not a better way to do this?
     basic_reward_net.predict_processed = mock.Mock(  # type: ignore[assignment]
-        return_value=np.zeros((10,))
+        return_value=np.zeros((10,)),
     )
     wrapped_reward_net = make_wrapper(
         basic_reward_net,

From 18a533f0d6dc87b6547308d03a2fb273dfa9f948 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 13:09:02 +0200
Subject: [PATCH 033/187] Fixed types: imitation/policies/

---
 src/imitation/policies/base.py                | 13 ++++---
 .../policies/replay_buffer_wrapper.py         | 31 +++++++++--------
 src/imitation/policies/serialize.py           | 34 +++++++++++--------
 3 files changed, 45 insertions(+), 33 deletions(-)

diff --git a/src/imitation/policies/base.py b/src/imitation/policies/base.py
index 3f9b0d919..5f15906d1 100644
--- a/src/imitation/policies/base.py
+++ b/src/imitation/policies/base.py
@@ -92,9 +92,9 @@ class NormalizeFeaturesExtractor(torch_layers.FlattenExtractor):
     """Feature extractor that flattens then normalizes input."""
 
     def __init__(
-        self,
-        observation_space: gym.Space,
-        normalize_class: Type[nn.Module] = networks.RunningNorm,
+            self,
+            observation_space: gym.Space,
+            normalize_class: Type[nn.Module] = networks.RunningNorm,
     ):
         """Builds NormalizeFeaturesExtractor.
 
@@ -105,7 +105,12 @@ def __init__(
                 e.g. `nn.BatchNorm*` or `nn.LayerNorm`.
         """
         super().__init__(observation_space)
-        self.normalize = normalize_class(self.features_dim)
+        # Below we have to ignore the type error when initializing the class because
+        # there is no simple way of specifying a protocol that admits one positional
+        # argument for the number of features while being compatible with nn.Module.
+        # (it would require defining a base class and forcing all the subclasses
+        # to inherit from it).
+        self.normalize = normalize_class(self.features_dim)  # type: ignore[call-arg]
 
     def forward(self, observations: th.Tensor) -> th.Tensor:
         flattened = super().forward(observations)
diff --git a/src/imitation/policies/replay_buffer_wrapper.py b/src/imitation/policies/replay_buffer_wrapper.py
index 9bb011063..590df9972 100644
--- a/src/imitation/policies/replay_buffer_wrapper.py
+++ b/src/imitation/policies/replay_buffer_wrapper.py
@@ -1,6 +1,5 @@
 """Wrapper for reward labeling for transitions sampled from a replay buffer."""
 
-
 from typing import Mapping, Type
 
 import numpy as np
@@ -13,7 +12,7 @@
 
 
 def _samples_to_reward_fn_input(
-    samples: ReplayBufferSamples,
+        samples: ReplayBufferSamples,
 ) -> Mapping[str, np.ndarray]:
     """Convert a sample from a replay buffer to a numpy array."""
     return dict(
@@ -28,14 +27,14 @@ class ReplayBufferRewardWrapper(BaseBuffer):
     """Relabel the rewards in transitions sampled from a ReplayBuffer."""
 
     def __init__(
-        self,
-        buffer_size: int,
-        observation_space: spaces.Space,
-        action_space: spaces.Space,
-        *,
-        replay_buffer_class: Type[ReplayBuffer],
-        reward_fn: RewardFn,
-        **kwargs,
+            self,
+            buffer_size: int,
+            observation_space: spaces.Space,
+            action_space: spaces.Space,
+            *,
+            replay_buffer_class: Type[ReplayBuffer],
+            reward_fn: RewardFn,
+            **kwargs,
     ):
         """Builds ReplayBufferRewardWrapper.
 
@@ -63,16 +62,20 @@ def __init__(
         _base_kwargs = {k: v for k, v in kwargs.items() if k in ["device", "n_envs"]}
         super().__init__(buffer_size, observation_space, action_space, **_base_kwargs)
 
-    @property
-    def pos(self) -> int:
+    # TODO(juan) remove the type ignore once the merged PR
+    #  https://github.com/python/mypy/pull/13475
+    #  is released into a mypy version on pypi.
+
+    @property  # type: ignore[override]
+    def pos(self) -> int:  # type: ignore[override]
         return self.replay_buffer.pos
 
     @pos.setter
     def pos(self, pos: int):
         self.replay_buffer.pos = pos
 
-    @property
-    def full(self) -> bool:
+    @property  # type: ignore[override]
+    def full(self) -> bool:  # type: ignore[override]
         return self.replay_buffer.full
 
     @full.setter
diff --git a/src/imitation/policies/serialize.py b/src/imitation/policies/serialize.py
index e43651a0e..232dc7ac0 100644
--- a/src/imitation/policies/serialize.py
+++ b/src/imitation/policies/serialize.py
@@ -21,10 +21,10 @@
 
 
 def load_stable_baselines_model(
-    cls: Type[Algorithm],
-    path: str,
-    venv: vec_env.VecEnv,
-    **kwargs,
+        cls: Type[Algorithm],
+        path: str,
+        venv: vec_env.VecEnv,
+        **kwargs,
 ) -> Algorithm:
     """Helper method to load RL models from Stable Baselines.
 
@@ -57,11 +57,14 @@ def load_stable_baselines_model(
         )
 
     model_path = policy_dir / "model.zip"
-    return cls.load(model_path, env=venv, **kwargs)
+    # TODO(juan) remove the type ignore when this SB3 PR gets merged
+    #  and released:
+    #  https://github.com/DLR-RM/stable-baselines3/pull/1043
+    return cls.load(model_path, env=venv, **kwargs)  # type: ignore[return-value]
 
 
 def _load_stable_baselines(
-    cls: Type[base_class.BaseAlgorithm],
+        cls: Type[base_class.BaseAlgorithm],
 ) -> PolicyLoaderFn:
     """Higher-order function, returning a policy loading function.
 
@@ -105,9 +108,9 @@ def _add_stable_baselines_policies(classes):
 
 
 def load_policy(
-    policy_type: str,
-    policy_path: str,
-    venv: vec_env.VecEnv,
+        policy_type: str,
+        policy_path: str,
+        venv: vec_env.VecEnv,
 ) -> policies.BasePolicy:
     """Load serialized policy.
 
@@ -124,8 +127,8 @@ def load_policy(
 
 
 def save_stable_model(
-    output_dir: str,
-    model: base_class.BaseAlgorithm,
+        output_dir: str,
+        model: base_class.BaseAlgorithm,
 ) -> None:
     """Serialize Stable Baselines model.
 
@@ -151,10 +154,10 @@ class SavePolicyCallback(callbacks.EventCallback):
     """
 
     def __init__(
-        self,
-        policy_dir: str,
-        *args,
-        **kwargs,
+            self,
+            policy_dir: str,
+            *args,
+            **kwargs,
     ):
         """Builds SavePolicyCallback.
 
@@ -167,6 +170,7 @@ def __init__(
         self.policy_dir = policy_dir
 
     def _on_step(self) -> bool:
+        assert self.model is not None
         output_dir = os.path.join(self.policy_dir, f"{self.num_timesteps:012d}")
         save_stable_model(output_dir, self.model)
         return True

From c2786cf310b9fbec6d0b2cf26e9f9a1af6a8004a Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 13:17:23 +0200
Subject: [PATCH 034/187] Formatting

---
 src/imitation/policies/base.py                |  6 ++--
 .../policies/replay_buffer_wrapper.py         | 18 ++++++------
 src/imitation/policies/serialize.py           | 28 +++++++++----------
 3 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/src/imitation/policies/base.py b/src/imitation/policies/base.py
index 5f15906d1..f0ee6d588 100644
--- a/src/imitation/policies/base.py
+++ b/src/imitation/policies/base.py
@@ -92,9 +92,9 @@ class NormalizeFeaturesExtractor(torch_layers.FlattenExtractor):
     """Feature extractor that flattens then normalizes input."""
 
     def __init__(
-            self,
-            observation_space: gym.Space,
-            normalize_class: Type[nn.Module] = networks.RunningNorm,
+        self,
+        observation_space: gym.Space,
+        normalize_class: Type[nn.Module] = networks.RunningNorm,
     ):
         """Builds NormalizeFeaturesExtractor.
 
diff --git a/src/imitation/policies/replay_buffer_wrapper.py b/src/imitation/policies/replay_buffer_wrapper.py
index 590df9972..2624be335 100644
--- a/src/imitation/policies/replay_buffer_wrapper.py
+++ b/src/imitation/policies/replay_buffer_wrapper.py
@@ -12,7 +12,7 @@
 
 
 def _samples_to_reward_fn_input(
-        samples: ReplayBufferSamples,
+    samples: ReplayBufferSamples,
 ) -> Mapping[str, np.ndarray]:
     """Convert a sample from a replay buffer to a numpy array."""
     return dict(
@@ -27,14 +27,14 @@ class ReplayBufferRewardWrapper(BaseBuffer):
     """Relabel the rewards in transitions sampled from a ReplayBuffer."""
 
     def __init__(
-            self,
-            buffer_size: int,
-            observation_space: spaces.Space,
-            action_space: spaces.Space,
-            *,
-            replay_buffer_class: Type[ReplayBuffer],
-            reward_fn: RewardFn,
-            **kwargs,
+        self,
+        buffer_size: int,
+        observation_space: spaces.Space,
+        action_space: spaces.Space,
+        *,
+        replay_buffer_class: Type[ReplayBuffer],
+        reward_fn: RewardFn,
+        **kwargs,
     ):
         """Builds ReplayBufferRewardWrapper.
 
diff --git a/src/imitation/policies/serialize.py b/src/imitation/policies/serialize.py
index 232dc7ac0..d03513c43 100644
--- a/src/imitation/policies/serialize.py
+++ b/src/imitation/policies/serialize.py
@@ -21,10 +21,10 @@
 
 
 def load_stable_baselines_model(
-        cls: Type[Algorithm],
-        path: str,
-        venv: vec_env.VecEnv,
-        **kwargs,
+    cls: Type[Algorithm],
+    path: str,
+    venv: vec_env.VecEnv,
+    **kwargs,
 ) -> Algorithm:
     """Helper method to load RL models from Stable Baselines.
 
@@ -64,7 +64,7 @@ def load_stable_baselines_model(
 
 
 def _load_stable_baselines(
-        cls: Type[base_class.BaseAlgorithm],
+    cls: Type[base_class.BaseAlgorithm],
 ) -> PolicyLoaderFn:
     """Higher-order function, returning a policy loading function.
 
@@ -108,9 +108,9 @@ def _add_stable_baselines_policies(classes):
 
 
 def load_policy(
-        policy_type: str,
-        policy_path: str,
-        venv: vec_env.VecEnv,
+    policy_type: str,
+    policy_path: str,
+    venv: vec_env.VecEnv,
 ) -> policies.BasePolicy:
     """Load serialized policy.
 
@@ -127,8 +127,8 @@ def load_policy(
 
 
 def save_stable_model(
-        output_dir: str,
-        model: base_class.BaseAlgorithm,
+    output_dir: str,
+    model: base_class.BaseAlgorithm,
 ) -> None:
     """Serialize Stable Baselines model.
 
@@ -154,10 +154,10 @@ class SavePolicyCallback(callbacks.EventCallback):
     """
 
     def __init__(
-            self,
-            policy_dir: str,
-            *args,
-            **kwargs,
+        self,
+        policy_dir: str,
+        *args,
+        **kwargs,
     ):
         """Builds SavePolicyCallback.
 

From 4c0a0738cdd5687faacafeb1554c9bb11941dad1 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 13:17:45 +0200
Subject: [PATCH 035/187] Fixed types: imitation/rewards/

---
 src/imitation/rewards/serialize.py | 56 ++++++++++++++++--------------
 1 file changed, 29 insertions(+), 27 deletions(-)

diff --git a/src/imitation/rewards/serialize.py b/src/imitation/rewards/serialize.py
index 54b674e6f..d580cfa6a 100644
--- a/src/imitation/rewards/serialize.py
+++ b/src/imitation/rewards/serialize.py
@@ -1,19 +1,12 @@
 """Load serialized reward functions of different types."""
 
-from typing import Any, Callable, Iterable, Sequence, Type, Union
+from typing import Any, Callable, Iterable, Sequence, Type, Union, cast
 
 import numpy as np
 import torch as th
 from stable_baselines3.common.vec_env import VecEnv
 
-from imitation.rewards import reward_function
-from imitation.rewards.reward_nets import (
-    AddSTDRewardWrapper,
-    NormalizedRewardNet,
-    RewardNet,
-    RewardNetWrapper,
-    ShapedRewardNet,
-)
+from imitation.rewards import reward_function, reward_nets
 from imitation.util import registry, util
 
 # TODO(sam): I suspect this whole file can be replaced with th.load calls. Try
@@ -55,9 +48,9 @@ def __call__(
 
 
 def _strip_wrappers(
-    reward_net: RewardNet,
-    wrapper_types: Iterable[Type[RewardNetWrapper]],
-) -> RewardNet:
+    reward_net: reward_nets.RewardNet,
+    wrapper_types: Iterable[Type[reward_nets.RewardNetWrapper]],
+) -> reward_nets.RewardNet:
     """Attempts to remove provided wrappers.
 
     Strips wrappers of type `wrapper_type` from `reward_net` in order until either the
@@ -74,7 +67,7 @@ def _strip_wrappers(
     for wrapper_type in wrapper_types:
         assert issubclass(
             wrapper_type,
-            RewardNetWrapper,
+            reward_nets.RewardNetWrapper,
         ), f"trying to remove non-wrapper type {wrapper_type}"
 
         if isinstance(reward_net, wrapper_type):
@@ -86,7 +79,7 @@ def _strip_wrappers(
 
 
 def _make_functional(
-    net: RewardNet,
+    net: reward_nets.RewardNet,
     attr: str = "predict",
     default_kwargs=None,
     **kwargs,
@@ -97,7 +90,7 @@ def _make_functional(
     return lambda *args: getattr(net, attr)(*args, **default_kwargs)
 
 
-WrapperPrefix = Sequence[Type[RewardNet]]
+WrapperPrefix = Sequence[Type[reward_nets.RewardNet]]
 
 
 def _prefix_matches(wrappers: Sequence[Type[Any]], prefix: Sequence[Type[Any]]):
@@ -120,9 +113,9 @@ def _prefix_matches(wrappers: Sequence[Type[Any]], prefix: Sequence[Type[Any]]):
 
 
 def _validate_wrapper_structure(
-    reward_net: Union[RewardNet, RewardNetWrapper],
+    reward_net: Union[reward_nets.RewardNet, reward_nets.RewardNetWrapper],
     prefixes: Iterable[WrapperPrefix],
-) -> RewardNet:
+) -> reward_nets.RewardNet:
     """Reward net if it has a valid structure.
 
     A wrapper prefix specifies, from outermost to innermost, which wrappers must
@@ -155,7 +148,7 @@ def _validate_wrapper_structure(
     wrappers = []
     while hasattr(wrapper, "base"):
         wrappers.append(wrapper.__class__)
-        wrapper = wrapper.base
+        wrapper = cast(reward_nets.RewardNet, wrapper.base)
     wrappers.append(wrapper.__class__)  # append the final reward net
 
     if any(_prefix_matches(wrappers, prefix) for prefix in prefixes):
@@ -198,16 +191,19 @@ def f(
     key="RewardNet_shaped",
     value=lambda path, _, **kwargs: ValidateRewardFn(
         _make_functional(
-            _validate_wrapper_structure(th.load(str(path)), {(ShapedRewardNet,)}),
+            _validate_wrapper_structure(
+                th.load(str(path)), {(reward_nets.ShapedRewardNet,)}
+            ),
         ),
     ),
 )
 
-
 reward_registry.register(
     key="RewardNet_unshaped",
     value=lambda path, _, **kwargs: ValidateRewardFn(
-        _make_functional(_strip_wrappers(th.load(str(path)), (ShapedRewardNet,))),
+        _make_functional(
+            _strip_wrappers(th.load(str(path)), (reward_nets.ShapedRewardNet,))
+        ),
     ),
 )
 
@@ -215,7 +211,9 @@ def f(
     key="RewardNet_normalized",
     value=lambda path, _, **kwargs: ValidateRewardFn(
         _make_functional(
-            _validate_wrapper_structure(th.load(str(path)), {(NormalizedRewardNet,)}),
+            _validate_wrapper_structure(
+                th.load(str(path)), {(reward_nets.NormalizedRewardNet,)}
+            ),
             attr="predict_processed",
             default_kwargs={"update_stats": False},
             **kwargs,
@@ -226,7 +224,9 @@ def f(
 reward_registry.register(
     key="RewardNet_unnormalized",
     value=lambda path, _, **kwargs: ValidateRewardFn(
-        _make_functional(_strip_wrappers(th.load(str(path)), (NormalizedRewardNet,))),
+        _make_functional(
+            _strip_wrappers(th.load(str(path)), (reward_nets.NormalizedRewardNet,))
+        ),
     ),
 )
 
@@ -238,11 +238,14 @@ def f(
                 _validate_wrapper_structure(
                     th.load(str(path)),
                     {
-                        (AddSTDRewardWrapper,),
-                        (NormalizedRewardNet, AddSTDRewardWrapper),
+                        (reward_nets.AddSTDRewardWrapper,),
+                        (
+                            reward_nets.NormalizedRewardNet,
+                            reward_nets.AddSTDRewardWrapper,
+                        ),
                     },
                 ),
-                (NormalizedRewardNet,),
+                (reward_nets.NormalizedRewardNet,),
             ),
             attr="predict_processed",
             default_kwargs={},
@@ -251,7 +254,6 @@ def f(
     ),
 )
 
-
 reward_registry.register(key="zero", value=load_zero)
 
 

From 44d7130b890713aa05ed42c2678e770cd7892cb3 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 13:59:45 +0200
Subject: [PATCH 036/187] Fixed types: imitation/rewards/

---
 src/imitation/rewards/reward_nets.py    | 332 ++++++++++++------------
 src/imitation/rewards/reward_wrapper.py |  19 +-
 2 files changed, 177 insertions(+), 174 deletions(-)

diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index 85e04cce9..27490d25e 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -1,7 +1,7 @@
 """Constructs deep network reward models."""
 
 import abc
-from typing import Callable, Iterable, Optional, Sequence, Tuple, Type
+from typing import Callable, Iterable, Optional, Sequence, Tuple, Type, cast
 
 import gym
 import numpy as np
@@ -10,6 +10,7 @@
 from torch import nn
 
 from imitation.util import networks, util
+from imitation.util.networks import BaseNorm
 
 
 class RewardNet(nn.Module, abc.ABC):
@@ -20,10 +21,10 @@ class RewardNet(nn.Module, abc.ABC):
     """
 
     def __init__(
-        self,
-        observation_space: gym.Space,
-        action_space: gym.Space,
-        normalize_images: bool = True,
+            self,
+            observation_space: gym.Space,
+            action_space: gym.Space,
+            normalize_images: bool = True,
     ):
         """Initialize the RewardNet.
 
@@ -40,20 +41,20 @@ def __init__(
 
     @abc.abstractmethod
     def forward(
-        self,
-        state: th.Tensor,
-        action: th.Tensor,
-        next_state: th.Tensor,
-        done: th.Tensor,
+            self,
+            state: th.Tensor,
+            action: th.Tensor,
+            next_state: th.Tensor,
+            done: th.Tensor,
     ) -> th.Tensor:
         """Compute rewards for a batch of transitions and keep gradients."""
 
     def preprocess(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
     ) -> Tuple[th.Tensor, th.Tensor, th.Tensor, th.Tensor]:
         """Preprocess a batch of input transitions and convert it to PyTorch tensors.
 
@@ -82,21 +83,21 @@ def preprocess(
         del state, action, next_state, done  # unused
 
         # preprocess
-        state_th = preprocessing.preprocess_obs(
+        state_th = cast(th.Tensor, preprocessing.preprocess_obs(
             state_th,
             self.observation_space,
             self.normalize_images,
-        )
-        action_th = preprocessing.preprocess_obs(
+        ))
+        action_th = cast(th.Tensor, preprocessing.preprocess_obs(
             action_th,
             self.action_space,
             self.normalize_images,
-        )
-        next_state_th = preprocessing.preprocess_obs(
+        ))
+        next_state_th = cast(th.Tensor, preprocessing.preprocess_obs(
             next_state_th,
             self.observation_space,
             self.normalize_images,
-        )
+        ))
         done_th = done_th.to(th.float32)
 
         n_gen = len(state_th)
@@ -106,11 +107,11 @@ def preprocess(
         return state_th, action_th, next_state_th, done_th
 
     def predict_th(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
     ) -> th.Tensor:
         """Compute th.Tensor rewards for a batch of transitions without gradients.
 
@@ -141,11 +142,11 @@ def predict_th(
             return rew_th
 
     def predict(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
     ) -> np.ndarray:
         """Compute rewards for a batch of transitions without gradients.
 
@@ -164,12 +165,12 @@ def predict(
         return rew_th.detach().cpu().numpy().flatten()
 
     def predict_processed(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
-        **kwargs,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
+            **kwargs,
     ) -> np.ndarray:
         """Compute the processed rewards for a batch of transitions without gradients.
 
@@ -221,8 +222,8 @@ class RewardNetWrapper(RewardNet):
     """
 
     def __init__(
-        self,
-        base: RewardNet,
+            self,
+            base: RewardNet,
     ):
         """Initialize a RewardNet wrapper.
 
@@ -241,52 +242,52 @@ def base(self) -> RewardNet:
         return self._base
 
     def forward(
-        self,
-        state: th.Tensor,
-        action: th.Tensor,
-        next_state: th.Tensor,
-        done: th.Tensor,
+            self,
+            state: th.Tensor,
+            action: th.Tensor,
+            next_state: th.Tensor,
+            done: th.Tensor,
     ) -> th.Tensor:
         __doc__ = super().forward.__doc__  # noqa: F841
         return self.base.forward(state, action, next_state, done)
 
     def predict_processed(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
-        **kwargs,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
+            **kwargs,
     ) -> np.ndarray:
         __doc__ = super().predict_processed.__doc__  # noqa: F841
         return self.base.predict_processed(state, action, next_state, done, **kwargs)
 
     def predict(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
     ) -> np.ndarray:
         __doc__ = super().predict.__doc__  # noqa: F841
         return self.base.predict(state, action, next_state, done)
 
     def predict_th(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
     ) -> th.Tensor:
         __doc__ = super().predict_th.__doc__  # noqa: F841
         return self.base.predict_th(state, action, next_state, done)
 
     def preprocess(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
     ) -> Tuple[th.Tensor, th.Tensor, th.Tensor, th.Tensor]:
         __doc__ = super().preprocess.__doc__  # noqa: F841
         return self.base.preprocess(state, action, next_state, done)
@@ -306,12 +307,12 @@ class RewardNetWithVariance(RewardNet):
 
     @abc.abstractmethod
     def predict_reward_moments(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
-        **kwargs,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
+            **kwargs,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Compute the mean and variance of the reward distribution.
 
@@ -336,14 +337,14 @@ class BasicRewardNet(RewardNet):
     """
 
     def __init__(
-        self,
-        observation_space: gym.Space,
-        action_space: gym.Space,
-        use_state: bool = True,
-        use_action: bool = True,
-        use_next_state: bool = False,
-        use_done: bool = False,
-        **kwargs,
+            self,
+            observation_space: gym.Space,
+            action_space: gym.Space,
+            use_state: bool = True,
+            use_action: bool = True,
+            use_next_state: bool = False,
+            use_done: bool = False,
+            **kwargs,
     ):
         """Builds reward MLP.
 
@@ -375,21 +376,20 @@ def __init__(
         if self.use_done:
             combined_size += 1
 
-        full_build_mlp_kwargs = {
-            "hid_sizes": (32, 32),
+        # kwargs except for in_size, out_size, squeeze_output keys,
+        # so they are not overriden.
+        kwargs = {
+            k: v for k, v in kwargs.items()
+            if k not in ("in_size", "out_size", "squeeze_output")
         }
-        full_build_mlp_kwargs.update(kwargs)
-        full_build_mlp_kwargs.update(
-            {
-                # we do not want these overridden
-                "in_size": combined_size,
-                "out_size": 1,
-                "squeeze_output": True,
-            },
+        self.mlp = networks.build_mlp(
+            hid_sizes=(32, 32),
+            **kwargs,
+            in_size=combined_size,
+            out_size=1,
+            squeeze_output=True
         )
 
-        self.mlp = networks.build_mlp(**full_build_mlp_kwargs)
-
     def forward(self, state, action, next_state, done):
         inputs = []
         if self.use_state:
@@ -413,9 +413,9 @@ class NormalizedRewardNet(RewardNetWrapper):
     """A reward net that normalizes the output of its base network."""
 
     def __init__(
-        self,
-        base: RewardNet,
-        normalize_output_layer: Type[nn.Module],
+            self,
+            base: RewardNet,
+            normalize_output_layer: Type[BaseNorm],
     ):
         """Initialize the NormalizedRewardNet.
 
@@ -434,13 +434,13 @@ def __init__(
         self.normalize_output_layer = normalize_output_layer(1)
 
     def predict_processed(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
-        update_stats: bool = True,
-        **kwargs,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
+            update_stats: bool = True,
+            **kwargs,
     ) -> np.ndarray:
         """Compute normalized rewards for a batch of transitions without gradients.
 
@@ -474,10 +474,10 @@ class ShapedRewardNet(RewardNetWrapper):
     """A RewardNet consisting of a base network and a potential shaping."""
 
     def __init__(
-        self,
-        base: RewardNet,
-        potential: Callable[[th.Tensor], th.Tensor],
-        discount_factor: float,
+            self,
+            base: RewardNet,
+            potential: Callable[[th.Tensor], th.Tensor],
+            discount_factor: float,
     ):
         """Setup a ShapedRewardNet instance.
 
@@ -497,11 +497,11 @@ def __init__(
         self.discount_factor = discount_factor
 
     def forward(
-        self,
-        state: th.Tensor,
-        action: th.Tensor,
-        next_state: th.Tensor,
-        done: th.Tensor,
+            self,
+            state: th.Tensor,
+            action: th.Tensor,
+            next_state: th.Tensor,
+            done: th.Tensor,
     ):
         base_reward_net_output = self.base(state, action, next_state, done)
         new_shaping_output = self.potential(next_state).flatten()
@@ -526,9 +526,9 @@ def forward(
         # length!
         new_shaping = (1 - done.float()) * new_shaping_output
         final_rew = (
-            base_reward_net_output
-            + self.discount_factor * new_shaping
-            - old_shaping_output
+                base_reward_net_output
+                + self.discount_factor * new_shaping
+                - old_shaping_output
         )
         assert final_rew.shape == state.shape[:1]
         return final_rew
@@ -551,18 +551,18 @@ class BasicShapedRewardNet(ShapedRewardNet):
     """
 
     def __init__(
-        self,
-        observation_space: gym.Space,
-        action_space: gym.Space,
-        *,
-        reward_hid_sizes: Sequence[int] = (32,),
-        potential_hid_sizes: Sequence[int] = (32, 32),
-        use_state: bool = True,
-        use_action: bool = True,
-        use_next_state: bool = False,
-        use_done: bool = False,
-        discount_factor: float = 0.99,
-        **kwargs,
+            self,
+            observation_space: gym.Space,
+            action_space: gym.Space,
+            *,
+            reward_hid_sizes: Sequence[int] = (32,),
+            potential_hid_sizes: Sequence[int] = (32, 32),
+            use_state: bool = True,
+            use_action: bool = True,
+            use_next_state: bool = False,
+            use_done: bool = False,
+            discount_factor: float = 0.99,
+            **kwargs,
     ):
         """Builds a simple shaped reward network.
 
@@ -611,10 +611,10 @@ class BasicPotentialMLP(nn.Module):
     """Simple implementation of a potential using an MLP."""
 
     def __init__(
-        self,
-        observation_space: gym.Space,
-        hid_sizes: Iterable[int],
-        **kwargs,
+            self,
+            observation_space: gym.Space,
+            hid_sizes: Iterable[int],
+            **kwargs,
     ):
         """Initialize the potential.
 
@@ -649,10 +649,10 @@ class RewardEnsemble(RewardNetWithVariance):
     members: nn.ModuleList
 
     def __init__(
-        self,
-        observation_space: gym.Space,
-        action_space: gym.Space,
-        members: Iterable[RewardNet],
+            self,
+            observation_space: gym.Space,
+            action_space: gym.Space,
+            members: Iterable[RewardNet],
     ):
         """Initialize the RewardEnsemble.
 
@@ -680,12 +680,12 @@ def num_members(self):
         return len(self.members)
 
     def predict_processed_all(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
-        **kwargs,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
+            **kwargs,
     ) -> np.ndarray:
         """Get the results of predict processed on all of the members.
 
@@ -701,22 +701,22 @@ def predict_processed_all(
                 shape `(batch_size, num_members)`.
         """
         batch_size = state.shape[0]
-        rewards = [
+        rewards_list = [
             member.predict_processed(state, action, next_state, done, **kwargs)
             for member in self.members
         ]
-        rewards: np.ndarray = np.stack(rewards, axis=-1)
+        rewards: np.ndarray = np.stack(rewards_list, axis=-1)
         assert rewards.shape == (batch_size, self.num_members)
         return rewards
 
     @th.no_grad()
     def predict_reward_moments(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
-        **kwargs,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
+            **kwargs,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Compute the standard deviation of the reward distribution for a batch.
 
@@ -749,23 +749,23 @@ def forward(self, *args) -> th.Tensor:
         raise NotImplementedError()
 
     def predict_processed(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
-        **kwargs,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
+            **kwargs,
     ) -> np.ndarray:
         """Return the mean of the ensemble members."""
         return self.predict(state, action, next_state, done, **kwargs)
 
     def predict(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
-        **kwargs,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
+            **kwargs,
     ):
         """Return the mean of the ensemble members."""
         mean, _ = self.predict_reward_moments(state, action, next_state, done, **kwargs)
@@ -799,13 +799,13 @@ def __init__(self, base: RewardNetWithVariance, default_alpha: float = 0.0):
         self.default_alpha = default_alpha
 
     def predict_processed(
-        self,
-        state: np.ndarray,
-        action: np.ndarray,
-        next_state: np.ndarray,
-        done: np.ndarray,
-        alpha: Optional[float] = None,
-        **kwargs,
+            self,
+            state: np.ndarray,
+            action: np.ndarray,
+            next_state: np.ndarray,
+            done: np.ndarray,
+            alpha: Optional[float] = None,
+            **kwargs,
     ) -> np.ndarray:
         """Compute a lower/upper confidence bound on the reward without gradients.
 
diff --git a/src/imitation/rewards/reward_wrapper.py b/src/imitation/rewards/reward_wrapper.py
index b333cfbb8..6a9243664 100644
--- a/src/imitation/rewards/reward_wrapper.py
+++ b/src/imitation/rewards/reward_wrapper.py
@@ -1,10 +1,10 @@
 """Common wrapper for adding custom reward values to an environment."""
 
 import collections
-from typing import Deque
+from typing import Deque, Optional
 
 import numpy as np
-from stable_baselines3.common import callbacks, vec_env
+from stable_baselines3.common import callbacks, vec_env, logger as sb_logger
 
 from imitation.rewards import reward_function
 
@@ -12,6 +12,8 @@
 class WrappedRewardCallback(callbacks.BaseCallback):
     """Logs mean wrapped reward as part of RL (or other) training."""
 
+    logger: Optional[sb_logger.Logger]  # type: ignore[assignment]
+
     def __init__(self, episode_rewards: Deque[float], *args, **kwargs):
         """Builds WrappedRewardCallback.
 
@@ -21,7 +23,7 @@ def __init__(self, episode_rewards: Deque[float], *args, **kwargs):
             **kwargs: Passed through to `callbacks.BaseCallback`.
         """
         self.episode_rewards = episode_rewards
-        super().__init__(self, *args, **kwargs)
+        super().__init__(*args, **kwargs)
 
     def _on_step(self) -> bool:
         return True
@@ -30,6 +32,7 @@ def _on_rollout_start(self) -> None:
         if len(self.episode_rewards) == 0:
             return
         mean = sum(self.episode_rewards) / len(self.episode_rewards)
+        assert self.logger is not None
         self.logger.record("rollout/ep_rew_wrapped_mean", mean)
 
 
@@ -45,10 +48,10 @@ class RewardVecEnvWrapper(vec_env.VecEnvWrapper):
     """
 
     def __init__(
-        self,
-        venv: vec_env.VecEnv,
-        reward_fn: reward_function.RewardFn,
-        ep_history: int = 100,
+            self,
+            venv: vec_env.VecEnv,
+            reward_fn: reward_function.RewardFn,
+            ep_history: int = 100,
     ):
         """Builds RewardVecEnvWrapper.
 
@@ -62,7 +65,7 @@ def __init__(
         """
         assert not isinstance(venv, RewardVecEnvWrapper)
         super().__init__(venv)
-        self.episode_rewards = collections.deque(maxlen=ep_history)
+        self.episode_rewards: Deque = collections.deque(maxlen=ep_history)
         self._cumulative_rew = np.zeros((venv.num_envs,))
         self.reward_fn = reward_fn
         self._old_obs = None

From af55ab257604e14cba37e0de08c25830e495a0c6 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 19:13:23 +0200
Subject: [PATCH 037/187] Fixed types: imitation/scripts/

---
 src/imitation/algorithms/dagger.py            | 66 ++++++++++---------
 src/imitation/scripts/analyze.py              | 26 ++++----
 src/imitation/scripts/common/common.py        | 30 ++++-----
 src/imitation/scripts/common/reward.py        | 27 ++++----
 src/imitation/scripts/common/wb.py            | 27 ++++----
 src/imitation/scripts/eval_policy.py          | 27 ++++----
 src/imitation/scripts/parallel.py             | 42 ++++++------
 src/imitation/scripts/train_adversarial.py    | 18 ++---
 src/imitation/scripts/train_imitation.py      | 54 +++++++++------
 .../scripts/train_preference_comparisons.py   | 60 ++++++++---------
 src/imitation/scripts/train_rl.py             | 28 ++++----
 11 files changed, 209 insertions(+), 196 deletions(-)

diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index 8700e0b12..122962024 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -67,10 +67,10 @@ def __call__(self, round_num: int) -> float:
 
 
 def reconstruct_trainer(
-    scratch_dir: types.AnyPath,
-    venv: vec_env.VecEnv,
-    custom_logger: Optional[logger.HierarchicalLogger] = None,
-    device: Union[th.device, str] = "auto",
+        scratch_dir: types.AnyPath,
+        venv: vec_env.VecEnv,
+        custom_logger: Optional[logger.HierarchicalLogger] = None,
+        device: Union[th.device, str] = "auto",
 ) -> "DAggerTrainer":
     """Reconstruct trainer from the latest snapshot in some working directory.
 
@@ -97,9 +97,9 @@ def reconstruct_trainer(
 
 
 def _save_dagger_demo(
-    trajectory: types.Trajectory,
-    save_dir: types.AnyPath,
-    prefix: str = "",
+        trajectory: types.Trajectory,
+        save_dir: types.AnyPath,
+        prefix: str = "",
 ) -> None:
     # TODO(shwang): This is possibly redundant with types.save(). Note
     #   however that NPZ save here is likely more space efficient than
@@ -147,11 +147,11 @@ class InteractiveTrajectoryCollector(vec_env.VecEnvWrapper):
     """
 
     def __init__(
-        self,
-        venv: vec_env.VecEnv,
-        get_robot_acts: Callable[[np.ndarray], np.ndarray],
-        beta: float,
-        save_dir: types.AnyPath,
+            self,
+            venv: vec_env.VecEnv,
+            get_robot_acts: Callable[[np.ndarray], np.ndarray],
+            beta: float,
+            save_dir: types.AnyPath,
     ):
         """Builds InteractiveTrajectoryCollector.
 
@@ -296,17 +296,19 @@ class DAggerTrainer(base.BaseImitationAlgorithm):
                    …
     """
 
+    _all_demos: List[types.Trajectory]
+
     DEFAULT_N_EPOCHS: int = 4
     """The default number of BC training epochs in `extend_and_update`."""
 
     def __init__(
-        self,
-        *,
-        venv: vec_env.VecEnv,
-        scratch_dir: types.AnyPath,
-        beta_schedule: Optional[Callable[[int], float]] = None,
-        bc_trainer: bc.BC,
-        custom_logger: Optional[logger.HierarchicalLogger] = None,
+            self,
+            *,
+            venv: vec_env.VecEnv,
+            scratch_dir: types.AnyPath,
+            beta_schedule: Optional[Callable[[int], float]] = None,
+            bc_trainer: bc.BC,
+            custom_logger: Optional[logger.HierarchicalLogger] = None,
     ):
         """Builds DAggerTrainer.
 
@@ -520,13 +522,13 @@ class SimpleDAggerTrainer(DAggerTrainer):
     """Simpler subclass of DAggerTrainer for training with synthetic feedback."""
 
     def __init__(
-        self,
-        *,
-        venv: vec_env.VecEnv,
-        scratch_dir: types.AnyPath,
-        expert_policy: policies.BasePolicy,
-        expert_trajs: Optional[Sequence[types.Trajectory]] = None,
-        **dagger_trainer_kwargs,
+            self,
+            *,
+            venv: vec_env.VecEnv,
+            scratch_dir: types.AnyPath,
+            expert_policy: policies.BasePolicy,
+            expert_trajs: Optional[Sequence[types.Trajectory]] = None,
+            **dagger_trainer_kwargs,
     ):
         """Builds SimpleDAggerTrainer.
 
@@ -571,12 +573,12 @@ def __init__(
                 )
 
     def train(
-        self,
-        total_timesteps: int,
-        *,
-        rollout_round_min_episodes: int = 3,
-        rollout_round_min_timesteps: int = 500,
-        bc_train_kwargs: Optional[dict] = None,
+            self,
+            total_timesteps: int,
+            *,
+            rollout_round_min_episodes: int = 3,
+            rollout_round_min_timesteps: int = 500,
+            bc_train_kwargs: Optional[dict] = None,
     ) -> None:
         """Train the DAgger agent.
 
diff --git a/src/imitation/scripts/analyze.py b/src/imitation/scripts/analyze.py
index 74617891c..b3654f9c5 100644
--- a/src/imitation/scripts/analyze.py
+++ b/src/imitation/scripts/analyze.py
@@ -9,7 +9,7 @@
 import tempfile
 import warnings
 from collections import OrderedDict
-from typing import Any, Callable, List, Mapping, Optional, Sequence, Set
+from typing import Any, Callable, List, Mapping, Optional, Sequence, Set, Iterable
 
 import pandas as pd
 from sacred.observers import FileStorageObserver
@@ -21,10 +21,10 @@
 
 @analysis_ex.capture
 def _gather_sacred_dicts(
-    source_dirs: Sequence[str],
-    run_name: Optional[str],
-    env_name: Optional[str],
-    skip_failed_runs: bool,
+        source_dirs: Sequence[str],
+        run_name: Optional[str],
+        env_name: Optional[str],
+        skip_failed_runs: bool,
 ) -> List[sacred_util.SacredDicts]:
     """Helper function for parsing and selecting Sacred experiment JSON files.
 
@@ -49,14 +49,15 @@ def _gather_sacred_dicts(
     sacred_dirs = itertools.chain.from_iterable(
         sacred_util.filter_subdirs(source_dir) for source_dir in source_dirs
     )
-    sacred_dicts = []
+    sacred_dicts_list = []
 
     for sacred_dir in sacred_dirs:
         try:
-            sacred_dicts.append(sacred_util.SacredDicts.load_from_dir(sacred_dir))
+            sacred_dicts_list.append(sacred_util.SacredDicts.load_from_dir(sacred_dir))
         except json.JSONDecodeError:
             warnings.warn(f"Invalid JSON file in {sacred_dir}", RuntimeWarning)
 
+    sacred_dicts: Iterable = sacred_dicts_list
     if run_name is not None:
         sacred_dicts = filter(
             lambda sd: get(sd.run, "experiment.name") == run_name,
@@ -179,7 +180,7 @@ def _return_summaries(sd: sacred_util.SacredDicts) -> dict:
         # Assuming here that `result.imit_stats` and `result.expert_stats` are
         # formatted correctly.
         imit_expert_ratio = (
-            imit_stats["monitor_return_mean"] / expert_stats["return_mean"]
+                imit_stats["monitor_return_mean"] / expert_stats["return_mean"]
         )
     else:
         imit_expert_ratio = None
@@ -217,7 +218,6 @@ def _return_summaries(sd: sacred_util.SacredDicts) -> dict:
     ],
 )
 
-
 # If `verbosity` is at least the length of this list, then we use all table_entry_fns
 # as columns of table.
 # Otherwise, use only the subset at index `verbosity`. The subset of columns is
@@ -260,10 +260,10 @@ def _get_table_entry_fns_subset(table_verbosity: int) -> sd_to_table_entry_type:
 
 @analysis_ex.command
 def analyze_imitation(
-    csv_output_path: Optional[str],
-    tex_output_path: Optional[str],
-    print_table: bool,
-    table_verbosity: int,
+        csv_output_path: Optional[str],
+        tex_output_path: Optional[str],
+        print_table: bool,
+        table_verbosity: int,
 ) -> pd.DataFrame:
     """Parse Sacred logs and generate a DataFrame for imitation learning results.
 
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index c08eda766..9878c6f63 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -3,7 +3,7 @@
 import contextlib
 import logging
 import os
-from typing import Any, Mapping, Sequence, Tuple, Union
+from typing import Any, Mapping, Sequence, Tuple, Union, Generator
 
 import sacred
 from stable_baselines3.common import vec_env
@@ -76,9 +76,9 @@ def fast():
 
 @common_ingredient.capture
 def make_log_dir(
-    _run,
-    log_dir: str,
-    log_level: Union[int, str],
+        _run,
+        log_dir: str,
+        log_level: Union[int, str],
 ) -> str:
     """Creates log directory and sets up symlink to Sacred logs.
 
@@ -105,8 +105,8 @@ def make_log_dir(
 
 @common_ingredient.capture
 def setup_logging(
-    _run,
-    log_format_strs: Sequence[str],
+        _run,
+        log_format_strs: Sequence[str],
 ) -> Tuple[imit_logger.HierarchicalLogger, str]:
     """Builds the imitation logger.
 
@@ -130,15 +130,15 @@ def setup_logging(
 @contextlib.contextmanager
 @common_ingredient.capture
 def make_venv(
-    _seed,
-    env_name: str,
-    num_vec: int,
-    parallel: bool,
-    log_dir: str,
-    max_episode_steps: int,
-    env_make_kwargs: Mapping[str, Any],
-    **kwargs,
-) -> vec_env.VecEnv:
+        _seed,
+        env_name: str,
+        num_vec: int,
+        parallel: bool,
+        log_dir: str,
+        max_episode_steps: int,
+        env_make_kwargs: Mapping[str, Any],
+        **kwargs,
+) -> Generator[vec_env.VecEnv, None, None]:
     """Builds the vector environment.
 
     Args:
diff --git a/src/imitation/scripts/common/reward.py b/src/imitation/scripts/common/reward.py
index 548bca855..6b26e970b 100644
--- a/src/imitation/scripts/common/reward.py
+++ b/src/imitation/scripts/common/reward.py
@@ -6,7 +6,6 @@
 
 import sacred
 from stable_baselines3.common import vec_env
-from torch import nn
 
 from imitation.rewards import reward_nets
 from imitation.util import networks
@@ -82,10 +81,10 @@ def config_hook(config, command_name, logger):
 
 
 def _make_reward_net(
-    venv: vec_env.VecEnv,
-    net_cls: Type[reward_nets.RewardNet],
-    net_kwargs: Mapping[str, Any],
-    normalize_output_layer: Optional[Type[nn.Module]],
+        venv: vec_env.VecEnv,
+        net_cls: Type[reward_nets.RewardNet],
+        net_kwargs: Mapping[str, Any],
+        normalize_output_layer: Optional[Type[networks.BaseNorm]],
 ):
     """Helper function for creating reward nets."""
     reward_net = net_cls(
@@ -105,13 +104,13 @@ def _make_reward_net(
 
 @reward_ingredient.capture
 def make_reward_net(
-    venv: vec_env.VecEnv,
-    net_cls: Type[reward_nets.RewardNet],
-    net_kwargs: Mapping[str, Any],
-    normalize_output_layer: Optional[Type[nn.Module]],
-    add_std_alpha: Optional[float],
-    ensemble_size: Optional[int],
-    ensemble_member_config: Optional[Mapping[str, Any]],
+        venv: vec_env.VecEnv,
+        net_cls: Type[reward_nets.RewardNet],
+        net_kwargs: Mapping[str, Any],
+        normalize_output_layer: Optional[Type[networks.BaseNorm]],
+        add_std_alpha: Optional[float],
+        ensemble_size: Optional[int],
+        ensemble_member_config: Optional[Mapping[str, Any]],
 ) -> reward_nets.RewardNet:
     """Builds a reward network.
 
@@ -150,9 +149,11 @@ def make_reward_net(
             for _ in range(ensemble_size)
         ]
 
-        reward_net = net_cls(venv.observation_space, venv.action_space, members)
+        reward_net: reward_nets.RewardNet = net_cls(venv.observation_space, venv.action_space, members)
 
         if add_std_alpha is not None:
+            if not isinstance(reward_net, reward_nets.RewardNetWithVariance):
+                raise ValueError("add_std_alpha is only supported for reward nets with variance tracking.")
             reward_net = reward_nets.AddSTDRewardWrapper(
                 reward_net,
                 default_alpha=add_std_alpha,
diff --git a/src/imitation/scripts/common/wb.py b/src/imitation/scripts/common/wb.py
index 2689e16a2..47b5e3fe5 100644
--- a/src/imitation/scripts/common/wb.py
+++ b/src/imitation/scripts/common/wb.py
@@ -26,12 +26,12 @@ def wandb_config():
 
 @wandb_ingredient.capture
 def wandb_init(
-    _run,
-    wandb_name_prefix: str,
-    wandb_tag: Optional[str],
-    wandb_kwargs: Mapping[str, Any],
-    wandb_additional_info: Mapping[str, Any],
-    log_dir: str,
+        _run,
+        wandb_name_prefix: str,
+        wandb_tag: Optional[str],
+        wandb_kwargs: Mapping[str, Any],
+        wandb_additional_info: Mapping[str, Any],
+        log_dir: str,
 ) -> None:
     """Putting everything together to get the W&B kwargs for wandb.init().
 
@@ -49,15 +49,12 @@ def wandb_init(
     env_name = _run.config["common"]["env_name"]
     root_seed = _run.config["seed"]
 
-    updated_wandb_kwargs = {}
-    updated_wandb_kwargs.update(wandb_kwargs)
-    updated_wandb_kwargs.update(
-        dict(
-            name="-".join([wandb_name_prefix, env_name, f"seed{root_seed}"]),
-            tags=[env_name, f"seed{root_seed}"] + ([wandb_tag] if wandb_tag else []),
-            dir=log_dir,
-        ),
-    )
+    updated_wandb_kwargs: Mapping[str, Any] = {
+        **wandb_kwargs,
+        "name": f"{wandb_name_prefix}-{env_name}-seed{root_seed}",
+        "tags": [env_name, f"seed{root_seed}"] + ([wandb_tag] if wandb_tag else []),
+        "dir": log_dir,
+    }
     try:
         import wandb
     except ModuleNotFoundError as e:
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 8a667df36..6b943532d 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -53,19 +53,19 @@ def f(env: gym.Env, i: int) -> gym.Env:
 
 @eval_policy_ex.main
 def eval_policy(
-    _run,
-    _seed: int,
-    eval_n_timesteps: Optional[int],
-    eval_n_episodes: Optional[int],
-    render: bool,
-    render_fps: int,
-    videos: bool,
-    video_kwargs: Mapping[str, Any],
-    policy_type: Optional[str],
-    policy_path: Optional[str],
-    reward_type: Optional[str] = None,
-    reward_path: Optional[str] = None,
-    rollout_save_path: Optional[str] = None,
+        _run,
+        _seed: int,
+        eval_n_timesteps: Optional[int],
+        eval_n_episodes: Optional[int],
+        render: bool,
+        render_fps: int,
+        videos: bool,
+        video_kwargs: Mapping[str, Any],
+        policy_type: Optional[str],
+        policy_path: Optional[str],
+        reward_type: Optional[str] = None,
+        reward_path: Optional[str] = None,
+        rollout_save_path: Optional[str] = None,
 ):
     """Rolls a policy out in an environment, collecting statistics.
 
@@ -106,6 +106,7 @@ def eval_policy(
 
         policy = None
         if policy_type is not None:
+            assert policy_path is not None
             policy = serialize.load_policy(policy_type, policy_path, venv)
         trajs = rollout.generate_trajectories(policy, venv, sample_until)
 
diff --git a/src/imitation/scripts/parallel.py b/src/imitation/scripts/parallel.py
index 2bc62c745..fc2c3625a 100644
--- a/src/imitation/scripts/parallel.py
+++ b/src/imitation/scripts/parallel.py
@@ -3,7 +3,7 @@
 import collections.abc
 import copy
 import os
-from typing import Any, Callable, Mapping, Optional, Sequence
+from typing import Any, Callable, Mapping, Optional, Sequence, Dict
 
 import ray
 import ray.tune
@@ -15,15 +15,15 @@
 
 @parallel_ex.main
 def parallel(
-    sacred_ex_name: str,
-    run_name: str,
-    search_space: Mapping[str, Any],
-    base_named_configs: Sequence[str],
-    base_config_updates: Mapping[str, Any],
-    resources_per_trial: Mapping[str, Any],
-    init_kwargs: Mapping[str, Any],
-    local_dir: Optional[str],
-    upload_dir: Optional[str],
+        sacred_ex_name: str,
+        run_name: str,
+        search_space: Mapping[str, Any],
+        base_named_configs: Sequence[str],
+        base_config_updates: Mapping[str, Any],
+        resources_per_trial: Mapping[str, Any],
+        init_kwargs: Mapping[str, Any],
+        local_dir: Optional[str],
+        upload_dir: Optional[str],
 ) -> None:
     """Parallelize multiple runs of another Sacred Experiment using Ray Tune.
 
@@ -92,8 +92,8 @@ def parallel(
     # each Raylet.
     if sacred_ex_name == "train_adversarial":
         no_data_dir = (
-            "demonstrations.data_dir" not in base_config_updates
-            and "data_dir" not in base_config_updates.get("demonstrations", {})
+                "demonstrations.data_dir" not in base_config_updates
+                and "data_dir" not in base_config_updates.get("demonstrations", {})
         )
         if no_data_dir:
             data_dir = os.path.join(os.getcwd(), "data/")
@@ -122,10 +122,10 @@ def parallel(
 
 
 def _ray_tune_sacred_wrapper(
-    sacred_ex_name: str,
-    run_name: str,
-    base_named_configs: list,
-    base_config_updates: Mapping[str, Any],
+        sacred_ex_name: str,
+        run_name: str,
+        base_named_configs: list,
+        base_config_updates: Mapping[str, Any],
 ) -> Callable[[Mapping[str, Any], Any], Mapping[str, Any]]:
     """From an Experiment build a wrapped run function suitable for Ray Tune.
 
@@ -178,7 +178,7 @@ def inner(config: Mapping[str, Any], reporter) -> Mapping[str, Any]:
         sacred.SETTINGS.CAPTURE_MODE = "sys"
 
         run_kwargs = config
-        updated_run_kwargs = {}
+        updated_run_kwargs: Dict[str, Any] = {}
         # Import inside function rather than in module because Sacred experiments
         # are not picklable, and Ray requires this function to be picklable.
         from imitation.scripts.train_adversarial import train_adversarial_ex
@@ -192,14 +192,10 @@ def inner(config: Mapping[str, Any], reporter) -> Mapping[str, Any]:
         ex.observers = [FileStorageObserver("sacred")]
 
         # Apply base configs to get modified `named_configs` and `config_updates`.
-        named_configs = []
-        named_configs.extend(base_named_configs)
-        named_configs.extend(run_kwargs["named_configs"])
+        named_configs = [*base_named_configs, *run_kwargs["named_configs"]]
         updated_run_kwargs["named_configs"] = named_configs
 
-        config_updates = {}
-        config_updates.update(base_config_updates)
-        config_updates.update(run_kwargs["config_updates"])
+        config_updates = {**base_config_updates, **run_kwargs["config_updates"]}
         updated_run_kwargs["config_updates"] = config_updates
 
         # Add other run_kwargs items to updated_run_kwargs.
diff --git a/src/imitation/scripts/train_adversarial.py b/src/imitation/scripts/train_adversarial.py
index 1d24955ad..291800d0f 100644
--- a/src/imitation/scripts/train_adversarial.py
+++ b/src/imitation/scripts/train_adversarial.py
@@ -60,20 +60,20 @@ def dummy_config():
         algorithm_specific = {}  # noqa: F841
 
 
-for ingredient in [train_adversarial_ex] + train_adversarial_ex.ingredients:
+for ingredient in [train_adversarial_ex, *train_adversarial_ex.ingredients]:
     _add_hook(ingredient)
 
 
 @train_adversarial_ex.capture
 def train_adversarial(
-    _run,
-    _seed: int,
-    show_config: bool,
-    algo_cls: Type[common.AdversarialTrainer],
-    algorithm_kwargs: Mapping[str, Any],
-    total_timesteps: int,
-    checkpoint_interval: int,
-    agent_path: Optional[str],
+        _run,
+        _seed: int,
+        show_config: bool,
+        algo_cls: Type[common.AdversarialTrainer],
+        algorithm_kwargs: Mapping[str, Any],
+        total_timesteps: int,
+        checkpoint_interval: int,
+        agent_path: Optional[str],
 ) -> Mapping[str, Mapping[str, float]]:
     """Train an adversarial-network-based imitation learning algorithm.
 
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index 043268550..8a87d928c 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -3,14 +3,14 @@
 import logging
 import os.path as osp
 import warnings
-from typing import Any, Mapping, Optional, Type
+from typing import Any, Mapping, Optional, Type, Sequence, List, cast
 
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import policies, utils, vec_env
 
 from imitation.algorithms import bc as bc_algorithm
 from imitation.algorithms.dagger import SimpleDAggerTrainer
-from imitation.data import rollout
+from imitation.data import rollout, types
 from imitation.policies import serialize
 from imitation.scripts.common import common, demonstrations, train
 from imitation.scripts.config.train_imitation import train_imitation_ex
@@ -20,10 +20,10 @@
 
 @train_imitation_ex.capture(prefix="train")
 def make_policy(
-    venv: vec_env.VecEnv,
-    policy_cls: Type[policies.BasePolicy],
-    policy_kwargs: Mapping[str, Any],
-    agent_path: Optional[str],
+        venv: vec_env.VecEnv,
+        policy_cls: Type[policies.BasePolicy],
+        policy_kwargs: Mapping[str, Any],
+        agent_path: Optional[str],
 ) -> policies.BasePolicy:
     """Makes policy.
 
@@ -50,6 +50,7 @@ def make_policy(
                 "lr_schedule": utils.get_schedule_fn(1),
             },
         )
+    policy: policies.BasePolicy
     if agent_path is not None:
         warnings.warn(
             "When agent_path is specified, policy_cls and policy_kwargs are ignored.",
@@ -64,9 +65,9 @@ def make_policy(
 
 @train_imitation_ex.capture(prefix="dagger")
 def load_expert_policy(
-    venv: vec_env.VecEnv,
-    expert_policy_type: Optional[str],
-    expert_policy_path: Optional[str],
+        venv: vec_env.VecEnv,
+        expert_policy_type: Optional[str],
+        expert_policy_path: Optional[str],
 ) -> policies.BasePolicy:
     """Loads expert policy from `expert_policy_path`.
 
@@ -88,6 +89,10 @@ def load_expert_policy(
     """
     if expert_policy_path is None:
         raise ValueError("expert_policy_path cannot be None")
+
+    if expert_policy_type is None:
+        raise ValueError("expert_policy_type cannot be None")
+
     # TODO(shwang): Add support for directly loading a BasePolicy `*.th` file.
     expert_policy = serialize.load_policy(expert_policy_type, expert_policy_path, venv)
     if not isinstance(expert_policy, policies.BasePolicy):
@@ -97,12 +102,12 @@ def load_expert_policy(
 
 @train_imitation_ex.capture
 def train_imitation(
-    _run,
-    bc_kwargs: Mapping[str, Any],
-    bc_train_kwargs: Mapping[str, Any],
-    dagger: Mapping[str, Any],
-    use_dagger: bool,
-    agent_path: Optional[str],
+        _run,
+        bc_kwargs: Mapping[str, Any],
+        bc_train_kwargs: Mapping[str, Any],
+        dagger: Mapping[str, Any],
+        use_dagger: bool,
+        agent_path: Optional[str],
 ) -> Mapping[str, Mapping[str, float]]:
     """Runs DAgger (if `use_dagger`) or BC (otherwise) training.
 
@@ -123,7 +128,7 @@ def train_imitation(
     with common.make_venv() as venv:
         imit_policy = make_policy(venv, agent_path=agent_path)
 
-        expert_trajs = None
+        expert_trajs: Optional[Sequence[types.Trajectory]] = None
         if not use_dagger or dagger["use_offline_rollouts"]:
             expert_trajs = demonstrations.load_expert_trajs()
 
@@ -166,11 +171,22 @@ def train_imitation(
 
         imit_stats = train.eval_policy(imit_policy, venv)
 
+    # TODO(juan): I'm not super happy with this solution for the type system.
+    #  is model._all_demos always Sequence[TrajectoryWithRew]? We can change
+    #  the type in the class definition. Same goes for demonstrations.load_expert_trajs.
+    #  using assert doesn't work because we'd have to loop over all the trajectories
+    #  and check that each one is a TrajectoryWithRew, which seems inefficient
+    #  just for adding type annotations.
+    trajectories: List[types.TrajectoryWithRew]
+    if use_dagger:
+        assert expert_trajs is not None
+        trajectories = cast(List[types.TrajectoryWithRew], expert_trajs)
+    else:
+        trajectories = cast(List[types.TrajectoryWithRew], model._all_demos)
+
     return {
         "imit_stats": imit_stats,
-        "expert_stats": rollout.rollout_stats(
-            model._all_demos if use_dagger else expert_trajs,
-        ),
+        "expert_stats": rollout.rollout_stats(trajectories),
     }
 
 
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 704bed568..5aa861a2a 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -24,8 +24,8 @@
 
 
 def save_model(
-    agent_trainer: preference_comparisons.AgentTrainer,
-    save_path: str,
+        agent_trainer: preference_comparisons.AgentTrainer,
+        save_path: str,
 ):
     """Save the model as model.pkl."""
     serialize.save_stable_model(
@@ -35,9 +35,9 @@ def save_model(
 
 
 def save_checkpoint(
-    trainer: preference_comparisons.PreferenceComparisons,
-    save_path: str,
-    allow_save_policy: Optional[bool],
+        trainer: preference_comparisons.PreferenceComparisons,
+        save_path: str,
+        allow_save_policy: Optional[bool],
 ):
     """Save reward model and optionally policy."""
     os.makedirs(save_path, exist_ok=True)
@@ -56,30 +56,30 @@ def save_checkpoint(
 
 @train_preference_comparisons_ex.main
 def train_preference_comparisons(
-    _seed: int,
-    total_timesteps: int,
-    total_comparisons: int,
-    num_iterations: int,
-    comparison_queue_size: Optional[int],
-    fragment_length: int,
-    transition_oversampling: float,
-    initial_comparison_frac: float,
-    exploration_frac: float,
-    trajectory_path: Optional[str],
-    trajectory_generator_kwargs: Mapping[str, Any],
-    save_preferences: bool,
-    agent_path: Optional[str],
-    preference_model_kwargs: Mapping[str, Any],
-    reward_trainer_kwargs: Mapping[str, Any],
-    gatherer_cls: Type[preference_comparisons.PreferenceGatherer],
-    gatherer_kwargs: Mapping[str, Any],
-    active_selection: bool,
-    active_selection_oversampling: int,
-    uncertainty_on: str,
-    fragmenter_kwargs: Mapping[str, Any],
-    allow_variable_horizon: bool,
-    checkpoint_interval: int,
-    query_schedule: Union[str, type_aliases.Schedule],
+        _seed: int,
+        total_timesteps: int,
+        total_comparisons: int,
+        num_iterations: int,
+        comparison_queue_size: Optional[int],
+        fragment_length: int,
+        transition_oversampling: float,
+        initial_comparison_frac: float,
+        exploration_frac: float,
+        trajectory_path: Optional[str],
+        trajectory_generator_kwargs: Mapping[str, Any],
+        save_preferences: bool,
+        agent_path: Optional[str],
+        preference_model_kwargs: Mapping[str, Any],
+        reward_trainer_kwargs: Mapping[str, Any],
+        gatherer_cls: Type[preference_comparisons.PreferenceGatherer],
+        gatherer_kwargs: Mapping[str, Any],
+        active_selection: bool,
+        active_selection_oversampling: int,
+        uncertainty_on: str,
+        fragmenter_kwargs: Mapping[str, Any],
+        allow_variable_horizon: bool,
+        checkpoint_interval: int,
+        query_schedule: Union[str, type_aliases.Schedule],
 ) -> Mapping[str, Any]:
     """Train a reward model using preference comparisons.
 
@@ -191,7 +191,7 @@ def train_preference_comparisons(
                 **trajectory_generator_kwargs,
             )
 
-        fragmenter = preference_comparisons.RandomFragmenter(
+        fragmenter: preference_comparisons.Fragmenter = preference_comparisons.RandomFragmenter(
             **fragmenter_kwargs,
             seed=_seed,
             custom_logger=custom_logger,
diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
index 2540ac79c..fa222dad2 100644
--- a/src/imitation/scripts/train_rl.py
+++ b/src/imitation/scripts/train_rl.py
@@ -28,19 +28,19 @@
 
 @train_rl_ex.main
 def train_rl(
-    *,
-    total_timesteps: int,
-    normalize_reward: bool,
-    normalize_kwargs: dict,
-    reward_type: Optional[str],
-    reward_path: Optional[str],
-    load_reward_kwargs: Optional[Mapping[str, Any]],
-    rollout_save_final: bool,
-    rollout_save_n_timesteps: Optional[int],
-    rollout_save_n_episodes: Optional[int],
-    policy_save_interval: int,
-    policy_save_final: bool,
-    agent_path: Optional[str],
+        *,
+        total_timesteps: int,
+        normalize_reward: bool,
+        normalize_kwargs: dict,
+        reward_type: Optional[str],
+        reward_path: Optional[str],
+        load_reward_kwargs: Optional[Mapping[str, Any]],
+        rollout_save_final: bool,
+        rollout_save_n_timesteps: Optional[int],
+        rollout_save_n_episodes: Optional[int],
+        policy_save_interval: int,
+        policy_save_final: bool,
+        agent_path: Optional[str],
 ) -> Mapping[str, float]:
     """Trains an expert policy from scratch and saves the rollouts and policy.
 
@@ -120,7 +120,7 @@ def train_rl(
                 )
 
         if policy_save_interval > 0:
-            save_policy_callback = serialize.SavePolicyCallback(policy_dir)
+            save_policy_callback: callbacks.EventCallback = serialize.SavePolicyCallback(policy_dir)
             save_policy_callback = callbacks.EveryNTimesteps(
                 policy_save_interval,
                 save_policy_callback,

From 9d320bd439bd2e424e3fa8244f8a8657e6da0f42 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 19:27:41 +0200
Subject: [PATCH 038/187] Fixed types: imitation/util/ and formatting

---
 src/imitation/algorithms/dagger.py            |  64 ++--
 src/imitation/rewards/reward_nets.py          | 334 +++++++++---------
 src/imitation/rewards/reward_wrapper.py       |  12 +-
 src/imitation/scripts/analyze.py              |  20 +-
 src/imitation/scripts/common/common.py        |  28 +-
 src/imitation/scripts/common/reward.py        |  30 +-
 src/imitation/scripts/common/wb.py            |  12 +-
 src/imitation/scripts/eval_policy.py          |  26 +-
 src/imitation/scripts/parallel.py             |  32 +-
 src/imitation/scripts/train_adversarial.py    |  16 +-
 src/imitation/scripts/train_imitation.py      |  28 +-
 .../scripts/train_preference_comparisons.py   |  68 ++--
 src/imitation/scripts/train_rl.py             |  30 +-
 src/imitation/util/logger.py                  |  13 +-
 src/imitation/util/networks.py                |   6 +-
 src/imitation/util/sacred.py                  |   4 +
 src/imitation/util/util.py                    |   7 +-
 src/imitation/util/video_wrapper.py           |   8 +-
 18 files changed, 389 insertions(+), 349 deletions(-)

diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index 122962024..4ad0b6ebc 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -67,10 +67,10 @@ def __call__(self, round_num: int) -> float:
 
 
 def reconstruct_trainer(
-        scratch_dir: types.AnyPath,
-        venv: vec_env.VecEnv,
-        custom_logger: Optional[logger.HierarchicalLogger] = None,
-        device: Union[th.device, str] = "auto",
+    scratch_dir: types.AnyPath,
+    venv: vec_env.VecEnv,
+    custom_logger: Optional[logger.HierarchicalLogger] = None,
+    device: Union[th.device, str] = "auto",
 ) -> "DAggerTrainer":
     """Reconstruct trainer from the latest snapshot in some working directory.
 
@@ -97,9 +97,9 @@ def reconstruct_trainer(
 
 
 def _save_dagger_demo(
-        trajectory: types.Trajectory,
-        save_dir: types.AnyPath,
-        prefix: str = "",
+    trajectory: types.Trajectory,
+    save_dir: types.AnyPath,
+    prefix: str = "",
 ) -> None:
     # TODO(shwang): This is possibly redundant with types.save(). Note
     #   however that NPZ save here is likely more space efficient than
@@ -147,11 +147,11 @@ class InteractiveTrajectoryCollector(vec_env.VecEnvWrapper):
     """
 
     def __init__(
-            self,
-            venv: vec_env.VecEnv,
-            get_robot_acts: Callable[[np.ndarray], np.ndarray],
-            beta: float,
-            save_dir: types.AnyPath,
+        self,
+        venv: vec_env.VecEnv,
+        get_robot_acts: Callable[[np.ndarray], np.ndarray],
+        beta: float,
+        save_dir: types.AnyPath,
     ):
         """Builds InteractiveTrajectoryCollector.
 
@@ -302,13 +302,13 @@ class DAggerTrainer(base.BaseImitationAlgorithm):
     """The default number of BC training epochs in `extend_and_update`."""
 
     def __init__(
-            self,
-            *,
-            venv: vec_env.VecEnv,
-            scratch_dir: types.AnyPath,
-            beta_schedule: Optional[Callable[[int], float]] = None,
-            bc_trainer: bc.BC,
-            custom_logger: Optional[logger.HierarchicalLogger] = None,
+        self,
+        *,
+        venv: vec_env.VecEnv,
+        scratch_dir: types.AnyPath,
+        beta_schedule: Optional[Callable[[int], float]] = None,
+        bc_trainer: bc.BC,
+        custom_logger: Optional[logger.HierarchicalLogger] = None,
     ):
         """Builds DAggerTrainer.
 
@@ -522,13 +522,13 @@ class SimpleDAggerTrainer(DAggerTrainer):
     """Simpler subclass of DAggerTrainer for training with synthetic feedback."""
 
     def __init__(
-            self,
-            *,
-            venv: vec_env.VecEnv,
-            scratch_dir: types.AnyPath,
-            expert_policy: policies.BasePolicy,
-            expert_trajs: Optional[Sequence[types.Trajectory]] = None,
-            **dagger_trainer_kwargs,
+        self,
+        *,
+        venv: vec_env.VecEnv,
+        scratch_dir: types.AnyPath,
+        expert_policy: policies.BasePolicy,
+        expert_trajs: Optional[Sequence[types.Trajectory]] = None,
+        **dagger_trainer_kwargs,
     ):
         """Builds SimpleDAggerTrainer.
 
@@ -573,12 +573,12 @@ def __init__(
                 )
 
     def train(
-            self,
-            total_timesteps: int,
-            *,
-            rollout_round_min_episodes: int = 3,
-            rollout_round_min_timesteps: int = 500,
-            bc_train_kwargs: Optional[dict] = None,
+        self,
+        total_timesteps: int,
+        *,
+        rollout_round_min_episodes: int = 3,
+        rollout_round_min_timesteps: int = 500,
+        bc_train_kwargs: Optional[dict] = None,
     ) -> None:
         """Train the DAgger agent.
 
diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index 27490d25e..60f7853cd 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -21,10 +21,10 @@ class RewardNet(nn.Module, abc.ABC):
     """
 
     def __init__(
-            self,
-            observation_space: gym.Space,
-            action_space: gym.Space,
-            normalize_images: bool = True,
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        normalize_images: bool = True,
     ):
         """Initialize the RewardNet.
 
@@ -41,20 +41,20 @@ def __init__(
 
     @abc.abstractmethod
     def forward(
-            self,
-            state: th.Tensor,
-            action: th.Tensor,
-            next_state: th.Tensor,
-            done: th.Tensor,
+        self,
+        state: th.Tensor,
+        action: th.Tensor,
+        next_state: th.Tensor,
+        done: th.Tensor,
     ) -> th.Tensor:
         """Compute rewards for a batch of transitions and keep gradients."""
 
     def preprocess(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
     ) -> Tuple[th.Tensor, th.Tensor, th.Tensor, th.Tensor]:
         """Preprocess a batch of input transitions and convert it to PyTorch tensors.
 
@@ -83,21 +83,30 @@ def preprocess(
         del state, action, next_state, done  # unused
 
         # preprocess
-        state_th = cast(th.Tensor, preprocessing.preprocess_obs(
-            state_th,
-            self.observation_space,
-            self.normalize_images,
-        ))
-        action_th = cast(th.Tensor, preprocessing.preprocess_obs(
-            action_th,
-            self.action_space,
-            self.normalize_images,
-        ))
-        next_state_th = cast(th.Tensor, preprocessing.preprocess_obs(
-            next_state_th,
-            self.observation_space,
-            self.normalize_images,
-        ))
+        state_th = cast(
+            th.Tensor,
+            preprocessing.preprocess_obs(
+                state_th,
+                self.observation_space,
+                self.normalize_images,
+            ),
+        )
+        action_th = cast(
+            th.Tensor,
+            preprocessing.preprocess_obs(
+                action_th,
+                self.action_space,
+                self.normalize_images,
+            ),
+        )
+        next_state_th = cast(
+            th.Tensor,
+            preprocessing.preprocess_obs(
+                next_state_th,
+                self.observation_space,
+                self.normalize_images,
+            ),
+        )
         done_th = done_th.to(th.float32)
 
         n_gen = len(state_th)
@@ -107,11 +116,11 @@ def preprocess(
         return state_th, action_th, next_state_th, done_th
 
     def predict_th(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
     ) -> th.Tensor:
         """Compute th.Tensor rewards for a batch of transitions without gradients.
 
@@ -142,11 +151,11 @@ def predict_th(
             return rew_th
 
     def predict(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
     ) -> np.ndarray:
         """Compute rewards for a batch of transitions without gradients.
 
@@ -165,12 +174,12 @@ def predict(
         return rew_th.detach().cpu().numpy().flatten()
 
     def predict_processed(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
-            **kwargs,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
+        **kwargs,
     ) -> np.ndarray:
         """Compute the processed rewards for a batch of transitions without gradients.
 
@@ -222,8 +231,8 @@ class RewardNetWrapper(RewardNet):
     """
 
     def __init__(
-            self,
-            base: RewardNet,
+        self,
+        base: RewardNet,
     ):
         """Initialize a RewardNet wrapper.
 
@@ -242,52 +251,52 @@ def base(self) -> RewardNet:
         return self._base
 
     def forward(
-            self,
-            state: th.Tensor,
-            action: th.Tensor,
-            next_state: th.Tensor,
-            done: th.Tensor,
+        self,
+        state: th.Tensor,
+        action: th.Tensor,
+        next_state: th.Tensor,
+        done: th.Tensor,
     ) -> th.Tensor:
         __doc__ = super().forward.__doc__  # noqa: F841
         return self.base.forward(state, action, next_state, done)
 
     def predict_processed(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
-            **kwargs,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
+        **kwargs,
     ) -> np.ndarray:
         __doc__ = super().predict_processed.__doc__  # noqa: F841
         return self.base.predict_processed(state, action, next_state, done, **kwargs)
 
     def predict(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
     ) -> np.ndarray:
         __doc__ = super().predict.__doc__  # noqa: F841
         return self.base.predict(state, action, next_state, done)
 
     def predict_th(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
     ) -> th.Tensor:
         __doc__ = super().predict_th.__doc__  # noqa: F841
         return self.base.predict_th(state, action, next_state, done)
 
     def preprocess(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
     ) -> Tuple[th.Tensor, th.Tensor, th.Tensor, th.Tensor]:
         __doc__ = super().preprocess.__doc__  # noqa: F841
         return self.base.preprocess(state, action, next_state, done)
@@ -307,12 +316,12 @@ class RewardNetWithVariance(RewardNet):
 
     @abc.abstractmethod
     def predict_reward_moments(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
-            **kwargs,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
+        **kwargs,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Compute the mean and variance of the reward distribution.
 
@@ -337,14 +346,14 @@ class BasicRewardNet(RewardNet):
     """
 
     def __init__(
-            self,
-            observation_space: gym.Space,
-            action_space: gym.Space,
-            use_state: bool = True,
-            use_action: bool = True,
-            use_next_state: bool = False,
-            use_done: bool = False,
-            **kwargs,
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        use_state: bool = True,
+        use_action: bool = True,
+        use_next_state: bool = False,
+        use_done: bool = False,
+        **kwargs,
     ):
         """Builds reward MLP.
 
@@ -379,7 +388,8 @@ def __init__(
         # kwargs except for in_size, out_size, squeeze_output keys,
         # so they are not overriden.
         kwargs = {
-            k: v for k, v in kwargs.items()
+            k: v
+            for k, v in kwargs.items()
             if k not in ("in_size", "out_size", "squeeze_output")
         }
         self.mlp = networks.build_mlp(
@@ -387,7 +397,7 @@ def __init__(
             **kwargs,
             in_size=combined_size,
             out_size=1,
-            squeeze_output=True
+            squeeze_output=True,
         )
 
     def forward(self, state, action, next_state, done):
@@ -413,9 +423,9 @@ class NormalizedRewardNet(RewardNetWrapper):
     """A reward net that normalizes the output of its base network."""
 
     def __init__(
-            self,
-            base: RewardNet,
-            normalize_output_layer: Type[BaseNorm],
+        self,
+        base: RewardNet,
+        normalize_output_layer: Type[BaseNorm],
     ):
         """Initialize the NormalizedRewardNet.
 
@@ -434,13 +444,13 @@ def __init__(
         self.normalize_output_layer = normalize_output_layer(1)
 
     def predict_processed(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
-            update_stats: bool = True,
-            **kwargs,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
+        update_stats: bool = True,
+        **kwargs,
     ) -> np.ndarray:
         """Compute normalized rewards for a batch of transitions without gradients.
 
@@ -474,10 +484,10 @@ class ShapedRewardNet(RewardNetWrapper):
     """A RewardNet consisting of a base network and a potential shaping."""
 
     def __init__(
-            self,
-            base: RewardNet,
-            potential: Callable[[th.Tensor], th.Tensor],
-            discount_factor: float,
+        self,
+        base: RewardNet,
+        potential: Callable[[th.Tensor], th.Tensor],
+        discount_factor: float,
     ):
         """Setup a ShapedRewardNet instance.
 
@@ -497,11 +507,11 @@ def __init__(
         self.discount_factor = discount_factor
 
     def forward(
-            self,
-            state: th.Tensor,
-            action: th.Tensor,
-            next_state: th.Tensor,
-            done: th.Tensor,
+        self,
+        state: th.Tensor,
+        action: th.Tensor,
+        next_state: th.Tensor,
+        done: th.Tensor,
     ):
         base_reward_net_output = self.base(state, action, next_state, done)
         new_shaping_output = self.potential(next_state).flatten()
@@ -526,9 +536,9 @@ def forward(
         # length!
         new_shaping = (1 - done.float()) * new_shaping_output
         final_rew = (
-                base_reward_net_output
-                + self.discount_factor * new_shaping
-                - old_shaping_output
+            base_reward_net_output
+            + self.discount_factor * new_shaping
+            - old_shaping_output
         )
         assert final_rew.shape == state.shape[:1]
         return final_rew
@@ -551,18 +561,18 @@ class BasicShapedRewardNet(ShapedRewardNet):
     """
 
     def __init__(
-            self,
-            observation_space: gym.Space,
-            action_space: gym.Space,
-            *,
-            reward_hid_sizes: Sequence[int] = (32,),
-            potential_hid_sizes: Sequence[int] = (32, 32),
-            use_state: bool = True,
-            use_action: bool = True,
-            use_next_state: bool = False,
-            use_done: bool = False,
-            discount_factor: float = 0.99,
-            **kwargs,
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        *,
+        reward_hid_sizes: Sequence[int] = (32,),
+        potential_hid_sizes: Sequence[int] = (32, 32),
+        use_state: bool = True,
+        use_action: bool = True,
+        use_next_state: bool = False,
+        use_done: bool = False,
+        discount_factor: float = 0.99,
+        **kwargs,
     ):
         """Builds a simple shaped reward network.
 
@@ -611,10 +621,10 @@ class BasicPotentialMLP(nn.Module):
     """Simple implementation of a potential using an MLP."""
 
     def __init__(
-            self,
-            observation_space: gym.Space,
-            hid_sizes: Iterable[int],
-            **kwargs,
+        self,
+        observation_space: gym.Space,
+        hid_sizes: Iterable[int],
+        **kwargs,
     ):
         """Initialize the potential.
 
@@ -649,10 +659,10 @@ class RewardEnsemble(RewardNetWithVariance):
     members: nn.ModuleList
 
     def __init__(
-            self,
-            observation_space: gym.Space,
-            action_space: gym.Space,
-            members: Iterable[RewardNet],
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        members: Iterable[RewardNet],
     ):
         """Initialize the RewardEnsemble.
 
@@ -680,12 +690,12 @@ def num_members(self):
         return len(self.members)
 
     def predict_processed_all(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
-            **kwargs,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
+        **kwargs,
     ) -> np.ndarray:
         """Get the results of predict processed on all of the members.
 
@@ -711,12 +721,12 @@ def predict_processed_all(
 
     @th.no_grad()
     def predict_reward_moments(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
-            **kwargs,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
+        **kwargs,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Compute the standard deviation of the reward distribution for a batch.
 
@@ -749,23 +759,23 @@ def forward(self, *args) -> th.Tensor:
         raise NotImplementedError()
 
     def predict_processed(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
-            **kwargs,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
+        **kwargs,
     ) -> np.ndarray:
         """Return the mean of the ensemble members."""
         return self.predict(state, action, next_state, done, **kwargs)
 
     def predict(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
-            **kwargs,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
+        **kwargs,
     ):
         """Return the mean of the ensemble members."""
         mean, _ = self.predict_reward_moments(state, action, next_state, done, **kwargs)
@@ -799,13 +809,13 @@ def __init__(self, base: RewardNetWithVariance, default_alpha: float = 0.0):
         self.default_alpha = default_alpha
 
     def predict_processed(
-            self,
-            state: np.ndarray,
-            action: np.ndarray,
-            next_state: np.ndarray,
-            done: np.ndarray,
-            alpha: Optional[float] = None,
-            **kwargs,
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        next_state: np.ndarray,
+        done: np.ndarray,
+        alpha: Optional[float] = None,
+        **kwargs,
     ) -> np.ndarray:
         """Compute a lower/upper confidence bound on the reward without gradients.
 
diff --git a/src/imitation/rewards/reward_wrapper.py b/src/imitation/rewards/reward_wrapper.py
index 6a9243664..412234527 100644
--- a/src/imitation/rewards/reward_wrapper.py
+++ b/src/imitation/rewards/reward_wrapper.py
@@ -4,7 +4,9 @@
 from typing import Deque, Optional
 
 import numpy as np
-from stable_baselines3.common import callbacks, vec_env, logger as sb_logger
+from stable_baselines3.common import callbacks
+from stable_baselines3.common import logger as sb_logger
+from stable_baselines3.common import vec_env
 
 from imitation.rewards import reward_function
 
@@ -48,10 +50,10 @@ class RewardVecEnvWrapper(vec_env.VecEnvWrapper):
     """
 
     def __init__(
-            self,
-            venv: vec_env.VecEnv,
-            reward_fn: reward_function.RewardFn,
-            ep_history: int = 100,
+        self,
+        venv: vec_env.VecEnv,
+        reward_fn: reward_function.RewardFn,
+        ep_history: int = 100,
     ):
         """Builds RewardVecEnvWrapper.
 
diff --git a/src/imitation/scripts/analyze.py b/src/imitation/scripts/analyze.py
index b3654f9c5..3787ceb88 100644
--- a/src/imitation/scripts/analyze.py
+++ b/src/imitation/scripts/analyze.py
@@ -9,7 +9,7 @@
 import tempfile
 import warnings
 from collections import OrderedDict
-from typing import Any, Callable, List, Mapping, Optional, Sequence, Set, Iterable
+from typing import Any, Callable, Iterable, List, Mapping, Optional, Sequence, Set
 
 import pandas as pd
 from sacred.observers import FileStorageObserver
@@ -21,10 +21,10 @@
 
 @analysis_ex.capture
 def _gather_sacred_dicts(
-        source_dirs: Sequence[str],
-        run_name: Optional[str],
-        env_name: Optional[str],
-        skip_failed_runs: bool,
+    source_dirs: Sequence[str],
+    run_name: Optional[str],
+    env_name: Optional[str],
+    skip_failed_runs: bool,
 ) -> List[sacred_util.SacredDicts]:
     """Helper function for parsing and selecting Sacred experiment JSON files.
 
@@ -180,7 +180,7 @@ def _return_summaries(sd: sacred_util.SacredDicts) -> dict:
         # Assuming here that `result.imit_stats` and `result.expert_stats` are
         # formatted correctly.
         imit_expert_ratio = (
-                imit_stats["monitor_return_mean"] / expert_stats["return_mean"]
+            imit_stats["monitor_return_mean"] / expert_stats["return_mean"]
         )
     else:
         imit_expert_ratio = None
@@ -260,10 +260,10 @@ def _get_table_entry_fns_subset(table_verbosity: int) -> sd_to_table_entry_type:
 
 @analysis_ex.command
 def analyze_imitation(
-        csv_output_path: Optional[str],
-        tex_output_path: Optional[str],
-        print_table: bool,
-        table_verbosity: int,
+    csv_output_path: Optional[str],
+    tex_output_path: Optional[str],
+    print_table: bool,
+    table_verbosity: int,
 ) -> pd.DataFrame:
     """Parse Sacred logs and generate a DataFrame for imitation learning results.
 
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index 9878c6f63..ac5888817 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -3,7 +3,7 @@
 import contextlib
 import logging
 import os
-from typing import Any, Mapping, Sequence, Tuple, Union, Generator
+from typing import Any, Generator, Mapping, Sequence, Tuple, Union
 
 import sacred
 from stable_baselines3.common import vec_env
@@ -76,9 +76,9 @@ def fast():
 
 @common_ingredient.capture
 def make_log_dir(
-        _run,
-        log_dir: str,
-        log_level: Union[int, str],
+    _run,
+    log_dir: str,
+    log_level: Union[int, str],
 ) -> str:
     """Creates log directory and sets up symlink to Sacred logs.
 
@@ -105,8 +105,8 @@ def make_log_dir(
 
 @common_ingredient.capture
 def setup_logging(
-        _run,
-        log_format_strs: Sequence[str],
+    _run,
+    log_format_strs: Sequence[str],
 ) -> Tuple[imit_logger.HierarchicalLogger, str]:
     """Builds the imitation logger.
 
@@ -130,14 +130,14 @@ def setup_logging(
 @contextlib.contextmanager
 @common_ingredient.capture
 def make_venv(
-        _seed,
-        env_name: str,
-        num_vec: int,
-        parallel: bool,
-        log_dir: str,
-        max_episode_steps: int,
-        env_make_kwargs: Mapping[str, Any],
-        **kwargs,
+    _seed,
+    env_name: str,
+    num_vec: int,
+    parallel: bool,
+    log_dir: str,
+    max_episode_steps: int,
+    env_make_kwargs: Mapping[str, Any],
+    **kwargs,
 ) -> Generator[vec_env.VecEnv, None, None]:
     """Builds the vector environment.
 
diff --git a/src/imitation/scripts/common/reward.py b/src/imitation/scripts/common/reward.py
index 6b26e970b..78fc3e0a4 100644
--- a/src/imitation/scripts/common/reward.py
+++ b/src/imitation/scripts/common/reward.py
@@ -81,10 +81,10 @@ def config_hook(config, command_name, logger):
 
 
 def _make_reward_net(
-        venv: vec_env.VecEnv,
-        net_cls: Type[reward_nets.RewardNet],
-        net_kwargs: Mapping[str, Any],
-        normalize_output_layer: Optional[Type[networks.BaseNorm]],
+    venv: vec_env.VecEnv,
+    net_cls: Type[reward_nets.RewardNet],
+    net_kwargs: Mapping[str, Any],
+    normalize_output_layer: Optional[Type[networks.BaseNorm]],
 ):
     """Helper function for creating reward nets."""
     reward_net = net_cls(
@@ -104,13 +104,13 @@ def _make_reward_net(
 
 @reward_ingredient.capture
 def make_reward_net(
-        venv: vec_env.VecEnv,
-        net_cls: Type[reward_nets.RewardNet],
-        net_kwargs: Mapping[str, Any],
-        normalize_output_layer: Optional[Type[networks.BaseNorm]],
-        add_std_alpha: Optional[float],
-        ensemble_size: Optional[int],
-        ensemble_member_config: Optional[Mapping[str, Any]],
+    venv: vec_env.VecEnv,
+    net_cls: Type[reward_nets.RewardNet],
+    net_kwargs: Mapping[str, Any],
+    normalize_output_layer: Optional[Type[networks.BaseNorm]],
+    add_std_alpha: Optional[float],
+    ensemble_size: Optional[int],
+    ensemble_member_config: Optional[Mapping[str, Any]],
 ) -> reward_nets.RewardNet:
     """Builds a reward network.
 
@@ -149,11 +149,15 @@ def make_reward_net(
             for _ in range(ensemble_size)
         ]
 
-        reward_net: reward_nets.RewardNet = net_cls(venv.observation_space, venv.action_space, members)
+        reward_net: reward_nets.RewardNet = net_cls(
+            venv.observation_space, venv.action_space, members
+        )
 
         if add_std_alpha is not None:
             if not isinstance(reward_net, reward_nets.RewardNetWithVariance):
-                raise ValueError("add_std_alpha is only supported for reward nets with variance tracking.")
+                raise ValueError(
+                    "add_std_alpha is only supported for reward nets with variance tracking."
+                )
             reward_net = reward_nets.AddSTDRewardWrapper(
                 reward_net,
                 default_alpha=add_std_alpha,
diff --git a/src/imitation/scripts/common/wb.py b/src/imitation/scripts/common/wb.py
index 47b5e3fe5..e64edc978 100644
--- a/src/imitation/scripts/common/wb.py
+++ b/src/imitation/scripts/common/wb.py
@@ -26,12 +26,12 @@ def wandb_config():
 
 @wandb_ingredient.capture
 def wandb_init(
-        _run,
-        wandb_name_prefix: str,
-        wandb_tag: Optional[str],
-        wandb_kwargs: Mapping[str, Any],
-        wandb_additional_info: Mapping[str, Any],
-        log_dir: str,
+    _run,
+    wandb_name_prefix: str,
+    wandb_tag: Optional[str],
+    wandb_kwargs: Mapping[str, Any],
+    wandb_additional_info: Mapping[str, Any],
+    log_dir: str,
 ) -> None:
     """Putting everything together to get the W&B kwargs for wandb.init().
 
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 6b943532d..47e6220f3 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -53,19 +53,19 @@ def f(env: gym.Env, i: int) -> gym.Env:
 
 @eval_policy_ex.main
 def eval_policy(
-        _run,
-        _seed: int,
-        eval_n_timesteps: Optional[int],
-        eval_n_episodes: Optional[int],
-        render: bool,
-        render_fps: int,
-        videos: bool,
-        video_kwargs: Mapping[str, Any],
-        policy_type: Optional[str],
-        policy_path: Optional[str],
-        reward_type: Optional[str] = None,
-        reward_path: Optional[str] = None,
-        rollout_save_path: Optional[str] = None,
+    _run,
+    _seed: int,
+    eval_n_timesteps: Optional[int],
+    eval_n_episodes: Optional[int],
+    render: bool,
+    render_fps: int,
+    videos: bool,
+    video_kwargs: Mapping[str, Any],
+    policy_type: Optional[str],
+    policy_path: Optional[str],
+    reward_type: Optional[str] = None,
+    reward_path: Optional[str] = None,
+    rollout_save_path: Optional[str] = None,
 ):
     """Rolls a policy out in an environment, collecting statistics.
 
diff --git a/src/imitation/scripts/parallel.py b/src/imitation/scripts/parallel.py
index fc2c3625a..6e938d43f 100644
--- a/src/imitation/scripts/parallel.py
+++ b/src/imitation/scripts/parallel.py
@@ -3,7 +3,7 @@
 import collections.abc
 import copy
 import os
-from typing import Any, Callable, Mapping, Optional, Sequence, Dict
+from typing import Any, Callable, Dict, Mapping, Optional, Sequence
 
 import ray
 import ray.tune
@@ -15,15 +15,15 @@
 
 @parallel_ex.main
 def parallel(
-        sacred_ex_name: str,
-        run_name: str,
-        search_space: Mapping[str, Any],
-        base_named_configs: Sequence[str],
-        base_config_updates: Mapping[str, Any],
-        resources_per_trial: Mapping[str, Any],
-        init_kwargs: Mapping[str, Any],
-        local_dir: Optional[str],
-        upload_dir: Optional[str],
+    sacred_ex_name: str,
+    run_name: str,
+    search_space: Mapping[str, Any],
+    base_named_configs: Sequence[str],
+    base_config_updates: Mapping[str, Any],
+    resources_per_trial: Mapping[str, Any],
+    init_kwargs: Mapping[str, Any],
+    local_dir: Optional[str],
+    upload_dir: Optional[str],
 ) -> None:
     """Parallelize multiple runs of another Sacred Experiment using Ray Tune.
 
@@ -92,8 +92,8 @@ def parallel(
     # each Raylet.
     if sacred_ex_name == "train_adversarial":
         no_data_dir = (
-                "demonstrations.data_dir" not in base_config_updates
-                and "data_dir" not in base_config_updates.get("demonstrations", {})
+            "demonstrations.data_dir" not in base_config_updates
+            and "data_dir" not in base_config_updates.get("demonstrations", {})
         )
         if no_data_dir:
             data_dir = os.path.join(os.getcwd(), "data/")
@@ -122,10 +122,10 @@ def parallel(
 
 
 def _ray_tune_sacred_wrapper(
-        sacred_ex_name: str,
-        run_name: str,
-        base_named_configs: list,
-        base_config_updates: Mapping[str, Any],
+    sacred_ex_name: str,
+    run_name: str,
+    base_named_configs: list,
+    base_config_updates: Mapping[str, Any],
 ) -> Callable[[Mapping[str, Any], Any], Mapping[str, Any]]:
     """From an Experiment build a wrapped run function suitable for Ray Tune.
 
diff --git a/src/imitation/scripts/train_adversarial.py b/src/imitation/scripts/train_adversarial.py
index 291800d0f..35320e12c 100644
--- a/src/imitation/scripts/train_adversarial.py
+++ b/src/imitation/scripts/train_adversarial.py
@@ -66,14 +66,14 @@ def dummy_config():
 
 @train_adversarial_ex.capture
 def train_adversarial(
-        _run,
-        _seed: int,
-        show_config: bool,
-        algo_cls: Type[common.AdversarialTrainer],
-        algorithm_kwargs: Mapping[str, Any],
-        total_timesteps: int,
-        checkpoint_interval: int,
-        agent_path: Optional[str],
+    _run,
+    _seed: int,
+    show_config: bool,
+    algo_cls: Type[common.AdversarialTrainer],
+    algorithm_kwargs: Mapping[str, Any],
+    total_timesteps: int,
+    checkpoint_interval: int,
+    agent_path: Optional[str],
 ) -> Mapping[str, Mapping[str, float]]:
     """Train an adversarial-network-based imitation learning algorithm.
 
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index 8a87d928c..68a9ade02 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -3,7 +3,7 @@
 import logging
 import os.path as osp
 import warnings
-from typing import Any, Mapping, Optional, Type, Sequence, List, cast
+from typing import Any, List, Mapping, Optional, Sequence, Type, cast
 
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import policies, utils, vec_env
@@ -20,10 +20,10 @@
 
 @train_imitation_ex.capture(prefix="train")
 def make_policy(
-        venv: vec_env.VecEnv,
-        policy_cls: Type[policies.BasePolicy],
-        policy_kwargs: Mapping[str, Any],
-        agent_path: Optional[str],
+    venv: vec_env.VecEnv,
+    policy_cls: Type[policies.BasePolicy],
+    policy_kwargs: Mapping[str, Any],
+    agent_path: Optional[str],
 ) -> policies.BasePolicy:
     """Makes policy.
 
@@ -65,9 +65,9 @@ def make_policy(
 
 @train_imitation_ex.capture(prefix="dagger")
 def load_expert_policy(
-        venv: vec_env.VecEnv,
-        expert_policy_type: Optional[str],
-        expert_policy_path: Optional[str],
+    venv: vec_env.VecEnv,
+    expert_policy_type: Optional[str],
+    expert_policy_path: Optional[str],
 ) -> policies.BasePolicy:
     """Loads expert policy from `expert_policy_path`.
 
@@ -102,12 +102,12 @@ def load_expert_policy(
 
 @train_imitation_ex.capture
 def train_imitation(
-        _run,
-        bc_kwargs: Mapping[str, Any],
-        bc_train_kwargs: Mapping[str, Any],
-        dagger: Mapping[str, Any],
-        use_dagger: bool,
-        agent_path: Optional[str],
+    _run,
+    bc_kwargs: Mapping[str, Any],
+    bc_train_kwargs: Mapping[str, Any],
+    dagger: Mapping[str, Any],
+    use_dagger: bool,
+    agent_path: Optional[str],
 ) -> Mapping[str, Mapping[str, float]]:
     """Runs DAgger (if `use_dagger`) or BC (otherwise) training.
 
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 5aa861a2a..6011ca7f4 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -24,8 +24,8 @@
 
 
 def save_model(
-        agent_trainer: preference_comparisons.AgentTrainer,
-        save_path: str,
+    agent_trainer: preference_comparisons.AgentTrainer,
+    save_path: str,
 ):
     """Save the model as model.pkl."""
     serialize.save_stable_model(
@@ -35,9 +35,9 @@ def save_model(
 
 
 def save_checkpoint(
-        trainer: preference_comparisons.PreferenceComparisons,
-        save_path: str,
-        allow_save_policy: Optional[bool],
+    trainer: preference_comparisons.PreferenceComparisons,
+    save_path: str,
+    allow_save_policy: Optional[bool],
 ):
     """Save reward model and optionally policy."""
     os.makedirs(save_path, exist_ok=True)
@@ -56,30 +56,30 @@ def save_checkpoint(
 
 @train_preference_comparisons_ex.main
 def train_preference_comparisons(
-        _seed: int,
-        total_timesteps: int,
-        total_comparisons: int,
-        num_iterations: int,
-        comparison_queue_size: Optional[int],
-        fragment_length: int,
-        transition_oversampling: float,
-        initial_comparison_frac: float,
-        exploration_frac: float,
-        trajectory_path: Optional[str],
-        trajectory_generator_kwargs: Mapping[str, Any],
-        save_preferences: bool,
-        agent_path: Optional[str],
-        preference_model_kwargs: Mapping[str, Any],
-        reward_trainer_kwargs: Mapping[str, Any],
-        gatherer_cls: Type[preference_comparisons.PreferenceGatherer],
-        gatherer_kwargs: Mapping[str, Any],
-        active_selection: bool,
-        active_selection_oversampling: int,
-        uncertainty_on: str,
-        fragmenter_kwargs: Mapping[str, Any],
-        allow_variable_horizon: bool,
-        checkpoint_interval: int,
-        query_schedule: Union[str, type_aliases.Schedule],
+    _seed: int,
+    total_timesteps: int,
+    total_comparisons: int,
+    num_iterations: int,
+    comparison_queue_size: Optional[int],
+    fragment_length: int,
+    transition_oversampling: float,
+    initial_comparison_frac: float,
+    exploration_frac: float,
+    trajectory_path: Optional[str],
+    trajectory_generator_kwargs: Mapping[str, Any],
+    save_preferences: bool,
+    agent_path: Optional[str],
+    preference_model_kwargs: Mapping[str, Any],
+    reward_trainer_kwargs: Mapping[str, Any],
+    gatherer_cls: Type[preference_comparisons.PreferenceGatherer],
+    gatherer_kwargs: Mapping[str, Any],
+    active_selection: bool,
+    active_selection_oversampling: int,
+    uncertainty_on: str,
+    fragmenter_kwargs: Mapping[str, Any],
+    allow_variable_horizon: bool,
+    checkpoint_interval: int,
+    query_schedule: Union[str, type_aliases.Schedule],
 ) -> Mapping[str, Any]:
     """Train a reward model using preference comparisons.
 
@@ -191,10 +191,12 @@ def train_preference_comparisons(
                 **trajectory_generator_kwargs,
             )
 
-        fragmenter: preference_comparisons.Fragmenter = preference_comparisons.RandomFragmenter(
-            **fragmenter_kwargs,
-            seed=_seed,
-            custom_logger=custom_logger,
+        fragmenter: preference_comparisons.Fragmenter = (
+            preference_comparisons.RandomFragmenter(
+                **fragmenter_kwargs,
+                seed=_seed,
+                custom_logger=custom_logger,
+            )
         )
         preference_model = preference_comparisons.PreferenceModel(
             **preference_model_kwargs,
diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
index fa222dad2..b36f9518b 100644
--- a/src/imitation/scripts/train_rl.py
+++ b/src/imitation/scripts/train_rl.py
@@ -28,19 +28,19 @@
 
 @train_rl_ex.main
 def train_rl(
-        *,
-        total_timesteps: int,
-        normalize_reward: bool,
-        normalize_kwargs: dict,
-        reward_type: Optional[str],
-        reward_path: Optional[str],
-        load_reward_kwargs: Optional[Mapping[str, Any]],
-        rollout_save_final: bool,
-        rollout_save_n_timesteps: Optional[int],
-        rollout_save_n_episodes: Optional[int],
-        policy_save_interval: int,
-        policy_save_final: bool,
-        agent_path: Optional[str],
+    *,
+    total_timesteps: int,
+    normalize_reward: bool,
+    normalize_kwargs: dict,
+    reward_type: Optional[str],
+    reward_path: Optional[str],
+    load_reward_kwargs: Optional[Mapping[str, Any]],
+    rollout_save_final: bool,
+    rollout_save_n_timesteps: Optional[int],
+    rollout_save_n_episodes: Optional[int],
+    policy_save_interval: int,
+    policy_save_final: bool,
+    agent_path: Optional[str],
 ) -> Mapping[str, float]:
     """Trains an expert policy from scratch and saves the rollouts and policy.
 
@@ -120,7 +120,9 @@ def train_rl(
                 )
 
         if policy_save_interval > 0:
-            save_policy_callback: callbacks.EventCallback = serialize.SavePolicyCallback(policy_dir)
+            save_policy_callback: callbacks.EventCallback = (
+                serialize.SavePolicyCallback(policy_dir)
+            )
             save_policy_callback = callbacks.EveryNTimesteps(
                 policy_save_interval,
                 save_policy_callback,
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 8875cb211..5e45525bd 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -4,7 +4,7 @@
 import datetime
 import os
 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
 
@@ -26,7 +26,7 @@ def _build_output_formats(
         A list of output formats, one corresponding to each `format_strs`.
     """
     os.makedirs(folder, exist_ok=True)
-    output_formats = []
+    output_formats: List[sb_logger.KVWriter] = []
     for f in format_strs:
         if f == "wandb":
             output_formats.append(WandbOutputFormat())
@@ -43,6 +43,12 @@ class HierarchicalLogger(sb_logger.Logger):
     top-level (root) logger.
     """
 
+    default_logger: sb_logger.Logger
+    current_logger: Optional[sb_logger.Logger]
+    _cached_loggers: Dict[str, sb_logger.Logger]
+    _subdir: Optional[str]
+    format_strs: Sequence[str]
+
     def __init__(
         self,
         default_logger: sb_logger.Logger,
@@ -108,10 +114,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 = types.path_to_str(subdir)
         if subdir in self._cached_loggers:
             logger = self._cached_loggers[subdir]
         else:
-            subdir = types.path_to_str(subdir)
+            assert self.default_logger.dir is not None
             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)
diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index ba9680967..50fb5350f 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -3,7 +3,7 @@
 import collections
 import contextlib
 import functools
-from typing import Iterable, Optional, Type
+from typing import Iterable, Optional, OrderedDict, Type
 
 import torch as th
 from torch import nn
@@ -205,7 +205,7 @@ def build_mlp(
     dropout_prob: float = 0.0,
     squeeze_output: bool = False,
     flatten_input: bool = False,
-    normalize_input_layer: Optional[Type[nn.Module]] = None,
+    normalize_input_layer: Optional[Type[BaseNorm]] = None,
 ) -> nn.Module:
     """Constructs a Torch MLP.
 
@@ -234,7 +234,7 @@ def build_mlp(
     Raises:
         ValueError: if squeeze_output was supplied with out_size!=1.
     """
-    layers = collections.OrderedDict()
+    layers: OrderedDict[str, nn.Module] = collections.OrderedDict()
 
     if name is None:
         prefix = ""
diff --git a/src/imitation/util/sacred.py b/src/imitation/util/sacred.py
index d6d7fe8bd..b96f872df 100644
--- a/src/imitation/util/sacred.py
+++ b/src/imitation/util/sacred.py
@@ -7,6 +7,8 @@
 from typing import Any, Callable, NamedTuple, Sequence, Union
 
 import sacred
+import sacred.observers
+import sacred.run
 
 from imitation.data import types
 
@@ -78,6 +80,8 @@ def filter_subdirs(
 
 def build_sacred_symlink(log_dir: types.AnyPath, run: sacred.run.Run) -> None:
     """Constructs a symlink "{log_dir}/sacred" => "${SACRED_PATH}"."""
+    if isinstance(log_dir, bytes):
+        log_dir = log_dir.decode("utf-8")
     log_dir = pathlib.Path(log_dir)
 
     sacred_dir = get_sacred_dir_from_run(run)
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index f9d4e83bc..a59829bdc 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -10,6 +10,7 @@
     Callable,
     Iterable,
     Iterator,
+    List,
     Mapping,
     Optional,
     Sequence,
@@ -99,7 +100,7 @@ def make_vec_env(
     spec = gym.spec(env_name)
     env_make_kwargs = env_make_kwargs or {}
 
-    def make_env(i, this_seed):
+    def make_env(i, this_seed) -> gym.Env:
         # Previously, we directly called `gym.make(env_name)`, but running
         # `imitation.scripts.train_adversarial` within `imitation.scripts.parallel`
         # created a weird interaction between Gym and Ray -- `gym.make` would fail
@@ -138,7 +139,9 @@ def make_env(i, this_seed):
 
     rng = np.random.RandomState(seed)
     env_seeds = rng.randint(0, (1 << 31) - 1, (n_envs,))
-    env_fns = [functools.partial(make_env, i, s) for i, s in enumerate(env_seeds)]
+    env_fns: List[Callable[[], gym.Env]] = [
+        functools.partial(make_env, i, s) for i, s in enumerate(env_seeds)
+    ]
     if parallel:
         # See GH hill-a/stable-baselines issue #217
         return SubprocVecEnv(env_fns, start_method="forkserver")
diff --git a/src/imitation/util/video_wrapper.py b/src/imitation/util/video_wrapper.py
index 5fc2ae4f5..6da02aa70 100644
--- a/src/imitation/util/video_wrapper.py
+++ b/src/imitation/util/video_wrapper.py
@@ -1,6 +1,7 @@
 """Wrapper to record rendered video frames from an environment."""
 
 import os
+from typing import Optional
 
 import gym
 from gym.wrappers.monitoring import video_recorder
@@ -11,6 +12,11 @@
 class VideoWrapper(gym.Wrapper):
     """Creates videos from wrapped environment by calling render after each timestep."""
 
+    episode_id: int
+    video_recorder: Optional[video_recorder.VideoRecorder]
+    single_video: bool
+    directory: str
+
     def __init__(
         self,
         env: gym.Env,
@@ -33,7 +39,7 @@ def __init__(
         self.video_recorder = None
         self.single_video = single_video
 
-        self.directory = os.path.abspath(directory)
+        self.directory = str(os.path.abspath(directory))
         os.makedirs(self.directory)
 
     def _reset_video_recorder(self) -> None:

From e68e32cd1622c99e956b16933d17a3d3132f34ab Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 19:29:32 +0200
Subject: [PATCH 039/187] Linting and formatting

---
 src/imitation/rewards/reward_nets.py   |  2 +-
 src/imitation/rewards/serialize.py     | 10 ++++++----
 src/imitation/scripts/common/reward.py |  7 +++++--
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index 60f7853cd..cae52e8e9 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -386,7 +386,7 @@ def __init__(
             combined_size += 1
 
         # kwargs except for in_size, out_size, squeeze_output keys,
-        # so they are not overriden.
+        # so they are not overridden.
         kwargs = {
             k: v
             for k, v in kwargs.items()
diff --git a/src/imitation/rewards/serialize.py b/src/imitation/rewards/serialize.py
index d580cfa6a..07817291b 100644
--- a/src/imitation/rewards/serialize.py
+++ b/src/imitation/rewards/serialize.py
@@ -192,7 +192,8 @@ def f(
     value=lambda path, _, **kwargs: ValidateRewardFn(
         _make_functional(
             _validate_wrapper_structure(
-                th.load(str(path)), {(reward_nets.ShapedRewardNet,)}
+                th.load(str(path)),
+                {(reward_nets.ShapedRewardNet,)},
             ),
         ),
     ),
@@ -202,7 +203,7 @@ def f(
     key="RewardNet_unshaped",
     value=lambda path, _, **kwargs: ValidateRewardFn(
         _make_functional(
-            _strip_wrappers(th.load(str(path)), (reward_nets.ShapedRewardNet,))
+            _strip_wrappers(th.load(str(path)), (reward_nets.ShapedRewardNet,)),
         ),
     ),
 )
@@ -212,7 +213,8 @@ def f(
     value=lambda path, _, **kwargs: ValidateRewardFn(
         _make_functional(
             _validate_wrapper_structure(
-                th.load(str(path)), {(reward_nets.NormalizedRewardNet,)}
+                th.load(str(path)),
+                {(reward_nets.NormalizedRewardNet,)},
             ),
             attr="predict_processed",
             default_kwargs={"update_stats": False},
@@ -225,7 +227,7 @@ def f(
     key="RewardNet_unnormalized",
     value=lambda path, _, **kwargs: ValidateRewardFn(
         _make_functional(
-            _strip_wrappers(th.load(str(path)), (reward_nets.NormalizedRewardNet,))
+            _strip_wrappers(th.load(str(path)), (reward_nets.NormalizedRewardNet,)),
         ),
     ),
 )
diff --git a/src/imitation/scripts/common/reward.py b/src/imitation/scripts/common/reward.py
index 78fc3e0a4..a00498138 100644
--- a/src/imitation/scripts/common/reward.py
+++ b/src/imitation/scripts/common/reward.py
@@ -150,13 +150,16 @@ def make_reward_net(
         ]
 
         reward_net: reward_nets.RewardNet = net_cls(
-            venv.observation_space, venv.action_space, members
+            venv.observation_space,
+            venv.action_space,
+            members,
         )
 
         if add_std_alpha is not None:
             if not isinstance(reward_net, reward_nets.RewardNetWithVariance):
                 raise ValueError(
-                    "add_std_alpha is only supported for reward nets with variance tracking."
+                    "add_std_alpha is only supported for "
+                    "reward nets with variance tracking.",
                 )
             reward_net = reward_nets.AddSTDRewardWrapper(
                 reward_net,

From b062d47fea3b549e836c11fbdb72cd029a23174f Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 19:54:09 +0200
Subject: [PATCH 040/187] Bug fixes for test errors

---
 src/imitation/rewards/reward_nets.py     | 2 +-
 src/imitation/scripts/train_imitation.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index cae52e8e9..b607af1ea 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -390,7 +390,7 @@ def __init__(
         kwargs = {
             k: v
             for k, v in kwargs.items()
-            if k not in ("in_size", "out_size", "squeeze_output")
+            if k not in ("in_size", "out_size", "squeeze_output", "hid_sizes")
         }
         self.mlp = networks.build_mlp(
             hid_sizes=(32, 32),
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index 68a9ade02..e206f6b9a 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -179,10 +179,10 @@ def train_imitation(
     #  just for adding type annotations.
     trajectories: List[types.TrajectoryWithRew]
     if use_dagger:
+        trajectories = cast(List[types.TrajectoryWithRew], model._all_demos)
+    else:
         assert expert_trajs is not None
         trajectories = cast(List[types.TrajectoryWithRew], expert_trajs)
-    else:
-        trajectories = cast(List[types.TrajectoryWithRew], model._all_demos)
 
     return {
         "imit_stats": imit_stats,

From 49cb0144e1252927ab0c4e0a6ca05e3d637505e3 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 20:12:28 +0200
Subject: [PATCH 041/187] Linting and typing

---
 src/imitation/data/rollout.py           | 3 +++
 src/imitation/rewards/reward_wrapper.py | 6 ++----
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index a9ded8bcf..5d070d855 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -28,6 +28,9 @@ def unwrap_traj(traj: types.TrajectoryWithRew) -> types.TrajectoryWithRew:
 
     Returns:
         A copy of `traj` with replaced `obs` and `rews` fields.
+
+    Raises:
+        ValueError: If `traj.infos` is None
     """
     if traj.infos is None:
         raise ValueError("Trajectory must have infos to unwrap")
diff --git a/src/imitation/rewards/reward_wrapper.py b/src/imitation/rewards/reward_wrapper.py
index 412234527..7afa551b3 100644
--- a/src/imitation/rewards/reward_wrapper.py
+++ b/src/imitation/rewards/reward_wrapper.py
@@ -1,7 +1,7 @@
 """Common wrapper for adding custom reward values to an environment."""
 
 import collections
-from typing import Deque, Optional
+from typing import Deque
 
 import numpy as np
 from stable_baselines3.common import callbacks
@@ -14,8 +14,6 @@
 class WrappedRewardCallback(callbacks.BaseCallback):
     """Logs mean wrapped reward as part of RL (or other) training."""
 
-    logger: Optional[sb_logger.Logger]  # type: ignore[assignment]
-
     def __init__(self, episode_rewards: Deque[float], *args, **kwargs):
         """Builds WrappedRewardCallback.
 
@@ -34,7 +32,7 @@ def _on_rollout_start(self) -> None:
         if len(self.episode_rewards) == 0:
             return
         mean = sum(self.episode_rewards) / len(self.episode_rewards)
-        assert self.logger is not None
+        assert isinstance(self.logger, sb_logger.Logger)
         self.logger.record("rollout/ep_rew_wrapped_mean", mean)
 
 

From a332025ffe87350105ae30b6ba01b2515bb28387 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 21:00:24 +0200
Subject: [PATCH 042/187] Improve typing in algorithms

---
 src/imitation/algorithms/adversarial/airl.py  |  29 +-
 .../algorithms/adversarial/common.py          | 113 ++++----
 src/imitation/algorithms/base.py              |  37 ++-
 .../algorithms/preference_comparisons.py      | 268 +++++++++---------
 src/imitation/util/util.py                    |  24 +-
 5 files changed, 247 insertions(+), 224 deletions(-)

diff --git a/src/imitation/algorithms/adversarial/airl.py b/src/imitation/algorithms/adversarial/airl.py
index 7676b4548..f8ae5f78d 100644
--- a/src/imitation/algorithms/adversarial/airl.py
+++ b/src/imitation/algorithms/adversarial/airl.py
@@ -1,4 +1,5 @@
 """Adversarial Inverse Reinforcement Learning (AIRL)."""
+from typing import Optional
 
 import torch as th
 from stable_baselines3.common import base_class, policies, vec_env
@@ -18,14 +19,14 @@ class AIRL(common.AdversarialTrainer):
     """
 
     def __init__(
-        self,
-        *,
-        demonstrations: base.AnyTransitions,
-        demo_batch_size: int,
-        venv: vec_env.VecEnv,
-        gen_algo: base_class.BaseAlgorithm,
-        reward_net: reward_nets.RewardNet,
-        **kwargs,
+            self,
+            *,
+            demonstrations: base.AnyTransitions,
+            demo_batch_size: int,
+            venv: vec_env.VecEnv,
+            gen_algo: base_class.BaseAlgorithm,
+            reward_net: reward_nets.RewardNet,
+            **kwargs,
     ):
         """Builds an AIRL trainer.
 
@@ -64,12 +65,12 @@ def __init__(
             )
 
     def logits_expert_is_high(
-        self,
-        state: th.Tensor,
-        action: th.Tensor,
-        next_state: th.Tensor,
-        done: th.Tensor,
-        log_policy_act_prob: th.Tensor,
+            self,
+            state: th.Tensor,
+            action: th.Tensor,
+            next_state: th.Tensor,
+            done: th.Tensor,
+            log_policy_act_prob: Optional[th.Tensor] = None,
     ) -> th.Tensor:
         r"""Compute the discriminator's logits for each state-action sample.
 
diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
index eee6937e4..d08a55c4f 100644
--- a/src/imitation/algorithms/adversarial/common.py
+++ b/src/imitation/algorithms/adversarial/common.py
@@ -4,7 +4,7 @@
 import dataclasses
 import logging
 import os
-from typing import Callable, Mapping, Optional, Sequence, Tuple, Type
+from typing import Callable, Mapping, Optional, Sequence, Tuple, Type, Iterable, Iterator, overload
 
 import numpy as np
 import torch as th
@@ -21,9 +21,9 @@
 
 
 def compute_train_stats(
-    disc_logits_expert_is_high: th.Tensor,
-    labels_expert_is_one: th.Tensor,
-    disc_loss: th.Tensor,
+        disc_logits_expert_is_high: th.Tensor,
+        labels_expert_is_one: th.Tensor,
+        disc_loss: th.Tensor,
 ) -> Mapping[str, float]:
     """Train statistics for GAIL/AIRL discriminator.
 
@@ -63,7 +63,7 @@ def compute_train_stats(
         else:
             # float() is defensive, since we cannot divide Torch tensors by
             # Python ints
-            expert_acc = _n_pred_expert / float(n_expert)
+            expert_acc = _n_pred_expert.item() / float(n_expert)
 
         _n_pred_gen = th.sum(th.logical_and(bin_is_generated_true, correct_vec))
         _n_gen_or_1 = max(1, n_generated)
@@ -103,25 +103,28 @@ class AdversarialTrainer(base.DemonstrationAlgorithm[types.Transitions]):
     If `debug_use_ground_truth=True` was passed into the initializer then
     `self.venv_train` is the same as `self.venv`."""
 
+    _demo_data_loader: Optional[Iterable[base.TransitionMapping]]
+    _endless_expert_iterator: Optional[Iterator[base.TransitionMapping]]
+
     def __init__(
-        self,
-        *,
-        demonstrations: base.AnyTransitions,
-        demo_batch_size: int,
-        venv: vec_env.VecEnv,
-        gen_algo: base_class.BaseAlgorithm,
-        reward_net: reward_nets.RewardNet,
-        n_disc_updates_per_round: int = 2,
-        log_dir: str = "output/",
-        disc_opt_cls: Type[th.optim.Optimizer] = th.optim.Adam,
-        disc_opt_kwargs: Optional[Mapping] = None,
-        gen_train_timesteps: Optional[int] = None,
-        gen_replay_buffer_capacity: Optional[int] = None,
-        custom_logger: Optional[logger.HierarchicalLogger] = None,
-        init_tensorboard: bool = False,
-        init_tensorboard_graph: bool = False,
-        debug_use_ground_truth: bool = False,
-        allow_variable_horizon: bool = False,
+            self,
+            *,
+            demonstrations: base.AnyTransitions,
+            demo_batch_size: int,
+            venv: vec_env.VecEnv,
+            gen_algo: base_class.BaseAlgorithm,
+            reward_net: reward_nets.RewardNet,
+            n_disc_updates_per_round: int = 2,
+            log_dir: str = "output/",
+            disc_opt_cls: Type[th.optim.Optimizer] = th.optim.Adam,
+            disc_opt_kwargs: Optional[Mapping] = None,
+            gen_train_timesteps: Optional[int] = None,
+            gen_replay_buffer_capacity: Optional[int] = None,
+            custom_logger: Optional[logger.HierarchicalLogger] = None,
+            init_tensorboard: bool = False,
+            init_tensorboard_graph: bool = False,
+            debug_use_ground_truth: bool = False,
+            allow_variable_horizon: bool = False,
     ):
         """Builds AdversarialTrainer.
 
@@ -212,6 +215,7 @@ def __init__(
             self.venv_wrapped = venv
             self.gen_callback = None
         else:
+            # TODO(juan) this is a bug; venv is going unused.
             venv = self.venv_wrapped = reward_wrapper.RewardVecEnvWrapper(
                 venv,
                 reward_fn=self.reward_train.predict_processed,
@@ -227,6 +231,8 @@ def __init__(
             assert gen_algo_env is not None
             self.gen_train_timesteps = gen_algo_env.num_envs
             if hasattr(self.gen_algo, "n_steps"):  # on policy
+                # TODO(juan) this looks like a bug; could not find n_steps
+                #  defined anywhere in the codebase.
                 self.gen_train_timesteps *= self.gen_algo.n_steps
         else:
             self.gen_train_timesteps = gen_train_timesteps
@@ -240,16 +246,17 @@ def __init__(
 
     @property
     def policy(self) -> policies.BasePolicy:
+        # TODO(juan) either change the return type to optional or add an assertion.
         return self.gen_algo.policy
 
     @abc.abstractmethod
     def logits_expert_is_high(
-        self,
-        state: th.Tensor,
-        action: th.Tensor,
-        next_state: th.Tensor,
-        done: th.Tensor,
-        log_policy_act_prob: Optional[th.Tensor] = None,
+            self,
+            state: th.Tensor,
+            action: th.Tensor,
+            next_state: th.Tensor,
+            done: th.Tensor,
+            log_policy_act_prob: Optional[th.Tensor] = None,
     ) -> th.Tensor:
         """Compute the discriminator's logits for each state-action sample.
 
@@ -288,13 +295,14 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
         self._endless_expert_iterator = util.endless_iter(self._demo_data_loader)
 
     def _next_expert_batch(self) -> Mapping:
+        assert self._endless_expert_iterator is not None
         return next(self._endless_expert_iterator)
 
     def train_disc(
-        self,
-        *,
-        expert_samples: Optional[Mapping] = None,
-        gen_samples: Optional[Mapping] = None,
+            self,
+            *,
+            expert_samples: Optional[Mapping] = None,
+            gen_samples: Optional[Mapping] = None,
     ) -> Optional[Mapping[str, float]]:
         """Perform a single discriminator update, optionally using provided samples.
 
@@ -358,9 +366,9 @@ def train_disc(
         return train_stats
 
     def train_gen(
-        self,
-        total_timesteps: Optional[int] = None,
-        learn_kwargs: Optional[Mapping] = None,
+            self,
+            total_timesteps: Optional[int] = None,
+            learn_kwargs: Optional[Mapping] = None,
     ) -> None:
         """Trains the generator to maximize the discriminator loss.
 
@@ -394,9 +402,9 @@ def train_gen(
         self._gen_replay_buffer.store(gen_samples)
 
     def train(
-        self,
-        total_timesteps: int,
-        callback: Optional[Callable[[int], None]] = None,
+            self,
+            total_timesteps: int,
+            callback: Optional[Callable[[int], None]] = None,
     ) -> None:
         """Alternates between training the generator and discriminator.
 
@@ -429,14 +437,23 @@ def train(
                 callback(r)
             self.logger.dump(self._global_step)
 
+    @overload
+    def _torchify_array(self, ndarray: np.ndarray) -> th.Tensor:
+        ...
+
+    @overload
+    def _torchify_array(self, ndarray: None) -> None:
+        ...
+    
     def _torchify_array(self, ndarray: Optional[np.ndarray]) -> Optional[th.Tensor]:
         if ndarray is not None:
             return th.as_tensor(ndarray, device=self.reward_train.device)
+        return None
 
     def _get_log_policy_act_prob(
-        self,
-        obs_th: th.Tensor,
-        acts_th: th.Tensor,
+            self,
+            obs_th: th.Tensor,
+            acts_th: th.Tensor,
     ) -> Optional[th.Tensor]:
         """Evaluates the given actions on the given observations.
 
@@ -474,10 +491,10 @@ def _get_log_policy_act_prob(
         return log_policy_act_prob_th
 
     def _make_disc_train_batch(
-        self,
-        *,
-        gen_samples: Optional[Mapping] = None,
-        expert_samples: Optional[Mapping] = None,
+            self,
+            *,
+            gen_samples: Optional[Mapping] = None,
+            expert_samples: Optional[Mapping] = None,
     ) -> Mapping[str, th.Tensor]:
         """Build and return training batch for the next discriminator update.
 
@@ -502,8 +519,8 @@ def _make_disc_train_batch(
                 raise RuntimeError(
                     "No generator samples for training. " "Call `train_gen()` first.",
                 )
-            gen_samples = self._gen_replay_buffer.sample(self.demo_batch_size)
-            gen_samples = types.dataclass_quick_asdict(gen_samples)
+            gen_samples_dataclass = self._gen_replay_buffer.sample(self.demo_batch_size)
+            gen_samples = types.dataclass_quick_asdict(gen_samples_dataclass)
 
         n_gen = len(gen_samples["obs"])
         n_expert = len(expert_samples["obs"])
diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index ebad52125..2d8ca04ea 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -25,10 +25,10 @@ class BaseImitationAlgorithm(abc.ABC):
     """Horizon of trajectories seen so far (None if no trajectories seen)."""
 
     def __init__(
-        self,
-        *,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-        allow_variable_horizon: bool = False,
+            self,
+            *,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            allow_variable_horizon: bool = False,
     ):
         """Creates an imitation learning algorithm.
 
@@ -124,11 +124,11 @@ class DemonstrationAlgorithm(BaseImitationAlgorithm, Generic[TransitionKind]):
     """An algorithm that learns from demonstration: BC, IRL, etc."""
 
     def __init__(
-        self,
-        *,
-        demonstrations: Optional[AnyTransitions],
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-        allow_variable_horizon: bool = False,
+            self,
+            *,
+            demonstrations: Optional[AnyTransitions],
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            allow_variable_horizon: bool = False,
     ):
         """Creates an algorithm that learns from demonstrations.
 
@@ -177,9 +177,9 @@ class _WrappedDataLoader:
     """Wraps a data loader (batch iterable) and checks for specified batch size."""
 
     def __init__(
-        self,
-        data_loader: Iterable[TransitionMapping],
-        expected_batch_size: int,
+            self,
+            data_loader: Iterable[TransitionMapping],
+            expected_batch_size: int,
     ):
         """Builds _WrapedDataLoader.
 
@@ -215,9 +215,9 @@ def __iter__(self):
 
 
 def make_data_loader(
-    transitions: AnyTransitions,
-    batch_size: int,
-    data_loader_kwargs: Optional[Mapping[str, Any]] = None,
+        transitions: AnyTransitions,
+        batch_size: int,
+        data_loader_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> Iterable[TransitionMapping]:
     """Converts demonstration data to Torch data loader.
 
@@ -256,14 +256,13 @@ def make_data_loader(
                 f"is smaller than batch size {batch_size}.",
             )
 
-        extra_kwargs = dict(shuffle=True, drop_last=True)
-        if data_loader_kwargs is not None:
-            extra_kwargs.update(data_loader_kwargs)
         return th_data.DataLoader(
             transitions,
             batch_size=batch_size,
             collate_fn=types.transitions_collate_fn,
-            **extra_kwargs,
+            shuffle=True,
+            drop_last=True,
+            **(data_loader_kwargs or {}),
         )
     elif isinstance(transitions, Iterable):
         return _WrappedDataLoader(transitions, batch_size)
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 532a83237..18d4f3390 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -96,10 +96,10 @@ class TrajectoryDataset(TrajectoryGenerator):
     """A fixed dataset of trajectories."""
 
     def __init__(
-        self,
-        trajectories: Sequence[TrajectoryWithRew],
-        seed: Optional[int] = None,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            trajectories: Sequence[TrajectoryWithRew],
+            seed: Optional[int] = None,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Creates a dataset loaded from `path`.
 
@@ -123,15 +123,15 @@ class AgentTrainer(TrajectoryGenerator):
     """Wrapper for training an SB3 algorithm on an arbitrary reward function."""
 
     def __init__(
-        self,
-        algorithm: base_class.BaseAlgorithm,
-        reward_fn: Union[reward_function.RewardFn, reward_nets.RewardNet],
-        venv: vec_env.VecEnv,
-        exploration_frac: float = 0.0,
-        switch_prob: float = 0.5,
-        random_prob: float = 0.5,
-        seed: Optional[int] = None,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            algorithm: base_class.BaseAlgorithm,
+            reward_fn: Union[reward_function.RewardFn, reward_nets.RewardNet],
+            venv: vec_env.VecEnv,
+            exploration_frac: float = 0.0,
+            switch_prob: float = 0.5,
+            random_prob: float = 0.5,
+            seed: Optional[int] = None,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the agent trainer.
 
@@ -310,8 +310,8 @@ def logger(self, value: imit_logger.HierarchicalLogger):
 
 
 def _get_trajectories(
-    trajectories: Sequence[TrajectoryWithRew],
-    steps: int,
+        trajectories: Sequence[TrajectoryWithRew],
+        steps: int,
 ) -> Sequence[TrajectoryWithRew]:
     """Get enough trajectories to have at least `steps` transitions in total."""
     if steps == 0:
@@ -339,11 +339,11 @@ class PreferenceModel(nn.Module):
     """Class to convert two fragments' rewards into preference probability."""
 
     def __init__(
-        self,
-        model: reward_nets.RewardNet,
-        noise_prob: float = 0.0,
-        discount_factor: float = 1.0,
-        threshold: float = 50,
+            self,
+            model: reward_nets.RewardNet,
+            noise_prob: float = 0.0,
+            discount_factor: float = 1.0,
+            threshold: float = 50,
     ):
         """Create Preference Prediction Model.
 
@@ -388,9 +388,9 @@ def __init__(
                 self.member_pref_models.append(member_pref_model)
 
     def forward(
-        self,
-        fragment_pairs: Sequence[TrajectoryPair],
-        ensemble_member_index: Optional[int] = None,
+            self,
+            fragment_pairs: Sequence[TrajectoryPair],
+            ensemble_member_index: Optional[int] = None,
     ) -> Tuple[th.Tensor, Optional[th.Tensor]]:
         """Computes the preference probability of the first fragment for all pairs.
 
@@ -451,8 +451,8 @@ def forward(
         return probs, (gt_probs if gt_reward_available else None)
 
     def rewards(
-        self,
-        transitions: Transitions,
+            self,
+            transitions: Transitions,
     ) -> th.Tensor:
         """Computes the reward for all transitions.
 
@@ -469,6 +469,8 @@ def rewards(
         next_state = transitions.next_obs
         done = transitions.dones
         if self.is_ensemble:
+            # TODO(juan) for some super strange reason mypy thinks
+            #  self.model.predict_processed_all is a tensor?
             rews = self.model.predict_processed_all(state, action, next_state, done)
             assert rews.shape == (len(state), self.model.num_members)
             return util.safe_to_tensor(rews).to(self.model.device)
@@ -479,9 +481,9 @@ def rewards(
             return rews
 
     def probability(
-        self,
-        rews1: th.Tensor,
-        rews2: th.Tensor,
+            self,
+            rews1: th.Tensor,
+            rews2: th.Tensor,
     ) -> th.Tensor:
         """Computes the Boltzmann rational probability that the first trajectory is best.
 
@@ -539,10 +541,10 @@ def __init__(self, custom_logger: Optional[imit_logger.HierarchicalLogger] = Non
 
     @abc.abstractmethod
     def __call__(
-        self,
-        trajectories: Sequence[TrajectoryWithRew],
-        fragment_length: int,
-        num_pairs: int,
+            self,
+            trajectories: Sequence[TrajectoryWithRew],
+            fragment_length: int,
+            num_pairs: int,
     ) -> Sequence[TrajectoryWithRewPair]:
         """Create fragment pairs out of a sequence of trajectories.
 
@@ -570,10 +572,10 @@ class RandomFragmenter(Fragmenter):
     """
 
     def __init__(
-        self,
-        seed: Optional[float] = None,
-        warning_threshold: int = 10,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            seed: Optional[float] = None,
+            warning_threshold: int = 10,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the fragmenter.
 
@@ -589,10 +591,10 @@ def __init__(
         self.warning_threshold = warning_threshold
 
     def __call__(
-        self,
-        trajectories: Sequence[TrajectoryWithRew],
-        fragment_length: int,
-        num_pairs: int,
+            self,
+            trajectories: Sequence[TrajectoryWithRew],
+            fragment_length: int,
+            num_pairs: int,
     ) -> Sequence[TrajectoryWithRewPair]:
         fragments: List[TrajectoryWithRew] = []
 
@@ -622,8 +624,8 @@ def __call__(
                 "of fragment pairs. Some transitions will appear multiple times.",
             )
         elif (
-            self.warning_threshold
-            and sum(weights) < self.warning_threshold * num_transitions
+                self.warning_threshold
+                and sum(weights) < self.warning_threshold * num_transitions
         ):
             # If the number of available transitions is not much larger
             # than the number of requires ones, we already give a warning.
@@ -643,7 +645,7 @@ def __call__(
             end = start + fragment_length
             terminal = (end == n) and traj.terminal
             fragment = TrajectoryWithRew(
-                obs=traj.obs[start : end + 1],
+                obs=traj.obs[start: end + 1],
                 acts=traj.acts[start:end],
                 infos=traj.infos[start:end] if traj.infos is not None else None,
                 rews=traj.rews[start:end],
@@ -665,12 +667,12 @@ class ActiveSelectionFragmenter(Fragmenter):
     """
 
     def __init__(
-        self,
-        preference_model: PreferenceModel,
-        base_fragmenter: Fragmenter,
-        fragment_sample_factor: float,
-        uncertainty_on: str = "logits",
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            preference_model: PreferenceModel,
+            base_fragmenter: Fragmenter,
+            fragment_sample_factor: float,
+            uncertainty_on: str = "logits",
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the active selection fragmenter.
 
@@ -711,10 +713,10 @@ def raise_uncertainty_on_not_supported(self):
         )
 
     def __call__(
-        self,
-        trajectories: Sequence[TrajectoryWithRew],
-        fragment_length: int,
-        num_pairs: int,
+            self,
+            trajectories: Sequence[TrajectoryWithRew],
+            fragment_length: int,
+            num_pairs: int,
     ) -> Sequence[TrajectoryWithRewPair]:
         # sample a large number (self.fragment_sample_factor*num_pairs)
         # of fragments from all the trajectories
@@ -774,9 +776,9 @@ class PreferenceGatherer(abc.ABC):
     """Base class for gathering preference comparisons between trajectory fragments."""
 
     def __init__(
-        self,
-        seed: Optional[int] = None,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            seed: Optional[int] = None,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initializes the preference gatherer.
 
@@ -814,13 +816,13 @@ class SyntheticGatherer(PreferenceGatherer):
     """Computes synthetic preferences using ground-truth environment rewards."""
 
     def __init__(
-        self,
-        temperature: float = 1,
-        discount_factor: float = 1,
-        sample: bool = True,
-        seed: Optional[int] = None,
-        threshold: float = 50,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            temperature: float = 1,
+            discount_factor: float = 1,
+            sample: bool = True,
+            seed: Optional[int] = None,
+            threshold: float = 50,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the synthetic preference gatherer.
 
@@ -868,8 +870,8 @@ def __call__(self, fragment_pairs: Sequence[TrajectoryWithRewPair]) -> np.ndarra
         # how good we can expect the performance of the learned reward
         # model to be at predicting preferences.
         entropy = -(
-            special.xlogy(model_probs, model_probs)
-            + special.xlogy(1 - model_probs, 1 - model_probs)
+                special.xlogy(model_probs, model_probs)
+                + special.xlogy(1 - model_probs, 1 - model_probs)
         ).mean()
         self.logger.record("entropy", entropy)
 
@@ -916,9 +918,9 @@ def __init__(self, max_size: Optional[int] = None):
         self.preferences = np.array([])
 
     def push(
-        self,
-        fragments: Sequence[TrajectoryWithRewPair],
-        preferences: np.ndarray,
+            self,
+            fragments: Sequence[TrajectoryWithRewPair],
+            preferences: np.ndarray,
     ):
         """Add more samples to the dataset.
 
@@ -970,7 +972,7 @@ def load(path: AnyPath) -> "PreferenceDataset":
 
 
 def preference_collate_fn(
-    batch: Sequence[Tuple[TrajectoryWithRewPair, float]],
+        batch: Sequence[Tuple[TrajectoryWithRewPair, float]],
 ) -> Tuple[Sequence[TrajectoryWithRewPair], np.ndarray]:
     fragment_pairs, preferences = zip(*batch)
     return list(fragment_pairs), np.array(preferences)
@@ -988,10 +990,10 @@ class RewardLoss(nn.Module, abc.ABC):
 
     @abc.abstractmethod
     def forward(
-        self,
-        fragment_pairs: Sequence[TrajectoryPair],
-        preferences: np.ndarray,
-        ensemble_member_index: Optional[int] = None,
+            self,
+            fragment_pairs: Sequence[TrajectoryPair],
+            preferences: np.ndarray,
+            ensemble_member_index: Optional[int] = None,
     ) -> LossAndMetrics:
         """Computes the loss.
 
@@ -1017,8 +1019,8 @@ class CrossEntropyRewardLoss(RewardLoss):
     """Compute the cross entropy reward loss."""
 
     def __init__(
-        self,
-        preference_model: PreferenceModel,
+            self,
+            preference_model: PreferenceModel,
     ):
         """Create cross entropy reward loss.
 
@@ -1029,10 +1031,10 @@ def __init__(
         self.preference_model = preference_model
 
     def forward(
-        self,
-        fragment_pairs: Sequence[TrajectoryPair],
-        preferences: np.ndarray,
-        ensemble_member_index: Optional[int] = None,
+            self,
+            fragment_pairs: Sequence[TrajectoryPair],
+            preferences: np.ndarray,
+            ensemble_member_index: Optional[int] = None,
     ) -> LossAndMetrics:
         """Computes the loss.
 
@@ -1080,9 +1082,9 @@ class RewardTrainer(abc.ABC):
     """
 
     def __init__(
-        self,
-        model: reward_nets.RewardNet,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            model: reward_nets.RewardNet,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the reward trainer.
 
@@ -1113,14 +1115,14 @@ class BasicRewardTrainer(RewardTrainer):
     """Train a basic reward model."""
 
     def __init__(
-        self,
-        model: reward_nets.RewardNet,
-        loss: RewardLoss,
-        batch_size: int = 32,
-        epochs: int = 1,
-        lr: float = 1e-3,
-        weight_decay: float = 0.0,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            model: reward_nets.RewardNet,
+            loss: RewardLoss,
+            batch_size: int = 32,
+            epochs: int = 1,
+            lr: float = 1e-3,
+            weight_decay: float = 0.0,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the reward model trainer.
 
@@ -1169,9 +1171,9 @@ def _train(self, dataset: PreferenceDataset, epoch_multiplier: float = 1.0) -> N
                 self.optim.step()
 
     def _training_inner_loop(
-        self,
-        fragment_pairs: Sequence[TrajectoryPair],
-        preferences: np.ndarray,
+            self,
+            fragment_pairs: Sequence[TrajectoryPair],
+            preferences: np.ndarray,
     ) -> th.Tensor:
         output = self.loss.forward(fragment_pairs, preferences)
         loss = output.loss
@@ -1187,15 +1189,15 @@ class EnsembleTrainer(BasicRewardTrainer):
     _model: reward_nets.RewardEnsemble
 
     def __init__(
-        self,
-        model: reward_nets.RewardEnsemble,
-        loss: RewardLoss,
-        batch_size: int = 32,
-        epochs: int = 1,
-        lr: float = 1e-3,
-        weight_decay: float = 0.0,
-        seed: Optional[int] = None,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            self,
+            model: reward_nets.RewardEnsemble,
+            loss: RewardLoss,
+            batch_size: int = 32,
+            epochs: int = 1,
+            lr: float = 1e-3,
+            weight_decay: float = 0.0,
+            seed: Optional[int] = None,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the reward model trainer.
 
@@ -1233,9 +1235,9 @@ def __init__(
         self.rng = np.random.default_rng(seed=seed)
 
     def _training_inner_loop(
-        self,
-        fragment_pairs: Sequence[TrajectoryPair],
-        preferences: np.ndarray,
+            self,
+            fragment_pairs: Sequence[TrajectoryPair],
+            preferences: np.ndarray,
     ) -> th.Tensor:
         assert len(fragment_pairs) == preferences.shape[0]
         losses_list = []
@@ -1268,7 +1270,7 @@ def _training_inner_loop(
 
 
 def get_base_model(
-    reward_model: reward_nets.RewardNet,
+        reward_model: reward_nets.RewardNet,
 ) -> Union[reward_nets.RewardNet, reward_nets.RewardEnsemble]:
     base_model = reward_model
     while hasattr(base_model, "base"):
@@ -1278,10 +1280,10 @@ def get_base_model(
 
 
 def _make_reward_trainer(
-    reward_model: reward_nets.RewardNet,
-    loss: RewardLoss,
-    reward_trainer_kwargs: Optional[Mapping[str, Any]] = None,
-    seed: Optional[int] = None,
+        reward_model: reward_nets.RewardNet,
+        loss: RewardLoss,
+        reward_trainer_kwargs: Optional[Mapping[str, Any]] = None,
+        seed: Optional[int] = None,
 ) -> RewardTrainer:
     """Construct the correct type of reward trainer for this reward function."""
     if reward_trainer_kwargs is None:
@@ -1296,8 +1298,8 @@ def _make_reward_trainer(
         # must train directly on the base model for reward model training.
         is_base = reward_model is base_model
         is_std_wrapper = (
-            isinstance(reward_model, reward_nets.AddSTDRewardWrapper)
-            and reward_model.base is base_model
+                isinstance(reward_model, reward_nets.AddSTDRewardWrapper)
+                and reward_model.base is base_model
         )
 
         if is_base or is_std_wrapper:
@@ -1318,7 +1320,7 @@ def _make_reward_trainer(
 QUERY_SCHEDULES: Dict[str, type_aliases.Schedule] = {
     "constant": lambda t: 1.0,
     "hyperbolic": lambda t: 1.0 / (1.0 + t),
-    "inverse_quadratic": lambda t: 1.0 / (1.0 + t**2),
+    "inverse_quadratic": lambda t: 1.0 / (1.0 + t ** 2),
 }
 
 
@@ -1326,22 +1328,22 @@ class PreferenceComparisons(base.BaseImitationAlgorithm):
     """Main interface for reward learning using preference comparisons."""
 
     def __init__(
-        self,
-        trajectory_generator: TrajectoryGenerator,
-        reward_model: reward_nets.RewardNet,
-        num_iterations: int,
-        fragmenter: Optional[Fragmenter] = None,
-        preference_gatherer: Optional[PreferenceGatherer] = None,
-        reward_trainer: Optional[RewardTrainer] = None,
-        comparison_queue_size: Optional[int] = None,
-        fragment_length: int = 100,
-        transition_oversampling: float = 1,
-        initial_comparison_frac: float = 0.1,
-        initial_epoch_multiplier: float = 200.0,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-        allow_variable_horizon: bool = False,
-        seed: Optional[int] = None,
-        query_schedule: Union[str, type_aliases.Schedule] = "hyperbolic",
+            self,
+            trajectory_generator: TrajectoryGenerator,
+            reward_model: reward_nets.RewardNet,
+            num_iterations: int,
+            fragmenter: Optional[Fragmenter] = None,
+            preference_gatherer: Optional[PreferenceGatherer] = None,
+            reward_trainer: Optional[RewardTrainer] = None,
+            comparison_queue_size: Optional[int] = None,
+            fragment_length: int = 100,
+            transition_oversampling: float = 1,
+            initial_comparison_frac: float = 0.1,
+            initial_epoch_multiplier: float = 200.0,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            allow_variable_horizon: bool = False,
+            seed: Optional[int] = None,
+            query_schedule: Union[str, type_aliases.Schedule] = "hyperbolic",
     ):
         """Initialize the preference comparison trainer.
 
@@ -1462,10 +1464,10 @@ def __init__(
         self.dataset = PreferenceDataset(max_size=comparison_queue_size)
 
     def train(
-        self,
-        total_timesteps: int,
-        total_comparisons: int,
-        callback: Optional[Callable[[int], None]] = None,
+            self,
+            total_timesteps: int,
+            total_comparisons: int,
+            callback: Optional[Callable[[int], None]] = None,
     ) -> Mapping[str, Any]:
         """Train the reward model and the policy if applicable.
 
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index a59829bdc..b97bf7cba 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -63,14 +63,14 @@ def make_unique_timestamp() -> str:
 
 
 def make_vec_env(
-    env_name: str,
-    n_envs: int = 8,
-    seed: int = 0,
-    parallel: bool = False,
-    log_dir: Optional[str] = None,
-    max_episode_steps: Optional[int] = None,
-    post_wrappers: Optional[Sequence[Callable[[gym.Env, int], gym.Env]]] = None,
-    env_make_kwargs: Optional[Mapping[str, Any]] = None,
+        env_name: str,
+        n_envs: int = 8,
+        seed: int = 0,
+        parallel: bool = False,
+        log_dir: Optional[str] = None,
+        max_episode_steps: Optional[int] = None,
+        post_wrappers: Optional[Sequence[Callable[[gym.Env, int], gym.Env]]] = None,
+        env_make_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> VecEnv:
     """Makes a vectorized environment.
 
@@ -183,6 +183,10 @@ def endless_iter(iterable: Iterable[T]) -> Iterator[T]:
     Raises:
         ValueError: `iterable` is empty -- the first call it to returns no elements.
     """
+    # TODO(juan) this is wrong; if the iterable is not a container then the first
+    #  element will be wasted if it's not stored. The sensible solution is to
+    #  restrict the type of `iterable` to `Sequence` (this same issue is present
+    #  in a few other places in the codebase).
     try:
         next(iter(iterable))
     except StopIteration:
@@ -212,8 +216,8 @@ def safe_to_tensor(numpy_array: np.ndarray, **kwargs) -> th.Tensor:
 
 
 def tensor_iter_norm(
-    tensor_iter: Iterable[th.Tensor],
-    ord: Union[int, float] = 2,  # noqa: A002
+        tensor_iter: Iterable[th.Tensor],
+        ord: Union[int, float] = 2,  # noqa: A002
 ) -> th.Tensor:
     """Compute the norm of a big vector that is produced one tensor chunk at a time.
 

From 2de6d8a94d834abd853fdaea0e413034181f716a Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 21:51:58 +0200
Subject: [PATCH 043/187] Formatting

---
 src/imitation/algorithms/adversarial/airl.py  |  28 +-
 .../algorithms/adversarial/common.py          | 102 ++++---
 src/imitation/algorithms/base.py              |  30 +-
 .../algorithms/preference_comparisons.py      | 266 +++++++++---------
 src/imitation/util/util.py                    |  20 +-
 5 files changed, 228 insertions(+), 218 deletions(-)

diff --git a/src/imitation/algorithms/adversarial/airl.py b/src/imitation/algorithms/adversarial/airl.py
index f8ae5f78d..7f0a9d9d6 100644
--- a/src/imitation/algorithms/adversarial/airl.py
+++ b/src/imitation/algorithms/adversarial/airl.py
@@ -19,14 +19,14 @@ class AIRL(common.AdversarialTrainer):
     """
 
     def __init__(
-            self,
-            *,
-            demonstrations: base.AnyTransitions,
-            demo_batch_size: int,
-            venv: vec_env.VecEnv,
-            gen_algo: base_class.BaseAlgorithm,
-            reward_net: reward_nets.RewardNet,
-            **kwargs,
+        self,
+        *,
+        demonstrations: base.AnyTransitions,
+        demo_batch_size: int,
+        venv: vec_env.VecEnv,
+        gen_algo: base_class.BaseAlgorithm,
+        reward_net: reward_nets.RewardNet,
+        **kwargs,
     ):
         """Builds an AIRL trainer.
 
@@ -65,12 +65,12 @@ def __init__(
             )
 
     def logits_expert_is_high(
-            self,
-            state: th.Tensor,
-            action: th.Tensor,
-            next_state: th.Tensor,
-            done: th.Tensor,
-            log_policy_act_prob: Optional[th.Tensor] = None,
+        self,
+        state: th.Tensor,
+        action: th.Tensor,
+        next_state: th.Tensor,
+        done: th.Tensor,
+        log_policy_act_prob: Optional[th.Tensor] = None,
     ) -> th.Tensor:
         r"""Compute the discriminator's logits for each state-action sample.
 
diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
index d08a55c4f..00367349a 100644
--- a/src/imitation/algorithms/adversarial/common.py
+++ b/src/imitation/algorithms/adversarial/common.py
@@ -4,7 +4,17 @@
 import dataclasses
 import logging
 import os
-from typing import Callable, Mapping, Optional, Sequence, Tuple, Type, Iterable, Iterator, overload
+from typing import (
+    Callable,
+    Iterable,
+    Iterator,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    overload,
+)
 
 import numpy as np
 import torch as th
@@ -21,9 +31,9 @@
 
 
 def compute_train_stats(
-        disc_logits_expert_is_high: th.Tensor,
-        labels_expert_is_one: th.Tensor,
-        disc_loss: th.Tensor,
+    disc_logits_expert_is_high: th.Tensor,
+    labels_expert_is_one: th.Tensor,
+    disc_loss: th.Tensor,
 ) -> Mapping[str, float]:
     """Train statistics for GAIL/AIRL discriminator.
 
@@ -107,24 +117,24 @@ class AdversarialTrainer(base.DemonstrationAlgorithm[types.Transitions]):
     _endless_expert_iterator: Optional[Iterator[base.TransitionMapping]]
 
     def __init__(
-            self,
-            *,
-            demonstrations: base.AnyTransitions,
-            demo_batch_size: int,
-            venv: vec_env.VecEnv,
-            gen_algo: base_class.BaseAlgorithm,
-            reward_net: reward_nets.RewardNet,
-            n_disc_updates_per_round: int = 2,
-            log_dir: str = "output/",
-            disc_opt_cls: Type[th.optim.Optimizer] = th.optim.Adam,
-            disc_opt_kwargs: Optional[Mapping] = None,
-            gen_train_timesteps: Optional[int] = None,
-            gen_replay_buffer_capacity: Optional[int] = None,
-            custom_logger: Optional[logger.HierarchicalLogger] = None,
-            init_tensorboard: bool = False,
-            init_tensorboard_graph: bool = False,
-            debug_use_ground_truth: bool = False,
-            allow_variable_horizon: bool = False,
+        self,
+        *,
+        demonstrations: base.AnyTransitions,
+        demo_batch_size: int,
+        venv: vec_env.VecEnv,
+        gen_algo: base_class.BaseAlgorithm,
+        reward_net: reward_nets.RewardNet,
+        n_disc_updates_per_round: int = 2,
+        log_dir: str = "output/",
+        disc_opt_cls: Type[th.optim.Optimizer] = th.optim.Adam,
+        disc_opt_kwargs: Optional[Mapping] = None,
+        gen_train_timesteps: Optional[int] = None,
+        gen_replay_buffer_capacity: Optional[int] = None,
+        custom_logger: Optional[logger.HierarchicalLogger] = None,
+        init_tensorboard: bool = False,
+        init_tensorboard_graph: bool = False,
+        debug_use_ground_truth: bool = False,
+        allow_variable_horizon: bool = False,
     ):
         """Builds AdversarialTrainer.
 
@@ -251,12 +261,12 @@ def policy(self) -> policies.BasePolicy:
 
     @abc.abstractmethod
     def logits_expert_is_high(
-            self,
-            state: th.Tensor,
-            action: th.Tensor,
-            next_state: th.Tensor,
-            done: th.Tensor,
-            log_policy_act_prob: Optional[th.Tensor] = None,
+        self,
+        state: th.Tensor,
+        action: th.Tensor,
+        next_state: th.Tensor,
+        done: th.Tensor,
+        log_policy_act_prob: Optional[th.Tensor] = None,
     ) -> th.Tensor:
         """Compute the discriminator's logits for each state-action sample.
 
@@ -299,10 +309,10 @@ def _next_expert_batch(self) -> Mapping:
         return next(self._endless_expert_iterator)
 
     def train_disc(
-            self,
-            *,
-            expert_samples: Optional[Mapping] = None,
-            gen_samples: Optional[Mapping] = None,
+        self,
+        *,
+        expert_samples: Optional[Mapping] = None,
+        gen_samples: Optional[Mapping] = None,
     ) -> Optional[Mapping[str, float]]:
         """Perform a single discriminator update, optionally using provided samples.
 
@@ -366,9 +376,9 @@ def train_disc(
         return train_stats
 
     def train_gen(
-            self,
-            total_timesteps: Optional[int] = None,
-            learn_kwargs: Optional[Mapping] = None,
+        self,
+        total_timesteps: Optional[int] = None,
+        learn_kwargs: Optional[Mapping] = None,
     ) -> None:
         """Trains the generator to maximize the discriminator loss.
 
@@ -402,9 +412,9 @@ def train_gen(
         self._gen_replay_buffer.store(gen_samples)
 
     def train(
-            self,
-            total_timesteps: int,
-            callback: Optional[Callable[[int], None]] = None,
+        self,
+        total_timesteps: int,
+        callback: Optional[Callable[[int], None]] = None,
     ) -> None:
         """Alternates between training the generator and discriminator.
 
@@ -444,16 +454,16 @@ def _torchify_array(self, ndarray: np.ndarray) -> th.Tensor:
     @overload
     def _torchify_array(self, ndarray: None) -> None:
         ...
-    
+
     def _torchify_array(self, ndarray: Optional[np.ndarray]) -> Optional[th.Tensor]:
         if ndarray is not None:
             return th.as_tensor(ndarray, device=self.reward_train.device)
         return None
 
     def _get_log_policy_act_prob(
-            self,
-            obs_th: th.Tensor,
-            acts_th: th.Tensor,
+        self,
+        obs_th: th.Tensor,
+        acts_th: th.Tensor,
     ) -> Optional[th.Tensor]:
         """Evaluates the given actions on the given observations.
 
@@ -491,10 +501,10 @@ def _get_log_policy_act_prob(
         return log_policy_act_prob_th
 
     def _make_disc_train_batch(
-            self,
-            *,
-            gen_samples: Optional[Mapping] = None,
-            expert_samples: Optional[Mapping] = None,
+        self,
+        *,
+        gen_samples: Optional[Mapping] = None,
+        expert_samples: Optional[Mapping] = None,
     ) -> Mapping[str, th.Tensor]:
         """Build and return training batch for the next discriminator update.
 
diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index 2d8ca04ea..3ca0e7888 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -25,10 +25,10 @@ class BaseImitationAlgorithm(abc.ABC):
     """Horizon of trajectories seen so far (None if no trajectories seen)."""
 
     def __init__(
-            self,
-            *,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-            allow_variable_horizon: bool = False,
+        self,
+        *,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        allow_variable_horizon: bool = False,
     ):
         """Creates an imitation learning algorithm.
 
@@ -124,11 +124,11 @@ class DemonstrationAlgorithm(BaseImitationAlgorithm, Generic[TransitionKind]):
     """An algorithm that learns from demonstration: BC, IRL, etc."""
 
     def __init__(
-            self,
-            *,
-            demonstrations: Optional[AnyTransitions],
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-            allow_variable_horizon: bool = False,
+        self,
+        *,
+        demonstrations: Optional[AnyTransitions],
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        allow_variable_horizon: bool = False,
     ):
         """Creates an algorithm that learns from demonstrations.
 
@@ -177,9 +177,9 @@ class _WrappedDataLoader:
     """Wraps a data loader (batch iterable) and checks for specified batch size."""
 
     def __init__(
-            self,
-            data_loader: Iterable[TransitionMapping],
-            expected_batch_size: int,
+        self,
+        data_loader: Iterable[TransitionMapping],
+        expected_batch_size: int,
     ):
         """Builds _WrapedDataLoader.
 
@@ -215,9 +215,9 @@ def __iter__(self):
 
 
 def make_data_loader(
-        transitions: AnyTransitions,
-        batch_size: int,
-        data_loader_kwargs: Optional[Mapping[str, Any]] = None,
+    transitions: AnyTransitions,
+    batch_size: int,
+    data_loader_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> Iterable[TransitionMapping]:
     """Converts demonstration data to Torch data loader.
 
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 18d4f3390..23d367175 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -96,10 +96,10 @@ class TrajectoryDataset(TrajectoryGenerator):
     """A fixed dataset of trajectories."""
 
     def __init__(
-            self,
-            trajectories: Sequence[TrajectoryWithRew],
-            seed: Optional[int] = None,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        trajectories: Sequence[TrajectoryWithRew],
+        seed: Optional[int] = None,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Creates a dataset loaded from `path`.
 
@@ -123,15 +123,15 @@ class AgentTrainer(TrajectoryGenerator):
     """Wrapper for training an SB3 algorithm on an arbitrary reward function."""
 
     def __init__(
-            self,
-            algorithm: base_class.BaseAlgorithm,
-            reward_fn: Union[reward_function.RewardFn, reward_nets.RewardNet],
-            venv: vec_env.VecEnv,
-            exploration_frac: float = 0.0,
-            switch_prob: float = 0.5,
-            random_prob: float = 0.5,
-            seed: Optional[int] = None,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        algorithm: base_class.BaseAlgorithm,
+        reward_fn: Union[reward_function.RewardFn, reward_nets.RewardNet],
+        venv: vec_env.VecEnv,
+        exploration_frac: float = 0.0,
+        switch_prob: float = 0.5,
+        random_prob: float = 0.5,
+        seed: Optional[int] = None,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the agent trainer.
 
@@ -310,8 +310,8 @@ def logger(self, value: imit_logger.HierarchicalLogger):
 
 
 def _get_trajectories(
-        trajectories: Sequence[TrajectoryWithRew],
-        steps: int,
+    trajectories: Sequence[TrajectoryWithRew],
+    steps: int,
 ) -> Sequence[TrajectoryWithRew]:
     """Get enough trajectories to have at least `steps` transitions in total."""
     if steps == 0:
@@ -339,11 +339,11 @@ class PreferenceModel(nn.Module):
     """Class to convert two fragments' rewards into preference probability."""
 
     def __init__(
-            self,
-            model: reward_nets.RewardNet,
-            noise_prob: float = 0.0,
-            discount_factor: float = 1.0,
-            threshold: float = 50,
+        self,
+        model: reward_nets.RewardNet,
+        noise_prob: float = 0.0,
+        discount_factor: float = 1.0,
+        threshold: float = 50,
     ):
         """Create Preference Prediction Model.
 
@@ -388,9 +388,9 @@ def __init__(
                 self.member_pref_models.append(member_pref_model)
 
     def forward(
-            self,
-            fragment_pairs: Sequence[TrajectoryPair],
-            ensemble_member_index: Optional[int] = None,
+        self,
+        fragment_pairs: Sequence[TrajectoryPair],
+        ensemble_member_index: Optional[int] = None,
     ) -> Tuple[th.Tensor, Optional[th.Tensor]]:
         """Computes the preference probability of the first fragment for all pairs.
 
@@ -451,8 +451,8 @@ def forward(
         return probs, (gt_probs if gt_reward_available else None)
 
     def rewards(
-            self,
-            transitions: Transitions,
+        self,
+        transitions: Transitions,
     ) -> th.Tensor:
         """Computes the reward for all transitions.
 
@@ -481,9 +481,9 @@ def rewards(
             return rews
 
     def probability(
-            self,
-            rews1: th.Tensor,
-            rews2: th.Tensor,
+        self,
+        rews1: th.Tensor,
+        rews2: th.Tensor,
     ) -> th.Tensor:
         """Computes the Boltzmann rational probability that the first trajectory is best.
 
@@ -541,10 +541,10 @@ def __init__(self, custom_logger: Optional[imit_logger.HierarchicalLogger] = Non
 
     @abc.abstractmethod
     def __call__(
-            self,
-            trajectories: Sequence[TrajectoryWithRew],
-            fragment_length: int,
-            num_pairs: int,
+        self,
+        trajectories: Sequence[TrajectoryWithRew],
+        fragment_length: int,
+        num_pairs: int,
     ) -> Sequence[TrajectoryWithRewPair]:
         """Create fragment pairs out of a sequence of trajectories.
 
@@ -572,10 +572,10 @@ class RandomFragmenter(Fragmenter):
     """
 
     def __init__(
-            self,
-            seed: Optional[float] = None,
-            warning_threshold: int = 10,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        seed: Optional[float] = None,
+        warning_threshold: int = 10,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the fragmenter.
 
@@ -591,10 +591,10 @@ def __init__(
         self.warning_threshold = warning_threshold
 
     def __call__(
-            self,
-            trajectories: Sequence[TrajectoryWithRew],
-            fragment_length: int,
-            num_pairs: int,
+        self,
+        trajectories: Sequence[TrajectoryWithRew],
+        fragment_length: int,
+        num_pairs: int,
     ) -> Sequence[TrajectoryWithRewPair]:
         fragments: List[TrajectoryWithRew] = []
 
@@ -624,8 +624,8 @@ def __call__(
                 "of fragment pairs. Some transitions will appear multiple times.",
             )
         elif (
-                self.warning_threshold
-                and sum(weights) < self.warning_threshold * num_transitions
+            self.warning_threshold
+            and sum(weights) < self.warning_threshold * num_transitions
         ):
             # If the number of available transitions is not much larger
             # than the number of requires ones, we already give a warning.
@@ -645,7 +645,7 @@ def __call__(
             end = start + fragment_length
             terminal = (end == n) and traj.terminal
             fragment = TrajectoryWithRew(
-                obs=traj.obs[start: end + 1],
+                obs=traj.obs[start : end + 1],
                 acts=traj.acts[start:end],
                 infos=traj.infos[start:end] if traj.infos is not None else None,
                 rews=traj.rews[start:end],
@@ -667,12 +667,12 @@ class ActiveSelectionFragmenter(Fragmenter):
     """
 
     def __init__(
-            self,
-            preference_model: PreferenceModel,
-            base_fragmenter: Fragmenter,
-            fragment_sample_factor: float,
-            uncertainty_on: str = "logits",
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        preference_model: PreferenceModel,
+        base_fragmenter: Fragmenter,
+        fragment_sample_factor: float,
+        uncertainty_on: str = "logits",
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the active selection fragmenter.
 
@@ -713,10 +713,10 @@ def raise_uncertainty_on_not_supported(self):
         )
 
     def __call__(
-            self,
-            trajectories: Sequence[TrajectoryWithRew],
-            fragment_length: int,
-            num_pairs: int,
+        self,
+        trajectories: Sequence[TrajectoryWithRew],
+        fragment_length: int,
+        num_pairs: int,
     ) -> Sequence[TrajectoryWithRewPair]:
         # sample a large number (self.fragment_sample_factor*num_pairs)
         # of fragments from all the trajectories
@@ -776,9 +776,9 @@ class PreferenceGatherer(abc.ABC):
     """Base class for gathering preference comparisons between trajectory fragments."""
 
     def __init__(
-            self,
-            seed: Optional[int] = None,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        seed: Optional[int] = None,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initializes the preference gatherer.
 
@@ -816,13 +816,13 @@ class SyntheticGatherer(PreferenceGatherer):
     """Computes synthetic preferences using ground-truth environment rewards."""
 
     def __init__(
-            self,
-            temperature: float = 1,
-            discount_factor: float = 1,
-            sample: bool = True,
-            seed: Optional[int] = None,
-            threshold: float = 50,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        temperature: float = 1,
+        discount_factor: float = 1,
+        sample: bool = True,
+        seed: Optional[int] = None,
+        threshold: float = 50,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the synthetic preference gatherer.
 
@@ -870,8 +870,8 @@ def __call__(self, fragment_pairs: Sequence[TrajectoryWithRewPair]) -> np.ndarra
         # how good we can expect the performance of the learned reward
         # model to be at predicting preferences.
         entropy = -(
-                special.xlogy(model_probs, model_probs)
-                + special.xlogy(1 - model_probs, 1 - model_probs)
+            special.xlogy(model_probs, model_probs)
+            + special.xlogy(1 - model_probs, 1 - model_probs)
         ).mean()
         self.logger.record("entropy", entropy)
 
@@ -918,9 +918,9 @@ def __init__(self, max_size: Optional[int] = None):
         self.preferences = np.array([])
 
     def push(
-            self,
-            fragments: Sequence[TrajectoryWithRewPair],
-            preferences: np.ndarray,
+        self,
+        fragments: Sequence[TrajectoryWithRewPair],
+        preferences: np.ndarray,
     ):
         """Add more samples to the dataset.
 
@@ -972,7 +972,7 @@ def load(path: AnyPath) -> "PreferenceDataset":
 
 
 def preference_collate_fn(
-        batch: Sequence[Tuple[TrajectoryWithRewPair, float]],
+    batch: Sequence[Tuple[TrajectoryWithRewPair, float]],
 ) -> Tuple[Sequence[TrajectoryWithRewPair], np.ndarray]:
     fragment_pairs, preferences = zip(*batch)
     return list(fragment_pairs), np.array(preferences)
@@ -990,10 +990,10 @@ class RewardLoss(nn.Module, abc.ABC):
 
     @abc.abstractmethod
     def forward(
-            self,
-            fragment_pairs: Sequence[TrajectoryPair],
-            preferences: np.ndarray,
-            ensemble_member_index: Optional[int] = None,
+        self,
+        fragment_pairs: Sequence[TrajectoryPair],
+        preferences: np.ndarray,
+        ensemble_member_index: Optional[int] = None,
     ) -> LossAndMetrics:
         """Computes the loss.
 
@@ -1019,8 +1019,8 @@ class CrossEntropyRewardLoss(RewardLoss):
     """Compute the cross entropy reward loss."""
 
     def __init__(
-            self,
-            preference_model: PreferenceModel,
+        self,
+        preference_model: PreferenceModel,
     ):
         """Create cross entropy reward loss.
 
@@ -1031,10 +1031,10 @@ def __init__(
         self.preference_model = preference_model
 
     def forward(
-            self,
-            fragment_pairs: Sequence[TrajectoryPair],
-            preferences: np.ndarray,
-            ensemble_member_index: Optional[int] = None,
+        self,
+        fragment_pairs: Sequence[TrajectoryPair],
+        preferences: np.ndarray,
+        ensemble_member_index: Optional[int] = None,
     ) -> LossAndMetrics:
         """Computes the loss.
 
@@ -1082,9 +1082,9 @@ class RewardTrainer(abc.ABC):
     """
 
     def __init__(
-            self,
-            model: reward_nets.RewardNet,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        model: reward_nets.RewardNet,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the reward trainer.
 
@@ -1115,14 +1115,14 @@ class BasicRewardTrainer(RewardTrainer):
     """Train a basic reward model."""
 
     def __init__(
-            self,
-            model: reward_nets.RewardNet,
-            loss: RewardLoss,
-            batch_size: int = 32,
-            epochs: int = 1,
-            lr: float = 1e-3,
-            weight_decay: float = 0.0,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        model: reward_nets.RewardNet,
+        loss: RewardLoss,
+        batch_size: int = 32,
+        epochs: int = 1,
+        lr: float = 1e-3,
+        weight_decay: float = 0.0,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the reward model trainer.
 
@@ -1171,9 +1171,9 @@ def _train(self, dataset: PreferenceDataset, epoch_multiplier: float = 1.0) -> N
                 self.optim.step()
 
     def _training_inner_loop(
-            self,
-            fragment_pairs: Sequence[TrajectoryPair],
-            preferences: np.ndarray,
+        self,
+        fragment_pairs: Sequence[TrajectoryPair],
+        preferences: np.ndarray,
     ) -> th.Tensor:
         output = self.loss.forward(fragment_pairs, preferences)
         loss = output.loss
@@ -1189,15 +1189,15 @@ class EnsembleTrainer(BasicRewardTrainer):
     _model: reward_nets.RewardEnsemble
 
     def __init__(
-            self,
-            model: reward_nets.RewardEnsemble,
-            loss: RewardLoss,
-            batch_size: int = 32,
-            epochs: int = 1,
-            lr: float = 1e-3,
-            weight_decay: float = 0.0,
-            seed: Optional[int] = None,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        self,
+        model: reward_nets.RewardEnsemble,
+        loss: RewardLoss,
+        batch_size: int = 32,
+        epochs: int = 1,
+        lr: float = 1e-3,
+        weight_decay: float = 0.0,
+        seed: Optional[int] = None,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the reward model trainer.
 
@@ -1235,9 +1235,9 @@ def __init__(
         self.rng = np.random.default_rng(seed=seed)
 
     def _training_inner_loop(
-            self,
-            fragment_pairs: Sequence[TrajectoryPair],
-            preferences: np.ndarray,
+        self,
+        fragment_pairs: Sequence[TrajectoryPair],
+        preferences: np.ndarray,
     ) -> th.Tensor:
         assert len(fragment_pairs) == preferences.shape[0]
         losses_list = []
@@ -1270,7 +1270,7 @@ def _training_inner_loop(
 
 
 def get_base_model(
-        reward_model: reward_nets.RewardNet,
+    reward_model: reward_nets.RewardNet,
 ) -> Union[reward_nets.RewardNet, reward_nets.RewardEnsemble]:
     base_model = reward_model
     while hasattr(base_model, "base"):
@@ -1280,10 +1280,10 @@ def get_base_model(
 
 
 def _make_reward_trainer(
-        reward_model: reward_nets.RewardNet,
-        loss: RewardLoss,
-        reward_trainer_kwargs: Optional[Mapping[str, Any]] = None,
-        seed: Optional[int] = None,
+    reward_model: reward_nets.RewardNet,
+    loss: RewardLoss,
+    reward_trainer_kwargs: Optional[Mapping[str, Any]] = None,
+    seed: Optional[int] = None,
 ) -> RewardTrainer:
     """Construct the correct type of reward trainer for this reward function."""
     if reward_trainer_kwargs is None:
@@ -1298,8 +1298,8 @@ def _make_reward_trainer(
         # must train directly on the base model for reward model training.
         is_base = reward_model is base_model
         is_std_wrapper = (
-                isinstance(reward_model, reward_nets.AddSTDRewardWrapper)
-                and reward_model.base is base_model
+            isinstance(reward_model, reward_nets.AddSTDRewardWrapper)
+            and reward_model.base is base_model
         )
 
         if is_base or is_std_wrapper:
@@ -1320,7 +1320,7 @@ def _make_reward_trainer(
 QUERY_SCHEDULES: Dict[str, type_aliases.Schedule] = {
     "constant": lambda t: 1.0,
     "hyperbolic": lambda t: 1.0 / (1.0 + t),
-    "inverse_quadratic": lambda t: 1.0 / (1.0 + t ** 2),
+    "inverse_quadratic": lambda t: 1.0 / (1.0 + t**2),
 }
 
 
@@ -1328,22 +1328,22 @@ class PreferenceComparisons(base.BaseImitationAlgorithm):
     """Main interface for reward learning using preference comparisons."""
 
     def __init__(
-            self,
-            trajectory_generator: TrajectoryGenerator,
-            reward_model: reward_nets.RewardNet,
-            num_iterations: int,
-            fragmenter: Optional[Fragmenter] = None,
-            preference_gatherer: Optional[PreferenceGatherer] = None,
-            reward_trainer: Optional[RewardTrainer] = None,
-            comparison_queue_size: Optional[int] = None,
-            fragment_length: int = 100,
-            transition_oversampling: float = 1,
-            initial_comparison_frac: float = 0.1,
-            initial_epoch_multiplier: float = 200.0,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-            allow_variable_horizon: bool = False,
-            seed: Optional[int] = None,
-            query_schedule: Union[str, type_aliases.Schedule] = "hyperbolic",
+        self,
+        trajectory_generator: TrajectoryGenerator,
+        reward_model: reward_nets.RewardNet,
+        num_iterations: int,
+        fragmenter: Optional[Fragmenter] = None,
+        preference_gatherer: Optional[PreferenceGatherer] = None,
+        reward_trainer: Optional[RewardTrainer] = None,
+        comparison_queue_size: Optional[int] = None,
+        fragment_length: int = 100,
+        transition_oversampling: float = 1,
+        initial_comparison_frac: float = 0.1,
+        initial_epoch_multiplier: float = 200.0,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        allow_variable_horizon: bool = False,
+        seed: Optional[int] = None,
+        query_schedule: Union[str, type_aliases.Schedule] = "hyperbolic",
     ):
         """Initialize the preference comparison trainer.
 
@@ -1464,10 +1464,10 @@ def __init__(
         self.dataset = PreferenceDataset(max_size=comparison_queue_size)
 
     def train(
-            self,
-            total_timesteps: int,
-            total_comparisons: int,
-            callback: Optional[Callable[[int], None]] = None,
+        self,
+        total_timesteps: int,
+        total_comparisons: int,
+        callback: Optional[Callable[[int], None]] = None,
     ) -> Mapping[str, Any]:
         """Train the reward model and the policy if applicable.
 
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index b97bf7cba..351e3ac9f 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -63,14 +63,14 @@ def make_unique_timestamp() -> str:
 
 
 def make_vec_env(
-        env_name: str,
-        n_envs: int = 8,
-        seed: int = 0,
-        parallel: bool = False,
-        log_dir: Optional[str] = None,
-        max_episode_steps: Optional[int] = None,
-        post_wrappers: Optional[Sequence[Callable[[gym.Env, int], gym.Env]]] = None,
-        env_make_kwargs: Optional[Mapping[str, Any]] = None,
+    env_name: str,
+    n_envs: int = 8,
+    seed: int = 0,
+    parallel: bool = False,
+    log_dir: Optional[str] = None,
+    max_episode_steps: Optional[int] = None,
+    post_wrappers: Optional[Sequence[Callable[[gym.Env, int], gym.Env]]] = None,
+    env_make_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> VecEnv:
     """Makes a vectorized environment.
 
@@ -216,8 +216,8 @@ def safe_to_tensor(numpy_array: np.ndarray, **kwargs) -> th.Tensor:
 
 
 def tensor_iter_norm(
-        tensor_iter: Iterable[th.Tensor],
-        ord: Union[int, float] = 2,  # noqa: A002
+    tensor_iter: Iterable[th.Tensor],
+    ord: Union[int, float] = 2,  # noqa: A002
 ) -> th.Tensor:
     """Compute the norm of a big vector that is produced one tensor chunk at a time.
 

From baa65120de1f1afffe7a359b3ae98ed57243d6a6 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 22:08:35 +0200
Subject: [PATCH 044/187] Bug fix

---
 src/imitation/algorithms/base.py | 39 +++++++++++++++++---------------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index 3ca0e7888..46b022aa1 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -25,10 +25,10 @@ class BaseImitationAlgorithm(abc.ABC):
     """Horizon of trajectories seen so far (None if no trajectories seen)."""
 
     def __init__(
-        self,
-        *,
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-        allow_variable_horizon: bool = False,
+            self,
+            *,
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            allow_variable_horizon: bool = False,
     ):
         """Creates an imitation learning algorithm.
 
@@ -124,11 +124,11 @@ class DemonstrationAlgorithm(BaseImitationAlgorithm, Generic[TransitionKind]):
     """An algorithm that learns from demonstration: BC, IRL, etc."""
 
     def __init__(
-        self,
-        *,
-        demonstrations: Optional[AnyTransitions],
-        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-        allow_variable_horizon: bool = False,
+            self,
+            *,
+            demonstrations: Optional[AnyTransitions],
+            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+            allow_variable_horizon: bool = False,
     ):
         """Creates an algorithm that learns from demonstrations.
 
@@ -177,9 +177,9 @@ class _WrappedDataLoader:
     """Wraps a data loader (batch iterable) and checks for specified batch size."""
 
     def __init__(
-        self,
-        data_loader: Iterable[TransitionMapping],
-        expected_batch_size: int,
+            self,
+            data_loader: Iterable[TransitionMapping],
+            expected_batch_size: int,
     ):
         """Builds _WrapedDataLoader.
 
@@ -215,9 +215,9 @@ def __iter__(self):
 
 
 def make_data_loader(
-    transitions: AnyTransitions,
-    batch_size: int,
-    data_loader_kwargs: Optional[Mapping[str, Any]] = None,
+        transitions: AnyTransitions,
+        batch_size: int,
+        data_loader_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> Iterable[TransitionMapping]:
     """Converts demonstration data to Torch data loader.
 
@@ -256,13 +256,16 @@ def make_data_loader(
                 f"is smaller than batch size {batch_size}.",
             )
 
+        kwargs: Mapping[str, Any] = {
+            "shuffle": True,
+            "drop_last": True,
+            **(data_loader_kwargs or {}),
+        }
         return th_data.DataLoader(
             transitions,
             batch_size=batch_size,
             collate_fn=types.transitions_collate_fn,
-            shuffle=True,
-            drop_last=True,
-            **(data_loader_kwargs or {}),
+            **kwargs,
         )
     elif isinstance(transitions, Iterable):
         return _WrappedDataLoader(transitions, batch_size)

From 4679dfd46f832c418daa60b4a22239da53364528 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 5 Sep 2022 22:08:49 +0200
Subject: [PATCH 045/187] Formatting

---
 src/imitation/algorithms/base.py | 30 +++++++++++++++---------------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index 46b022aa1..03dd600bb 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -25,10 +25,10 @@ class BaseImitationAlgorithm(abc.ABC):
     """Horizon of trajectories seen so far (None if no trajectories seen)."""
 
     def __init__(
-            self,
-            *,
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-            allow_variable_horizon: bool = False,
+        self,
+        *,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        allow_variable_horizon: bool = False,
     ):
         """Creates an imitation learning algorithm.
 
@@ -124,11 +124,11 @@ class DemonstrationAlgorithm(BaseImitationAlgorithm, Generic[TransitionKind]):
     """An algorithm that learns from demonstration: BC, IRL, etc."""
 
     def __init__(
-            self,
-            *,
-            demonstrations: Optional[AnyTransitions],
-            custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
-            allow_variable_horizon: bool = False,
+        self,
+        *,
+        demonstrations: Optional[AnyTransitions],
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        allow_variable_horizon: bool = False,
     ):
         """Creates an algorithm that learns from demonstrations.
 
@@ -177,9 +177,9 @@ class _WrappedDataLoader:
     """Wraps a data loader (batch iterable) and checks for specified batch size."""
 
     def __init__(
-            self,
-            data_loader: Iterable[TransitionMapping],
-            expected_batch_size: int,
+        self,
+        data_loader: Iterable[TransitionMapping],
+        expected_batch_size: int,
     ):
         """Builds _WrapedDataLoader.
 
@@ -215,9 +215,9 @@ def __iter__(self):
 
 
 def make_data_loader(
-        transitions: AnyTransitions,
-        batch_size: int,
-        data_loader_kwargs: Optional[Mapping[str, Any]] = None,
+    transitions: AnyTransitions,
+    batch_size: int,
+    data_loader_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> Iterable[TransitionMapping]:
     """Converts demonstration data to Torch data loader.
 

From 73bd0127f99b9d31169aa273c382ea0ee6d89d57 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 6 Sep 2022 10:46:00 +0200
Subject: [PATCH 046/187] Fixes suggested by Adam.

---
 setup.py                                      |  2 +-
 .../algorithms/adversarial/common.py          | 20 ++++++++++---------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/setup.py b/setup.py
index 9851e658c..890d82361 100644
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,7 @@
         # TODO: upgrade jupyter-client once
         #  https://github.com/jupyter/jupyter_client/issues/637 is fixed
         "jupyter-client~=6.1.12",
-        "mypy~=0.791",
+        "mypy==0.791",
         "pandas~=1.4.3",
         "pytest~=7.1.2",
         "pytest-cov~=3.0.0",
diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
index 00367349a..c1e638f15 100644
--- a/src/imitation/algorithms/adversarial/common.py
+++ b/src/imitation/algorithms/adversarial/common.py
@@ -20,7 +20,7 @@
 import torch as th
 import torch.utils.tensorboard as thboard
 import tqdm
-from stable_baselines3.common import base_class, policies, vec_env
+from stable_baselines3.common import on_policy_algorithm, policies, vec_env
 from stable_baselines3.sac import policies as sac_policies
 from torch.nn import functional as F
 
@@ -116,13 +116,15 @@ class AdversarialTrainer(base.DemonstrationAlgorithm[types.Transitions]):
     _demo_data_loader: Optional[Iterable[base.TransitionMapping]]
     _endless_expert_iterator: Optional[Iterator[base.TransitionMapping]]
 
+    venv_wrapped: vec_env.VecEnvWrapper
+
     def __init__(
         self,
         *,
         demonstrations: base.AnyTransitions,
         demo_batch_size: int,
         venv: vec_env.VecEnv,
-        gen_algo: base_class.BaseAlgorithm,
+        gen_algo: on_policy_algorithm.OnPolicyAlgorithm,
         reward_net: reward_nets.RewardNet,
         n_disc_updates_per_round: int = 2,
         log_dir: str = "output/",
@@ -218,16 +220,15 @@ def __init__(
             os.makedirs(summary_dir, exist_ok=True)
             self._summary_writer = thboard.SummaryWriter(summary_dir)
 
-        venv = self.venv_buffering = wrappers.BufferingWrapper(self.venv)
+        self.venv_buffering = wrappers.BufferingWrapper(self.venv)
 
         if debug_use_ground_truth:
             # Would use an identity reward fn here, but RewardFns can't see rewards.
-            self.venv_wrapped = venv
+            self.venv_wrapped = self.venv_buffering
             self.gen_callback = None
         else:
-            # TODO(juan) this is a bug; venv is going unused.
-            venv = self.venv_wrapped = reward_wrapper.RewardVecEnvWrapper(
-                venv,
+            self.venv_wrapped = reward_wrapper.RewardVecEnvWrapper(
+                self.venv_buffering,
                 reward_fn=self.reward_train.predict_processed,
             )
             self.gen_callback = self.venv_wrapped.make_log_callback()
@@ -256,8 +257,9 @@ def __init__(
 
     @property
     def policy(self) -> policies.BasePolicy:
-        # TODO(juan) either change the return type to optional or add an assertion.
-        return self.gen_algo.policy
+        policy = self.gen_algo.policy
+        assert policy is not None
+        return policy
 
     @abc.abstractmethod
     def logits_expert_is_high(

From 1ecc300f06ae2a7388efdd578449ddc253f997ea Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 6 Sep 2022 10:57:33 +0200
Subject: [PATCH 047/187] Fix mypy version.

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 890d82361..acdd7d790 100644
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,7 @@
         # TODO: upgrade jupyter-client once
         #  https://github.com/jupyter/jupyter_client/issues/637 is fixed
         "jupyter-client~=6.1.12",
-        "mypy==0.791",
+        "mypy==0.971",
         "pandas~=1.4.3",
         "pytest~=7.1.2",
         "pytest-cov~=3.0.0",

From e177c34774114fc77da0e5f883dd18f07b21ad1f Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 6 Sep 2022 16:51:06 +0200
Subject: [PATCH 048/187] Fix bugs

---
 src/imitation/data/types.py                   | 56 +++++++++++++++++++
 src/imitation/policies/serialize.py           |  2 -
 .../scripts/common/demonstrations.py          |  7 +--
 .../scripts/config/train_imitation.py         |  3 +-
 tests/policies/test_policies.py               |  4 +-
 5 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index 3ce236bb4..c27a9ed35 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -33,6 +33,60 @@ def dataclass_quick_asdict(obj) -> Dict[str, Any]:
     return d
 
 
+def parse_path(
+        path: AnyPath,
+        allow_relative: bool = True,
+        base_directory: Optional[pathlib.Path] = None
+) -> pathlib.Path:
+    """
+    Parse a path, and check that it is absolute. If `allow_relative` is True,
+    then relative paths are allowed as input, and are resolved relative to the
+    current working directory, or relative to `base_directory` if it is
+    specified.
+
+    Args:
+        path: The path to parse. Can be a string, bytes, or `os.PathLike`.
+        allow_relative: If True, then relative paths are allowed as input, and
+            are resolved relative to the current working directory. If false,
+            an error is raised if the path is not absolute.
+        base_directory: If specified, then relative paths are resolved relative
+            to this directory, instead of the current working directory.
+
+    Returns:
+        A `pathlib.Path` object.
+
+    Raises:
+        ValueError: If `allow_relative` is False and the path is not absolute.
+        ValueError: If `base_directory` is specified and `allow_relative` is
+            False.
+    """
+    if base_directory is not None and not allow_relative:
+        raise ValueError(
+            "If `base_directory` is specified, then `allow_relative` must be True."
+        )
+
+    parsed_path: pathlib.Path
+    if isinstance(path, pathlib.Path):
+        parsed_path = path
+    elif isinstance(path, str):
+        parsed_path = pathlib.Path(path)
+    elif isinstance(path, bytes):
+        parsed_path = pathlib.Path(path.decode())
+    else:
+        parsed_path = pathlib.Path(str(path))
+
+    parsed_path = parsed_path.resolve()
+    if parsed_path.is_absolute():
+        return parsed_path
+    else:
+        if allow_relative:
+            base_directory = base_directory or pathlib.Path.cwd()
+            # relative to current working directory
+            return base_directory / parsed_path
+        else:
+            raise ValueError(f"Path {str(parsed_path)} is not absolute")
+
+
 def path_to_str(path: AnyPath) -> str:
     if isinstance(path, bytes):
         return path.decode()
@@ -42,6 +96,8 @@ def path_to_str(path: AnyPath) -> str:
 
 def path_to_pathlib(path: AnyPath) -> pathlib.Path:
     """Converts a path to a pathlib.Path."""
+    if isinstance(path, pathlib.Path):
+        return path
     return pathlib.Path(path_to_str(path))
 
 
diff --git a/src/imitation/policies/serialize.py b/src/imitation/policies/serialize.py
index ca925213a..197a3beff 100644
--- a/src/imitation/policies/serialize.py
+++ b/src/imitation/policies/serialize.py
@@ -4,7 +4,6 @@
 # torch.load() and torch.save() calls
 
 import logging
-import os
 import pathlib
 from typing import Callable, Type, TypeVar
 
@@ -138,7 +137,6 @@ def save_stable_model(
     # Save each model in new directory in case we want to add metadata or other
     # information in future. (E.g. we used to save `VecNormalize` statistics here,
     # although that is no longer necessary.)
-    output_dir = output_dir.resolve()
     output_dir.mkdir(parents=True, exist_ok=True)
     model.save(str(output_dir / "model.zip"))
     logging.info(f"Saved policy to {output_dir}")
diff --git a/src/imitation/scripts/common/demonstrations.py b/src/imitation/scripts/common/demonstrations.py
index c97eec58c..a0522a315 100644
--- a/src/imitation/scripts/common/demonstrations.py
+++ b/src/imitation/scripts/common/demonstrations.py
@@ -1,7 +1,6 @@
 """Common configuration element for scripts learning from demonstrations."""
 
 import logging
-import os
 import pathlib
 from typing import Dict, Optional, Sequence
 
@@ -28,9 +27,7 @@ def fast():
     n_expert_demos = 1  # noqa: F841
 
 
-def guess_expert_dir(data_dir: str, env_name: str) -> pathlib.Path:
-    data_dir = pathlib.Path(data_dir)
-    assert data_dir.is_absolute()
+def guess_expert_dir(data_dir: pathlib.Path, env_name: str) -> pathlib.Path:
     rollout_hint = env_name.rsplit("-", 1)[0].replace("/", "_").lower()
     return data_dir / "expert_models" / f"{rollout_hint}_0"
 
@@ -41,7 +38,7 @@ def hook(config, command_name, logger):
     del command_name, logger
     updates: Dict[str, str] = {}
     if config["demonstrations"]["rollout_path"] is None:
-        data_dir = config["demonstrations"]["data_dir"]
+        data_dir = types.parse_path(config["demonstrations"]["data_dir"])
         env_name = config["common"]["env_name"].replace("/", "_")
         updates["rollout_path"] = str(
             guess_expert_dir(data_dir, env_name) / "rollouts" / "final.pkl"
diff --git a/src/imitation/scripts/config/train_imitation.py b/src/imitation/scripts/config/train_imitation.py
index bb5a8c0c6..33bf767de 100644
--- a/src/imitation/scripts/config/train_imitation.py
+++ b/src/imitation/scripts/config/train_imitation.py
@@ -5,6 +5,7 @@
 import sacred
 import torch as th
 
+from imitation.data import types
 from imitation.scripts.common import common
 from imitation.scripts.common import demonstrations as demos_common
 from imitation.scripts.common import train
@@ -54,7 +55,7 @@ def defaults(
             expert_policy_type="ppo",
             expert_policy_path=os.path.join(
                 demos_common.guess_expert_dir(
-                    demonstrations["data_dir"],
+                    types.parse_path(demonstrations["data_dir"]),
                     common["env_name"],
                 ),
                 "policies",
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index f095edb7c..54a74a4b6 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -10,7 +10,7 @@
 from stable_baselines3.common import preprocessing
 from torch import nn
 
-from imitation.data import rollout
+from imitation.data import rollout, types
 from imitation.policies import base, serialize
 from imitation.util import registry, util
 
@@ -84,7 +84,7 @@ def _test_serialize_identity(env_name, model_cfg, tmpdir):
         rng=np.random.RandomState(0),
     )
 
-    serialize.save_stable_model(tmpdir, model)
+    serialize.save_stable_model(types.parse_path(tmpdir), model)
     loaded = serialize.load_policy(model_name, tmpdir, venv)
     venv.env_method("seed", 0)
     venv.reset()

From 3c45d8cdd117a8a831f8e85204f5264dcb3fc3f0 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 6 Sep 2022 17:12:32 +0200
Subject: [PATCH 049/187] Remove unused imports

---
 src/imitation/algorithms/adversarial/common.py        | 1 -
 src/imitation/data/types.py                           | 6 +++---
 src/imitation/scripts/common/common.py                | 1 -
 src/imitation/scripts/convert_trajs.py                | 1 -
 src/imitation/scripts/eval_policy.py                  | 1 -
 src/imitation/scripts/parallel.py                     | 1 -
 src/imitation/scripts/train_adversarial.py            | 2 --
 src/imitation/scripts/train_preference_comparisons.py | 1 -
 src/imitation/scripts/train_rl.py                     | 2 --
 src/imitation/util/logger.py                          | 1 -
 src/imitation/util/video_wrapper.py                   | 1 -
 11 files changed, 3 insertions(+), 15 deletions(-)

diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
index 3632a3883..80093d346 100644
--- a/src/imitation/algorithms/adversarial/common.py
+++ b/src/imitation/algorithms/adversarial/common.py
@@ -3,7 +3,6 @@
 import collections
 import dataclasses
 import logging
-import os
 import pathlib
 from typing import (
     Callable,
diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index 90ef73f66..d08f1aa51 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -34,9 +34,9 @@ def dataclass_quick_asdict(obj) -> Dict[str, Any]:
 
 
 def parse_path(
-        path: AnyPath,
-        allow_relative: bool = True,
-        base_directory: Optional[pathlib.Path] = None
+    path: AnyPath,
+    allow_relative: bool = True,
+    base_directory: Optional[pathlib.Path] = None,
 ) -> pathlib.Path:
     """
     Parse a path, and check that it is absolute. If `allow_relative` is True,
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index ffaccb453..287fd5cc6 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -2,7 +2,6 @@
 
 import contextlib
 import logging
-import os
 import pathlib
 from typing import Any, Generator, Mapping, Sequence, Tuple, Union
 
diff --git a/src/imitation/scripts/convert_trajs.py b/src/imitation/scripts/convert_trajs.py
index 8eabc7edd..23fce0e11 100644
--- a/src/imitation/scripts/convert_trajs.py
+++ b/src/imitation/scripts/convert_trajs.py
@@ -9,7 +9,6 @@
 (i.e. "A.pkl" -> "A.npz", "A.npz" -> "A.npz", "A" -> "A.npz", "A.foo" -> "A.foo.npz").
 """
 
-import os
 import pathlib
 import warnings
 
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 5b95dfc80..654e76582 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -2,7 +2,6 @@
 
 import logging
 import pathlib
-import re
 import time
 from typing import Any, Mapping, Optional
 
diff --git a/src/imitation/scripts/parallel.py b/src/imitation/scripts/parallel.py
index d3185d132..e808d2f7d 100644
--- a/src/imitation/scripts/parallel.py
+++ b/src/imitation/scripts/parallel.py
@@ -2,7 +2,6 @@
 
 import collections.abc
 import copy
-import os
 import pathlib
 from typing import Any, Callable, Dict, Mapping, Optional, Sequence
 
diff --git a/src/imitation/scripts/train_adversarial.py b/src/imitation/scripts/train_adversarial.py
index 46d030c3a..abebe23f8 100644
--- a/src/imitation/scripts/train_adversarial.py
+++ b/src/imitation/scripts/train_adversarial.py
@@ -2,8 +2,6 @@
 
 import functools
 import logging
-import os
-import os.path as osp
 import pathlib
 from typing import Any, Mapping, Optional, Type
 
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 142a4f1b3..7c61b76fe 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -5,7 +5,6 @@
 """
 
 import functools
-import os
 import pathlib
 from typing import Any, Mapping, Optional, Type, Union
 
diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
index 8b35ab9fb..86aa235de 100644
--- a/src/imitation/scripts/train_rl.py
+++ b/src/imitation/scripts/train_rl.py
@@ -9,8 +9,6 @@
 """
 
 import logging
-import os
-import os.path as osp
 import pathlib
 import warnings
 from typing import Any, Mapping, Optional
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index abad01059..759549500 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -2,7 +2,6 @@
 
 import contextlib
 import datetime
-import os
 import pathlib
 import tempfile
 from typing import Any, Dict, Generator, List, Optional, Sequence, Tuple, Union
diff --git a/src/imitation/util/video_wrapper.py b/src/imitation/util/video_wrapper.py
index 7334bc0d4..6788f927e 100644
--- a/src/imitation/util/video_wrapper.py
+++ b/src/imitation/util/video_wrapper.py
@@ -1,6 +1,5 @@
 """Wrapper to record rendered video frames from an environment."""
 
-import os
 import pathlib
 from typing import Optional
 

From 27581913313f08987810adef7309c1888d5a2335 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 6 Sep 2022 17:14:16 +0200
Subject: [PATCH 050/187] Formatting

---
 src/imitation/data/types.py                    | 7 ++++---
 src/imitation/scripts/common/demonstrations.py | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index d08f1aa51..3e4799969 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -38,8 +38,9 @@ def parse_path(
     allow_relative: bool = True,
     base_directory: Optional[pathlib.Path] = None,
 ) -> pathlib.Path:
-    """
-    Parse a path, and check that it is absolute. If `allow_relative` is True,
+    """Parse a path to a `pathlib.Path` object.
+
+    All resulting paths are resolved, absolute paths. If `allow_relative` is True,
     then relative paths are allowed as input, and are resolved relative to the
     current working directory, or relative to `base_directory` if it is
     specified.
@@ -62,7 +63,7 @@ def parse_path(
     """
     if base_directory is not None and not allow_relative:
         raise ValueError(
-            "If `base_directory` is specified, then `allow_relative` must be True."
+            "If `base_directory` is specified, then `allow_relative` must be True.",
         )
 
     parsed_path: pathlib.Path
diff --git a/src/imitation/scripts/common/demonstrations.py b/src/imitation/scripts/common/demonstrations.py
index a0522a315..d4a8cde88 100644
--- a/src/imitation/scripts/common/demonstrations.py
+++ b/src/imitation/scripts/common/demonstrations.py
@@ -41,7 +41,7 @@ def hook(config, command_name, logger):
         data_dir = types.parse_path(config["demonstrations"]["data_dir"])
         env_name = config["common"]["env_name"].replace("/", "_")
         updates["rollout_path"] = str(
-            guess_expert_dir(data_dir, env_name) / "rollouts" / "final.pkl"
+            guess_expert_dir(data_dir, env_name) / "rollouts" / "final.pkl",
         )
     return updates
 

From 141e638a056626d3fc47efaf302ffefa3c3da346 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 6 Sep 2022 17:52:06 +0200
Subject: [PATCH 051/187] Added parse_path func and refactored code to use it

---
 .../algorithms/adversarial/common.py          |  5 +--
 src/imitation/algorithms/dagger.py            |  9 +++--
 src/imitation/data/types.py                   | 40 ++++++++++++++-----
 src/imitation/policies/serialize.py           |  3 +-
 src/imitation/scripts/analyze.py              |  3 +-
 src/imitation/scripts/common/common.py        | 17 ++++----
 src/imitation/scripts/convert_trajs.py        |  5 +--
 src/imitation/scripts/train_adversarial.py    |  3 +-
 .../scripts/train_preference_comparisons.py   |  3 +-
 src/imitation/scripts/train_rl.py             |  3 +-
 src/imitation/util/logger.py                  |  6 +--
 src/imitation/util/sacred.py                  | 14 +++----
 src/imitation/util/video_wrapper.py           |  8 ++--
 tests/data/test_types.py                      |  7 ++--
 tests/policies/test_policies.py               |  3 +-
 tests/scripts/test_scripts.py                 | 16 ++++----
 16 files changed, 80 insertions(+), 65 deletions(-)

diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
index 80093d346..6829175c8 100644
--- a/src/imitation/algorithms/adversarial/common.py
+++ b/src/imitation/algorithms/adversarial/common.py
@@ -3,7 +3,6 @@
 import collections
 import dataclasses
 import logging
-import pathlib
 from typing import (
     Callable,
     Iterable,
@@ -127,7 +126,7 @@ def __init__(
         gen_algo: on_policy_algorithm.OnPolicyAlgorithm,
         reward_net: reward_nets.RewardNet,
         n_disc_updates_per_round: int = 2,
-        log_dir: str = "output/",
+        log_dir: types.AnyPath = "output/",
         disc_opt_cls: Type[th.optim.Optimizer] = th.optim.Adam,
         disc_opt_kwargs: Optional[Mapping] = None,
         gen_train_timesteps: Optional[int] = None,
@@ -202,7 +201,7 @@ def __init__(
         self.venv = venv
         self.gen_algo = gen_algo
         self._reward_net = reward_net.to(gen_algo.device)
-        self._log_dir = pathlib.Path(log_dir).resolve()
+        self._log_dir = types.parse_path(log_dir)
 
         # Create graph for optimising/recording stats on discriminator
         self._disc_opt_cls = disc_opt_cls
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index c345b1628..3a8e3c83b 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -89,7 +89,8 @@ def reconstruct_trainer(
         A deserialized `DAggerTrainer`.
     """
     custom_logger = custom_logger or logger.configure()
-    checkpoint_path = pathlib.Path(scratch_dir, "checkpoint-latest.pt")
+    scratch_dir = types.parse_path(scratch_dir)
+    checkpoint_path = scratch_dir / "checkpoint-latest.pt"
     trainer = th.load(checkpoint_path, map_location=utils.get_device(device))
     trainer.venv = venv
     trainer._logger = custom_logger
@@ -105,14 +106,14 @@ def _save_dagger_demo(
     #   however that NPZ save here is likely more space efficient than
     #   pickle from types.save(), and types.save only accepts
     #   TrajectoryWithRew right now (subclass of Trajectory).
-    save_dir = pathlib.Path(save_dir)
+    save_dir = types.parse_path(save_dir)
     assert isinstance(trajectory, types.Trajectory)
     actual_prefix = f"{prefix}-" if prefix else ""
     timestamp = util.make_unique_timestamp()
     filename = f"{actual_prefix}dagger-demo-{timestamp}.npz"
 
     save_dir.mkdir(parents=True, exist_ok=True)
-    npz_path = pathlib.Path(save_dir, filename)
+    npz_path = save_dir / filename
     np.savez_compressed(npz_path, **dataclasses.asdict(trajectory))
     logging.info(f"Saved demo at '{npz_path}'")
 
@@ -327,7 +328,7 @@ def __init__(
         if beta_schedule is None:
             beta_schedule = LinearBetaSchedule(15)
         self.beta_schedule = beta_schedule
-        self.scratch_dir = pathlib.Path(scratch_dir)
+        self.scratch_dir = types.parse_path(scratch_dir)
         self.venv = venv
         self.round_num = 0
         self._last_loaded_round = -1
diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index 3e4799969..03d937240 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -88,18 +88,38 @@ def parse_path(
             raise ValueError(f"Path {str(parsed_path)} is not absolute")
 
 
-def path_to_str(path: AnyPath) -> str:
-    if isinstance(path, bytes):
-        return path.decode()
-    else:
-        return str(path)
+def parse_optional_path(
+    path: Optional[AnyPath],
+    allow_relative: bool = True,
+    base_directory: Optional[pathlib.Path] = None,
+) -> Optional[pathlib.Path]:
+    """Parse an optional path to a `pathlib.Path` object.
 
+    All resulting paths are resolved, absolute paths. If `allow_relative` is True,
+    then relative paths are allowed as input, and are resolved relative to the
+    current working directory, or relative to `base_directory` if it is
+    specified.
 
-def path_to_pathlib(path: AnyPath) -> pathlib.Path:
-    """Converts a path to a pathlib.Path."""
-    if isinstance(path, pathlib.Path):
-        return path
-    return pathlib.Path(path_to_str(path))
+    Args:
+        path: The path to parse. Can be a string, bytes, or `os.PathLike`.
+        allow_relative: If True, then relative paths are allowed as input, and
+            are resolved relative to the current working directory. If false,
+            an error is raised if the path is not absolute.
+        base_directory: If specified, then relative paths are resolved relative
+            to this directory, instead of the current working directory.
+
+    Returns:
+        A `pathlib.Path` object, or None if `path` is None.
+
+    Raises:
+        ValueError: If `allow_relative` is False and the path is not absolute.
+        ValueError: If `base_directory` is specified and `allow_relative` is
+            False.
+    """
+    if path is None:
+        return None
+    else:
+        return parse_path(path, allow_relative, base_directory)
 
 
 @dataclasses.dataclass(frozen=True)
diff --git a/src/imitation/policies/serialize.py b/src/imitation/policies/serialize.py
index 046738323..96d04cba1 100644
--- a/src/imitation/policies/serialize.py
+++ b/src/imitation/policies/serialize.py
@@ -9,6 +9,7 @@
 
 from stable_baselines3.common import base_class, callbacks, policies, vec_env
 
+from imitation.data import types
 from imitation.policies import base
 from imitation.util import registry
 
@@ -41,7 +42,7 @@ def load_stable_baselines_model(
         The deserialized RL algorithm.
     """
     logging.info(f"Loading Stable Baselines policy for '{cls}' from '{path}'")
-    policy_dir = pathlib.Path(path)
+    policy_dir = types.parse_path(path)
     if not policy_dir.is_dir():
         raise FileNotFoundError(
             f"path={path} needs to be a directory containing model.zip.",
diff --git a/src/imitation/scripts/analyze.py b/src/imitation/scripts/analyze.py
index 6553aaf78..a3a2bf683 100644
--- a/src/imitation/scripts/analyze.py
+++ b/src/imitation/scripts/analyze.py
@@ -16,6 +16,7 @@
 from sacred.observers import FileStorageObserver
 
 import imitation.util.sacred as sacred_util
+from imitation.data import types
 from imitation.scripts.config.analyze import analysis_ex
 from imitation.util.sacred import dict_get_nested as get
 
@@ -108,7 +109,7 @@ def gather_tb_directories() -> dict:
         # Expecting a path like "~/ray_results/{run_name}/sacred/1".
         # Want to search for all Tensorboard dirs inside
         # "~/ray_results/{run_name}".
-        sacred_dir = pathlib.Path(sd.sacred_dir)
+        sacred_dir = types.parse_path(sd.sacred_dir)
         run_dir = sacred_dir.parent.parent
         run_name = run_dir.name
 
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index 287fd5cc6..e12f57edc 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -8,6 +8,7 @@
 import sacred
 from stable_baselines3.common import vec_env
 
+from imitation.data import types
 from imitation.scripts.common import wb
 from imitation.util import logger as imit_logger
 from imitation.util import sacred as sacred_util
@@ -50,11 +51,11 @@ def hook(config, command_name: str, logger):
     if config["common"]["log_dir"] is None:
         env_sanitized = config["common"]["env_name"].replace("/", "_")
         assert isinstance(env_sanitized, str)
-        config_log_root = config["common"]["log_root"]
+        config_log_root = types.parse_optional_path(config["common"]["log_root"])
         log_root = (
-            pathlib.Path(config_log_root)
+            config_log_root
             if config_log_root
-            else pathlib.Path.cwd() / "output"
+            else types.parse_path("output")  # relative to cwd
         )
         log_dir = log_root / command_name / env_sanitized / util.make_unique_timestamp()
         updates["log_dir"] = str(log_dir)
@@ -92,17 +93,17 @@ def make_log_dir(
     Returns:
         The `log_dir`. This avoids the caller needing to capture this argument.
     """
-    _log_dir = pathlib.Path(log_dir)
-    _log_dir.mkdir(parents=True, exist_ok=True)
+    parsed_log_dir = types.parse_path(log_dir)
+    parsed_log_dir.mkdir(parents=True, exist_ok=True)
     # convert strings of digits to numbers; but leave levels like 'INFO' unmodified
     try:
         log_level = int(log_level)
     except ValueError:
         pass
     logging.basicConfig(level=log_level)
-    logger.info("Logging to %s", _log_dir)
-    sacred_util.build_sacred_symlink(_log_dir, _run)
-    return _log_dir
+    logger.info("Logging to %s", parsed_log_dir)
+    sacred_util.build_sacred_symlink(parsed_log_dir, _run)
+    return parsed_log_dir
 
 
 @common_ingredient.capture
diff --git a/src/imitation/scripts/convert_trajs.py b/src/imitation/scripts/convert_trajs.py
index 23fce0e11..2e3c7c016 100644
--- a/src/imitation/scripts/convert_trajs.py
+++ b/src/imitation/scripts/convert_trajs.py
@@ -9,13 +9,12 @@
 (i.e. "A.pkl" -> "A.npz", "A.npz" -> "A.npz", "A" -> "A.npz", "A.foo" -> "A.foo.npz").
 """
 
-import pathlib
 import warnings
 
 from imitation.data import types
 
 
-def update_traj_file_in_place(path: str) -> None:
+def update_traj_file_in_place(path_str: str, /) -> None:
     """Modifies trajectories pickle file in-place to update data to new format.
 
     The new data is saved as `Sequence[imitation.types.TrajectoryWithRew]`.
@@ -24,6 +23,7 @@ def update_traj_file_in_place(path: str) -> None:
         path: Path to a pickle file containing `Sequence[imitation.types.Trajectory]`
             or `Sequence[imitation.old_types.TrajectoryWithRew]`.
     """
+    path = types.parse_path(path_str)
     with warnings.catch_warnings():
         # Filter out DeprecationWarning because we expect to load old trajectories here.
         warnings.filterwarnings(
@@ -33,7 +33,6 @@ def update_traj_file_in_place(path: str) -> None:
         )
         trajs = types.load(path)
 
-    path = pathlib.Path(path)
     ext = path.suffix
     new_ext = ".npz" if ext in (".pkl", ".npz") else ext + ".npz"
     types.save(path.with_suffix(new_ext), trajs)
diff --git a/src/imitation/scripts/train_adversarial.py b/src/imitation/scripts/train_adversarial.py
index abebe23f8..3ba305951 100644
--- a/src/imitation/scripts/train_adversarial.py
+++ b/src/imitation/scripts/train_adversarial.py
@@ -112,8 +112,7 @@ def train_adversarial(
         # So, support showing merged config from `train_adversarial {airl,gail}`.
         sacred.commands.print_config(_run)
 
-    custom_logger, _log_dir = common_config.setup_logging()
-    log_dir = pathlib.Path(_log_dir)
+    custom_logger, log_dir = common_config.setup_logging()
     expert_trajs = demonstrations.load_expert_trajs()
 
     with common_config.make_venv() as venv:
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 7c61b76fe..f74b0d468 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -147,8 +147,7 @@ def train_preference_comparisons(
     Raises:
         ValueError: Inconsistency between config and deserialized policy normalization.
     """
-    custom_logger, _log_dir = common.setup_logging()
-    log_dir = pathlib.Path(_log_dir)
+    custom_logger, log_dir = common.setup_logging()
 
     with common.make_venv() as venv:
         reward_net = reward.make_reward_net(venv)
diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
index 86aa235de..288804a76 100644
--- a/src/imitation/scripts/train_rl.py
+++ b/src/imitation/scripts/train_rl.py
@@ -86,8 +86,7 @@ def train_rl(
     Returns:
         The return value of `rollout_stats()` using the final policy.
     """
-    custom_logger, _log_dir = common.setup_logging()
-    log_dir = pathlib.Path(_log_dir)
+    custom_logger, log_dir = common.setup_logging()
     rollout_dir = log_dir / "rollouts"
     policy_dir = log_dir / "policies"
     rollout_dir.mkdir(parents=True, exist_ok=True)
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 759549500..9620d2835 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -120,7 +120,7 @@ def accumulate_means(self, subdir: types.AnyPath) -> Generator[None, None, None]
         else:
             default_logger_dir = self.default_logger.dir
             assert default_logger_dir is not None
-            folder = pathlib.Path(default_logger_dir) / "raw" / subdir_str
+            folder = types.parse_path(default_logger_dir) / "raw" / subdir_str
             folder.mkdir(exist_ok=True, parents=True)
             output_formats = _build_output_formats(folder, self.format_strs)
             logger = sb_logger.Logger(str(folder), list(output_formats))
@@ -235,12 +235,12 @@ def configure(
         The configured HierarchicalLogger instance.
     """
     if folder is None:
-        tempdir = pathlib.Path(tempfile.gettempdir())
+        tempdir = types.parse_path(tempfile.gettempdir())
         now = datetime.datetime.now()
         timestamp = now.strftime("imitation-%Y-%m-%d-%H-%M-%S-%f")
         folder = tempdir / timestamp
     else:
-        folder = types.path_to_pathlib(folder)
+        folder = types.parse_path(folder)
     if format_strs is None:
         format_strs = ["stdout", "log", "csv"]
     output_formats = _build_output_formats(folder, format_strs)
diff --git a/src/imitation/util/sacred.py b/src/imitation/util/sacred.py
index b96f872df..3d223e633 100644
--- a/src/imitation/util/sacred.py
+++ b/src/imitation/util/sacred.py
@@ -4,7 +4,7 @@
 import os
 import pathlib
 import warnings
-from typing import Any, Callable, NamedTuple, Sequence, Union
+from typing import Any, Callable, NamedTuple, Sequence, Optional
 
 import sacred
 import sacred.observers
@@ -80,16 +80,14 @@ def filter_subdirs(
 
 def build_sacred_symlink(log_dir: types.AnyPath, run: sacred.run.Run) -> None:
     """Constructs a symlink "{log_dir}/sacred" => "${SACRED_PATH}"."""
-    if isinstance(log_dir, bytes):
-        log_dir = log_dir.decode("utf-8")
-    log_dir = pathlib.Path(log_dir)
+    log_dir = types.parse_path(log_dir)
 
     sacred_dir = get_sacred_dir_from_run(run)
     if sacred_dir is None:
         warnings.warn(RuntimeWarning("Couldn't find sacred directory."))
         return
-    symlink_path = pathlib.Path(log_dir, "sacred")
-    target_path = pathlib.Path(os.path.relpath(sacred_dir, start=log_dir))
+    symlink_path = log_dir / "sacred"
+    target_path = sacred_dir.relative_to(log_dir)
 
     # Path.symlink_to errors if the symlink already exists. In our case, we actually
     # want to override the symlink to point to the most recent Sacred dir. The
@@ -116,11 +114,11 @@ def build_sacred_symlink(log_dir: types.AnyPath, run: sacred.run.Run) -> None:
             raise e
 
 
-def get_sacred_dir_from_run(run: sacred.run.Run) -> Union[pathlib.Path, None]:
+def get_sacred_dir_from_run(run: sacred.run.Run) -> Optional[pathlib.Path]:
     """Returns path to the sacred directory, or None if not found."""
     for obs in run.observers:
         if isinstance(obs, sacred.observers.FileStorageObserver):
-            return pathlib.Path(obs.dir)
+            return types.parse_path(obs.dir)
     return None
 
 
diff --git a/src/imitation/util/video_wrapper.py b/src/imitation/util/video_wrapper.py
index 6788f927e..a59641aa1 100644
--- a/src/imitation/util/video_wrapper.py
+++ b/src/imitation/util/video_wrapper.py
@@ -6,8 +6,6 @@
 import gym
 from gym.wrappers.monitoring import video_recorder
 
-from imitation.data import types
-
 
 class VideoWrapper(gym.Wrapper):
     """Creates videos from wrapped environment by calling render after each timestep."""
@@ -15,12 +13,12 @@ class VideoWrapper(gym.Wrapper):
     episode_id: int
     video_recorder: Optional[video_recorder.VideoRecorder]
     single_video: bool
-    directory: str
+    directory: pathlib.Path
 
     def __init__(
         self,
         env: gym.Env,
-        directory: types.AnyPath,
+        directory: pathlib.Path,
         single_video: bool = True,
     ):
         """Builds a VideoWrapper.
@@ -39,7 +37,7 @@ def __init__(
         self.video_recorder = None
         self.single_video = single_video
 
-        self.directory = pathlib.Path(types.path_to_str(directory)).resolve()
+        self.directory = directory
         self.directory.mkdir(parents=True, exist_ok=True)
 
     def _reset_video_recorder(self) -> None:
diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index 851626404..b9511e067 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -213,13 +213,14 @@ def test_save_trajectories(
         if use_chdir:
             # Test no relative path without directory edge-case.
             chdir_context = pushd(tmpdir)
-            save_dir = ""
+            save_dir_str = ""
         else:
             chdir_context = contextlib.nullcontext()
-            save_dir = tmpdir
+            save_dir_str = tmpdir
+        save_dir = types.parse_path(save_dir_str)
 
         trajs = [trajectory_rew if use_rewards else trajectory]
-        save_path = pathlib.Path(save_dir, "trajs")
+        save_path = save_dir / "trajs"
 
         with chdir_context:
             if use_pickle:
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index 3d7da00f8..0ee80c947 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -1,7 +1,6 @@
 """Tests `imitation.policies.*`."""
 
 import functools
-import pathlib
 from typing import cast
 
 import gym
@@ -45,7 +44,7 @@ def test_actions_valid(env_name, policy_type):
 def test_save_stable_model_errors_and_warnings(tmpdir, policy_env_name_pair):
     """Check errors and warnings in `save_stable_model()`."""
     policy, env_name = policy_env_name_pair
-    tmpdir = pathlib.Path(tmpdir)
+    tmpdir = types.parse_path(tmpdir)
     venv = util.make_vec_env(env_name)
 
     # Trigger FileNotFoundError for no model.{zip,pkl}
diff --git a/tests/scripts/test_scripts.py b/tests/scripts/test_scripts.py
index 7b592f0c4..cae7b599f 100644
--- a/tests/scripts/test_scripts.py
+++ b/tests/scripts/test_scripts.py
@@ -53,7 +53,7 @@
     train_rl,
 ]
 
-TEST_DATA_PATH = pathlib.Path("tests/testdata")
+TEST_DATA_PATH = types.parse_path("tests/testdata")
 CARTPOLE_TEST_DATA_PATH = TEST_DATA_PATH / "expert_models/cartpole_0/"
 CARTPOLE_TEST_ROLLOUT_PATH = CARTPOLE_TEST_DATA_PATH / "rollouts/final.pkl"
 CARTPOLE_TEST_POLICY_PATH = CARTPOLE_TEST_DATA_PATH / "policies/final"
@@ -297,7 +297,7 @@ def test_train_dagger_warmstart(tmpdir):
     )
     assert run.status == "COMPLETED"
 
-    log_dir = pathlib.Path(run.config["common"]["log_dir"])
+    log_dir = types.parse_path(run.config["common"]["log_dir"])
     policy_path = log_dir / "scratch" / "policy-latest.pt"
     run_warmstart = train_imitation.train_imitation_ex.run(
         command_name="dagger",
@@ -356,7 +356,7 @@ def test_train_bc_warmstart(tmpdir):
     )
     assert run.status == "COMPLETED"
 
-    policy_path = pathlib.Path(run.config["common"]["log_dir"]) / "final.th"
+    policy_path = types.parse_path(run.config["common"]["log_dir"]) / "final.th"
     run_warmstart = train_imitation.train_imitation_ex.run(
         command_name="bc",
         named_configs=["cartpole"] + ALGO_FAST_CONFIGS["imitation"],
@@ -508,7 +508,7 @@ def test_train_adversarial_warmstart(tmpdir, command):
         config_updates=config_updates,
     )
 
-    log_dir = pathlib.Path(run.config["common"]["log_dir"])
+    log_dir = types.parse_path(run.config["common"]["log_dir"])
     policy_path = log_dir / "checkpoints" / "final" / "gen_policy"
 
     run_warmstart = train_adversarial.train_adversarial_ex.run(
@@ -593,7 +593,7 @@ def test_transfer_learning(tmpdir: str) -> None:
     Args:
         tmpdir: Temporary directory to save results to.
     """
-    tmpdir_path = pathlib.Path(tmpdir)
+    tmpdir_path = types.parse_path(tmpdir)
     log_dir_train = tmpdir_path / "train"
     run = train_adversarial.train_adversarial_ex.run(
         command_name="airl",
@@ -642,7 +642,7 @@ def test_preference_comparisons_transfer_learning(
         tmpdir: Temporary directory to save results to.
         named_configs_dict: Named configs for preference_comparisons and rl.
     """
-    tmpdir_path = pathlib.Path(tmpdir)
+    tmpdir_path = types.parse_path(tmpdir)
 
     log_dir_train = tmpdir_path / "train"
     run = train_preference_comparisons.train_preference_comparisons_ex.run(
@@ -778,7 +778,7 @@ def test_parallel_arg_errors(tmpdir):
 
 
 def _generate_test_rollouts(tmpdir: str, env_named_config: str) -> pathlib.Path:
-    tmpdir_path = pathlib.Path(tmpdir)
+    tmpdir_path = types.parse_path(tmpdir)
     train_rl.train_rl_ex.run(
         named_configs=[env_named_config] + ALGO_FAST_CONFIGS["rl"],
         config_updates=dict(
@@ -850,7 +850,7 @@ def _run_train_bc_for_test_analyze_imit(run_name, sacred_logs_dir, log_dir):
     ),
 )
 def test_analyze_imitation(tmpdir: str, run_names: List[str], run_sacred_fn):
-    sacred_logs_dir = tmpdir_path = pathlib.Path(tmpdir)
+    sacred_logs_dir = tmpdir_path = types.parse_path(tmpdir)
 
     # Generate sacred logs (other logs are put in separate tmpdir for deletion).
     for i, run_name in enumerate(run_names):

From aa2a7e464a05788947de500826d17367216faf3e Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 6 Sep 2022 18:00:59 +0200
Subject: [PATCH 052/187] Fix typing, linting

---
 src/imitation/util/logger.py | 13 ++++++-------
 src/imitation/util/sacred.py |  2 +-
 tests/data/test_types.py     | 11 -----------
 3 files changed, 7 insertions(+), 19 deletions(-)

diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 9620d2835..c49fbdeb6 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -78,7 +78,7 @@ 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 accumulate_means(self, subdir: 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
@@ -114,21 +114,20 @@ 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_str = types.path_to_str(subdir)
-        if subdir_str in self._cached_loggers:
-            logger = self._cached_loggers[subdir_str]
+        if subdir in self._cached_loggers:
+            logger = self._cached_loggers[subdir]
         else:
             default_logger_dir = self.default_logger.dir
             assert default_logger_dir is not None
-            folder = types.parse_path(default_logger_dir) / "raw" / subdir_str
+            folder = types.parse_path(default_logger_dir) / "raw" / subdir
             folder.mkdir(exist_ok=True, parents=True)
             output_formats = _build_output_formats(folder, self.format_strs)
             logger = sb_logger.Logger(str(folder), list(output_formats))
-            self._cached_loggers[subdir_str] = logger
+            self._cached_loggers[subdir] = logger
 
         try:
             self.current_logger = logger
-            self._subdir = subdir_str
+            self._subdir = subdir
             self._update_name_to_maps()
             yield
         finally:
diff --git a/src/imitation/util/sacred.py b/src/imitation/util/sacred.py
index 3d223e633..8a85d146e 100644
--- a/src/imitation/util/sacred.py
+++ b/src/imitation/util/sacred.py
@@ -4,7 +4,7 @@
 import os
 import pathlib
 import warnings
-from typing import Any, Callable, NamedTuple, Sequence, Optional
+from typing import Any, Callable, NamedTuple, Optional, Sequence
 
 import sacred
 import sacred.observers
diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index b9511e067..01a12e55a 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -388,14 +388,3 @@ def test_zero_length_fails():
     empty = np.array([])
     with pytest.raises(ValueError, match=r"Degenerate trajectory.*"):
         types.Trajectory(obs=np.array([42]), acts=empty, infos=None, terminal=True)
-
-
-def test_path_to_str():
-    assert types.path_to_str("") == ""
-    assert types.path_to_str(b"") == ""
-    assert types.path_to_str("foo") == "foo"
-    assert types.path_to_str(b"foo") == "foo"
-    assert types.path_to_str(pathlib.Path("foo")) == "foo"
-    assert types.path_to_str("/foo/bar") == "/foo/bar"
-    assert types.path_to_str(b"/foo/bar") == "/foo/bar"
-    assert types.path_to_str(pathlib.Path("/foo", "bar"))

From 970931ae1c678ac5e1c538216010e13fb51490e1 Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Tue, 6 Sep 2022 14:53:57 -0700
Subject: [PATCH 053/187] Update TabularPolicy.predict to match base class

---
 src/imitation/algorithms/mce_irl.py | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index d8cf4e108..77d84872a 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -185,13 +185,13 @@ def forward(  # type: ignore[override]
     ):
         raise NotImplementedError("Should never be called.")
 
-    def predict(  # type: ignore[override]
+    def predict(
         self,
-        observation: np.ndarray,
-        state: Optional[np.ndarray] = None,
-        mask: Optional[np.ndarray] = None,
+        observation: Union[np.ndarray, Mapping[str, np.ndarray]],
+        state: Optional[Tuple[np.ndarray, ...]] = None,
+        episode_start: Optional[np.ndarray] = None,
         deterministic: bool = False,
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    ) -> Tuple[np.ndarray, Optional[Tuple[np.ndarray, ...]]]:
         """Predict action to take in given state.
 
         Arguments follow SB3 naming convention as this is an SB3 policy.
@@ -205,24 +205,22 @@ def predict(  # type: ignore[override]
         Args:
             observation: States in the underlying MDP.
             state: Hidden states of the policy -- used to represent timesteps by us.
-            mask: Has episode completed?
+            episode_start: Has episode completed?
             deterministic: If true, pick action with highest probability; otherwise,
                 sample.
 
         Returns:
             Tuple of the actions and new hidden states.
         """
-        timesteps = state  # rename to avoid confusion
-        del state
-
-        if timesteps is None:
+        if state is None:
             timesteps = np.zeros(len(observation), dtype=int)
         else:
-            timesteps = np.array(timesteps)
+            assert len(state) == 1
+            timesteps = state[0]
         assert len(timesteps) == len(observation), "timestep and obs batch size differ"
 
-        if mask is not None:
-            timesteps[mask] = 0
+        if episode_start is not None:
+            timesteps[episode_start] = 0
 
         actions: List[int] = []
         for obs, t in zip(observation, timesteps):
@@ -234,7 +232,8 @@ def predict(  # type: ignore[override]
                 actions.append(self.rng.choice(len(dist), p=dist))
 
         timesteps += 1  # increment timestep
-        return np.array(actions), timesteps
+        state = (timesteps,)
+        return np.array(actions), state
 
 
 MCEDemonstrations = Union[np.ndarray, base.AnyTransitions]

From 02a4aa574c36a834ba3d283dfae102f410bb6cc9 Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Tue, 6 Sep 2022 14:55:35 -0700
Subject: [PATCH 054/187] Fix not checking for dones

---
 src/imitation/algorithms/mce_irl.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index 77d84872a..6ad54872d 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -355,11 +355,7 @@ def _set_demo_from_obs(
         # then possibly shuffled. So add next observations for terminal states,
         # as they will not appear anywhere else; but ignore next observations
         # for all other states as they occur elsewhere in dataset.
-        if next_obses is not None:
-            # TODO(juan) this is wrong; dones is an optional
-            #  and zip() cannot handle this when it's None.
-            #  either require an np.ndarray or do something else
-            #  if the value is None.
+        if dones is not None and next_obses is not None:
             for done, obs in zip(dones, next_obses):
                 if isinstance(done, th.Tensor):
                     done = done.item()  # must be scalar

From f16f92424b51f31c155cf7666bf284d1a2ba4324 Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Tue, 6 Sep 2022 15:34:46 -0700
Subject: [PATCH 055/187] Change for loop to dict comprehension

---
 src/imitation/algorithms/mce_irl.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index 6ad54872d..df2079471 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -416,9 +416,7 @@ def set_demonstrations(self, demonstrations: MCEDemonstrations) -> None:
                 for k in ("obs", "dones", "next_obs"):
                     if k in batch:
                         collated_list[k].append(batch[k])
-            collated = {}
-            for k, v in collated_list.items():
-                collated[k] = np.concatenate(v)
+            collated = {k: np.concatenate(v) for k, v in collated_list.items()}
 
             assert "obs" in collated
             for k, v in collated.items():

From ba00b40db16b596e8555cf3b77b2ad4e9d4673d9 Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Tue, 6 Sep 2022 16:01:14 -0700
Subject: [PATCH 056/187] Remove is_ensemble to clear up type checking errors

---
 .../algorithms/preference_comparisons.py      | 44 +++++++++----------
 1 file changed, 21 insertions(+), 23 deletions(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 23d367175..6d5560d31 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -368,17 +368,14 @@ def __init__(
         self.noise_prob = noise_prob
         self.discount_factor = discount_factor
         self.threshold = threshold
-        self.is_ensemble: bool
         base_model = get_base_model(model)
-        self.is_ensemble = isinstance(base_model, reward_nets.RewardEnsemble)
+        self.ensemble_model = None
         # if the base model is an ensemble model, then keep the base model as
         # model to get rewards from all networks
-        if self.is_ensemble:
-            # For some reason, mypy is not able to infer the type of base_model
-            # when self.is_ensemble is True.
-            self.model = cast(reward_nets.RewardEnsemble, base_model)
+        if isinstance(base_model, reward_nets.RewardEnsemble):
+            self.ensemble_model = base_model
             self.member_pref_models = []
-            for member in self.model.members:
+            for member in self.ensemble_model.members:
                 member_pref_model = PreferenceModel(
                     cast(reward_nets.RewardNet, member),  # nn.ModuleList is not generic
                     self.noise_prob,
@@ -419,7 +416,7 @@ def forward(
             network and (num_fragment_pairs, num_networks) for an ensemble of networks.
 
         """
-        if self.is_ensemble:
+        if self.ensemble_model is not None:
             if ensemble_member_index is None:
                 raise ValueError(
                     "`ensemble_member_index` required for ensemble models.",
@@ -468,17 +465,20 @@ def rewards(
         action = transitions.acts
         next_state = transitions.next_obs
         done = transitions.dones
-        if self.is_ensemble:
-            # TODO(juan) for some super strange reason mypy thinks
-            #  self.model.predict_processed_all is a tensor?
-            rews = self.model.predict_processed_all(state, action, next_state, done)
-            assert rews.shape == (len(state), self.model.num_members)
-            return util.safe_to_tensor(rews).to(self.model.device)
+        if self.ensemble_model is not None:
+            rews_np = self.ensemble_model.predict_processed_all(
+                state,
+                action,
+                next_state,
+                done,
+            )
+            assert rews_np.shape == (len(state), self.ensemble_model.num_members)
+            rews = util.safe_to_tensor(rews_np).to(self.ensemble_model.device)
         else:
             preprocessed = self.model.preprocess(state, action, next_state, done)
             rews = self.model(*preprocessed)
             assert rews.shape == (len(state),)
-            return rews
+        return rews
 
     def probability(
         self,
@@ -500,7 +500,7 @@ def probability(
             () for non-ensemble model which is a torch scalar.
         """
         # check rews has correct shape based on the model
-        expected_dims = 2 if self.is_ensemble else 1
+        expected_dims = 2 if self.ensemble_model is not None else 1
         assert rews1.ndim == rews2.ndim == expected_dims
         # First, we compute the difference of the returns of
         # the two fragments. We have a special case for a discount
@@ -510,7 +510,7 @@ def probability(
             returns_diff = (rews2 - rews1).sum(axis=0)  # type: ignore[call-overload]
         else:
             discounts = self.discount_factor ** th.arange(len(rews1))
-            if self.is_ensemble:
+            if self.ensemble_model is not None:
                 discounts = discounts.reshape(-1, 1)
             returns_diff = (discounts * (rews2 - rews1)).sum(axis=0)
         # Clip to avoid overflows (which in particular may occur
@@ -521,7 +521,7 @@ def probability(
         # probability that fragment 1 is preferred.
         model_probability = 1 / (1 + returns_diff.exp())
         probability = self.noise_prob * 0.5 + (1 - self.noise_prob) * model_probability
-        if self.is_ensemble:
+        if self.ensemble_model is not None:
             assert probability.shape == (self.model.num_members,)
         else:
             assert probability.shape == ()
@@ -915,7 +915,7 @@ def __init__(self, max_size: Optional[int] = None):
         self.fragments1: List[TrajectoryWithRew] = []
         self.fragments2: List[TrajectoryWithRew] = []
         self.max_size = max_size
-        self.preferences = np.array([])
+        self.preferences: np.ndarray = np.array([])
 
     def push(
         self,
@@ -1271,7 +1271,7 @@ def _training_inner_loop(
 
 def get_base_model(
     reward_model: reward_nets.RewardNet,
-) -> Union[reward_nets.RewardNet, reward_nets.RewardEnsemble]:
+) -> reward_nets.RewardNet:
     base_model = reward_model
     while hasattr(base_model, "base"):
         base_model = cast(reward_nets.RewardNet, base_model.base)
@@ -1290,10 +1290,8 @@ def _make_reward_trainer(
         reward_trainer_kwargs = {}
 
     base_model = get_base_model(reward_model)
-    is_ensemble = isinstance(base_model, reward_nets.RewardEnsemble)
 
-    if is_ensemble:
-        base_model = cast(reward_nets.RewardEnsemble, base_model)
+    if isinstance(base_model, reward_nets.RewardEnsemble):
         # reward_model may include an AddSTDRewardWrapper for RL training; but we
         # must train directly on the base model for reward model training.
         is_base = reward_model is base_model

From b4fd47d293bb39c951c54950b1696da6f551e4e5 Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Tue, 6 Sep 2022 16:10:12 -0700
Subject: [PATCH 057/187] Reduce code duplication and general cleanup

---
 src/imitation/data/buffer.py | 43 ++++++++++++++++++++++--------------
 1 file changed, 27 insertions(+), 16 deletions(-)

diff --git a/src/imitation/data/buffer.py b/src/imitation/data/buffer.py
index 33bef6afd..5554f403a 100644
--- a/src/imitation/data/buffer.py
+++ b/src/imitation/data/buffer.py
@@ -1,7 +1,7 @@
 """Buffers to store NumPy arrays and transitions in."""
 
 import dataclasses
-from typing import Mapping, Optional, Tuple
+from typing import Any, Mapping, Optional, Tuple
 
 import numpy as np
 from stable_baselines3.common import vec_env
@@ -9,6 +9,25 @@
 from imitation.data import types
 
 
+def num_samples(data: Mapping[Any, np.ndarray]) -> int:
+    """Computes the number of samples contained in `data`.
+
+    Args:
+        data: A Mapping from keys to NumPy arrays.
+
+    Returns:
+        The unique length of the first dimension of arrays contained in `data`.
+
+    Raises:
+        ValueError: The length is not unique.
+    """
+    n_samples_list = [arr.shape[0] for arr in data.values()]
+    n_samples_np = np.unique(n_samples_list)
+    if len(n_samples_np) > 1:
+        raise ValueError("Keys map to different length values.")
+    return int(n_samples_np[0])
+
+
 class Buffer:
     """A FIFO ring buffer for NumPy arrays of a fixed shape and dtype.
 
@@ -150,12 +169,7 @@ def store(self, data: Mapping[str, np.ndarray], truncate_ok: bool = False) -> No
         if len(unexpected_keys) > 0:
             raise ValueError(f"Unexpected keys {unexpected_keys}")
 
-        n_samples_list = [arr.shape[0] for arr in data.values()]
-        n_samples_np = np.unique(n_samples_list)
-        if len(n_samples_np) > 1:
-            raise ValueError("Keys map to different length values.")
-        n_samples = int(n_samples_np[0])
-
+        n_samples = num_samples(data)
         if n_samples == 0:
             raise ValueError("Trying to store empty data.")
         if n_samples > self.capacity:
@@ -192,11 +206,7 @@ def _store_easy(self, data: Mapping[str, np.ndarray]) -> None:
             data: Same as in `self.store`'s docstring, except with the additional
                 constraint `size(data) <= self.capacity - self._idx`.
         """
-        n_samples_list = [arr.shape[0] for arr in data.values()]
-        n_samples_np = np.unique(n_samples_list)
-        assert len(n_samples_np) == 1
-        n_samples = int(n_samples_np[0])
-
+        n_samples = num_samples(data)
         assert n_samples <= self.capacity - self._idx
         idx_hi = self._idx + n_samples
         for k, arr in data.items():
@@ -250,7 +260,7 @@ def __init__(
             capacity: The number of samples that can be stored.
             venv: The environment whose action and observation
                 spaces can be used to determine the data shapes of the underlying
-                buffers. Overrides all the following arguments.
+                buffers. Mutually exclusive with shape and dtype arguments.
             obs_shape: The shape of the observation space.
             act_shape: The shape of the action space.
             obs_dtype: The dtype of the observation space.
@@ -259,10 +269,11 @@ def __init__(
         Raises:
             ValueError: Couldn't infer the observation and action shapes and dtypes
                 from the arguments.
+            ValueError: Specified both venv and shapes/dtypes.
         """
-        params = [obs_shape, act_shape, obs_dtype, act_dtype]
+        params = (obs_shape, act_shape, obs_dtype, act_dtype)
         if venv is not None:
-            if not all([x is None for x in params]):
+            if not all(x is None for x in params):
                 raise ValueError(
                     "Cannot specify both shape/dtype and also environment.",
                 )
@@ -271,7 +282,7 @@ def __init__(
             obs_dtype = venv.observation_space.dtype
             act_dtype = venv.action_space.dtype
         else:
-            if any([x is None for x in params]):
+            if any(x is None for x in params):
                 raise ValueError("Shape or dtype missing and no environment specified.")
 
         assert obs_shape is not None

From 19227c2e327e276f1bfa40d23ccb0cfdcf76ddb8 Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Tue, 6 Sep 2022 16:26:09 -0700
Subject: [PATCH 058/187] Fix type annotation of step_dict

---
 src/imitation/data/rollout.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index 5d070d855..1ad41d3b1 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -3,7 +3,17 @@
 import collections
 import dataclasses
 import logging
-from typing import Callable, Dict, Hashable, List, Mapping, Optional, Sequence, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Hashable,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Union,
+)
 
 import numpy as np
 from stable_baselines3.common.base_class import BaseAlgorithm
@@ -57,7 +67,7 @@ def __init__(self):
 
     def add_step(
         self,
-        step_dict: Mapping[str, np.ndarray],
+        step_dict: Mapping[str, Union[np.ndarray, Mapping[str, Any]]],
         key: Hashable = None,
     ) -> None:
         """Add a single step to the partial trajectory identified by `key`.
@@ -130,7 +140,7 @@ def add_steps_and_auto_finish(
             A list of completed trajectories. There should be one trajectory for
             each `True` in the `dones` argument.
         """
-        trajs = []
+        trajs: List[types.TrajectoryWithRew] = []
         for env_idx in range(len(obs)):
             assert env_idx in self.partial_trajectories
             assert list(self.partial_trajectories[env_idx][0].keys()) == ["obs"], (

From a601567628900412436850300acde0a6e175d8ca Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Tue, 6 Sep 2022 19:09:28 -0700
Subject: [PATCH 059/187] Change List to Sequence

---
 src/imitation/scripts/train_imitation.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index e206f6b9a..020ce9cd8 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -3,7 +3,7 @@
 import logging
 import os.path as osp
 import warnings
-from typing import Any, List, Mapping, Optional, Sequence, Type, cast
+from typing import Any, Mapping, Optional, Sequence, Type, cast
 
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import policies, utils, vec_env
@@ -177,12 +177,12 @@ def train_imitation(
     #  using assert doesn't work because we'd have to loop over all the trajectories
     #  and check that each one is a TrajectoryWithRew, which seems inefficient
     #  just for adding type annotations.
-    trajectories: List[types.TrajectoryWithRew]
+    trajectories: Sequence[types.TrajectoryWithRew]
     if use_dagger:
-        trajectories = cast(List[types.TrajectoryWithRew], model._all_demos)
+        trajectories = cast(Sequence[types.TrajectoryWithRew], model._all_demos)
     else:
         assert expert_trajs is not None
-        trajectories = cast(List[types.TrajectoryWithRew], expert_trajs)
+        trajectories = cast(Sequence[types.TrajectoryWithRew], expert_trajs)
 
     return {
         "imit_stats": imit_stats,

From b7d4364fdf5cdb63e4e2aa40496bbb24e8d822f2 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 10:31:45 +0200
Subject: [PATCH 060/187] Fix density.py::DensityAlgorithm._set_demo_from_batch

---
 src/imitation/algorithms/density.py | 26 ++++++++++++--------------
 1 file changed, 12 insertions(+), 14 deletions(-)

diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index e7e9464ad..f11046512 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -133,15 +133,16 @@ def _set_demo_from_batch(
         act_b: np.ndarray,
         next_obs_b: Optional[np.ndarray],
     ) -> None:
-        next_obs_b = next_obs_b or itertools.repeat(None)
-        # TODO(juan) I think this assumes that if next_obs_b is an array,
-        #  it has at least two axes, and zip maps along the first one.
-        #  This is not checked for, and if the array has only one dimension
-        #  it would raise an obscure error within _preprocess_transition.
-        #  _preprocess_transition also requires next_obs to be an ndarray
-        #  that is part of the observation space, so not sure how
-        #  this isn't failing when next_obs_b is None.
-        for obs, act, next_obs in zip(obs_b, act_b, next_obs_b):
+        if next_obs_b is None and self.density_type == DensityType.STATE_STATE_DENSITY:
+            raise ValueError(
+                "STATE_STATE_DENSITY requires next_obs_b "
+                "to be provided, but it was None",
+            )
+        if next_obs_b is not None:
+            assert next_obs_b.shape[1:] == self.venv.observation_space.shape
+        
+        next_obs_b_iterator = next_obs_b or itertools.repeat(None)
+        for obs, act, next_obs in zip(obs_b, act_b, next_obs_b_iterator):
             flat_trans = self._preprocess_transition(obs, act, next_obs)
             self.transitions.setdefault(None, []).append(flat_trans)
 
@@ -220,7 +221,7 @@ def _preprocess_transition(
         self,
         obs: np.ndarray,
         act: np.ndarray,
-        next_obs: np.ndarray,
+        next_obs: Optional[np.ndarray],
     ) -> np.ndarray:
         """Compute flattened transition on subset specified by `self.density_type`."""
         if self.density_type == DensityType.STATE_DENSITY:
@@ -233,6 +234,7 @@ def _preprocess_transition(
                 ],
             )
         elif self.density_type == DensityType.STATE_STATE_DENSITY:
+            assert next_obs is not None
             return np.concatenate(
                 [
                     flatten(self.venv.observation_space, obs),
@@ -242,10 +244,6 @@ def _preprocess_transition(
         else:
             raise ValueError(f"Unknown density type {self.density_type}")
 
-    # TODO(juan) I opted for renaming the function signature to match the
-    #  rewards.reward_function.RewardFn protocol, but changing the protocol
-    #  or making the arguments positional-only are two other valid approaches
-    #  if the previous names are better suited.
     def __call__(
         self,
         state: np.ndarray,

From 3ca3bf266d33eaaed79fe1a5e1f1742f057901db Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 11:11:47 +0200
Subject: [PATCH 061/187] Fixed n_steps (OnPolicyAlgorithm)

---
 src/imitation/algorithms/adversarial/common.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
index c1e638f15..7e01ab2fa 100644
--- a/src/imitation/algorithms/adversarial/common.py
+++ b/src/imitation/algorithms/adversarial/common.py
@@ -20,7 +20,7 @@
 import torch as th
 import torch.utils.tensorboard as thboard
 import tqdm
-from stable_baselines3.common import on_policy_algorithm, policies, vec_env
+from stable_baselines3.common import on_policy_algorithm, policies, vec_env, base_class
 from stable_baselines3.sac import policies as sac_policies
 from torch.nn import functional as F
 
@@ -124,7 +124,7 @@ def __init__(
         demonstrations: base.AnyTransitions,
         demo_batch_size: int,
         venv: vec_env.VecEnv,
-        gen_algo: on_policy_algorithm.OnPolicyAlgorithm,
+        gen_algo: base_class.BaseAlgorithm,
         reward_net: reward_nets.RewardNet,
         n_disc_updates_per_round: int = 2,
         log_dir: str = "output/",
@@ -241,9 +241,7 @@ def __init__(
             gen_algo_env = self.gen_algo.get_env()
             assert gen_algo_env is not None
             self.gen_train_timesteps = gen_algo_env.num_envs
-            if hasattr(self.gen_algo, "n_steps"):  # on policy
-                # TODO(juan) this looks like a bug; could not find n_steps
-                #  defined anywhere in the codebase.
+            if isinstance(self.gen_algo, on_policy_algorithm.OnPolicyAlgorithm):
                 self.gen_train_timesteps *= self.gen_algo.n_steps
         else:
             self.gen_train_timesteps = gen_train_timesteps

From bda08cb9af6c6434d6972eda5ab5af4a5c91c3b5 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 11:55:52 +0200
Subject: [PATCH 062/187] Fix errors in tests

---
 tests/algorithms/test_mce_irl.py                | 8 ++++----
 tests/algorithms/test_preference_comparisons.py | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/algorithms/test_mce_irl.py b/tests/algorithms/test_mce_irl.py
index 7bb9b7560..26d79fa16 100644
--- a/tests/algorithms/test_mce_irl.py
+++ b/tests/algorithms/test_mce_irl.py
@@ -251,22 +251,22 @@ def test_tabular_policy():
     states = np.array([0, 1, 1, 0, 1])
     actions, timesteps = tabular.predict(states)
     np.testing.assert_array_equal(states, actions)
-    np.testing.assert_equal(timesteps, 1)
+    np.testing.assert_equal(timesteps[0], 1)
 
     mask = np.zeros((5,), dtype=bool)
     actions, timesteps = tabular.predict(states, timesteps, mask)
     np.testing.assert_array_equal(1 - states, actions)
-    np.testing.assert_equal(timesteps, 2)
+    np.testing.assert_equal(timesteps[0], 2)
 
     mask = np.ones((5,), dtype=bool)
     actions, timesteps = tabular.predict(states, timesteps, mask)
     np.testing.assert_array_equal(states, actions)
-    np.testing.assert_equal(timesteps, 1)
+    np.testing.assert_equal(timesteps[0], 1)
 
     mask = (1 - states).astype(bool)
     actions, timesteps = tabular.predict(states, timesteps, mask)
     np.testing.assert_array_equal(np.zeros((5,)), actions)
-    np.testing.assert_equal(timesteps, 2 - mask.astype(int))
+    np.testing.assert_equal(timesteps[0], 2 - mask.astype(int))
 
 
 def test_tabular_policy_randomness():
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index ca85ce114..25c346de3 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -524,7 +524,7 @@ def preference_model(venv) -> preference_comparisons.PreferenceModel:
 def test_probability_model_raises_error_when_ensemble_member_index_not_provided(
     ensemble_preference_model,
 ):
-    assert ensemble_preference_model.is_ensemble
+    assert ensemble_preference_model.ensemble_model is not None
     with pytest.raises(
         ValueError,
         match="`ensemble_member_index` required for ensemble models",

From 276dfaa29869862c0fc7fb3896478f5a5b612d48 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 11:56:11 +0200
Subject: [PATCH 063/187] Include some suggestions into rollout.py and
 preference_comparisons.py

---
 src/imitation/algorithms/preference_comparisons.py | 8 +++-----
 src/imitation/data/rollout.py                      | 8 ++++----
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 6d5560d31..d48494e60 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -437,10 +437,8 @@ def forward(
             rews2 = self.rewards(trans2)
             probs[i] = self.probability(rews1, rews2)
             if gt_reward_available:
-                frag1, frag2 = cast(TrajectoryWithRew, frag1), cast(
-                    TrajectoryWithRew,
-                    frag2,
-                )
+                frag1 = cast(TrajectoryWithRew, frag1)
+                frag2 = cast(TrajectoryWithRew, frag2)
                 gt_rews_1 = th.from_numpy(frag1.rews)
                 gt_rews_2 = th.from_numpy(frag2.rews)
                 gt_probs[i] = self.probability(gt_rews_1, gt_rews_2)
@@ -691,7 +689,7 @@ def __init__(
             ValueError: Preference model not wrapped over an ensemble of networks.
         """
         super().__init__(custom_logger=custom_logger)
-        if not preference_model.is_ensemble:
+        if preference_model.ensemble_model is None:
             raise ValueError(
                 "Preference model not wrapped over an ensemble of networks.",
             )
diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index 1ad41d3b1..ab790e0dc 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -103,11 +103,11 @@ def finish_trajectory(
         del self.partial_trajectories[key]
         out_dict_unstacked = collections.defaultdict(list)
         for part_dict in part_dicts:
-            for key_, array in part_dict.items():
-                out_dict_unstacked[key_].append(array)
+            for k, array in part_dict.items():
+                out_dict_unstacked[k].append(array)
         out_dict_stacked = {
-            key_: np.stack(arr_list, axis=0)
-            for key_, arr_list in out_dict_unstacked.items()
+            k: np.stack(arr_list, axis=0)
+            for k, arr_list in out_dict_unstacked.items()
         }
         traj = types.TrajectoryWithRew(**out_dict_stacked, terminal=terminal)
         assert traj.rews.shape[0] == traj.acts.shape[0] == traj.obs.shape[0] - 1

From 3171a8b9704a5a7b778de6c0e03d8478ac642f01 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 11:57:18 +0200
Subject: [PATCH 064/187] Formatting

---
 src/imitation/algorithms/adversarial/common.py | 2 +-
 src/imitation/algorithms/density.py            | 2 +-
 src/imitation/data/rollout.py                  | 3 +--
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
index 7e01ab2fa..41c1b4129 100644
--- a/src/imitation/algorithms/adversarial/common.py
+++ b/src/imitation/algorithms/adversarial/common.py
@@ -20,7 +20,7 @@
 import torch as th
 import torch.utils.tensorboard as thboard
 import tqdm
-from stable_baselines3.common import on_policy_algorithm, policies, vec_env, base_class
+from stable_baselines3.common import base_class, on_policy_algorithm, policies, vec_env
 from stable_baselines3.sac import policies as sac_policies
 from torch.nn import functional as F
 
diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index f11046512..8dda39450 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -140,7 +140,7 @@ def _set_demo_from_batch(
             )
         if next_obs_b is not None:
             assert next_obs_b.shape[1:] == self.venv.observation_space.shape
-        
+
         next_obs_b_iterator = next_obs_b or itertools.repeat(None)
         for obs, act, next_obs in zip(obs_b, act_b, next_obs_b_iterator):
             flat_trans = self._preprocess_transition(obs, act, next_obs)
diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index ab790e0dc..03cb1f3d2 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -106,8 +106,7 @@ def finish_trajectory(
             for k, array in part_dict.items():
                 out_dict_unstacked[k].append(array)
         out_dict_stacked = {
-            k: np.stack(arr_list, axis=0)
-            for k, arr_list in out_dict_unstacked.items()
+            k: np.stack(arr_list, axis=0) for k, arr_list in out_dict_unstacked.items()
         }
         traj = types.TrajectoryWithRew(**out_dict_stacked, terminal=terminal)
         assert traj.rews.shape[0] == traj.acts.shape[0] == traj.obs.shape[0] - 1

From 479b2f154dd49870026ec95cc1298d0dd108cd14 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 12:27:08 +0200
Subject: [PATCH 065/187] Fix setter error as per
 https://github.com/python/mypy/issues/5936

---
 src/imitation/algorithms/preference_comparisons.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index d48494e60..88c9edabd 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -302,8 +302,11 @@ def sample(self, steps: int) -> Sequence[types.TrajectoryWithRew]:
             trajectories.extend(list(exploration_trajs))
         return trajectories
 
-    # Type ignore due to https://github.com/python/mypy/issues/5936
-    @TrajectoryGenerator.logger.setter  # type: ignore[attr-defined]
+    @property
+    def logger(self):
+        return super().logger
+
+    @logger.setter
     def logger(self, value: imit_logger.HierarchicalLogger):
         self._logger = value
         self.algorithm.set_logger(self.logger)

From a52c4a3b12ab846df9f2ac56981dd244821b2387 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 12:38:26 +0200
Subject: [PATCH 066/187] add reason for assertion.

---
 src/imitation/data/rollout.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index 03cb1f3d2..fdf84ae58 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -373,7 +373,7 @@ def generate_trajectories(
     #
     # To start with, all environments are active.
     active = np.ones(venv.num_envs, dtype=bool)
-    assert isinstance(obs, np.ndarray)
+    assert isinstance(obs, np.ndarray), "Dict/tuple observations are not supported."
     while np.any(active):
         acts = get_actions(obs)
         obs, rews, dones, infos = venv.step(acts)

From fb72cf6d5b02611ff51505880795482949abb824 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 12:40:14 +0200
Subject: [PATCH 067/187] Fix style guide violation:
 https://google.github.io/styleguide/pyguide.html#22-imports

---
 src/imitation/rewards/reward_nets.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index b607af1ea..c71dffc4e 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -10,7 +10,6 @@
 from torch import nn
 
 from imitation.util import networks, util
-from imitation.util.networks import BaseNorm
 
 
 class RewardNet(nn.Module, abc.ABC):
@@ -425,7 +424,7 @@ class NormalizedRewardNet(RewardNetWrapper):
     def __init__(
         self,
         base: RewardNet,
-        normalize_output_layer: Type[BaseNorm],
+        normalize_output_layer: Type[networks.BaseNorm],
     ):
         """Initialize the NormalizedRewardNet.
 

From 1cbf645c61db7bf666cff25cb5f3b38454d25e58 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 13:17:05 +0200
Subject: [PATCH 068/187] Update src/imitation/scripts/parallel.py

Co-authored-by: Adam Gleave <adam@gleave.me>
---
 src/imitation/scripts/parallel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/scripts/parallel.py b/src/imitation/scripts/parallel.py
index 6e938d43f..a1702d2b1 100644
--- a/src/imitation/scripts/parallel.py
+++ b/src/imitation/scripts/parallel.py
@@ -192,7 +192,7 @@ def inner(config: Mapping[str, Any], reporter) -> Mapping[str, Any]:
         ex.observers = [FileStorageObserver("sacred")]
 
         # Apply base configs to get modified `named_configs` and `config_updates`.
-        named_configs = [*base_named_configs, *run_kwargs["named_configs"]]
+        named_configs = base_named_configs + run_kwargs["named_configs"]
         updated_run_kwargs["named_configs"] = named_configs
 
         config_updates = {**base_config_updates, **run_kwargs["config_updates"]}

From f7443cb60a4de84f199c5a049d0608d4093f9668 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 13:18:52 +0200
Subject: [PATCH 069/187] Move kwargs to the end.

---
 src/imitation/rewards/reward_nets.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index c71dffc4e..a260fe10a 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -82,6 +82,8 @@ def preprocess(
         del state, action, next_state, done  # unused
 
         # preprocess
+        # we only support array spaces, so we cast
+        # the observation to torch tensors.
         state_th = cast(
             th.Tensor,
             preprocessing.preprocess_obs(
@@ -393,10 +395,10 @@ def __init__(
         }
         self.mlp = networks.build_mlp(
             hid_sizes=(32, 32),
-            **kwargs,
             in_size=combined_size,
             out_size=1,
             squeeze_output=True,
+            **kwargs,
         )
 
     def forward(self, state, action, next_state, done):

From 8dd12a75a67341041e07ecd75347fab1f96fb128 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 13:21:00 +0200
Subject: [PATCH 070/187] Swap order of expert_policy_type and
 expert_policy_path validation check

---
 src/imitation/scripts/train_imitation.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index e206f6b9a..a5047a3af 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -87,12 +87,12 @@ def load_expert_policy(
         ValueError: `expert_policy_path` is None.
         TypeError: The policy loaded from `expert_policy_path` is not a SB3 policy.
     """
-    if expert_policy_path is None:
-        raise ValueError("expert_policy_path cannot be None")
-
     if expert_policy_type is None:
         raise ValueError("expert_policy_type cannot be None")
 
+    if expert_policy_path is None:
+        raise ValueError("expert_policy_path cannot be None")
+
     # TODO(shwang): Add support for directly loading a BasePolicy `*.th` file.
     expert_policy = serialize.load_policy(expert_policy_type, expert_policy_path, venv)
     if not isinstance(expert_policy, policies.BasePolicy):

From 2bdc4b306feb82a02ea7fcadeb5bc351385972d7 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 13:32:05 +0200
Subject: [PATCH 071/187] Update src/imitation/util/util.py

Co-authored-by: Adam Gleave <adam@gleave.me>
---
 src/imitation/util/util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 351e3ac9f..2421dbc80 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -100,7 +100,7 @@ def make_vec_env(
     spec = gym.spec(env_name)
     env_make_kwargs = env_make_kwargs or {}
 
-    def make_env(i, this_seed) -> gym.Env:
+    def make_env(i: int, this_seed: int) -> gym.Env:
         # Previously, we directly called `gym.make(env_name)`, but running
         # `imitation.scripts.train_adversarial` within `imitation.scripts.parallel`
         # created a weird interaction between Gym and Ray -- `gym.make` would fail

From 8e46bae559fe6ec608022e229f0171f0d76cb2f3 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 7 Sep 2022 13:35:58 +0200
Subject: [PATCH 072/187] Update tests/rewards/test_reward_fn.py

Co-authored-by: Adam Gleave <adam@gleave.me>
---
 tests/rewards/test_reward_fn.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/rewards/test_reward_fn.py b/tests/rewards/test_reward_fn.py
index 98a3f130c..7a0e42410 100644
--- a/tests/rewards/test_reward_fn.py
+++ b/tests/rewards/test_reward_fn.py
@@ -7,7 +7,7 @@
 
 OBS = np.random.randint(0, 10, (64, 100))
 ACTS = NEXT_OBS = OBS
-DONES = np.zeros(64, dtype=np.bool)  # type: ignore
+DONES = np.zeros(64, dtype=np.bool_)
 
 
 def _funky_reward_fn(obs, act, next_obs, done):

From 3dcfae7b1469301256e174ab103ca2f6aa4586d1 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 00:05:49 +0200
Subject: [PATCH 073/187] Explicit random state setting and fix corresponding
 tests (except notebooks, sacred config, scripts)

---
 examples/quickstart.py                        |  4 +
 src/imitation/algorithms/bc.py                | 13 ++-
 src/imitation/algorithms/dagger.py            |  2 +-
 src/imitation/algorithms/density.py           |  4 +
 src/imitation/algorithms/mce_irl.py           | 16 +--
 .../algorithms/preference_comparisons.py      | 98 +++++++++++++------
 src/imitation/data/rollout.py                 | 15 ++-
 src/imitation/policies/exploration_wrapper.py |  9 +-
 src/imitation/scripts/common/common.py        |  7 +-
 .../scripts/train_preference_comparisons.py   | 20 ++--
 src/imitation/util/util.py                    | 39 +++++++-
 tests/algorithms/test_adversarial.py          | 44 +++++++--
 tests/algorithms/test_bc.py                   | 10 +-
 tests/algorithms/test_dagger.py               | 25 ++++-
 tests/algorithms/test_density_baselines.py    |  7 ++
 tests/algorithms/test_mce_irl.py              | 22 ++++-
 .../algorithms/test_preference_comparisons.py | 80 ++++++++++-----
 tests/conftest.py                             | 19 ++++
 tests/data/test_buffer.py                     |  2 +-
 tests/data/test_rollout.py                    | 42 +++++---
 tests/policies/test_exploration_wrapper.py    |  3 +-
 tests/policies/test_policies.py               | 38 ++++---
 tests/policies/test_replay_buffer_wrapper.py  | 12 ++-
 tests/rewards/test_reward_nets.py             | 51 +++++++---
 tests/rewards/test_reward_wrapper.py          |  9 +-
 tests/scripts/test_scripts.py                 |  7 +-
 26 files changed, 445 insertions(+), 153 deletions(-)

diff --git a/examples/quickstart.py b/examples/quickstart.py
index 98f68c53b..b255f971e 100644
--- a/examples/quickstart.py
+++ b/examples/quickstart.py
@@ -4,6 +4,7 @@
 """
 
 import gym
+import numpy as np
 from stable_baselines3 import PPO
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.vec_env import DummyVecEnv
@@ -14,6 +15,7 @@
 from imitation.data.wrappers import RolloutInfoWrapper
 
 env = gym.make("CartPole-v1")
+random_state = np.random.RandomState(0)
 
 
 def train_expert():
@@ -40,6 +42,7 @@ def sample_expert_transitions():
         expert,
         DummyVecEnv([lambda: RolloutInfoWrapper(env)]),
         rollout.make_sample_until(min_timesteps=None, min_episodes=50),
+        random_state=random_state,
     )
     return rollout.flatten_trajectories(rollouts)
 
@@ -49,6 +52,7 @@ def sample_expert_transitions():
     observation_space=env.observation_space,
     action_space=env.action_space,
     demonstrations=transitions,
+    random_state=random_state,
 )
 
 reward, _ = evaluate_policy(bc_trainer.policy, env, n_eval_episodes=3, render=True)
diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index 895c0cb55..434611365 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -185,12 +185,15 @@ class RolloutStatsComputer:
     #   EvalCallback could be a good fit:
     #   https://stable-baselines3.readthedocs.io/en/master/guide/callbacks.html#evalcallback
 
-    def __call__(self, policy: policies.ActorCriticPolicy) -> Mapping[str, float]:
+    def __call__(
+        self, policy: policies.ActorCriticPolicy, random_state: np.random.RandomState
+    ) -> Mapping[str, float]:
         if self.venv is not None and self.n_episodes > 0:
             trajs = rollout.generate_trajectories(
                 policy,
                 self.venv,
                 rollout.make_min_episodes(self.n_episodes),
+                random_state=random_state,
             )
             return rollout.rollout_stats(trajs)
         else:
@@ -272,6 +275,7 @@ def __init__(
         *,
         observation_space: gym.Space,
         action_space: gym.Space,
+        random_state: np.random.RandomState,
         policy: Optional[policies.ActorCriticPolicy] = None,
         demonstrations: Optional[algo_base.AnyTransitions] = None,
         batch_size: int = 32,
@@ -317,6 +321,8 @@ def __init__(
         self.action_space = action_space
         self.observation_space = observation_space
 
+        self.random_state = random_state
+
         if policy is None:
             policy = policy_base.FeedForward32Policy(
                 observation_space=observation_space,
@@ -401,6 +407,9 @@ def train(
             self._bc_logger.reset_tensorboard_steps()
         self._bc_logger.log_epoch(0)
 
+        # TODO(juan) docstrings above say that this can be none and that no rollouts
+        #  are generated. However initializing this requires passing a non-None
+        #  venv.
         compute_rollout_stats = RolloutStatsComputer(
             log_rollouts_venv,
             log_rollouts_n_episodes,
@@ -438,7 +447,7 @@ def _on_epoch_end(epoch_number: int):
             loss = self.trainer(batch)
 
             if batch_num % log_interval == 0:
-                rollout_stats = compute_rollout_stats(self.policy)
+                rollout_stats = compute_rollout_stats(self.policy, self.random_state)
 
                 self._bc_logger.log_batch(
                     batch_num,
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index 4ad0b6ebc..56b5ca8f9 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -631,7 +631,7 @@ def train(
                 venv=collector,
                 sample_until=sample_until,
                 deterministic_policy=True,
-                rng=collector.rng,
+                random_state=collector.rng,
             )
 
             for traj in trajectories:
diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index 8dda39450..492658331 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -59,6 +59,7 @@ def __init__(
         *,
         demonstrations: Optional[base.AnyTransitions],
         venv: vec_env.VecEnv,
+        random_state: np.random.RandomState,
         density_type: DensityType = DensityType.STATE_ACTION_DENSITY,
         kernel: str = "gaussian",
         kernel_bandwidth: float = 0.5,
@@ -82,6 +83,7 @@ def __init__(
                 any environment interaction to fit the reward model, but we use this
                 to extract the observation and action space, and to train the RL
                 algorithm `rl_algo` (if specified).
+            random_state: random state for sampling from demonstrations.
             rl_algo: An RL algorithm to train on the resulting reward model (optional).
             is_stationary: if True, share same density models for all timesteps;
                 if False, use a different density model for each timestep.
@@ -118,6 +120,7 @@ def __init__(
         self.standardise = standardise_inputs
         self._scaler = None
         self._density_models = dict()
+        self.random_state = random_state
 
         self.rl_algo = rl_algo
         self.buffering_wrapper = wrappers.BufferingWrapper(self.venv)
@@ -344,6 +347,7 @@ def test_policy(self, *, n_trajectories: int = 10, true_reward: bool = True):
             self.rl_algo,
             self.venv if true_reward else self.venv_wrapped,
             sample_until=rollout.make_min_episodes(n_trajectories),
+            random_state=self.random_state,
         )
         # We collect `trajs` above so disregard return value from `pop_trajectories`,
         # but still call it to clear out any saved trajectories.
diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index df2079471..8d3d4bbe5 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -149,7 +149,7 @@ def __init__(
         state_space: gym.Space,
         action_space: gym.Space,
         pi: np.ndarray,
-        rng: Optional[np.random.RandomState],
+        random_state: np.random.RandomState,
     ):
         """Builds TabularPolicy.
 
@@ -165,7 +165,7 @@ def __init__(
         assert isinstance(action_space, gym.spaces.Discrete), "action not tabular"
         # What we call state space here is observation space in SB3 nomenclature.
         super().__init__(observation_space=state_space, action_space=action_space)
-        self.rng = rng or np.random.RandomState()
+        self.rng = random_state
         self.set_pi(pi)
 
     def set_pi(self, pi: np.ndarray) -> None:
@@ -259,6 +259,7 @@ def __init__(
         demonstrations: Optional[MCEDemonstrations],
         env: resettable_env.TabularModelEnv,
         reward_net: reward_nets.RewardNet,
+        random_state: np.random.RandomState,
         optimizer_cls: Type[th.optim.Optimizer] = th.optim.Adam,
         optimizer_kwargs: Optional[Mapping[str, Any]] = None,
         discount: float = 1.0,
@@ -267,7 +268,6 @@ def __init__(
         # TODO(adam): do we need log_interval or can just use record_mean...?
         log_interval: Optional[int] = 100,
         *,
-        rng: Optional[np.random.RandomState] = None,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         r"""Creates MCE IRL.
@@ -279,10 +279,11 @@ def __init__(
                 The demonstrations must have observations one-hot coded unless
                 demonstrations is a state-occupancy measure.
             env: a tabular MDP.
+            random_state: random state used for sampling from policy.
             reward_net: a neural network that computes rewards for the supplied
                 observations.
-            optimizer_cls: optimiser to use for supervised training.
-            optimizer_kwargs: keyword arguments for optimiser construction.
+            optimizer_cls: optimizer to use for supervised training.
+            optimizer_kwargs: keyword arguments for optimizer construction.
             discount: the discount factor to use when computing occupancy measure.
                 If not 1.0 (undiscounted), then `demonstrations` must either be
                 a (discounted) state-occupancy measure, or trajectories. Transitions
@@ -295,7 +296,6 @@ def __init__(
                 MCE IRL gradient falls below this value.
             log_interval: how often to log current loss stats (using `logging`).
                 None to disable.
-            rng: random state used for sampling from policy.
             custom_logger: Where to log to; if None (default), creates a new logger.
         """
         self.discount = discount
@@ -313,7 +313,7 @@ def __init__(
         self.linf_eps = linf_eps
         self.grad_l2_eps = grad_l2_eps
         self.log_interval = log_interval
-        self.rng = rng
+        self.rng = random_state
 
         # Initialize policy to be uniform random. We don't use this for MCE IRL
         # training, but it gives us something to return at all times with `policy`
@@ -324,7 +324,7 @@ def __init__(
             state_space=self.env.pomdp_state_space,
             action_space=self.env.action_space,
             pi=uniform_pi,
-            rng=self.rng,
+            random_state=self.rng,
         )
 
     def _set_demo_from_trajectories(self, trajs: Iterable[types.Trajectory]) -> None:
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 88c9edabd..a07181ad8 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -98,19 +98,19 @@ class TrajectoryDataset(TrajectoryGenerator):
     def __init__(
         self,
         trajectories: Sequence[TrajectoryWithRew],
-        seed: Optional[int] = None,
+        random_state: random.Random,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Creates a dataset loaded from `path`.
 
         Args:
             trajectories: the dataset of rollouts.
-            seed: Seed for RNG used for shuffling dataset.
+            random_state: Random state for RNG used for shuffling dataset.
             custom_logger: Where to log to; if None (default), creates a new logger.
         """
         super().__init__(custom_logger=custom_logger)
         self._trajectories = trajectories
-        self.rng = random.Random(seed)
+        self.rng = random_state
 
     def sample(self, steps: int) -> Sequence[TrajectoryWithRew]:
         # make a copy before shuffling
@@ -127,10 +127,10 @@ def __init__(
         algorithm: base_class.BaseAlgorithm,
         reward_fn: Union[reward_function.RewardFn, reward_nets.RewardNet],
         venv: vec_env.VecEnv,
+        random_state: np.random.RandomState,
         exploration_frac: float = 0.0,
         switch_prob: float = 0.5,
         random_prob: float = 0.5,
-        seed: Optional[int] = None,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the agent trainer.
@@ -140,13 +140,13 @@ def __init__(
             reward_fn: either a RewardFn or a RewardNet instance that will supply
                 the rewards used for training the agent.
             venv: vectorized environment to train in.
+            random_state: random state used for exploration and for sampling.
             exploration_frac: fraction of the trajectories that will be generated
                 partially randomly rather than only by the agent when sampling.
             switch_prob: the probability of switching the current policy at each
                 step for the exploratory samples.
             random_prob: the probability of picking the random policy when switching
                 during exploration.
-            seed: random seed for exploratory trajectories.
             custom_logger: Where to log to; if None (default), creates a new logger.
         """
         self.algorithm = algorithm
@@ -162,6 +162,7 @@ def __init__(
             reward_fn = reward_fn.predict_processed
         self.reward_fn = reward_fn
         self.exploration_frac = exploration_frac
+        self.random_state = random_state
 
         # The BufferingWrapper records all trajectories, so we can return
         # them after training. This should come first (before the wrapper that
@@ -199,7 +200,7 @@ def __init__(
             venv=algo_venv,
             random_prob=random_prob,
             switch_prob=switch_prob,
-            seed=seed,
+            random_state=self.random_state,
         )
 
     def train(self, steps: int, **kwargs) -> None:
@@ -269,6 +270,7 @@ def sample(self, steps: int) -> Sequence[types.TrajectoryWithRew]:
                 # deterministic. If self.algorithm is stochastic, then policy_callable
                 # will also be stochastic.
                 deterministic_policy=False,
+                random_state=self.random_state,
             )
             additional_trajs, _ = self.buffering_wrapper.pop_finished_trajectories()
             agent_trajs = list(agent_trajs) + list(additional_trajs)
@@ -292,6 +294,7 @@ def sample(self, steps: int) -> Sequence[types.TrajectoryWithRew]:
                 # buffering_wrapper collects rollouts from a non-deterministic policy,
                 # so we do that here as well for consistency.
                 deterministic_policy=False,
+                random_state=self.random_state,
             )
             exploration_trajs, _ = self.buffering_wrapper.pop_finished_trajectories()
             exploration_trajs = _get_trajectories(exploration_trajs, exploration_steps)
@@ -574,21 +577,22 @@ class RandomFragmenter(Fragmenter):
 
     def __init__(
         self,
-        seed: Optional[float] = None,
+        random_state: random.Random,
         warning_threshold: int = 10,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the fragmenter.
 
         Args:
-            seed: an optional seed for the internal RNG
+            random_state: the random state for the random number generator
             warning_threshold: give a warning if the number of available
                 transitions is less than this many times the number of
                 required samples. Set to 0 to disable this warning.
             custom_logger: Where to log to; if None (default), creates a new logger.
         """
         super().__init__(custom_logger)
-        self.rng = random.Random(seed)
+        assert isinstance(random_state, random.Random)
+        self.rng = random_state
         self.warning_threshold = warning_threshold
 
     def __call__(
@@ -778,7 +782,7 @@ class PreferenceGatherer(abc.ABC):
 
     def __init__(
         self,
-        seed: Optional[int] = None,
+        random_state: Optional[np.random.RandomState] = None,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initializes the preference gatherer.
@@ -791,7 +795,7 @@ def __init__(
         # as an argument nevertheless because that means we can always
         # pass in a seed in training scripts (without worrying about whether
         # the PreferenceGatherer we use needs one).
-        del seed
+        del random_state
         self.logger = custom_logger or imit_logger.configure()
 
     @abc.abstractmethod
@@ -821,7 +825,7 @@ def __init__(
         temperature: float = 1,
         discount_factor: float = 1,
         sample: bool = True,
-        seed: Optional[int] = None,
+        random_state: Optional[np.random.RandomState] = None,
         threshold: float = 50,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
@@ -838,7 +842,8 @@ def __init__(
                 a Bernoulli distribution (or 0.5 in the case of ties with zero
                 temperature). If False, then the underlying Bernoulli probabilities
                 are returned instead.
-            seed: seed for the internal RNG (only used if temperature > 0 and sample)
+            random_state: random state for the internal RNG
+                (only used if temperature > 0 and sample)
             threshold: preferences are sampled from a softmax of returns.
                 To avoid overflows, we clip differences in returns that are
                 above this threshold (after multiplying with temperature).
@@ -850,9 +855,12 @@ def __init__(
         self.temperature = temperature
         self.discount_factor = discount_factor
         self.sample = sample
-        self.rng = np.random.default_rng(seed=seed)
+        self.rng = random_state
         self.threshold = threshold
 
+        if self.sample and self.rng is None:
+            raise ValueError("If sample is True, then random_state must be provided.")
+
     def __call__(self, fragment_pairs: Sequence[TrajectoryWithRewPair]) -> np.ndarray:
         """Computes probability fragment 1 is preferred over fragment 2."""
         returns1, returns2 = self._reward_sums(fragment_pairs)
@@ -1193,11 +1201,11 @@ def __init__(
         self,
         model: reward_nets.RewardEnsemble,
         loss: RewardLoss,
+        random_state: np.random.RandomState,
         batch_size: int = 32,
         epochs: int = 1,
         lr: float = 1e-3,
         weight_decay: float = 0.0,
-        seed: Optional[int] = None,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the reward model trainer.
@@ -1205,6 +1213,7 @@ def __init__(
         Args:
             model: the RewardNet instance to be trained
             loss: the loss to use
+            random_state: random state for the internal RNG used in bagging
             batch_size: number of fragment pairs per batch
             epochs: number of epochs in each training iteration (can be adjusted
                 on the fly by specifying an `epoch_multiplier` in `self.train()`
@@ -1213,7 +1222,6 @@ def __init__(
             weight_decay: the weight decay factor for the reward model's weights
                 to use with ``th.optim.AdamW``. This is similar to but not equivalent
                 to L2 regularization, see https://arxiv.org/abs/1711.05101
-            seed: seed for the internal RNG used in bagging
             custom_logger: Where to log to; if None (default), creates a new logger.
 
         Raises:
@@ -1233,7 +1241,7 @@ def __init__(
             weight_decay,
             custom_logger,
         )
-        self.rng = np.random.default_rng(seed=seed)
+        self.rng = random_state
 
     def _training_inner_loop(
         self,
@@ -1283,8 +1291,8 @@ def get_base_model(
 def _make_reward_trainer(
     reward_model: reward_nets.RewardNet,
     loss: RewardLoss,
+    random_state: np.random.RandomState,
     reward_trainer_kwargs: Optional[Mapping[str, Any]] = None,
-    seed: Optional[int] = None,
 ) -> RewardTrainer:
     """Construct the correct type of reward trainer for this reward function."""
     if reward_trainer_kwargs is None:
@@ -1302,7 +1310,9 @@ def _make_reward_trainer(
         )
 
         if is_base or is_std_wrapper:
-            return EnsembleTrainer(base_model, loss, seed=seed, **reward_trainer_kwargs)
+            return EnsembleTrainer(
+                base_model, loss, random_state=random_state, **reward_trainer_kwargs
+            )
         else:
             raise ValueError(
                 "RewardEnsemble can only be wrapped"
@@ -1341,7 +1351,7 @@ def __init__(
         initial_epoch_multiplier: float = 200.0,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
         allow_variable_horizon: bool = False,
-        seed: Optional[int] = None,
+        random_state: Optional[np.random.RandomState] = None,
         query_schedule: Union[str, type_aliases.Schedule] = "hyperbolic",
     ):
         """Initialize the preference comparison trainer.
@@ -1396,7 +1406,8 @@ def __init__(
                 condition, and can seriously confound evaluation. Read
                 https://imitation.readthedocs.io/en/latest/guide/variable_horizon.html
                 before overriding this.
-            seed: seed to use for initializing subcomponents such as fragmenter.
+            random_state: random state to use for initializing subcomponents such as
+                fragmenter.
                 Only used when default components are used; if you instantiate your
                 own fragmenter, preference gatherer, etc., you are responsible for
                 seeding them!
@@ -1423,11 +1434,29 @@ def __init__(
         self._iteration = 0
 
         self.model = reward_model
+        self.random_state = random_state
+
+        if self.random_state is None and None in (preference_gatherer, fragmenter):
+            raise ValueError(
+                "If you don't provide a random state, you must provide your own "
+                "seeded fragmenter and preference gatherer. You can initialize"
+                "a random state with `np.random.RandomState(seed)`."
+            )
+        elif self.random_state is not None and None not in (
+            preference_gatherer,
+            fragmenter,
+        ):
+            raise ValueError(
+                "If you provide your own fragmenter and preference gatherer, you "
+                "don't need to provide a random state."
+            )
 
         if reward_trainer is None:
             preference_model = PreferenceModel(reward_model)
             loss = CrossEntropyRewardLoss(preference_model)
-            self.reward_trainer = _make_reward_trainer(reward_model, loss, seed=seed)
+            self.reward_trainer = _make_reward_trainer(
+                reward_model, loss, random_state=self.random_state
+            )
         else:
             self.reward_trainer = reward_trainer
 
@@ -1437,15 +1466,24 @@ def __init__(
         self.reward_trainer.logger = self.logger
         self.trajectory_generator = trajectory_generator
         self.trajectory_generator.logger = self.logger
-        self.fragmenter = fragmenter or RandomFragmenter(
-            custom_logger=self.logger,
-            seed=seed,
-        )
+        if fragmenter:
+            self.fragmenter = fragmenter
+        else:
+            assert self.random_state is not None
+            self.fragmenter = RandomFragmenter(
+                custom_logger=self.logger,
+                random_state=random.Random(util.make_seeds(self.random_state)),
+            )
         self.fragmenter.logger = self.logger
-        self.preference_gatherer = preference_gatherer or SyntheticGatherer(
-            custom_logger=self.logger,
-            seed=seed,
-        )
+        if preference_gatherer:
+            self.preference_gatherer = preference_gatherer
+        else:
+            assert self.random_state is not None
+            self.preference_gatherer = SyntheticGatherer(
+                custom_logger=self.logger,
+                random_state=self.random_state,
+            )
+
         self.preference_gatherer.logger = self.logger
 
         self.fragment_length = fragment_length
diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index fdf84ae58..04e0876eb 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -323,9 +323,9 @@ def generate_trajectories(
     policy: AnyPolicy,
     venv: VecEnv,
     sample_until: GenTrajTerminationFn,
+    random_state: np.random.RandomState,
     *,
     deterministic_policy: bool = False,
-    rng: Optional[np.random.RandomState] = None,
 ) -> Sequence[types.TrajectoryWithRew]:
     """Generate trajectory dictionaries from a policy and an environment.
 
@@ -342,14 +342,13 @@ def generate_trajectories(
         deterministic_policy: If True, asks policy to deterministically return
             action. Note the trajectories might still be non-deterministic if the
             environment has non-determinism!
-        rng: used for shuffling trajectories.
+        random_state: used for shuffling trajectories.
 
     Returns:
         Sequence of trajectories, satisfying `sample_until`. Additional trajectories
         may be collected to avoid biasing process towards short episodes; the user
         should truncate if required.
     """
-    rng = rng or np.random.RandomState()
     get_actions = _policy_to_callable(policy, venv, deterministic_policy)
 
     # Collect rollout tuples.
@@ -406,7 +405,7 @@ def generate_trajectories(
     # `trajectories` sooner. Shuffle to avoid bias in order. This is important
     # when callees end up truncating the number of trajectories or transitions.
     # It is also cheap, since we're just shuffling pointers.
-    rng.shuffle(trajectories)  # type: ignore[arg-type]
+    random_state.shuffle(trajectories)  # type: ignore[arg-type]
 
     # Sanity checks.
     for trajectory in trajectories:
@@ -529,6 +528,7 @@ def generate_transitions(
     policy: AnyPolicy,
     venv: VecEnv,
     n_timesteps: int,
+    random_state: np.random.RandomState,
     *,
     truncate: bool = True,
     **kwargs,
@@ -556,6 +556,7 @@ def generate_transitions(
         policy,
         venv,
         sample_until=make_min_timesteps(n_timesteps),
+        random_state=random_state,
         **kwargs,
     )
     transitions = flatten_trajectories_with_rew(traj)
@@ -570,6 +571,7 @@ def rollout(
     policy: AnyPolicy,
     venv: VecEnv,
     sample_until: GenTrajTerminationFn,
+    random_state: np.random.RandomState,
     *,
     unwrap: bool = True,
     exclude_infos: bool = True,
@@ -588,6 +590,7 @@ def rollout(
             3) None, in which case actions will be sampled randomly.
         venv: The vectorized environments.
         sample_until: End condition for rollout sampling.
+        random_state: Random state to use for sampling.
         unwrap: If True, then save original observations and rewards (instead of
             potentially wrapped observations and rewards) by calling
             `unwrap_traj()`.
@@ -602,7 +605,9 @@ def rollout(
         may be collected to avoid biasing process towards short episodes; the user
         should truncate if required.
     """
-    trajs = generate_trajectories(policy, venv, sample_until, **kwargs)
+    trajs = generate_trajectories(
+        policy, venv, sample_until, random_state=random_state, **kwargs
+    )
     if unwrap:
         trajs = [unwrap_traj(traj) for traj in trajs]
     if exclude_infos:
diff --git a/src/imitation/policies/exploration_wrapper.py b/src/imitation/policies/exploration_wrapper.py
index 651476278..0ea5f8f9e 100644
--- a/src/imitation/policies/exploration_wrapper.py
+++ b/src/imitation/policies/exploration_wrapper.py
@@ -6,6 +6,7 @@
 from stable_baselines3.common import vec_env
 
 from imitation.data import rollout
+from imitation.util import util
 
 
 class ExplorationWrapper:
@@ -25,7 +26,7 @@ def __init__(
         venv: vec_env.VecEnv,
         random_prob: float,
         switch_prob: float,
-        seed: Optional[int] = None,
+        random_state: np.random.RandomState,
     ):
         """Initializes the ExplorationWrapper.
 
@@ -34,14 +35,16 @@ def __init__(
             venv: The environment to use (needed for sampling random actions).
             random_prob: The probability of picking the random policy when switching.
             switch_prob: The probability of switching away from the current policy.
-            seed: The random seed to use.
+            random_state: The random state to use for seeding the environment and for
+                switching policies.
         """
         self.wrapped_policy = policy_callable
         self.random_prob = random_prob
         self.switch_prob = switch_prob
         self.venv = venv
 
-        self.rng = np.random.RandomState(seed)
+        self.rng = random_state
+        seed = util.make_seeds(self.rng)
         self.venv.action_space.seed(seed)
 
         self.current_policy = policy_callable
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index ac5888817..b7c164eb3 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -5,6 +5,7 @@
 import os
 from typing import Any, Generator, Mapping, Sequence, Tuple, Union
 
+import numpy as np
 import sacred
 from stable_baselines3.common import vec_env
 
@@ -130,7 +131,7 @@ def setup_logging(
 @contextlib.contextmanager
 @common_ingredient.capture
 def make_venv(
-    _seed,
+    random_state: np.random.RandomState,
     env_name: str,
     num_vec: int,
     parallel: bool,
@@ -159,8 +160,8 @@ def make_venv(
     try:
         venv = util.make_vec_env(
             env_name,
-            num_vec,
-            seed=_seed,
+            random_state=random_state,
+            n_envs=num_vec,
             parallel=parallel,
             max_episode_steps=max_episode_steps,
             log_dir=log_dir,
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 6011ca7f4..2f2e89759 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -6,8 +6,10 @@
 
 import functools
 import os
+import random
 from typing import Any, Mapping, Optional, Type, Union
 
+import numpy as np
 import torch as th
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import type_aliases
@@ -21,6 +23,7 @@
 from imitation.scripts.config.train_preference_comparisons import (
     train_preference_comparisons_ex,
 )
+from imitation.util import util
 
 
 def save_model(
@@ -56,7 +59,7 @@ def save_checkpoint(
 
 @train_preference_comparisons_ex.main
 def train_preference_comparisons(
-    _seed: int,
+    random_state: np.random.RandomState,
     total_timesteps: int,
     total_comparisons: int,
     num_iterations: int,
@@ -84,7 +87,7 @@ def train_preference_comparisons(
     """Train a reward model using preference comparisons.
 
     Args:
-        _seed: Random seed.
+        random_state: Random state for the internal random number generation.
         total_timesteps: number of environment interaction steps
         total_comparisons: number of preferences to gather in total
         num_iterations: number of times to train the agent against the reward model
@@ -148,6 +151,7 @@ def train_preference_comparisons(
         ValueError: Inconsistency between config and deserialized policy normalization.
     """
     custom_logger, log_dir = common.setup_logging()
+    seed = util.make_seeds(random_state)
 
     with common.make_venv() as venv:
         reward_net = reward.make_reward_net(venv)
@@ -172,7 +176,7 @@ def train_preference_comparisons(
                 reward_fn=reward_net,
                 venv=venv,
                 exploration_frac=exploration_frac,
-                seed=_seed,
+                random_state=random_state,
                 custom_logger=custom_logger,
                 **trajectory_generator_kwargs,
             )
@@ -186,7 +190,7 @@ def train_preference_comparisons(
                 )
             trajectory_generator = preference_comparisons.TrajectoryDataset(
                 trajectories=types.load_with_rewards(trajectory_path),
-                seed=_seed,
+                random_state=random.Random(seed),
                 custom_logger=custom_logger,
                 **trajectory_generator_kwargs,
             )
@@ -194,7 +198,7 @@ def train_preference_comparisons(
         fragmenter: preference_comparisons.Fragmenter = (
             preference_comparisons.RandomFragmenter(
                 **fragmenter_kwargs,
-                seed=_seed,
+                random_state=random.Random(seed),
                 custom_logger=custom_logger,
             )
         )
@@ -212,7 +216,7 @@ def train_preference_comparisons(
             )
         gatherer = gatherer_cls(
             **gatherer_kwargs,
-            seed=_seed,
+            random_state=random_state,
             custom_logger=custom_logger,
         )
 
@@ -223,8 +227,8 @@ def train_preference_comparisons(
         reward_trainer = preference_comparisons._make_reward_trainer(
             reward_net,
             loss,
+            random_state,
             reward_trainer_kwargs,
-            seed=_seed,
         )
 
         main_trainer = preference_comparisons.PreferenceComparisons(
@@ -240,7 +244,7 @@ def train_preference_comparisons(
             initial_comparison_frac=initial_comparison_frac,
             custom_logger=custom_logger,
             allow_variable_horizon=allow_variable_horizon,
-            seed=_seed,
+            random_state=random_state,
             query_schedule=query_schedule,
         )
 
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 2421dbc80..e79139956 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -16,6 +16,7 @@
     Sequence,
     TypeVar,
     Union,
+    overload,
 )
 
 import gym
@@ -64,8 +65,8 @@ def make_unique_timestamp() -> str:
 
 def make_vec_env(
     env_name: str,
+    random_state: np.random.RandomState,
     n_envs: int = 8,
-    seed: int = 0,
     parallel: bool = False,
     log_dir: Optional[str] = None,
     max_episode_steps: Optional[int] = None,
@@ -76,6 +77,7 @@ def make_vec_env(
 
     Args:
         env_name: The Env's string id in Gym.
+        random_state: The random state to use to seed the environment.
         n_envs: The number of duplicate environments.
         seed: The environment seed.
         parallel: If True, uses SubprocVecEnv; otherwise, DummyVecEnv.
@@ -137,8 +139,7 @@ def make_env(i: int, this_seed: int) -> gym.Env:
 
         return env
 
-    rng = np.random.RandomState(seed)
-    env_seeds = rng.randint(0, (1 << 31) - 1, (n_envs,))
+    env_seeds = make_seeds(random_state, n_envs)
     env_fns: List[Callable[[], gym.Env]] = [
         functools.partial(make_env, i, s) for i, s in enumerate(env_seeds)
     ]
@@ -149,6 +150,38 @@ def make_env(i: int, this_seed: int) -> gym.Env:
         return DummyVecEnv(env_fns)
 
 
+@overload
+def make_seeds(
+    random_state: np.random.RandomState,
+) -> int:
+    ...
+
+
+@overload
+def make_seeds(random_state: np.random.RandomState, n: int) -> List[int]:
+    ...
+
+
+def make_seeds(
+    random_state: np.random.RandomState, n: Optional[int] = None
+) -> Union[List[int], int]:
+    """Generate n random seeds from a random state.
+
+    Args:
+        random_state: The random state to use to generate seeds.
+        n: The number of seeds to generate.
+
+    Returns:
+        A list of n random seeds.
+    """
+
+    seeds: List[int] = random_state.randint(0, (1 << 31) - 1, (n or 1,)).tolist()
+    if n is None:
+        return seeds[0]
+    else:
+        return seeds
+
+
 def docstring_parameter(*args, **kwargs):
     """Treats the docstring as a format string, substituting in the arguments."""
 
diff --git a/tests/algorithms/test_adversarial.py b/tests/algorithms/test_adversarial.py
index d8a5712f2..08f93855e 100644
--- a/tests/algorithms/test_adversarial.py
+++ b/tests/algorithms/test_adversarial.py
@@ -58,6 +58,7 @@ def make_trainer(
     algorithm_kwargs: Mapping[str, Any],
     tmpdir: str,
     expert_transitions: types.Transitions,
+    random_state: np.random.RandomState,
     expert_batch_size: int = 1,
     env_name: str = "seals/CartPole-v0",
     num_envs: int = 1,
@@ -76,7 +77,9 @@ def make_trainer(
     else:
         expert_data = expert_transitions
 
-    venv = util.make_vec_env(env_name, n_envs=num_envs, parallel=parallel)
+    venv = util.make_vec_env(
+        env_name, n_envs=num_envs, parallel=parallel, random_state=random_state
+    )
     model_cls = algorithm_kwargs["model_class"]
     gen_algo = model_cls(algorithm_kwargs["policy_class"], venv)
     reward_net_cls: Type[reward_nets.RewardNet] = reward_nets.BasicRewardNet
@@ -101,15 +104,19 @@ def make_trainer(
         venv.close()
 
 
-def test_airl_fail_fast(custom_logger, tmpdir):
+def test_airl_fail_fast(custom_logger, tmpdir, random_state_fixed):
+    random_state = random_state_fixed
     venv = util.make_vec_env(
         "seals/CartPole-v0",
         n_envs=1,
         parallel=False,
+        random_state=random_state,
     )
 
     gen_algo = stable_baselines3.DQN(stable_baselines3.dqn.MlpPolicy, venv)
-    small_data = rollout.generate_transitions(gen_algo, venv, n_timesteps=20)
+    small_data = rollout.generate_transitions(
+        gen_algo, venv, n_timesteps=20, random_state=random_state
+    )
     reward_net = reward_nets.BasicShapedRewardNet(
         observation_space=venv.observation_space,
         action_space=venv.action_space,
@@ -128,8 +135,11 @@ def test_airl_fail_fast(custom_logger, tmpdir):
 
 
 @pytest.fixture(params=ALGORITHM_KWARGS.values(), ids=list(ALGORITHM_KWARGS.keys()))
-def trainer(request, tmpdir, expert_transitions):
-    with make_trainer(request.param, tmpdir, expert_transitions) as trainer:
+def trainer(request, tmpdir, expert_transitions, random_state_fixed):
+    random_state = random_state_fixed
+    with make_trainer(
+        request.param, tmpdir, expert_transitions, random_state
+    ) as trainer:
         yield trainer
 
 
@@ -177,11 +187,14 @@ def trainer_parametrized(
     _expert_batch_size,
     tmpdir,
     expert_transitions,
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     with make_trainer(
         _algorithm_kwargs,
         tmpdir,
         expert_transitions,
+        random_state=random_state,
         parallel=_parallel,
         convert_dataset=_convert_dataset,
         expert_batch_size=_expert_batch_size,
@@ -189,12 +202,16 @@ def trainer_parametrized(
         yield trainer
 
 
-def test_train_disc_step_no_crash(trainer_parametrized, _expert_batch_size):
+def test_train_disc_step_no_crash(
+    trainer_parametrized, _expert_batch_size, random_state_fixed
+):
+    random_state = random_state_fixed
     transitions = rollout.generate_transitions(
         trainer_parametrized.gen_algo,
         trainer_parametrized.venv,
         n_timesteps=_expert_batch_size,
         truncate=True,
+        random_state=random_state,
     )
     trainer_parametrized.train_disc(
         gen_samples=types.dataclass_quick_asdict(transitions),
@@ -215,12 +232,15 @@ def trainer_batch_sizes(
     _expert_batch_size,
     tmpdir,
     expert_transitions,
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     with make_trainer(
         _algorithm_kwargs,
         tmpdir,
         expert_transitions,
         expert_batch_size=_expert_batch_size,
+        random_state=random_state,
     ) as trainer:
         yield trainer
 
@@ -230,8 +250,10 @@ def test_train_disc_improve_D(
     tmpdir,
     expert_transitions,
     _expert_batch_size,
+    random_state_fixed,
     n_steps=3,
 ):
+    random_state = random_state_fixed
     expert_samples = expert_transitions[:_expert_batch_size]
     expert_samples = types.dataclass_quick_asdict(expert_samples)
     gen_samples = rollout.generate_transitions(
@@ -239,6 +261,7 @@ def test_train_disc_improve_D(
         trainer_batch_sizes.venv_train,
         n_timesteps=_expert_batch_size,
         truncate=True,
+        random_state=random_state,
     )
     gen_samples = types.dataclass_quick_asdict(gen_samples)
     init_stats = final_stats = None
@@ -259,13 +282,17 @@ def _env_name(request):
 
 
 @pytest.fixture
-def trainer_diverse_env(_algorithm_kwargs, _env_name, tmpdir, expert_transitions):
+def trainer_diverse_env(
+    _algorithm_kwargs, _env_name, tmpdir, expert_transitions, random_state_fixed
+):
+    random_state = random_state_fixed
     if _algorithm_kwargs["model_class"] == stable_baselines3.DQN:
         pytest.skip("DQN does not support all environments.")
     with make_trainer(
         _algorithm_kwargs,
         tmpdir,
         expert_transitions,
+        random_state=random_state,
         env_name=_env_name,
     ) as trainer:
         yield trainer
@@ -275,6 +302,7 @@ def trainer_diverse_env(_algorithm_kwargs, _env_name, tmpdir, expert_transitions
 def test_logits_expert_is_high_log_policy_act_prob(
     trainer_diverse_env: common.AdversarialTrainer,
     n_timesteps: int,
+    random_state_fixed,
 ):
     """Smoke test calling `logits_expert_is_high` on `AdversarialTrainer`.
 
@@ -285,10 +313,12 @@ def test_logits_expert_is_high_log_policy_act_prob(
         trainer_diverse_env: The trainer to test.
         n_timesteps: The number of timesteps of rollouts to collect.
     """
+    random_state = random_state_fixed
     trans = rollout.generate_transitions(
         policy=None,
         venv=trainer_diverse_env.venv,
         n_timesteps=n_timesteps,
+        random_state=random_state,
     )
 
     obs, acts, next_obs, dones = trainer_diverse_env.reward_train.preprocess(
diff --git a/tests/algorithms/test_bc.py b/tests/algorithms/test_bc.py
index 200f2101b..ebf19ea1e 100644
--- a/tests/algorithms/test_bc.py
+++ b/tests/algorithms/test_bc.py
@@ -47,7 +47,9 @@ def trainer(
     expert_data_type,
     custom_logger,
     cartpole_expert_trajectories,
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     trans = rollout.flatten_trajectories(cartpole_expert_trajectories)
     if expert_data_type == "data_loader":
         expert_data = th_data.DataLoader(
@@ -70,10 +72,12 @@ def trainer(
         batch_size=batch_size,
         demonstrations=expert_data,
         custom_logger=custom_logger,
+        random_state=random_state,
     )
 
 
-def test_weight_decay_init_error(cartpole_venv, custom_logger):
+def test_weight_decay_init_error(cartpole_venv, custom_logger, random_state_fixed):
+    random_state = random_state_fixed
     with pytest.raises(ValueError, match=".*weight_decay.*"):
         bc.BC(
             observation_space=cartpole_venv.observation_space,
@@ -81,6 +85,7 @@ def test_weight_decay_init_error(cartpole_venv, custom_logger):
             demonstrations=None,
             optimizer_kwargs=dict(weight_decay=1e-4),
             custom_logger=custom_logger,
+            random_state=random_state,
         )
 
 
@@ -165,6 +170,7 @@ def test_bc_data_loader_empty_iter_error(
     no_yield_after_iter: bool,
     custom_logger: logger.HierarchicalLogger,
     cartpole_expert_trajectories,
+    random_state_fixed,
 ) -> None:
     """Check that we error out if the DataLoader suddenly stops yielding any batches.
 
@@ -177,6 +183,7 @@ def test_bc_data_loader_empty_iter_error(
         cartpole_expert_trajectories: The expert trajectories to use.
     """
     batch_size = 32
+    random_state = random_state_fixed
     trans = rollout.flatten_trajectories(cartpole_expert_trajectories)
     dummy_yield_value = dataclasses.asdict(trans[:batch_size])
 
@@ -189,6 +196,7 @@ def test_bc_data_loader_empty_iter_error(
         action_space=cartpole_venv.action_space,
         batch_size=batch_size,
         custom_logger=custom_logger,
+        random_state=random_state,
     )
     trainer.set_demonstrations(bad_data_loader)
     with pytest.raises(AssertionError, match=".*no data.*"):
diff --git a/tests/algorithms/test_dagger.py b/tests/algorithms/test_dagger.py
index f31d2edd3..5a54dc9ac 100644
--- a/tests/algorithms/test_dagger.py
+++ b/tests/algorithms/test_dagger.py
@@ -110,6 +110,7 @@ def _build_dagger_trainer(
     expert_policy,
     pendulum_expert_rollouts: List[TrajectoryWithRew],
     custom_logger,
+    random_state: np.random.RandomState,
 ):
     del expert_policy
     if pendulum_expert_rollouts is not None:
@@ -122,6 +123,7 @@ def _build_dagger_trainer(
         action_space=venv.action_space,
         optimizer_kwargs=dict(lr=1e-3),
         custom_logger=custom_logger,
+        random_state=random_state,
     )
     return dagger.DAggerTrainer(
         venv=venv,
@@ -139,12 +141,14 @@ def _build_simple_dagger_trainer(
     expert_policy,
     pendulum_expert_rollouts: Optional[List[TrajectoryWithRew]],
     custom_logger,
+    random_state,
 ):
     bc_trainer = bc.BC(
         observation_space=venv.observation_space,
         action_space=venv.action_space,
         optimizer_kwargs=dict(lr=1e-3),
         custom_logger=custom_logger,
+        random_state=random_state,
     )
     return dagger.SimpleDAggerTrainer(
         venv=venv,
@@ -171,6 +175,7 @@ def init_trainer_fn(
     pendulum_expert_policy,
     maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
     custom_logger,
+    random_state_fixed,
 ):
     # Provide a trainer initialization fixture in addition `trainer` fixture below
     # for tests that want to initialize multiple DAggerTrainer.
@@ -182,6 +187,7 @@ def init_trainer_fn(
         pendulum_expert_policy,
         maybe_pendulum_expert_trajectories,
         custom_logger,
+        random_state_fixed,
     )
 
 
@@ -198,6 +204,7 @@ def simple_dagger_trainer(
     pendulum_expert_policy,
     maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
     custom_logger,
+    random_state_fixed,
 ):
     return _build_simple_dagger_trainer(
         tmpdir,
@@ -206,6 +213,7 @@ def simple_dagger_trainer(
         pendulum_expert_policy,
         maybe_pendulum_expert_trajectories,
         custom_logger,
+        random_state_fixed,
     )
 
 
@@ -238,13 +246,16 @@ def test_trainer_needs_demos_exception_error(
         trainer.extend_and_update(dict(n_epochs=1))
 
 
-def test_trainer_train_arguments(trainer, pendulum_expert_policy):
+def test_trainer_train_arguments(trainer, pendulum_expert_policy, random_state_fixed):
+    random_state = random_state_fixed
+
     def add_samples():
         collector = trainer.create_trajectory_collector()
         rollout.generate_trajectories(
             pendulum_expert_policy,
             collector,
             sample_until=rollout.make_min_timesteps(40),
+            random_state=random_state,
         )
 
     # Lower default number of epochs for the no-arguments call that follows.
@@ -372,6 +383,7 @@ def test_simple_dagger_space_mismatch_error(
     pendulum_expert_policy,
     maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
     custom_logger,
+    random_state_fixed,
 ):
     class MismatchedSpace(gym.spaces.Space):
         """Dummy space that is not equal to any other space."""
@@ -389,17 +401,20 @@ class MismatchedSpace(gym.spaces.Space):
                     pendulum_expert_policy,
                     maybe_pendulum_expert_trajectories,
                     custom_logger,
+                    random_state_fixed,
                 )
 
 
-def test_dagger_not_enough_transitions_error(tmpdir, custom_logger):
-    venv = util.make_vec_env("CartPole-v0")
+def test_dagger_not_enough_transitions_error(tmpdir, custom_logger, random_state_fixed):
+    random_state = random_state_fixed
+    venv = util.make_vec_env("CartPole-v0", random_state=random_state)
     # Initialize with large batch size to ensure error down the line.
     bc_trainer = bc.BC(
         observation_space=venv.observation_space,
         action_space=venv.action_space,
         batch_size=100_000,
         custom_logger=custom_logger,
+        random_state=random_state,
     )
     trainer = dagger.DAggerTrainer(
         venv=venv,
@@ -409,6 +424,8 @@ def test_dagger_not_enough_transitions_error(tmpdir, custom_logger):
     )
     collector = trainer.create_trajectory_collector()
     policy = base.RandomPolicy(venv.observation_space, venv.action_space)
-    rollout.generate_trajectories(policy, collector, rollout.make_min_episodes(1))
+    rollout.generate_trajectories(
+        policy, collector, rollout.make_min_episodes(1), random_state=random_state
+    )
     with pytest.raises(ValueError, match="Not enough transitions.*"):
         trainer.extend_and_update()
diff --git a/tests/algorithms/test_density_baselines.py b/tests/algorithms/test_density_baselines.py
index 481d8c493..4a10ca057 100644
--- a/tests/algorithms/test_density_baselines.py
+++ b/tests/algorithms/test_density_baselines.py
@@ -44,7 +44,9 @@ def test_density_reward(
     is_stationary,
     pendulum_venv,
     pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     # use only a subset of trajectories
     expert_trajectories_all = pendulum_expert_trajectories[:8]
     n_experts = len(expert_trajectories_all)
@@ -57,6 +59,7 @@ def test_density_reward(
         is_stationary=is_stationary,
         kernel_bandwidth=0.2,
         standardise_inputs=True,
+        random_state=random_state,
     )
     reward_fn.train()
 
@@ -71,6 +74,7 @@ def test_density_reward(
         random_policy,
         pendulum_venv,
         sample_until=sample_until,
+        random_state=random_state,
     )
     expert_trajectories_test = expert_trajectories_all[n_experts // 2 :]
     random_returns = score_trajectories(random_trajectories, reward_fn)
@@ -85,7 +89,9 @@ def test_density_reward(
 def test_density_trainer_smoke(
     pendulum_venv,
     pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     # tests whether density trainer runs, not whether it's good
     # (it's actually really poor)
     rollouts = pendulum_expert_trajectories[:2]
@@ -94,6 +100,7 @@ def test_density_trainer_smoke(
         demonstrations=rollouts,
         venv=pendulum_venv,
         rl_algo=rl_algo,
+        random_state=random_state,
     )
     density_trainer.train()
     density_trainer.train_policy(n_timesteps=2)
diff --git a/tests/algorithms/test_mce_irl.py b/tests/algorithms/test_mce_irl.py
index 26d79fa16..fb374bfaf 100644
--- a/tests/algorithms/test_mce_irl.py
+++ b/tests/algorithms/test_mce_irl.py
@@ -269,7 +269,8 @@ def test_tabular_policy():
     np.testing.assert_equal(timesteps[0], 2 - mask.astype(int))
 
 
-def test_tabular_policy_randomness():
+def test_tabular_policy_randomness(random_state_fixed):
+    random_state = random_state_fixed
     state_space = gym.spaces.Discrete(2)
     action_space = gym.spaces.Discrete(2)
     pi = np.array(
@@ -280,12 +281,11 @@ def test_tabular_policy_randomness():
             ],
         ],
     )
-    rng = np.random.RandomState(42)
     tabular = TabularPolicy(
         state_space=state_space,
         action_space=action_space,
         pi=pi,
-        rng=rng,
+        random_state=random_state,
     )
 
     actions, _ = tabular.predict(np.zeros((100,), dtype=int))
@@ -297,7 +297,8 @@ def test_tabular_policy_randomness():
     np.testing.assert_equal(actions, 0)
 
 
-def test_mce_irl_demo_formats():
+def test_mce_irl_demo_formats(fixed_random_state):
+    random_state = fixed_random_state
     mdp = model_envs.RandomMDP(
         n_states=5,
         n_actions=3,
@@ -313,6 +314,7 @@ def test_mce_irl_demo_formats():
         policy=None,
         venv=state_venv,
         sample_until=rollout.make_min_timesteps(100),
+        random_state=random_state,
     )
     demonstrations = {
         "trajs": trajs,
@@ -357,7 +359,9 @@ def test_mce_irl_demo_formats():
 def test_mce_irl_reasonable_mdp(
     model_kwargs: Mapping[str, Any],
     discount: float,
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     with th.random.fork_rng():
         th.random.manual_seed(715298)
 
@@ -377,7 +381,14 @@ def test_mce_irl_reasonable_mdp(
             use_done=False,
             **model_kwargs,
         )
-        mce_irl = MCEIRL(D, mdp, reward_net, linf_eps=1e-3, discount=discount)
+        mce_irl = MCEIRL(
+            D,
+            mdp,
+            reward_net,
+            linf_eps=1e-3,
+            discount=discount,
+            random_state=random_state,
+        )
         final_counts = mce_irl.train()
 
         assert np.allclose(final_counts, D, atol=1e-3, rtol=1e-3)
@@ -390,6 +401,7 @@ def test_mce_irl_reasonable_mdp(
             mce_irl.policy,
             state_venv,
             sample_until=rollout.make_min_episodes(5),
+            random_state=random_state,
         )
         stats = rollout.rollout_stats(trajs)
         if discount > 0.0:  # skip check when discount==0.0 (random policy)
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 25c346de3..66c80653d 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -1,5 +1,5 @@
 """Tests for the preference comparisons reward learning implementation."""
-
+import random
 import re
 from typing import Sequence
 
@@ -22,10 +22,12 @@
 
 
 @pytest.fixture
-def venv():
+def venv(random_state_fixed):
+    random_state = random_state_fixed
     return util.make_vec_env(
         "seals/CartPole-v0",
         n_envs=1,
+        random_state=random_state,
     )
 
 
@@ -54,13 +56,17 @@ def agent(venv):
 
 
 @pytest.fixture
-def random_fragmenter():
-    return preference_comparisons.RandomFragmenter(seed=0, warning_threshold=0)
+def random_fragmenter(random_state_fixed):
+    random_state = random_state_fixed
+    return preference_comparisons.RandomFragmenter(
+        random_state=random.Random(util.make_seeds(random_state)), warning_threshold=0
+    )
 
 
 @pytest.fixture
-def agent_trainer(agent, reward_net, venv):
-    return preference_comparisons.AgentTrainer(agent, reward_net, venv)
+def agent_trainer(agent, reward_net, venv, random_state_fixed):
+    random_state = random_state_fixed
+    return preference_comparisons.AgentTrainer(agent, reward_net, venv, random_state)
 
 
 def _check_trajs_equal(
@@ -78,10 +84,12 @@ def _check_trajs_equal(
         assert traj1.terminal == traj2.terminal
 
 
-def test_mismatched_spaces(venv, agent):
+def test_mismatched_spaces(venv, agent, random_state_fixed):
+    random_state = random_state_fixed
     other_venv = util.make_vec_env(
         "seals/MountainCar-v0",
         n_envs=1,
+        random_state=random_state,
     )
     bad_reward_net = reward_nets.BasicRewardNet(
         other_venv.observation_space,
@@ -91,7 +99,7 @@ def test_mismatched_spaces(venv, agent):
         ValueError,
         match="spaces do not match",
     ):
-        preference_comparisons.AgentTrainer(agent, bad_reward_net, venv)
+        preference_comparisons.AgentTrainer(agent, bad_reward_net, venv, random_state=random_state)
 
 
 def test_trajectory_dataset_seeding(
@@ -100,12 +108,12 @@ def test_trajectory_dataset_seeding(
 ):
     dataset1 = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        seed=0,
+        random_state=random.Random(0),
     )
     sample1 = dataset1.sample(num_samples)
     dataset2 = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        seed=0,
+        random_state=random.Random(0),
     )
     sample2 = dataset2.sample(num_samples)
 
@@ -113,7 +121,7 @@ def test_trajectory_dataset_seeding(
 
     dataset3 = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        seed=42,
+        random_state=random.Random(42),
     )
     sample3 = dataset3.sample(num_samples)
     with pytest.raises(AssertionError):
@@ -128,7 +136,7 @@ def test_trajectory_dataset_len(
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        seed=0,
+        random_state=random.Random(0),
     )
     sample = dataset.sample(num_steps)
     lengths = [len(t) for t in sample]
@@ -142,7 +150,7 @@ def test_trajectory_dataset_too_long(
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        seed=0,
+        random_state=random.Random(0),
     )
     with pytest.raises(RuntimeError, match="Asked for.*but only.* available"):
         dataset.sample(100000)
@@ -154,7 +162,7 @@ def test_trajectory_dataset_shuffle(
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        seed=0,
+        random_state=random.Random(0),
     )
     sample = dataset.sample(num_steps)
     sample2 = dataset.sample(num_steps)
@@ -186,7 +194,9 @@ def test_trainer_no_crash(
     random_fragmenter,
     custom_logger,
     schedule,
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     main_trainer = preference_comparisons.PreferenceComparisons(
         agent_trainer,
         reward_net,
@@ -196,6 +206,7 @@ def test_trainer_no_crash(
         fragmenter=random_fragmenter,
         custom_logger=custom_logger,
         query_schedule=schedule,
+        random_state=random_state,
     )
     result = main_trainer.train(100, 10)
     # We don't expect good performance after training for 10 (!) timesteps,
@@ -204,7 +215,8 @@ def test_trainer_no_crash(
     assert 0.0 < result["reward_accuracy"] <= 1.0
 
 
-def test_reward_ensemble_trainer_raises_type_error(venv):
+def test_reward_ensemble_trainer_raises_type_error(venv, random_state_fixed):
+    random_state = random_state_fixed
     reward_net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     preference_model = preference_comparisons.PreferenceModel(
         model=reward_net,
@@ -221,6 +233,7 @@ def test_reward_ensemble_trainer_raises_type_error(venv):
         preference_comparisons.EnsembleTrainer(
             reward_net,
             loss,
+            random_state=random_state,
         )
 
 
@@ -229,7 +242,9 @@ def test_correct_reward_trainer_used_by_default(
     reward_net,
     random_fragmenter,
     custom_logger,
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     main_trainer = preference_comparisons.PreferenceComparisons(
         agent_trainer,
         reward_net,
@@ -237,6 +252,7 @@ def test_correct_reward_trainer_used_by_default(
         transition_oversampling=2,
         fragment_length=2,
         fragmenter=random_fragmenter,
+        random_state=random_state,
         custom_logger=custom_logger,
     )
 
@@ -258,7 +274,9 @@ def test_init_raises_error_when_trying_use_improperly_wrapped_ensemble(
     venv,
     random_fragmenter,
     custom_logger,
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     reward_net = testing_reward_nets.make_ensemble(
         venv.observation_space,
         venv.action_space,
@@ -279,11 +297,15 @@ def test_init_raises_error_when_trying_use_improperly_wrapped_ensemble(
             transition_oversampling=2,
             fragment_length=2,
             fragmenter=random_fragmenter,
+            random_state=random_state,
             custom_logger=custom_logger,
         )
 
 
-def test_discount_rate_no_crash(agent_trainer, venv, random_fragmenter, custom_logger):
+def test_discount_rate_no_crash(
+    agent_trainer, venv, random_fragmenter, custom_logger, random_state_fixed
+):
+    random_state = random_state_fixed
     # also use a non-zero noise probability to check that doesn't cause errors
     reward_net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     preference_model = preference_comparisons.PreferenceModel(
@@ -305,14 +327,16 @@ def test_discount_rate_no_crash(agent_trainer, venv, random_fragmenter, custom_l
         transition_oversampling=2,
         fragment_length=2,
         fragmenter=random_fragmenter,
+        random_state=random_state,
         reward_trainer=reward_trainer,
         custom_logger=custom_logger,
     )
     main_trainer.train(100, 10)
 
 
-def test_synthetic_gatherer_deterministic(agent_trainer, random_fragmenter):
-    gatherer = preference_comparisons.SyntheticGatherer(temperature=0)
+def test_synthetic_gatherer_deterministic(agent_trainer, random_fragmenter, random_state_fixed):
+    random_state = random_state_fixed
+    gatherer = preference_comparisons.SyntheticGatherer(temperature=0, random_state=random_state)
     trajectories = agent_trainer.sample(10)
     fragments = random_fragmenter(trajectories, fragment_length=2, num_pairs=2)
     preferences1 = gatherer(fragments)
@@ -346,7 +370,7 @@ def test_fragments_terminal(random_fragmenter):
 def test_fragments_too_short_error(agent_trainer):
     trajectories = agent_trainer.sample(2)
     random_fragmenter = preference_comparisons.RandomFragmenter(
-        seed=0,
+        random_state=random.Random(0),
         warning_threshold=0,
     )
     with pytest.raises(
@@ -373,11 +397,12 @@ def test_preference_dataset_errors(agent_trainer, random_fragmenter):
         dataset.push(fragments, preferences)
 
 
-def test_preference_dataset_queue(agent_trainer, random_fragmenter):
+def test_preference_dataset_queue(agent_trainer, random_fragmenter, random_state_fixed):
+    random_state = random_state_fixed
     dataset = preference_comparisons.PreferenceDataset(max_size=5)
     trajectories = agent_trainer.sample(10)
 
-    gatherer = preference_comparisons.SyntheticGatherer()
+    gatherer = preference_comparisons.SyntheticGatherer(random_state=random_state)
     for i in range(6):
         fragments = random_fragmenter(trajectories, fragment_length=2, num_pairs=1)
         preferences = gatherer(fragments)
@@ -389,11 +414,12 @@ def test_preference_dataset_queue(agent_trainer, random_fragmenter):
     assert len(dataset) == 5
 
 
-def test_store_and_load_preference_dataset(agent_trainer, random_fragmenter, tmp_path):
+def test_store_and_load_preference_dataset(agent_trainer, random_fragmenter, tmp_path, random_state_fixed):
+    random_state = random_state_fixed
     dataset = preference_comparisons.PreferenceDataset()
     trajectories = agent_trainer.sample(10)
     fragments = random_fragmenter(trajectories, fragment_length=2, num_pairs=2)
-    gatherer = preference_comparisons.SyntheticGatherer()
+    gatherer = preference_comparisons.SyntheticGatherer(random_state=random_state)
     preferences = gatherer(fragments)
     dataset.push(fragments, preferences)
 
@@ -415,11 +441,14 @@ def test_exploration_no_crash(
     venv,
     random_fragmenter,
     custom_logger,
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     agent_trainer = preference_comparisons.AgentTrainer(
         agent,
         reward_net,
         venv,
+        random_state=random_state,
         exploration_frac=0.5,
     )
     main_trainer = preference_comparisons.PreferenceComparisons(
@@ -429,6 +458,7 @@ def test_exploration_no_crash(
         transition_oversampling=2,
         fragment_length=5,
         fragmenter=random_fragmenter,
+        random_state=random_state,
         custom_logger=custom_logger,
     )
     main_trainer.train(100, 10)
@@ -441,7 +471,9 @@ def test_active_fragmenter_discount_rate_no_crash(
     random_fragmenter,
     uncertainty_on,
     custom_logger,
+    random_state_fixed,
 ):
+    random_state = random_state_fixed
     # also use a non-zero noise probability to check that doesn't cause errors
     reward_net = reward_nets.RewardEnsemble(
         venv.observation_space,
@@ -477,6 +509,7 @@ def test_active_fragmenter_discount_rate_no_crash(
     reward_trainer = preference_comparisons.EnsembleTrainer(
         reward_net,
         loss,
+        random_state=random_state,
     )
 
     main_trainer = preference_comparisons.PreferenceComparisons(
@@ -486,6 +519,7 @@ def test_active_fragmenter_discount_rate_no_crash(
         transition_oversampling=2,
         fragment_length=2,
         fragmenter=fragmenter,
+        random_state=random_state,
         reward_trainer=reward_trainer,
         custom_logger=custom_logger,
     )
diff --git a/tests/conftest.py b/tests/conftest.py
index b27244678..6bd429e35 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -6,6 +6,7 @@
 from typing import Callable, Optional, Sequence
 
 import gym
+import numpy as np
 import pytest
 import torch
 from filelock import FileLock
@@ -57,6 +58,7 @@ def load_or_rollout_trajectories(
     cache_path,
     policy,
     venv,
+    random_state,
 ) -> Sequence[TrajectoryWithRew]:
     os.makedirs(os.path.dirname(cache_path), exist_ok=True)
     with FileLock(cache_path + ".lock"):
@@ -71,6 +73,7 @@ def load_or_rollout_trajectories(
                 policy,
                 venv,
                 rollout.make_sample_until(min_timesteps=2000, min_episodes=57),
+                random_state=random_state,
             )
             types.save(cache_path, rollouts)
             return rollouts
@@ -138,7 +141,9 @@ def cartpole_expert_trajectories(
     cartpole_expert_policy,
     cartpole_venv,
     pytestconfig,
+    random_state_fixed,
 ) -> Sequence[TrajectoryWithRew]:
+    random_state = random_state_fixed
     rollouts_path = str(
         pytestconfig.cache.makedir("experts") / CARTPOLE_ENV_NAME / "rollout.npz",
     )
@@ -146,6 +151,7 @@ def cartpole_expert_trajectories(
         rollouts_path,
         cartpole_expert_policy,
         cartpole_venv,
+        random_state,
     )
 
 
@@ -198,14 +204,17 @@ def pendulum_expert_trajectories(
     pendulum_expert_policy,
     pendulum_venv,
     pytestconfig,
+    random_state_fixed,
 ) -> Sequence[TrajectoryWithRew]:
     rollouts_path = str(
         pytestconfig.cache.makedir("experts") / PENDULUM_ENV_NAME / "rollout.npz",
     )
+    random_state = random_state_fixed
     return load_or_rollout_trajectories(
         rollouts_path,
         pendulum_expert_policy,
         pendulum_venv,
+        random_state=random_state,
     )
 
 
@@ -227,3 +236,13 @@ def torch_single_threaded():
 @pytest.fixture()
 def custom_logger(tmpdir: str) -> logger.HierarchicalLogger:
     return logger.configure(tmpdir)
+
+
+@pytest.fixture()
+def random_state_fixed() -> np.random.RandomState:
+    return np.random.RandomState(0)
+
+
+@pytest.fixture()
+def random_state() -> np.random.RandomState:
+    return np.random.RandomState()
diff --git a/tests/data/test_buffer.py b/tests/data/test_buffer.py
index db6c52c3f..7f8caf4e0 100644
--- a/tests/data/test_buffer.py
+++ b/tests/data/test_buffer.py
@@ -53,7 +53,7 @@ def test_buffer(capacity, chunk_len, sample_shape) -> None:
     buf = Buffer(
         capacity,
         sample_shapes={"a": sample_shape, "b": sample_shape},
-        dtypes={"a": float, "b": float},
+        dtypes={"a": np.dtype(float), "b": np.dtype(float)},
     )
 
     to_insert = 3 * capacity
diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index 4b4f7ee4f..db845d3dd 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -38,6 +38,7 @@ def step(self, action):
 def _sample_fixed_length_trajectories(
     episode_lengths: Sequence[int],
     min_episodes: int,
+    random_state: np.random.RandomState,
     policy_type: str = "policy",
     **kwargs,
 ) -> Sequence[types.Trajectory]:
@@ -64,6 +65,7 @@ def policy(x):
         policy,
         venv,
         sample_until=sample_until,
+        random_state=random_state,
         **kwargs,
     )
     return trajectories
@@ -73,7 +75,7 @@ def policy(x):
     "policy_type",
     ["policy", "callable", "random"],
 )
-def test_complete_trajectories(policy_type) -> None:
+def test_complete_trajectories(policy_type, random_state_fixed) -> None:
     """Checks trajectories include the terminal observation.
 
     This is hidden by default by VecEnv's auto-reset; we add it back in using
@@ -82,6 +84,7 @@ def test_complete_trajectories(policy_type) -> None:
     Args:
         policy_type: Kind of policy to use when generating trajectories.
     """
+    random_state = random_state_fixed
     min_episodes = 13
     max_acts = 5
     num_envs = 4
@@ -89,6 +92,7 @@ def test_complete_trajectories(policy_type) -> None:
         [max_acts] * num_envs,
         min_episodes,
         policy_type=policy_type,
+        random_state=random_state,
     )
     assert len(trajectories) >= min_episodes
     expected_obs = np.array([[0]] * max_acts + [[1]])
@@ -118,6 +122,7 @@ def test_unbiased_trajectories(
     episode_lengths: Sequence[int],
     min_episodes: int,
     expected_counts: Mapping[int, int],
+    random_state_fixed,
 ) -> None:
     """Checks trajectories are sampled without bias towards shorter episodes.
 
@@ -138,7 +143,10 @@ def test_unbiased_trajectories(
         expected_counts: Mapping from episode length to expected number of episodes
             of that length (omit if 0 episodes of that length expected).
     """
-    trajectories = _sample_fixed_length_trajectories(episode_lengths, min_episodes)
+    random_state = random_state_fixed
+    trajectories = _sample_fixed_length_trajectories(
+        episode_lengths, min_episodes, random_state
+    )
     assert len(trajectories) == sum(expected_counts.values())
     traj_lens = np.array([len(traj) for traj in trajectories])
     for length, count in expected_counts.items():
@@ -157,9 +165,9 @@ def test_seed_trajectories():
     rng_a1 = np.random.RandomState(0)
     rng_a2 = np.random.RandomState(0)
     rng_b = np.random.RandomState(1)
-    traj_a1 = _sample_fixed_length_trajectories([3, 5], 2, rng=rng_a1)
-    traj_a2 = _sample_fixed_length_trajectories([3, 5], 2, rng=rng_a2)
-    traj_b = _sample_fixed_length_trajectories([3, 5], 2, rng=rng_b)
+    traj_a1 = _sample_fixed_length_trajectories([3, 5], 2, random_state=rng_a1)
+    traj_a2 = _sample_fixed_length_trajectories([3, 5], 2, random_state=rng_a2)
+    traj_b = _sample_fixed_length_trajectories([3, 5], 2, random_state=rng_b)
     assert [len(traj) for traj in traj_a1] == [len(traj) for traj in traj_a2]
     assert [len(traj) for traj in traj_a1] != [len(traj) for traj in traj_b]
 
@@ -176,18 +184,21 @@ def step(self, action):
         return obs / 2, rew / 2, done, info
 
 
-def test_rollout_stats():
+def test_rollout_stats(random_state_fixed):
     """Applying `ObsRewIncrementWrapper` halves the reward mean.
 
     `rollout_stats` should reflect this.
     """
+    random_state = random_state_fixed
     env = gym.make("CartPole-v1")
     env = monitor.Monitor(env, None)
     env = ObsRewHalveWrapper(env)
     venv = vec_env.DummyVecEnv([lambda: env])
 
     policy = serialize.load_policy("zero", "UNUSED", venv)
-    trajs = rollout.generate_trajectories(policy, venv, rollout.make_min_episodes(10))
+    trajs = rollout.generate_trajectories(
+        policy, venv, rollout.make_min_episodes(10), random_state=random_state
+    )
     s = rollout.rollout_stats(trajs)
 
     np.testing.assert_allclose(s["return_mean"], s["monitor_return_mean"] / 2)
@@ -196,18 +207,21 @@ def test_rollout_stats():
     np.testing.assert_allclose(s["return_max"], s["monitor_return_max"] / 2)
 
 
-def test_unwrap_traj():
+def test_unwrap_traj(random_state_fixed):
     """Check that unwrap_traj reverses `ObsRewIncrementWrapper`.
 
     Also check that unwrapping twice is a no-op.
     """
+    random_state = random_state_fixed
     env = gym.make("CartPole-v1")
     env = wrappers.RolloutInfoWrapper(env)
     env = ObsRewHalveWrapper(env)
     venv = vec_env.DummyVecEnv([lambda: env])
 
     policy = serialize.load_policy("zero", "UNUSED", venv)
-    trajs = rollout.generate_trajectories(policy, venv, rollout.make_min_episodes(10))
+    trajs = rollout.generate_trajectories(
+        policy, venv, rollout.make_min_episodes(10), random_state=random_state
+    )
     trajs_unwrapped = [rollout.unwrap_traj(t) for t in trajs]
     trajs_unwrapped_twice = [rollout.unwrap_traj(t) for t in trajs_unwrapped]
 
@@ -252,18 +266,21 @@ def test_compute_returns(gamma):
     assert abs(rollout.discounted_sum(rewards, gamma) - returns) < 1e-8
 
 
-def test_generate_trajectories_type_error():
+def test_generate_trajectories_type_error(random_state_fixed):
+    random_state = random_state_fixed
     venv = vec_env.DummyVecEnv([functools.partial(TerminalSentinelEnv, 1)])
     sample_until = rollout.make_min_episodes(1)
     with pytest.raises(TypeError, match="Policy must be.*got <class 'str'> instead"):
         rollout.generate_trajectories(
-            "strings_are_not_valid_policies",
+            "strings_are_not_valid_policies",  # type: ignore
             venv,
+            random_state=random_state,
             sample_until=sample_until,
         )
 
 
-def test_generate_trajectories_value_error():
+def test_generate_trajectories_value_error(random_state_fixed):
+    random_state = random_state_fixed
     venv = vec_env.DummyVecEnv([functools.partial(TerminalSentinelEnv, 1)])
     sample_until = rollout.make_min_episodes(1)
 
@@ -272,5 +289,6 @@ def test_generate_trajectories_value_error():
             lambda obs: np.zeros(len(obs), dtype=int),
             venv,
             sample_until=sample_until,
+            random_state=random_state,
             deterministic_policy=True,
         )
diff --git a/tests/policies/test_exploration_wrapper.py b/tests/policies/test_exploration_wrapper.py
index b2972e5e6..5b76f6f29 100644
--- a/tests/policies/test_exploration_wrapper.py
+++ b/tests/policies/test_exploration_wrapper.py
@@ -15,6 +15,7 @@ def make_wrapper(random_prob, switch_prob):
     venv = util.make_vec_env(
         "seals/CartPole-v0",
         n_envs=1,
+        random_state=np.random.RandomState(),
     )
     return (
         exploration_wrapper.ExplorationWrapper(
@@ -22,7 +23,7 @@ def make_wrapper(random_prob, switch_prob):
             venv=venv,
             random_prob=random_prob,
             switch_prob=switch_prob,
-            seed=0,
+            random_state=np.random.RandomState(0),
         ),
         venv,
     )
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index d48940042..36b71bea6 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -25,11 +25,16 @@
 
 @pytest.mark.parametrize("env_name", SIMPLE_ENVS)
 @pytest.mark.parametrize("policy_type", HARDCODED_TYPES)
-def test_actions_valid(env_name, policy_type):
+def test_actions_valid(env_name, policy_type, random_state_fixed):
     """Test output actions of our custom policies always lie in action space."""
-    venv = util.make_vec_env(env_name, n_envs=1, parallel=False)
+    random_state = random_state_fixed
+    venv = util.make_vec_env(
+        env_name, n_envs=1, parallel=False, random_state=random_state
+    )
     policy = serialize.load_policy(policy_type, "foobar", venv)
-    transitions = rollout.generate_transitions(policy, venv, n_timesteps=100)
+    transitions = rollout.generate_transitions(
+        policy, venv, n_timesteps=100, random_state=random_state
+    )
 
     for a in transitions.acts:
         assert venv.action_space.contains(a)
@@ -42,11 +47,14 @@ def test_actions_valid(env_name, policy_type):
         ("sac", SIMPLE_CONTINUOUS_ENV),
     ],
 )
-def test_save_stable_model_errors_and_warnings(tmpdir, policy_env_name_pair):
+def test_save_stable_model_errors_and_warnings(
+    tmpdir, policy_env_name_pair, random_state_fixed
+):
     """Check errors and warnings in `save_stable_model()`."""
+    random_state = random_state_fixed
     policy, env_name = policy_env_name_pair
     tmpdir = pathlib.Path(tmpdir)
-    venv = util.make_vec_env(env_name)
+    venv = util.make_vec_env(env_name, random_state=random_state)
 
     # Trigger FileNotFoundError for no model.{zip,pkl}
     dir_a = tmpdir / "a"
@@ -65,9 +73,11 @@ def test_save_stable_model_errors_and_warnings(tmpdir, policy_env_name_pair):
         serialize.load_policy(policy, str(dir_nonexistent), venv)
 
 
-def _test_serialize_identity(env_name, model_cfg, tmpdir):
+def _test_serialize_identity(env_name, model_cfg, tmpdir, random_state):
     """Test output actions of deserialized policy are same as original."""
-    venv = util.make_vec_env(env_name, n_envs=1, parallel=False)
+    venv = util.make_vec_env(
+        env_name, n_envs=1, parallel=False, random_state=random_state
+    )
 
     model_name, model_cls_name = model_cfg
     model_cls = registry.load_attr(model_cls_name)
@@ -82,7 +92,7 @@ def _test_serialize_identity(env_name, model_cfg, tmpdir):
         venv,
         n_timesteps=1000,
         deterministic_policy=True,
-        rng=np.random.RandomState(0),
+        random_state=np.random.RandomState(0),
     )
 
     serialize.save_stable_model(tmpdir, model)
@@ -94,7 +104,7 @@ def _test_serialize_identity(env_name, model_cfg, tmpdir):
         venv,
         n_timesteps=1000,
         deterministic_policy=True,
-        rng=np.random.RandomState(0),
+        random_state=np.random.RandomState(0),
     )
 
     assert np.allclose(orig_rollout.acts, new_rollout.acts)
@@ -108,16 +118,18 @@ def _test_serialize_identity(env_name, model_cfg, tmpdir):
 
 @pytest.mark.parametrize("env_name", SIMPLE_ENVS)
 @pytest.mark.parametrize("model_cfg", NORMAL_CONFIGS)
-def test_serialize_identity(env_name, model_cfg, tmpdir):
+def test_serialize_identity(env_name, model_cfg, tmpdir, random_state_fixed):
     """Test output actions of deserialized policy are same as original."""
-    _test_serialize_identity(env_name, model_cfg, tmpdir)
+    _test_serialize_identity(env_name, model_cfg, tmpdir, random_state_fixed)
 
 
 @pytest.mark.parametrize("env_name", [SIMPLE_CONTINUOUS_ENV])
 @pytest.mark.parametrize("model_cfg", CONTINUOUS_ONLY_CONFIGS)
-def test_serialize_identity_continuous_only(env_name, model_cfg, tmpdir):
+def test_serialize_identity_continuous_only(
+    env_name, model_cfg, tmpdir, random_state_fixed
+):
     """Test serialize identity for continuous_only algorithms."""
-    _test_serialize_identity(env_name, model_cfg, tmpdir)
+    _test_serialize_identity(env_name, model_cfg, tmpdir, random_state_fixed)
 
 
 class ZeroModule(nn.Module):
diff --git a/tests/policies/test_replay_buffer_wrapper.py b/tests/policies/test_replay_buffer_wrapper.py
index c08f5e46e..801b7ea8a 100644
--- a/tests/policies/test_replay_buffer_wrapper.py
+++ b/tests/policies/test_replay_buffer_wrapper.py
@@ -29,9 +29,10 @@ def make_algo_with_wrapped_buffer(
     rl_cls: Type[off_policy_algorithm.OffPolicyAlgorithm],
     policy_cls: Type[BasePolicy],
     replay_buffer_class: Type[buffers.ReplayBuffer],
+    random_state: np.random.RandomState,
     buffer_size: int = 100,
 ) -> off_policy_algorithm.OffPolicyAlgorithm:
-    venv = util.make_vec_env("Pendulum-v1", n_envs=1)
+    venv = util.make_vec_env("Pendulum-v1", n_envs=1, random_state=random_state)
     rl_algo = rl_cls(
         policy=policy_cls,
         policy_kwargs=dict(),
@@ -49,7 +50,8 @@ def make_algo_with_wrapped_buffer(
     return rl_algo
 
 
-def test_invalid_args():
+def test_invalid_args(random_state_fixed):
+    random_state = random_state_fixed
     with pytest.raises(
         TypeError,
         match=r".*unexpected keyword argument 'replay_buffer_class'.*",
@@ -60,6 +62,7 @@ def test_invalid_args():
             rl_cls=sb3.PPO,  # type: ignore
             policy_cls=policies.ActorCriticPolicy,
             replay_buffer_class=buffers.ReplayBuffer,
+            random_state=random_state,
         )
 
     with pytest.raises(AssertionError, match=r".*only ReplayBuffer is supported.*"):
@@ -67,10 +70,12 @@ def test_invalid_args():
             rl_cls=sb3.SAC,
             policy_cls=sb3.sac.policies.SACPolicy,
             replay_buffer_class=buffers.DictReplayBuffer,
+            random_state=random_state,
         )
 
 
-def test_wrapper_class(tmpdir):
+def test_wrapper_class(tmpdir, random_state_fixed):
+    random_state = random_state_fixed
     buffer_size = 15
     total_timesteps = 20
 
@@ -79,6 +84,7 @@ def test_wrapper_class(tmpdir):
         policy_cls=sb3.sac.policies.SACPolicy,
         replay_buffer_class=buffers.ReplayBuffer,
         buffer_size=buffer_size,
+        random_state=random_state,
     )
 
     rl_algo.learn(total_timesteps=total_timesteps)
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index 93bc6c79a..9e3717f37 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -114,8 +114,10 @@ def _sample(space, n):
     return np.array([space.sample() for _ in range(n)])
 
 
-def _make_env_and_save_reward_net(env_name, reward_type, tmpdir):
-    venv = util.make_vec_env(env_name, n_envs=1, parallel=False)
+def _make_env_and_save_reward_net(env_name, reward_type, tmpdir, random_state):
+    venv = util.make_vec_env(
+        env_name, n_envs=1, parallel=False, random_state=random_state
+    )
     save_path = os.path.join(tmpdir, "norm_reward.pt")
 
     assert reward_type in [
@@ -144,10 +146,14 @@ def _make_env_and_save_reward_net(env_name, reward_type, tmpdir):
 
 @pytest.mark.parametrize("env_name", ENVS)
 @pytest.mark.parametrize("reward_type", DESERIALIZATION_TYPES)
-def test_reward_valid(env_name, reward_type, tmpdir):
+def test_reward_valid(env_name, reward_type, tmpdir, random_state_fixed):
     """Test output of reward function is appropriate shape and type."""
-    venv = util.make_vec_env(env_name, n_envs=1, parallel=False)
-    venv, tmppath = _make_env_and_save_reward_net(env_name, reward_type, tmpdir)
+    random_state = random_state_fixed
+    # TODO(juan) the line below is not being used?
+    venv = util.make_vec_env(env_name, n_envs=1, parallel=False, random_state=random_state)
+    venv, tmppath = _make_env_and_save_reward_net(
+        env_name, reward_type, tmpdir, random_state
+    )
 
     TRAJECTORY_LEN = 10
     obs = _sample(venv.observation_space, TRAJECTORY_LEN)
@@ -190,8 +196,11 @@ def test_wrappers_default_to_passing_on_method_calls_to_base(
     assert wrapper.dtype is base.dtype
 
 
-def test_strip_wrappers_basic():
-    venv = util.make_vec_env("FrozenLake-v1", n_envs=1, parallel=False)
+def test_strip_wrappers_basic(random_state_fixed):
+    random_state = random_state_fixed
+    venv = util.make_vec_env(
+        "FrozenLake-v1", n_envs=1, parallel=False, random_state=random_state
+    )
     net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     net = reward_nets.NormalizedRewardNet(net, networks.RunningNorm)
     net = serialize._strip_wrappers(
@@ -204,8 +213,11 @@ def test_strip_wrappers_basic():
     assert isinstance(net, reward_nets.BasicRewardNet)
 
 
-def test_strip_wrappers_complex():
-    venv = util.make_vec_env("FrozenLake-v1", n_envs=1, parallel=False)
+def test_strip_wrappers_complex(random_state_fixed):
+    random_state = random_state_fixed
+    venv = util.make_vec_env(
+        "FrozenLake-v1", n_envs=1, parallel=False, random_state=random_state
+    )
     net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     net = reward_nets.ShapedRewardNet(net, _potential, discount_factor=0.99)
     net = reward_nets.NormalizedRewardNet(net, networks.RunningNorm)
@@ -279,11 +291,13 @@ def forward(*args):
 
 
 @pytest.mark.parametrize("env_name", ENVS)
-def test_cant_load_unnorm_as_norm(env_name, tmpdir):
+def test_cant_load_unnorm_as_norm(env_name, tmpdir, random_state_fixed):
+    random_state = random_state_fixed
     venv, tmppath = _make_env_and_save_reward_net(
         env_name,
         "RewardNet_unnormalized",
         tmpdir,
+        random_state=random_state,
     )
     with pytest.raises(TypeError):
         serialize.load_reward("RewardNet_normalized", tmppath, venv)
@@ -297,11 +311,15 @@ def test_serialize_identity(
     net_cls,
     normalize_rewards,
     tmpdir,
+    random_state_fixed,
 ):
     """Does output of deserialized reward network match that of original?"""
+    random_state = random_state_fixed
     logging.info(f"Testing {net_cls}")
 
-    venv = util.make_vec_env(env_name, n_envs=1, parallel=False)
+    venv = util.make_vec_env(
+        env_name, n_envs=1, parallel=False, random_state=random_state
+    )
     original = net_cls(venv.observation_space, venv.action_space)
     if normalize_rewards:
         original = reward_nets.NormalizedRewardNet(original, networks.RunningNorm)
@@ -314,7 +332,9 @@ def test_serialize_identity(
     assert original.observation_space == loaded.observation_space
     assert original.action_space == loaded.action_space
 
-    transitions = rollout.generate_transitions(random, venv, n_timesteps=100)
+    transitions = rollout.generate_transitions(
+        random, venv, n_timesteps=100, random_state=random_state
+    )
 
     if isinstance(original, reward_nets.NormalizedRewardNet):
         wrapped_rew_fn = serialize.load_reward("RewardNet_normalized", tmppath, venv)
@@ -591,8 +611,9 @@ def forward(self):
 
 
 @pytest.mark.parametrize("normalize_input_layer", [None, networks.RunningNorm])
-def test_training_regression(normalize_input_layer):
+def test_training_regression(normalize_input_layer, random_state_fixed):
     """Test reward_net normalization by training a regression model."""
+    random_state = random_state_fixed
     venv = DummyVecEnv([lambda: gym.make("CartPole-v0")] * 2)
     reward_net = reward_nets.BasicRewardNet(
         venv.observation_space,
@@ -611,7 +632,9 @@ def test_training_regression(normalize_input_layer):
     # Getting transitions from a random policy
     random = base.RandomPolicy(venv.observation_space, venv.action_space)
     for _ in range(2):
-        transitions = rollout.generate_transitions(random, venv, n_timesteps=100)
+        transitions = rollout.generate_transitions(
+            random, venv, n_timesteps=100, random_state=random_state
+        )
         trans_args = (
             transitions.obs,
             transitions.acts,
diff --git a/tests/rewards/test_reward_wrapper.py b/tests/rewards/test_reward_wrapper.py
index 8d8d8fc92..fee403880 100644
--- a/tests/rewards/test_reward_wrapper.py
+++ b/tests/rewards/test_reward_wrapper.py
@@ -17,20 +17,21 @@ def __call__(self, obs, act, next_obs, steps=None):
         return (np.arange(len(obs)) + 1).astype("float32")
 
 
-def test_reward_overwrite():
+def test_reward_overwrite(random_state_fixed):
     """Test that reward wrapper actually overwrites base rewards."""
+    random_state = random_state_fixed
     env_name = "Pendulum-v1"
     num_envs = 3
-    env = util.make_vec_env(env_name, num_envs)
+    env = util.make_vec_env(env_name, random_state=random_state, n_envs=num_envs)
     reward_fn = FunkyReward()
     wrapped_env = reward_wrapper.RewardVecEnvWrapper(env, reward_fn)
     policy = RandomPolicy(env.observation_space, env.action_space)
     sample_until = rollout.make_min_episodes(10)
     default_stats = rollout.rollout_stats(
-        rollout.generate_trajectories(policy, env, sample_until),
+        rollout.generate_trajectories(policy, env, sample_until, random_state),
     )
     wrapped_stats = rollout.rollout_stats(
-        rollout.generate_trajectories(policy, wrapped_env, sample_until),
+        rollout.generate_trajectories(policy, wrapped_env, sample_until, random_state),
     )
     # Pendulum-v1 always has negative rewards
     assert default_stats["return_max"] < 0
diff --git a/tests/scripts/test_scripts.py b/tests/scripts/test_scripts.py
index 7b592f0c4..9c257fcfb 100644
--- a/tests/scripts/test_scripts.py
+++ b/tests/scripts/test_scripts.py
@@ -679,8 +679,11 @@ def test_preference_comparisons_transfer_learning(
     _check_rollout_stats(run.result)
 
 
-def test_train_rl_double_normalization(tmpdir: str):
-    venv = util.make_vec_env("CartPole-v1", n_envs=1, parallel=False)
+def test_train_rl_double_normalization(tmpdir: str, random_state_fixed):
+    random_state = random_state_fixed
+    venv = util.make_vec_env(
+        "CartPole-v1", n_envs=1, parallel=False, random_state=random_state
+    )
     basic_reward_net = reward_nets.BasicRewardNet(
         venv.observation_space,
         venv.action_space,

From 8f5a13769ffd996141c21934a5b8b88dcbc3c377 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 00:19:08 +0200
Subject: [PATCH 074/187] Fix notebooks; add script to clean notebooks

---
 ci/clean_notebooks.py                         | 33 +++++++++++++++++++
 examples/1_train_bc.ipynb                     | 12 ++++---
 examples/2_train_dagger.ipynb                 | 10 +++---
 examples/3_train_gail.ipynb                   | 15 +++++----
 examples/4_train_airl.ipynb                   | 15 +++++----
 examples/5_train_preference_comparisons.ipynb | 18 ++++++----
 examples/6_train_mce.ipynb                    | 21 ++++++++----
 examples/7_train_density.ipynb                | 27 +++++----------
 8 files changed, 96 insertions(+), 55 deletions(-)
 create mode 100644 ci/clean_notebooks.py

diff --git a/ci/clean_notebooks.py b/ci/clean_notebooks.py
new file mode 100644
index 000000000..505739bd7
--- /dev/null
+++ b/ci/clean_notebooks.py
@@ -0,0 +1,33 @@
+import json
+import pathlib
+
+import nbformat
+
+
+def clean_notebook(file: pathlib.Path):
+    # Read the notebook
+    with open(file) as f:
+        nb = nbformat.read(f, as_version=4)
+
+    # Remove the output and metadata from each cell
+    # also reset the execution count
+    # if the cell has no code, remove it
+    for cell in nb.cells:
+        if 'outputs' in cell:
+            cell['outputs'] = []
+        if 'metadata' in cell:
+            cell['metadata'] = {}
+        if 'execution_count' in cell:
+            cell['execution_count'] = None
+        if cell['cell_type'] == 'code' and not cell['source']:
+            nb.cells.remove(cell)
+
+    # Write the notebook
+    with open(file, 'w') as f:
+        nbformat.write(nb, f)
+
+
+if __name__ == '__main__':
+    for file in pathlib.Path.cwd().glob('**/*.ipynb'):
+        print(f'Cleaning {file}')
+        clean_notebook(file)
\ No newline at end of file
diff --git a/examples/1_train_bc.ipynb b/examples/1_train_bc.ipynb
index 75a4659ee..7497c1dfc 100644
--- a/examples/1_train_bc.ipynb
+++ b/examples/1_train_bc.ipynb
@@ -26,11 +26,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import gym\n",
     "from stable_baselines3 import PPO\n",
     "from stable_baselines3.ppo import MlpPolicy\n",
     "\n",
-    "import gym\n",
-    "\n",
     "env = gym.make(\"CartPole-v1\")\n",
     "expert = PPO(\n",
     "    policy=MlpPolicy,\n",
@@ -84,11 +83,14 @@
     "from imitation.data import rollout\n",
     "from imitation.data.wrappers import RolloutInfoWrapper\n",
     "from stable_baselines3.common.vec_env import DummyVecEnv\n",
+    "import numpy as np\n",
     "\n",
+    "random_state = np.random.RandomState()\n",
     "rollouts = rollout.rollout(\n",
     "    expert,\n",
     "    DummyVecEnv([lambda: RolloutInfoWrapper(env)]),\n",
     "    rollout.make_sample_until(min_timesteps=None, min_episodes=50),\n",
+    "    random_state=random_state,\n",
     ")\n",
     "transitions = rollout.flatten_trajectories(rollouts)"
    ]
@@ -133,6 +135,7 @@
     "    observation_space=env.observation_space,\n",
     "    action_space=env.action_space,\n",
     "    demonstrations=transitions,\n",
+    "    random_state=random_state,\n",
     ")"
    ]
   },
@@ -192,9 +195,8 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.10"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/examples/2_train_dagger.ipynb b/examples/2_train_dagger.ipynb
index 0cea53f77..1920fc8af 100644
--- a/examples/2_train_dagger.ipynb
+++ b/examples/2_train_dagger.ipynb
@@ -25,11 +25,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import gym\n",
     "from stable_baselines3 import PPO\n",
     "from stable_baselines3.ppo import MlpPolicy\n",
     "\n",
-    "import gym\n",
-    "\n",
     "env = gym.make(\"CartPole-v1\")\n",
     "expert = PPO(\n",
     "    policy=MlpPolicy,\n",
@@ -59,6 +58,7 @@
    "source": [
     "import tempfile\n",
     "import gym\n",
+    "import numpy as np\n",
     "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "\n",
     "from imitation.algorithms import bc\n",
@@ -70,6 +70,7 @@
     "bc_trainer = bc.BC(\n",
     "    observation_space=env.observation_space,\n",
     "    action_space=env.action_space,\n",
+    "    random_state=np.random.RandomState(),\n",
     ")\n",
     "\n",
     "with tempfile.TemporaryDirectory(prefix=\"dagger_example_\") as tmpdir:\n",
@@ -121,9 +122,8 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.10"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/examples/3_train_gail.ipynb b/examples/3_train_gail.ipynb
index 944a61343..15eebedb0 100644
--- a/examples/3_train_gail.ipynb
+++ b/examples/3_train_gail.ipynb
@@ -24,10 +24,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import gym\n",
     "from stable_baselines3 import PPO\n",
     "from stable_baselines3.ppo import MlpPolicy\n",
-    "import gym\n",
-    "import seals\n",
+    "import seals  # needed to load environments\n",
     "\n",
     "env = gym.make(\"seals/CartPole-v0\")\n",
     "expert = PPO(\n",
@@ -59,11 +59,14 @@
     "from imitation.data import rollout\n",
     "from imitation.data.wrappers import RolloutInfoWrapper\n",
     "from stable_baselines3.common.vec_env import DummyVecEnv\n",
+    "import numpy as np\n",
     "\n",
+    "random_state = np.random.RandomState()\n",
     "rollouts = rollout.rollout(\n",
     "    expert,\n",
     "    DummyVecEnv([lambda: RolloutInfoWrapper(gym.make(\"seals/CartPole-v0\"))] * 5),\n",
     "    rollout.make_sample_until(min_timesteps=None, min_episodes=60),\n",
+    "    random_state=random_state,\n",
     ")"
    ]
   },
@@ -87,10 +90,9 @@
     "from imitation.util.networks import RunningNorm\n",
     "from stable_baselines3 import PPO\n",
     "from stable_baselines3.common.evaluation import evaluate_policy\n",
-    "from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv\n",
+    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "\n",
     "import gym\n",
-    "import seals\n",
     "\n",
     "\n",
     "venv = DummyVecEnv([lambda: gym.make(\"seals/CartPole-v0\")] * 8)\n",
@@ -173,9 +175,8 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.10"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/examples/4_train_airl.ipynb b/examples/4_train_airl.ipynb
index ccfc9a2a1..11987f716 100644
--- a/examples/4_train_airl.ipynb
+++ b/examples/4_train_airl.ipynb
@@ -21,10 +21,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import gym\n",
     "from stable_baselines3 import PPO\n",
     "from stable_baselines3.ppo import MlpPolicy\n",
-    "import gym\n",
-    "import seals\n",
+    "import seals  # needed to load environments\n",
     "\n",
     "env = gym.make(\"seals/CartPole-v0\")\n",
     "expert = PPO(\n",
@@ -56,11 +56,14 @@
     "from imitation.data import rollout\n",
     "from imitation.data.wrappers import RolloutInfoWrapper\n",
     "from stable_baselines3.common.vec_env import DummyVecEnv\n",
+    "import numpy as np\n",
     "\n",
+    "random_state = np.random.RandomState()\n",
     "rollouts = rollout.rollout(\n",
     "    expert,\n",
     "    DummyVecEnv([lambda: RolloutInfoWrapper(gym.make(\"seals/CartPole-v0\"))] * 5),\n",
     "    rollout.make_sample_until(min_timesteps=None, min_episodes=60),\n",
+    "    random_state=random_state,\n",
     ")"
    ]
   },
@@ -84,10 +87,9 @@
     "from imitation.util.networks import RunningNorm\n",
     "from stable_baselines3 import PPO\n",
     "from stable_baselines3.common.evaluation import evaluate_policy\n",
-    "from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv\n",
+    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "\n",
     "import gym\n",
-    "import seals\n",
     "\n",
     "\n",
     "venv = DummyVecEnv([lambda: gym.make(\"seals/CartPole-v0\")] * 8)\n",
@@ -170,9 +172,8 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.10"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/examples/5_train_preference_comparisons.ipynb b/examples/5_train_preference_comparisons.ipynb
index b3c382f2b..a3a7da64a 100644
--- a/examples/5_train_preference_comparisons.ipynb
+++ b/examples/5_train_preference_comparisons.ipynb
@@ -22,6 +22,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import random\n",
     "from imitation.algorithms import preference_comparisons\n",
     "from imitation.rewards.reward_nets import BasicRewardNet\n",
     "from imitation.util.networks import RunningNorm\n",
@@ -29,15 +30,20 @@
     "from imitation.policies.base import FeedForward32Policy, NormalizeFeaturesExtractor\n",
     "import gym\n",
     "from stable_baselines3 import PPO\n",
+    "import numpy as np\n",
     "\n",
-    "venv = make_vec_env(\"Pendulum-v1\")\n",
+    "random_state = np.random.RandomState(0)\n",
+    "\n",
+    "venv = make_vec_env(\"Pendulum-v1\", random_state=random_state)\n",
     "\n",
     "reward_net = BasicRewardNet(\n",
     "    venv.observation_space, venv.action_space, normalize_input_layer=RunningNorm\n",
     ")\n",
     "\n",
-    "fragmenter = preference_comparisons.RandomFragmenter(warning_threshold=0, seed=0)\n",
-    "gatherer = preference_comparisons.SyntheticGatherer(seed=0)\n",
+    "fragmenter = preference_comparisons.RandomFragmenter(\n",
+    "    warning_threshold=0, random_state=random.Random(0)\n",
+    ")\n",
+    "gatherer = preference_comparisons.SyntheticGatherer(random_state=random_state)\n",
     "preference_model = preference_comparisons.PreferenceModel(reward_net)\n",
     "reward_trainer = preference_comparisons.BasicRewardTrainer(\n",
     "    model=reward_net,\n",
@@ -65,7 +71,7 @@
     "    reward_fn=reward_net,\n",
     "    venv=venv,\n",
     "    exploration_frac=0.0,\n",
-    "    seed=0,\n",
+    "    random_state=random_state,\n",
     ")\n",
     "\n",
     "pref_comparisons = preference_comparisons.PreferenceComparisons(\n",
@@ -79,7 +85,7 @@
     "    transition_oversampling=1,\n",
     "    initial_comparison_frac=0.1,\n",
     "    allow_variable_horizon=False,\n",
-    "    seed=0,\n",
+    "    random_state=random_state,\n",
     "    initial_epoch_multiplier=1,\n",
     ")"
    ]
@@ -195,4 +201,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/examples/6_train_mce.ipynb b/examples/6_train_mce.ipynb
index e3aeb87bc..ff693ee7a 100644
--- a/examples/6_train_mce.ipynb
+++ b/examples/6_train_mce.ipynb
@@ -23,20 +23,21 @@
    "outputs": [],
    "source": [
     "from functools import partial\n",
+    "\n",
+    "import numpy as np\n",
+    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
+    "\n",
     "from imitation.algorithms.mce_irl import (\n",
     "    MCEIRL,\n",
     "    mce_occupancy_measures,\n",
     "    mce_partition_fh,\n",
     "    TabularPolicy,\n",
     ")\n",
-    "\n",
     "from imitation.data import rollout\n",
     "from imitation.envs import resettable_env\n",
     "from imitation.envs.examples.model_envs import CliffWorld\n",
-    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "from imitation.rewards import reward_nets\n",
     "\n",
-    "\n",
     "env_creator = partial(CliffWorld, height=4, horizon=8, width=7, use_xy_obs=True)\n",
     "env_single = env_creator()\n",
     "\n",
@@ -61,17 +62,19 @@
     "\n",
     "_, om = mce_occupancy_measures(env_single, pi=pi)\n",
     "\n",
+    "random_state = np.random.RandomState()\n",
     "expert = TabularPolicy(\n",
     "    state_space=env_single.pomdp_state_space,\n",
     "    action_space=env_single.action_space,\n",
     "    pi=pi,\n",
-    "    rng=None,\n",
+    "    random_state=random_state,\n",
     ")\n",
     "\n",
     "expert_trajs = rollout.generate_trajectories(\n",
     "    policy=expert,\n",
     "    venv=state_venv,\n",
     "    sample_until=rollout.make_min_timesteps(5000),\n",
+    "    random_state=random_state,\n",
     ")\n",
     "\n",
     "print(\"Expert stats: \", rollout.rollout_stats(expert_trajs))"
@@ -107,7 +110,12 @@
     "    )\n",
     "\n",
     "    mce_irl = MCEIRL(\n",
-    "        demos, env_single, reward_net, log_interval=250, optimizer_kwargs=dict(lr=lr)\n",
+    "        demos,\n",
+    "        env_single,\n",
+    "        reward_net,\n",
+    "        log_interval=250,\n",
+    "        optimizer_kwargs=dict(lr=lr),\n",
+    "        random_state=random_state,\n",
     "    )\n",
     "    occ_measure = mce_irl.train(**kwargs)\n",
     "\n",
@@ -115,6 +123,7 @@
     "        policy=mce_irl.policy,\n",
     "        venv=state_venv,\n",
     "        sample_until=rollout.make_min_timesteps(5000),\n",
+    "        random_state=random_state,\n",
     "    )\n",
     "    print(\"Imitation stats: \", rollout.rollout_stats(imitation_trajs))\n",
     "\n",
@@ -234,4 +243,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/examples/7_train_density.ipynb b/examples/7_train_density.ipynb
index 2f7a7420c..00070c50a 100644
--- a/examples/7_train_density.ipynb
+++ b/examples/7_train_density.ipynb
@@ -30,12 +30,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "is_executing": false,
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Set FAST = False for longer training. Use True for testing and CI.\n",
@@ -57,24 +52,23 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "is_executing": false
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "from stable_baselines3.common.policies import ActorCriticPolicy\n",
     "from stable_baselines3 import PPO\n",
+    "import numpy as np\n",
     "\n",
+    "random_state = np.random.RandomState()\n",
     "env_name = \"Pendulum-v1\"\n",
-    "env = util.make_vec_env(env_name, N_VEC)\n",
+    "env = util.make_vec_env(env_name, random_state=random_state, n_envs=N_VEC)\n",
     "rollouts = types.load(\"../tests/testdata/expert_models/pendulum_0/rollouts/final.pkl\")\n",
     "\n",
     "\n",
     "imitation_trainer = PPO(ActorCriticPolicy, env, learning_rate=3e-4, n_steps=2048)\n",
     "density_trainer = db.DensityAlgorithm(\n",
     "    venv=env,\n",
+    "    random_state=random_state,\n",
     "    demonstrations=rollouts,\n",
     "    rl_algo=imitation_trainer,\n",
     "    density_type=db.DensityType.STATE_ACTION_DENSITY,\n",
@@ -89,12 +83,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "is_executing": false,
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "novice_stats = density_trainer.test_policy()\n",
@@ -148,4 +137,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file

From 6af437d5e2bbc45e60cf60d360b2977e242a4caf Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 01:22:14 +0200
Subject: [PATCH 075/187] Fix all tests.

---
 examples/5_train_preference_comparisons.ipynb |  3 +--
 src/imitation/scripts/common/common.py        |  8 +++----
 src/imitation/scripts/common/rl.py            | 11 +++++----
 src/imitation/scripts/common/seeding.py       | 24 +++++++++++++++++++
 src/imitation/scripts/common/train.py         |  6 ++++-
 src/imitation/scripts/config/eval_policy.py   |  4 ++--
 .../config/train_preference_comparisons.py    |  3 ++-
 src/imitation/scripts/config/train_rl.py      |  4 ++--
 src/imitation/scripts/eval_policy.py          |  8 +++----
 src/imitation/scripts/train_adversarial.py    |  2 --
 src/imitation/scripts/train_imitation.py      |  5 +++-
 .../scripts/train_preference_comparisons.py   |  8 +++----
 src/imitation/scripts/train_rl.py             |  6 +++--
 tests/algorithms/test_mce_irl.py              | 10 ++++----
 .../algorithms/test_preference_comparisons.py |  4 +++-
 15 files changed, 69 insertions(+), 37 deletions(-)
 create mode 100644 src/imitation/scripts/common/seeding.py

diff --git a/examples/5_train_preference_comparisons.ipynb b/examples/5_train_preference_comparisons.ipynb
index a3a7da64a..511b29bad 100644
--- a/examples/5_train_preference_comparisons.ipynb
+++ b/examples/5_train_preference_comparisons.ipynb
@@ -85,7 +85,6 @@
     "    transition_oversampling=1,\n",
     "    initial_comparison_frac=0.1,\n",
     "    allow_variable_horizon=False,\n",
-    "    random_state=random_state,\n",
     "    initial_epoch_multiplier=1,\n",
     ")"
    ]
@@ -201,4 +200,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index b7c164eb3..eb8227ba4 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -9,12 +9,12 @@
 import sacred
 from stable_baselines3.common import vec_env
 
-from imitation.scripts.common import wb
+from imitation.scripts.common import wb, seeding
 from imitation.util import logger as imit_logger
 from imitation.util import sacred as sacred_util
 from imitation.util import util
 
-common_ingredient = sacred.Ingredient("common", ingredients=[wb.wandb_ingredient])
+common_ingredient = sacred.Ingredient("common", ingredients=[wb.wandb_ingredient, seeding.seeding_ingredient])
 logger = logging.getLogger(__name__)
 
 
@@ -131,7 +131,6 @@ def setup_logging(
 @contextlib.contextmanager
 @common_ingredient.capture
 def make_venv(
-    random_state: np.random.RandomState,
     env_name: str,
     num_vec: int,
     parallel: bool,
@@ -157,6 +156,7 @@ def make_venv(
     Yields:
         The constructed vector environment.
     """
+    random_state = seeding.make_random_state()
     try:
         venv = util.make_vec_env(
             env_name,
@@ -170,4 +170,4 @@ def make_venv(
         )
         yield venv
     finally:
-        venv.close()
+        venv.close()
\ No newline at end of file
diff --git a/src/imitation/scripts/common/rl.py b/src/imitation/scripts/common/rl.py
index 1551d20cf..8c28640a9 100644
--- a/src/imitation/scripts/common/rl.py
+++ b/src/imitation/scripts/common/rl.py
@@ -17,9 +17,10 @@
 from imitation.policies import serialize
 from imitation.policies.replay_buffer_wrapper import ReplayBufferRewardWrapper
 from imitation.rewards.reward_function import RewardFn
+from imitation.scripts.common import seeding
 from imitation.scripts.common.train import train_ingredient
 
-rl_ingredient = sacred.Ingredient("rl", ingredients=[train_ingredient])
+rl_ingredient = sacred.Ingredient("rl", ingredients=[train_ingredient, seeding.seeding_ingredient])
 logger = logging.getLogger(__name__)
 
 
@@ -102,7 +103,6 @@ def make_rl_algo(
     batch_size: int,
     rl_kwargs: Mapping[str, Any],
     train: Mapping[str, Any],
-    _seed: int,
     relabel_reward_fn: Optional[RewardFn] = None,
 ) -> base_class.BaseAlgorithm:
     """Instantiates a Stable Baselines3 RL algorithm.
@@ -124,6 +124,7 @@ def make_rl_algo(
         ValueError: `gen_batch_size` not divisible by `venv.num_envs`.
         TypeError: `rl_cls` is neither `OnPolicyAlgorithm` nor `OffPolicyAlgorithm`.
     """
+    seed = seeding.get_seed()
     if batch_size % venv.num_envs != 0:
         raise ValueError(
             f"num_envs={venv.num_envs} must evenly divide batch_size={batch_size}.",
@@ -157,7 +158,7 @@ def make_rl_algo(
         # https://github.com/DLR-RM/stable-baselines3/blob/30772aa9f53a4cf61571ee90046cdc454c1b11d7/sb3/common/off_policy_algorithm.py#L145
         policy_kwargs=dict(train["policy_kwargs"]),
         env=venv,
-        seed=_seed,
+        seed=seed,
         **rl_kwargs,
     )
     logger.info(f"RL algorithm: {type(rl_algo)}")
@@ -171,9 +172,9 @@ def load_rl_algo_from_path(
     venv: vec_env.VecEnv,
     rl_cls: Type[base_class.BaseAlgorithm],
     rl_kwargs: Mapping[str, Any],
-    _seed: int,
     relabel_reward_fn: Optional[RewardFn] = None,
 ) -> base_class.BaseAlgorithm:
+    seed = seeding.get_seed()
     rl_kwargs = dict(rl_kwargs)
     if issubclass(rl_cls, off_policy_algorithm.OffPolicyAlgorithm):
         rl_kwargs = _maybe_add_relabel_buffer(
@@ -184,7 +185,7 @@ def load_rl_algo_from_path(
         cls=rl_cls,
         path=agent_path,
         venv=venv,
-        seed=_seed,
+        seed=seed,
         **rl_kwargs,
     )
     logger.info(f"Warm starting agent from '{agent_path}'")
diff --git a/src/imitation/scripts/common/seeding.py b/src/imitation/scripts/common/seeding.py
new file mode 100644
index 000000000..0a2733f0c
--- /dev/null
+++ b/src/imitation/scripts/common/seeding.py
@@ -0,0 +1,24 @@
+from typing import Optional
+
+import numpy as np
+import sacred
+
+seeding_ingredient = sacred.Ingredient("seeding")
+
+
+@seeding_ingredient.config
+def config():
+    seed = 0
+    locals()  # quieten flake8
+
+
+@seeding_ingredient.capture
+def get_seed(seed: Optional[int]) -> Optional[int]:
+    """Returns the seed used by the seeding ingredient."""
+    return seed
+
+
+@seeding_ingredient.capture
+def make_random_state(seed: Optional[int]) -> np.random.RandomState:
+    """Creates a `np.random.RandomState` with the given seed."""
+    return np.random.RandomState(seed)
diff --git a/src/imitation/scripts/common/train.py b/src/imitation/scripts/common/train.py
index 88413fdf5..fa6002373 100644
--- a/src/imitation/scripts/common/train.py
+++ b/src/imitation/scripts/common/train.py
@@ -3,14 +3,16 @@
 import logging
 from typing import Any, Mapping, Union
 
+import numpy as np
 import sacred
 from stable_baselines3.common import base_class, policies, torch_layers, vec_env
 
 import imitation.util.networks
 from imitation.data import rollout
 from imitation.policies import base
+from imitation.scripts.common import seeding
 
-train_ingredient = sacred.Ingredient("train")
+train_ingredient = sacred.Ingredient("train", ingredients=[seeding.seeding_ingredient])
 logger = logging.getLogger(__name__)
 
 
@@ -89,11 +91,13 @@ def eval_policy(
         "monitor_return" key). "expert_stats" gives the return value of
         `rollout_stats()` on the expert demonstrations loaded from `rollout_path`.
     """
+    random_state = seeding.make_random_state()
     sample_until_eval = rollout.make_min_episodes(n_episodes_eval)
     trajs = rollout.generate_trajectories(
         rl_algo,
         venv,
         sample_until=sample_until_eval,
+        random_state=random_state,
     )
     return rollout.rollout_stats(trajs)
 
diff --git a/src/imitation/scripts/config/eval_policy.py b/src/imitation/scripts/config/eval_policy.py
index adde3bb53..4cd378a7a 100644
--- a/src/imitation/scripts/config/eval_policy.py
+++ b/src/imitation/scripts/config/eval_policy.py
@@ -2,11 +2,11 @@
 
 import sacred
 
-from imitation.scripts.common import common
+from imitation.scripts.common import common, seeding
 
 eval_policy_ex = sacred.Experiment(
     "eval_policy",
-    ingredients=[common.common_ingredient],
+    ingredients=[common.common_ingredient, seeding.seeding_ingredient],
 )
 
 
diff --git a/src/imitation/scripts/config/train_preference_comparisons.py b/src/imitation/scripts/config/train_preference_comparisons.py
index ba4e9483c..7792de8d9 100644
--- a/src/imitation/scripts/config/train_preference_comparisons.py
+++ b/src/imitation/scripts/config/train_preference_comparisons.py
@@ -3,7 +3,7 @@
 import sacred
 
 from imitation.algorithms import preference_comparisons
-from imitation.scripts.common import common, reward, rl, train
+from imitation.scripts.common import common, reward, rl, train, seeding
 
 train_preference_comparisons_ex = sacred.Experiment(
     "train_preference_comparisons",
@@ -12,6 +12,7 @@
         reward.reward_ingredient,
         rl.rl_ingredient,
         train.train_ingredient,
+        seeding.seeding_ingredient,
     ],
 )
 
diff --git a/src/imitation/scripts/config/train_rl.py b/src/imitation/scripts/config/train_rl.py
index ae3add76f..13edf0160 100644
--- a/src/imitation/scripts/config/train_rl.py
+++ b/src/imitation/scripts/config/train_rl.py
@@ -2,11 +2,11 @@
 
 import sacred
 
-from imitation.scripts.common import common, rl, train
+from imitation.scripts.common import common, rl, train, seeding
 
 train_rl_ex = sacred.Experiment(
     "train_rl",
-    ingredients=[common.common_ingredient, train.train_ingredient, rl.rl_ingredient],
+    ingredients=[common.common_ingredient, train.train_ingredient, rl.rl_ingredient, seeding.seeding_ingredient],
 )
 
 
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 47e6220f3..2ef2dac72 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -7,6 +7,7 @@
 from typing import Any, Mapping, Optional
 
 import gym
+import numpy as np
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common.vec_env import VecEnvWrapper
 
@@ -14,7 +15,7 @@
 from imitation.policies import serialize
 from imitation.rewards import reward_wrapper
 from imitation.rewards.serialize import load_reward
-from imitation.scripts.common import common
+from imitation.scripts.common import common, seeding
 from imitation.scripts.config.eval_policy import eval_policy_ex
 from imitation.util import video_wrapper
 
@@ -54,7 +55,6 @@ def f(env: gym.Env, i: int) -> gym.Env:
 @eval_policy_ex.main
 def eval_policy(
     _run,
-    _seed: int,
     eval_n_timesteps: Optional[int],
     eval_n_episodes: Optional[int],
     render: bool,
@@ -70,7 +70,6 @@ def eval_policy(
     """Rolls a policy out in an environment, collecting statistics.
 
     Args:
-        _seed: generated by Sacred.
         eval_n_timesteps: Minimum number of timesteps to evaluate for. Set exactly
             one of `eval_n_episodes` and `eval_n_timesteps`.
         eval_n_episodes: Minimum number of episodes to evaluate for. Set exactly
@@ -92,6 +91,7 @@ def eval_policy(
     Returns:
         Return value of `imitation.util.rollout.rollout_stats()`.
     """
+    random_state = seeding.make_random_state()
     log_dir = common.make_log_dir()
     sample_until = rollout.make_sample_until(eval_n_timesteps, eval_n_episodes)
     post_wrappers = [video_wrapper_factory(log_dir, **video_kwargs)] if videos else None
@@ -108,7 +108,7 @@ def eval_policy(
         if policy_type is not None:
             assert policy_path is not None
             policy = serialize.load_policy(policy_type, policy_path, venv)
-        trajs = rollout.generate_trajectories(policy, venv, sample_until)
+        trajs = rollout.generate_trajectories(policy, venv, sample_until, random_state=random_state)
 
     if rollout_save_path:
         types.save(rollout_save_path.replace("{log_dir}", log_dir), trajs)
diff --git a/src/imitation/scripts/train_adversarial.py b/src/imitation/scripts/train_adversarial.py
index 35320e12c..f4adc4be6 100644
--- a/src/imitation/scripts/train_adversarial.py
+++ b/src/imitation/scripts/train_adversarial.py
@@ -67,7 +67,6 @@ def dummy_config():
 @train_adversarial_ex.capture
 def train_adversarial(
     _run,
-    _seed: int,
     show_config: bool,
     algo_cls: Type[common.AdversarialTrainer],
     algorithm_kwargs: Mapping[str, Any],
@@ -84,7 +83,6 @@ def train_adversarial(
         - Generator policies are saved to `f"{log_dir}/checkpoints/{step}/gen_policy/"`.
 
     Args:
-        _seed: Random seed.
         show_config: Print the merged config before starting training. This is
             analogous to the print_config command, but will show config after
             rather than before merging `algorithm_specific` arguments.
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index 40341937e..1e28d2d51 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -5,6 +5,7 @@
 import warnings
 from typing import Any, Mapping, Optional, Sequence, Type, cast
 
+import numpy as np
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import policies, utils, vec_env
 
@@ -12,7 +13,7 @@
 from imitation.algorithms.dagger import SimpleDAggerTrainer
 from imitation.data import rollout, types
 from imitation.policies import serialize
-from imitation.scripts.common import common, demonstrations, train
+from imitation.scripts.common import common, demonstrations, train, seeding
 from imitation.scripts.config.train_imitation import train_imitation_ex
 
 logger = logging.getLogger(__name__)
@@ -123,6 +124,7 @@ def train_imitation(
     Returns:
         Statistics for rollouts from the trained policy and demonstration data.
     """
+    random_state = seeding.make_random_state()
     custom_logger, log_dir = common.setup_logging()
 
     with common.make_venv() as venv:
@@ -138,6 +140,7 @@ def train_imitation(
             policy=imit_policy,
             demonstrations=expert_trajs,
             custom_logger=custom_logger,
+            random_state=random_state,
             **bc_kwargs,
         )
         bc_train_kwargs = dict(log_rollouts_venv=venv, **bc_train_kwargs)
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 2f2e89759..61025f8d0 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -17,7 +17,7 @@
 from imitation.algorithms import preference_comparisons
 from imitation.data import types
 from imitation.policies import serialize
-from imitation.scripts.common import common, reward
+from imitation.scripts.common import common, reward, seeding
 from imitation.scripts.common import rl as rl_common
 from imitation.scripts.common import train
 from imitation.scripts.config.train_preference_comparisons import (
@@ -59,7 +59,6 @@ def save_checkpoint(
 
 @train_preference_comparisons_ex.main
 def train_preference_comparisons(
-    random_state: np.random.RandomState,
     total_timesteps: int,
     total_comparisons: int,
     num_iterations: int,
@@ -87,7 +86,6 @@ def train_preference_comparisons(
     """Train a reward model using preference comparisons.
 
     Args:
-        random_state: Random state for the internal random number generation.
         total_timesteps: number of environment interaction steps
         total_comparisons: number of preferences to gather in total
         num_iterations: number of times to train the agent against the reward model
@@ -151,7 +149,8 @@ def train_preference_comparisons(
         ValueError: Inconsistency between config and deserialized policy normalization.
     """
     custom_logger, log_dir = common.setup_logging()
-    seed = util.make_seeds(random_state)
+    random_state = seeding.make_random_state()
+    seed = seeding.get_seed()
 
     with common.make_venv() as venv:
         reward_net = reward.make_reward_net(venv)
@@ -244,7 +243,6 @@ def train_preference_comparisons(
             initial_comparison_frac=initial_comparison_frac,
             custom_logger=custom_logger,
             allow_variable_horizon=allow_variable_horizon,
-            random_state=random_state,
             query_schedule=query_schedule,
         )
 
diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
index b36f9518b..572737f70 100644
--- a/src/imitation/scripts/train_rl.py
+++ b/src/imitation/scripts/train_rl.py
@@ -14,6 +14,7 @@
 import warnings
 from typing import Any, Mapping, Optional
 
+import numpy as np
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import callbacks
 from stable_baselines3.common.vec_env import VecNormalize
@@ -22,7 +23,7 @@
 from imitation.policies import serialize
 from imitation.rewards.reward_wrapper import RewardVecEnvWrapper
 from imitation.rewards.serialize import load_reward
-from imitation.scripts.common import common, rl, train
+from imitation.scripts.common import common, rl, train, seeding
 from imitation.scripts.config.train_rl import train_rl_ex
 
 
@@ -87,6 +88,7 @@ def train_rl(
     Returns:
         The return value of `rollout_stats()` using the final policy.
     """
+    random_state = seeding.make_random_state()
     custom_logger, log_dir = common.setup_logging()
     rollout_dir = osp.join(log_dir, "rollouts")
     policy_dir = osp.join(log_dir, "policies")
@@ -144,7 +146,7 @@ def train_rl(
                 rollout_save_n_timesteps,
                 rollout_save_n_episodes,
             )
-            types.save(save_path, rollout.rollout(rl_algo, venv, sample_until))
+            types.save(save_path, rollout.rollout(rl_algo, venv, sample_until, random_state))
         if policy_save_final:
             output_dir = os.path.join(policy_dir, "final")
             serialize.save_stable_model(output_dir, rl_algo)
diff --git a/tests/algorithms/test_mce_irl.py b/tests/algorithms/test_mce_irl.py
index fb374bfaf..8f5ec705a 100644
--- a/tests/algorithms/test_mce_irl.py
+++ b/tests/algorithms/test_mce_irl.py
@@ -240,12 +240,12 @@ def test_tabular_policy():
     pi = np.stack(
         [np.eye(2), 1 - np.eye(2)],
     )
-    rng = np.random.RandomState(42)
+    random_state = np.random.RandomState(42)
     tabular = TabularPolicy(
         state_space=state_space,
         action_space=action_space,
         pi=pi,
-        rng=rng,
+        random_state=random_state,
     )
 
     states = np.array([0, 1, 1, 0, 1])
@@ -297,8 +297,8 @@ def test_tabular_policy_randomness(random_state_fixed):
     np.testing.assert_equal(actions, 0)
 
 
-def test_mce_irl_demo_formats(fixed_random_state):
-    random_state = fixed_random_state
+def test_mce_irl_demo_formats(random_state_fixed):
+    random_state = random_state_fixed
     mdp = model_envs.RandomMDP(
         n_states=5,
         n_actions=3,
@@ -339,7 +339,7 @@ def test_mce_irl_demo_formats(fixed_random_state):
                 use_done=False,
                 hid_sizes=[],
             )
-            mce_irl = MCEIRL(demo, mdp, reward_net, linf_eps=1e-3)
+            mce_irl = MCEIRL(demo, mdp, reward_net, linf_eps=1e-3, random_state=random_state)
             assert np.allclose(mce_irl.demo_state_om.sum(), mdp.horizon + 1)
             final_counts[kind] = mce_irl.train(max_iter=5)
 
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 66c80653d..220aaf561 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -621,13 +621,14 @@ def test_agent_trainer_sample(venv, agent_trainer):
     )
 
 
-def test_agent_trainer_sample_image_observations():
+def test_agent_trainer_sample_image_observations(random_state_fixed):
     """Test `AgentTrainer.sample()` in an image environment.
 
     SB3 algorithms may rearrange the channel dimension in environments with image
     observations, but `sample()` should return observations matching the original
     environment.
     """
+    random_state = random_state_fixed
     venv = DummyVecEnv([lambda: FakeImageEnv()])
     reward_net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     agent = stable_baselines3.PPO(
@@ -642,6 +643,7 @@ def test_agent_trainer_sample_image_observations():
         reward_net,
         venv,
         exploration_frac=0.5,
+        random_state=random_state,
     )
     trajectories = agent_trainer.sample(2)
     assert len(trajectories) > 0

From 852cfb1fe4191360822e684c957acbb90c0b4262 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 01:22:25 +0200
Subject: [PATCH 076/187] Formattting.

---
 src/imitation/scripts/common/common.py           |  8 +++++---
 src/imitation/scripts/common/rl.py               |  4 +++-
 .../config/train_preference_comparisons.py       |  2 +-
 src/imitation/scripts/config/train_rl.py         |  9 +++++++--
 src/imitation/scripts/eval_policy.py             |  4 +++-
 src/imitation/scripts/train_imitation.py         |  2 +-
 .../scripts/train_preference_comparisons.py      |  4 ++--
 src/imitation/scripts/train_rl.py                |  6 ++++--
 tests/algorithms/test_mce_irl.py                 |  4 +++-
 tests/algorithms/test_preference_comparisons.py  | 16 ++++++++++++----
 tests/rewards/test_reward_nets.py                |  4 +++-
 11 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index eb8227ba4..7a9ff7ddf 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -9,12 +9,14 @@
 import sacred
 from stable_baselines3.common import vec_env
 
-from imitation.scripts.common import wb, seeding
+from imitation.scripts.common import seeding, wb
 from imitation.util import logger as imit_logger
 from imitation.util import sacred as sacred_util
 from imitation.util import util
 
-common_ingredient = sacred.Ingredient("common", ingredients=[wb.wandb_ingredient, seeding.seeding_ingredient])
+common_ingredient = sacred.Ingredient(
+    "common", ingredients=[wb.wandb_ingredient, seeding.seeding_ingredient]
+)
 logger = logging.getLogger(__name__)
 
 
@@ -170,4 +172,4 @@ def make_venv(
         )
         yield venv
     finally:
-        venv.close()
\ No newline at end of file
+        venv.close()
diff --git a/src/imitation/scripts/common/rl.py b/src/imitation/scripts/common/rl.py
index 8c28640a9..46801b1b9 100644
--- a/src/imitation/scripts/common/rl.py
+++ b/src/imitation/scripts/common/rl.py
@@ -20,7 +20,9 @@
 from imitation.scripts.common import seeding
 from imitation.scripts.common.train import train_ingredient
 
-rl_ingredient = sacred.Ingredient("rl", ingredients=[train_ingredient, seeding.seeding_ingredient])
+rl_ingredient = sacred.Ingredient(
+    "rl", ingredients=[train_ingredient, seeding.seeding_ingredient]
+)
 logger = logging.getLogger(__name__)
 
 
diff --git a/src/imitation/scripts/config/train_preference_comparisons.py b/src/imitation/scripts/config/train_preference_comparisons.py
index 7792de8d9..64ee5bb07 100644
--- a/src/imitation/scripts/config/train_preference_comparisons.py
+++ b/src/imitation/scripts/config/train_preference_comparisons.py
@@ -3,7 +3,7 @@
 import sacred
 
 from imitation.algorithms import preference_comparisons
-from imitation.scripts.common import common, reward, rl, train, seeding
+from imitation.scripts.common import common, reward, rl, seeding, train
 
 train_preference_comparisons_ex = sacred.Experiment(
     "train_preference_comparisons",
diff --git a/src/imitation/scripts/config/train_rl.py b/src/imitation/scripts/config/train_rl.py
index 13edf0160..e71d7b14f 100644
--- a/src/imitation/scripts/config/train_rl.py
+++ b/src/imitation/scripts/config/train_rl.py
@@ -2,11 +2,16 @@
 
 import sacred
 
-from imitation.scripts.common import common, rl, train, seeding
+from imitation.scripts.common import common, rl, seeding, train
 
 train_rl_ex = sacred.Experiment(
     "train_rl",
-    ingredients=[common.common_ingredient, train.train_ingredient, rl.rl_ingredient, seeding.seeding_ingredient],
+    ingredients=[
+        common.common_ingredient,
+        train.train_ingredient,
+        rl.rl_ingredient,
+        seeding.seeding_ingredient,
+    ],
 )
 
 
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 2ef2dac72..70015ce75 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -108,7 +108,9 @@ def eval_policy(
         if policy_type is not None:
             assert policy_path is not None
             policy = serialize.load_policy(policy_type, policy_path, venv)
-        trajs = rollout.generate_trajectories(policy, venv, sample_until, random_state=random_state)
+        trajs = rollout.generate_trajectories(
+            policy, venv, sample_until, random_state=random_state
+        )
 
     if rollout_save_path:
         types.save(rollout_save_path.replace("{log_dir}", log_dir), trajs)
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index 1e28d2d51..ec1cdeebf 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -13,7 +13,7 @@
 from imitation.algorithms.dagger import SimpleDAggerTrainer
 from imitation.data import rollout, types
 from imitation.policies import serialize
-from imitation.scripts.common import common, demonstrations, train, seeding
+from imitation.scripts.common import common, demonstrations, seeding, train
 from imitation.scripts.config.train_imitation import train_imitation_ex
 
 logger = logging.getLogger(__name__)
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 61025f8d0..7b0866e60 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -17,9 +17,9 @@
 from imitation.algorithms import preference_comparisons
 from imitation.data import types
 from imitation.policies import serialize
-from imitation.scripts.common import common, reward, seeding
+from imitation.scripts.common import common, reward
 from imitation.scripts.common import rl as rl_common
-from imitation.scripts.common import train
+from imitation.scripts.common import seeding, train
 from imitation.scripts.config.train_preference_comparisons import (
     train_preference_comparisons_ex,
 )
diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
index 572737f70..1d17d5d7b 100644
--- a/src/imitation/scripts/train_rl.py
+++ b/src/imitation/scripts/train_rl.py
@@ -23,7 +23,7 @@
 from imitation.policies import serialize
 from imitation.rewards.reward_wrapper import RewardVecEnvWrapper
 from imitation.rewards.serialize import load_reward
-from imitation.scripts.common import common, rl, train, seeding
+from imitation.scripts.common import common, rl, seeding, train
 from imitation.scripts.config.train_rl import train_rl_ex
 
 
@@ -146,7 +146,9 @@ def train_rl(
                 rollout_save_n_timesteps,
                 rollout_save_n_episodes,
             )
-            types.save(save_path, rollout.rollout(rl_algo, venv, sample_until, random_state))
+            types.save(
+                save_path, rollout.rollout(rl_algo, venv, sample_until, random_state)
+            )
         if policy_save_final:
             output_dir = os.path.join(policy_dir, "final")
             serialize.save_stable_model(output_dir, rl_algo)
diff --git a/tests/algorithms/test_mce_irl.py b/tests/algorithms/test_mce_irl.py
index 8f5ec705a..73a70786d 100644
--- a/tests/algorithms/test_mce_irl.py
+++ b/tests/algorithms/test_mce_irl.py
@@ -339,7 +339,9 @@ def test_mce_irl_demo_formats(random_state_fixed):
                 use_done=False,
                 hid_sizes=[],
             )
-            mce_irl = MCEIRL(demo, mdp, reward_net, linf_eps=1e-3, random_state=random_state)
+            mce_irl = MCEIRL(
+                demo, mdp, reward_net, linf_eps=1e-3, random_state=random_state
+            )
             assert np.allclose(mce_irl.demo_state_om.sum(), mdp.horizon + 1)
             final_counts[kind] = mce_irl.train(max_iter=5)
 
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 220aaf561..6e488187f 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -99,7 +99,9 @@ def test_mismatched_spaces(venv, agent, random_state_fixed):
         ValueError,
         match="spaces do not match",
     ):
-        preference_comparisons.AgentTrainer(agent, bad_reward_net, venv, random_state=random_state)
+        preference_comparisons.AgentTrainer(
+            agent, bad_reward_net, venv, random_state=random_state
+        )
 
 
 def test_trajectory_dataset_seeding(
@@ -334,9 +336,13 @@ def test_discount_rate_no_crash(
     main_trainer.train(100, 10)
 
 
-def test_synthetic_gatherer_deterministic(agent_trainer, random_fragmenter, random_state_fixed):
+def test_synthetic_gatherer_deterministic(
+    agent_trainer, random_fragmenter, random_state_fixed
+):
     random_state = random_state_fixed
-    gatherer = preference_comparisons.SyntheticGatherer(temperature=0, random_state=random_state)
+    gatherer = preference_comparisons.SyntheticGatherer(
+        temperature=0, random_state=random_state
+    )
     trajectories = agent_trainer.sample(10)
     fragments = random_fragmenter(trajectories, fragment_length=2, num_pairs=2)
     preferences1 = gatherer(fragments)
@@ -414,7 +420,9 @@ def test_preference_dataset_queue(agent_trainer, random_fragmenter, random_state
     assert len(dataset) == 5
 
 
-def test_store_and_load_preference_dataset(agent_trainer, random_fragmenter, tmp_path, random_state_fixed):
+def test_store_and_load_preference_dataset(
+    agent_trainer, random_fragmenter, tmp_path, random_state_fixed
+):
     random_state = random_state_fixed
     dataset = preference_comparisons.PreferenceDataset()
     trajectories = agent_trainer.sample(10)
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index c1877d588..7e96b40f4 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -161,7 +161,9 @@ def test_reward_valid(env_name, reward_type, tmpdir, random_state_fixed):
     """Test output of reward function is appropriate shape and type."""
     random_state = random_state_fixed
     # TODO(juan) the line below is not being used?
-    venv = util.make_vec_env(env_name, n_envs=1, parallel=False, random_state=random_state)
+    venv = util.make_vec_env(
+        env_name, n_envs=1, parallel=False, random_state=random_state
+    )
     venv, tmppath = _make_env_and_save_reward_net(
         env_name, reward_type, tmpdir, random_state
     )

From 9e1622d4e8bf5158e73dee0bb745bd8bf3ffa60c Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 01:35:24 +0200
Subject: [PATCH 077/187] Additional fixes

---
 examples/2_train_dagger.ipynb            | 12 ++++++++++--
 src/imitation/algorithms/dagger.py       | 14 ++++++++++++--
 src/imitation/scripts/train_imitation.py |  1 +
 tests/algorithms/test_dagger.py          | 11 +++++++++--
 4 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/examples/2_train_dagger.ipynb b/examples/2_train_dagger.ipynb
index 1920fc8af..6bab7af54 100644
--- a/examples/2_train_dagger.ipynb
+++ b/examples/2_train_dagger.ipynb
@@ -53,7 +53,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "import tempfile\n",
@@ -76,7 +80,11 @@
     "with tempfile.TemporaryDirectory(prefix=\"dagger_example_\") as tmpdir:\n",
     "    print(tmpdir)\n",
     "    dagger_trainer = SimpleDAggerTrainer(\n",
-    "        venv=venv, scratch_dir=tmpdir, expert_policy=expert, bc_trainer=bc_trainer\n",
+    "        venv=venv,\n",
+    "        scratch_dir=tmpdir,\n",
+    "        expert_policy=expert,\n",
+    "        bc_trainer=bc_trainer,\n",
+    "        random_state=np.random.RandomState(),\n",
     "    )\n",
     "\n",
     "    dagger_trainer.train(2000)"
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index 56b5ca8f9..e95067956 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -152,6 +152,7 @@ def __init__(
         get_robot_acts: Callable[[np.ndarray], np.ndarray],
         beta: float,
         save_dir: types.AnyPath,
+        random_state: np.random.RandomState,
     ):
         """Builds InteractiveTrajectoryCollector.
 
@@ -175,7 +176,7 @@ def __init__(
         self._done_before = True
         self._is_reset = False
         self._last_user_actions = None
-        self.rng = np.random.RandomState()
+        self.rng = random_state
 
     def seed(self, seed=Optional[int]) -> List[Union[None, int]]:
         """Set the seed for the DAgger random number generator and wrapped VecEnv.
@@ -306,6 +307,7 @@ def __init__(
         *,
         venv: vec_env.VecEnv,
         scratch_dir: types.AnyPath,
+        random_state: np.random.RandomState,
         beta_schedule: Optional[Callable[[int], float]] = None,
         bc_trainer: bc.BC,
         custom_logger: Optional[logger.HierarchicalLogger] = None,
@@ -332,6 +334,7 @@ def __init__(
         self.round_num = 0
         self._last_loaded_round = -1
         self._all_demos = []
+        self.random_state = random_state
 
         utils.check_for_correct_spaces(
             self.venv,
@@ -473,6 +476,7 @@ def create_trajectory_collector(self) -> InteractiveTrajectoryCollector:
             get_robot_acts=lambda acts: self.bc_trainer.policy.predict(acts)[0],
             beta=beta,
             save_dir=save_dir,
+            random_state=self.random_state,
         )
         return collector
 
@@ -527,6 +531,7 @@ def __init__(
         venv: vec_env.VecEnv,
         scratch_dir: types.AnyPath,
         expert_policy: policies.BasePolicy,
+        random_state: np.random.RandomState,
         expert_trajs: Optional[Sequence[types.Trajectory]] = None,
         **dagger_trainer_kwargs,
     ):
@@ -549,7 +554,12 @@ def __init__(
             ValueError: The observation or action space does not match between
                 `venv` and `expert_policy`.
         """
-        super().__init__(venv=venv, scratch_dir=scratch_dir, **dagger_trainer_kwargs)
+        super().__init__(
+            venv=venv,
+            scratch_dir=scratch_dir,
+            random_state=random_state,
+            **dagger_trainer_kwargs,
+        )
         self.expert_policy = expert_policy
         if expert_policy.observation_space != self.venv.observation_space:
             raise ValueError(
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index ec1cdeebf..ed799ed3f 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -159,6 +159,7 @@ def train_imitation(
                 expert_policy=expert_policy,
                 custom_logger=custom_logger,
                 bc_trainer=bc_trainer,
+                random_state=random_state,
             )
             model.train(
                 total_timesteps=int(dagger["total_timesteps"]),
diff --git a/tests/algorithms/test_dagger.py b/tests/algorithms/test_dagger.py
index 5a54dc9ac..6b80e18a9 100644
--- a/tests/algorithms/test_dagger.py
+++ b/tests/algorithms/test_dagger.py
@@ -41,7 +41,8 @@ def test_beta_schedule():
         assert np.allclose(three_step_sched(i), (3 - i) / 3 if i <= 2 else 0)
 
 
-def test_traj_collector_seed(tmpdir, pendulum_venv):
+def test_traj_collector_seed(tmpdir, pendulum_venv, random_state_fixed):
+    random_state = random_state_fixed
     collector = dagger.InteractiveTrajectoryCollector(
         venv=pendulum_venv,
         get_robot_acts=lambda o: [
@@ -49,6 +50,7 @@ def test_traj_collector_seed(tmpdir, pendulum_venv):
         ],
         beta=0.5,
         save_dir=tmpdir,
+        random_state=random_state,
     )
     seeds1 = collector.seed(42)
     obs1 = collector.reset()
@@ -59,7 +61,8 @@ def test_traj_collector_seed(tmpdir, pendulum_venv):
     np.testing.assert_array_equal(obs1, obs2)
 
 
-def test_traj_collector(tmpdir, pendulum_venv):
+def test_traj_collector(tmpdir, pendulum_venv, random_state_fixed):
+    random_state = random_state_fixed
     robot_calls = 0
     num_episodes = 0
 
@@ -73,6 +76,7 @@ def get_random_acts(obs):
         get_robot_acts=get_random_acts,
         beta=0.5,
         save_dir=tmpdir,
+        random_state=random_state,
     )
     collector.reset()
     zero_acts = np.zeros(
@@ -131,6 +135,7 @@ def _build_dagger_trainer(
         beta_schedule=beta_schedule,
         bc_trainer=bc_trainer,
         custom_logger=custom_logger,
+        random_state=random_state,
     )
 
 
@@ -158,6 +163,7 @@ def _build_simple_dagger_trainer(
         expert_policy=expert_policy,
         expert_trajs=pendulum_expert_rollouts,
         custom_logger=custom_logger,
+        random_state=random_state,
     )
 
 
@@ -421,6 +427,7 @@ def test_dagger_not_enough_transitions_error(tmpdir, custom_logger, random_state
         scratch_dir=tmpdir,
         bc_trainer=bc_trainer,
         custom_logger=custom_logger,
+        random_state=random_state,
     )
     collector = trainer.create_trajectory_collector()
     policy = base.RandomPolicy(venv.observation_space, venv.action_space)

From 2e8b1f06c0153a6a03956dcfb27ec3fbbc1d2202 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 01:53:32 +0200
Subject: [PATCH 078/187] Linting

---
 src/imitation/algorithms/bc.py                |  4 +-
 src/imitation/algorithms/dagger.py            |  2 +
 src/imitation/algorithms/mce_irl.py           |  2 +-
 .../algorithms/preference_comparisons.py      | 15 ++++---
 src/imitation/data/rollout.py                 |  6 ++-
 src/imitation/policies/exploration_wrapper.py |  2 -
 src/imitation/scripts/common/common.py        |  4 +-
 src/imitation/scripts/common/rl.py            |  3 +-
 src/imitation/scripts/common/seeding.py       |  2 +
 src/imitation/scripts/common/train.py         |  1 -
 src/imitation/scripts/eval_policy.py          |  6 ++-
 src/imitation/scripts/train_imitation.py      |  1 -
 .../scripts/train_preference_comparisons.py   |  2 -
 src/imitation/scripts/train_rl.py             |  4 +-
 src/imitation/util/util.py                    |  4 +-
 tests/algorithms/test_adversarial.py          | 26 +++++++++---
 tests/algorithms/test_bc.py                   |  1 +
 tests/algorithms/test_dagger.py               |  5 ++-
 tests/algorithms/test_mce_irl.py              |  6 ++-
 .../algorithms/test_preference_comparisons.py | 26 +++++++++---
 tests/data/test_rollout.py                    | 16 ++++++--
 tests/policies/test_policies.py               | 24 ++++++++---
 tests/rewards/test_reward_nets.py             | 40 +++++++++++++++----
 tests/scripts/test_scripts.py                 |  5 ++-
 24 files changed, 154 insertions(+), 53 deletions(-)

diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index 434611365..41e95931d 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -186,7 +186,9 @@ class RolloutStatsComputer:
     #   https://stable-baselines3.readthedocs.io/en/master/guide/callbacks.html#evalcallback
 
     def __call__(
-        self, policy: policies.ActorCriticPolicy, random_state: np.random.RandomState
+        self,
+        policy: policies.ActorCriticPolicy,
+        random_state: np.random.RandomState,
     ) -> Mapping[str, float]:
         if self.venv is not None and self.n_episodes > 0:
             trajs = rollout.generate_trajectories(
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index e95067956..0c271a69c 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -165,6 +165,7 @@ def __init__(
                 robot action. The choice of robot or human action is independently
                 randomized for each individual `Env` at every timestep.
             save_dir: directory to save collected trajectories in.
+            random_state: random state for random number generation.
         """
         super().__init__(venv)
         self.get_robot_acts = get_robot_acts
@@ -318,6 +319,7 @@ def __init__(
             venv: Vectorized training environment.
             scratch_dir: Directory to use to store intermediate training
                 information (e.g. for resuming training).
+            random_state: random state for random number generation.
             beta_schedule: Provides a value of `beta` (the probability of taking
                 expert action in any given state) at each round of training. If
                 `None`, then `linear_beta_schedule` will be used instead.
diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index 8d3d4bbe5..b63ed8280 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -158,7 +158,7 @@ def __init__(
             action_space: The action space of the environment.
             pi: A tabular policy. Three-dimensional array, where pi[t,s,a]
                 is the probability of taking action a at state s at timestep t.
-            rng: Random state, used for sampling when `predict` is called with
+            random_state: Random state, used for sampling when `predict` is called with
                 `deterministic=False`.
         """
         assert isinstance(state_space, gym.spaces.Discrete), "state not tabular"
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index a07181ad8..c14eee9f7 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -788,7 +788,7 @@ def __init__(
         """Initializes the preference gatherer.
 
         Args:
-            seed: seed for the internal RNG, if applicable
+            random_state: random state for random number generation, if applicable.
             custom_logger: Where to log to; if None (default), creates a new logger.
         """
         # The random seed isn't used here, but it's useful to have this
@@ -1311,7 +1311,10 @@ def _make_reward_trainer(
 
         if is_base or is_std_wrapper:
             return EnsembleTrainer(
-                base_model, loss, random_state=random_state, **reward_trainer_kwargs
+                base_model,
+                loss,
+                random_state=random_state,
+                **reward_trainer_kwargs,
             )
         else:
             raise ValueError(
@@ -1440,7 +1443,7 @@ def __init__(
             raise ValueError(
                 "If you don't provide a random state, you must provide your own "
                 "seeded fragmenter and preference gatherer. You can initialize"
-                "a random state with `np.random.RandomState(seed)`."
+                "a random state with `np.random.RandomState(seed)`.",
             )
         elif self.random_state is not None and None not in (
             preference_gatherer,
@@ -1448,14 +1451,16 @@ def __init__(
         ):
             raise ValueError(
                 "If you provide your own fragmenter and preference gatherer, you "
-                "don't need to provide a random state."
+                "don't need to provide a random state.",
             )
 
         if reward_trainer is None:
             preference_model = PreferenceModel(reward_model)
             loss = CrossEntropyRewardLoss(preference_model)
             self.reward_trainer = _make_reward_trainer(
-                reward_model, loss, random_state=self.random_state
+                reward_model,
+                loss,
+                random_state=self.random_state,
             )
         else:
             self.reward_trainer = reward_trainer
diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index 04e0876eb..f8485c018 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -606,7 +606,11 @@ def rollout(
         should truncate if required.
     """
     trajs = generate_trajectories(
-        policy, venv, sample_until, random_state=random_state, **kwargs
+        policy,
+        venv,
+        sample_until,
+        random_state=random_state,
+        **kwargs,
     )
     if unwrap:
         trajs = [unwrap_traj(traj) for traj in trajs]
diff --git a/src/imitation/policies/exploration_wrapper.py b/src/imitation/policies/exploration_wrapper.py
index 0ea5f8f9e..65855dce2 100644
--- a/src/imitation/policies/exploration_wrapper.py
+++ b/src/imitation/policies/exploration_wrapper.py
@@ -1,7 +1,5 @@
 """Wrapper to turn a policy into a more exploratory version."""
 
-from typing import Optional
-
 import numpy as np
 from stable_baselines3.common import vec_env
 
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index 7a9ff7ddf..bc1bcdd5e 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -5,7 +5,6 @@
 import os
 from typing import Any, Generator, Mapping, Sequence, Tuple, Union
 
-import numpy as np
 import sacred
 from stable_baselines3.common import vec_env
 
@@ -15,7 +14,8 @@
 from imitation.util import util
 
 common_ingredient = sacred.Ingredient(
-    "common", ingredients=[wb.wandb_ingredient, seeding.seeding_ingredient]
+    "common",
+    ingredients=[wb.wandb_ingredient, seeding.seeding_ingredient],
 )
 logger = logging.getLogger(__name__)
 
diff --git a/src/imitation/scripts/common/rl.py b/src/imitation/scripts/common/rl.py
index 46801b1b9..ce687732b 100644
--- a/src/imitation/scripts/common/rl.py
+++ b/src/imitation/scripts/common/rl.py
@@ -21,7 +21,8 @@
 from imitation.scripts.common.train import train_ingredient
 
 rl_ingredient = sacred.Ingredient(
-    "rl", ingredients=[train_ingredient, seeding.seeding_ingredient]
+    "rl",
+    ingredients=[train_ingredient, seeding.seeding_ingredient],
 )
 logger = logging.getLogger(__name__)
 
diff --git a/src/imitation/scripts/common/seeding.py b/src/imitation/scripts/common/seeding.py
index 0a2733f0c..bcbbb1957 100644
--- a/src/imitation/scripts/common/seeding.py
+++ b/src/imitation/scripts/common/seeding.py
@@ -1,3 +1,5 @@
+"""Single source of truth for seeding random number generators in experiments."""
+
 from typing import Optional
 
 import numpy as np
diff --git a/src/imitation/scripts/common/train.py b/src/imitation/scripts/common/train.py
index fa6002373..384c930c9 100644
--- a/src/imitation/scripts/common/train.py
+++ b/src/imitation/scripts/common/train.py
@@ -3,7 +3,6 @@
 import logging
 from typing import Any, Mapping, Union
 
-import numpy as np
 import sacred
 from stable_baselines3.common import base_class, policies, torch_layers, vec_env
 
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 70015ce75..5a2b377c1 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -7,7 +7,6 @@
 from typing import Any, Mapping, Optional
 
 import gym
-import numpy as np
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common.vec_env import VecEnvWrapper
 
@@ -109,7 +108,10 @@ def eval_policy(
             assert policy_path is not None
             policy = serialize.load_policy(policy_type, policy_path, venv)
         trajs = rollout.generate_trajectories(
-            policy, venv, sample_until, random_state=random_state
+            policy,
+            venv,
+            sample_until,
+            random_state=random_state,
         )
 
     if rollout_save_path:
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index ed799ed3f..f18ee6d23 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -5,7 +5,6 @@
 import warnings
 from typing import Any, Mapping, Optional, Sequence, Type, cast
 
-import numpy as np
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import policies, utils, vec_env
 
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 7b0866e60..8bbbab0cf 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -9,7 +9,6 @@
 import random
 from typing import Any, Mapping, Optional, Type, Union
 
-import numpy as np
 import torch as th
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import type_aliases
@@ -23,7 +22,6 @@
 from imitation.scripts.config.train_preference_comparisons import (
     train_preference_comparisons_ex,
 )
-from imitation.util import util
 
 
 def save_model(
diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
index 1d17d5d7b..997dd1600 100644
--- a/src/imitation/scripts/train_rl.py
+++ b/src/imitation/scripts/train_rl.py
@@ -14,7 +14,6 @@
 import warnings
 from typing import Any, Mapping, Optional
 
-import numpy as np
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import callbacks
 from stable_baselines3.common.vec_env import VecNormalize
@@ -147,7 +146,8 @@ def train_rl(
                 rollout_save_n_episodes,
             )
             types.save(
-                save_path, rollout.rollout(rl_algo, venv, sample_until, random_state)
+                save_path,
+                rollout.rollout(rl_algo, venv, sample_until, random_state),
             )
         if policy_save_final:
             output_dir = os.path.join(policy_dir, "final")
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index e79139956..916a65b83 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -163,7 +163,8 @@ def make_seeds(random_state: np.random.RandomState, n: int) -> List[int]:
 
 
 def make_seeds(
-    random_state: np.random.RandomState, n: Optional[int] = None
+    random_state: np.random.RandomState,
+    n: Optional[int] = None,
 ) -> Union[List[int], int]:
     """Generate n random seeds from a random state.
 
@@ -174,7 +175,6 @@ def make_seeds(
     Returns:
         A list of n random seeds.
     """
-
     seeds: List[int] = random_state.randint(0, (1 << 31) - 1, (n or 1,)).tolist()
     if n is None:
         return seeds[0]
diff --git a/tests/algorithms/test_adversarial.py b/tests/algorithms/test_adversarial.py
index 08f93855e..81d37a30f 100644
--- a/tests/algorithms/test_adversarial.py
+++ b/tests/algorithms/test_adversarial.py
@@ -78,7 +78,10 @@ def make_trainer(
         expert_data = expert_transitions
 
     venv = util.make_vec_env(
-        env_name, n_envs=num_envs, parallel=parallel, random_state=random_state
+        env_name,
+        n_envs=num_envs,
+        parallel=parallel,
+        random_state=random_state,
     )
     model_cls = algorithm_kwargs["model_class"]
     gen_algo = model_cls(algorithm_kwargs["policy_class"], venv)
@@ -115,7 +118,10 @@ def test_airl_fail_fast(custom_logger, tmpdir, random_state_fixed):
 
     gen_algo = stable_baselines3.DQN(stable_baselines3.dqn.MlpPolicy, venv)
     small_data = rollout.generate_transitions(
-        gen_algo, venv, n_timesteps=20, random_state=random_state
+        gen_algo,
+        venv,
+        n_timesteps=20,
+        random_state=random_state,
     )
     reward_net = reward_nets.BasicShapedRewardNet(
         observation_space=venv.observation_space,
@@ -138,7 +144,10 @@ def test_airl_fail_fast(custom_logger, tmpdir, random_state_fixed):
 def trainer(request, tmpdir, expert_transitions, random_state_fixed):
     random_state = random_state_fixed
     with make_trainer(
-        request.param, tmpdir, expert_transitions, random_state
+        request.param,
+        tmpdir,
+        expert_transitions,
+        random_state,
     ) as trainer:
         yield trainer
 
@@ -203,7 +212,9 @@ def trainer_parametrized(
 
 
 def test_train_disc_step_no_crash(
-    trainer_parametrized, _expert_batch_size, random_state_fixed
+    trainer_parametrized,
+    _expert_batch_size,
+    random_state_fixed,
 ):
     random_state = random_state_fixed
     transitions = rollout.generate_transitions(
@@ -283,7 +294,11 @@ def _env_name(request):
 
 @pytest.fixture
 def trainer_diverse_env(
-    _algorithm_kwargs, _env_name, tmpdir, expert_transitions, random_state_fixed
+    _algorithm_kwargs,
+    _env_name,
+    tmpdir,
+    expert_transitions,
+    random_state_fixed,
 ):
     random_state = random_state_fixed
     if _algorithm_kwargs["model_class"] == stable_baselines3.DQN:
@@ -312,6 +327,7 @@ def test_logits_expert_is_high_log_policy_act_prob(
     Args:
         trainer_diverse_env: The trainer to test.
         n_timesteps: The number of timesteps of rollouts to collect.
+        random_state_fixed: The random state to use.
     """
     random_state = random_state_fixed
     trans = rollout.generate_transitions(
diff --git a/tests/algorithms/test_bc.py b/tests/algorithms/test_bc.py
index ebf19ea1e..87f6c18ed 100644
--- a/tests/algorithms/test_bc.py
+++ b/tests/algorithms/test_bc.py
@@ -181,6 +181,7 @@ def test_bc_data_loader_empty_iter_error(
         no_yield_after_iter: Data loader stops yielding after this many calls.
         custom_logger: Where to log to.
         cartpole_expert_trajectories: The expert trajectories to use.
+        random_state_fixed: Random state to use.
     """
     batch_size = 32
     random_state = random_state_fixed
diff --git a/tests/algorithms/test_dagger.py b/tests/algorithms/test_dagger.py
index 6b80e18a9..908c3a1ef 100644
--- a/tests/algorithms/test_dagger.py
+++ b/tests/algorithms/test_dagger.py
@@ -432,7 +432,10 @@ def test_dagger_not_enough_transitions_error(tmpdir, custom_logger, random_state
     collector = trainer.create_trajectory_collector()
     policy = base.RandomPolicy(venv.observation_space, venv.action_space)
     rollout.generate_trajectories(
-        policy, collector, rollout.make_min_episodes(1), random_state=random_state
+        policy,
+        collector,
+        rollout.make_min_episodes(1),
+        random_state=random_state,
     )
     with pytest.raises(ValueError, match="Not enough transitions.*"):
         trainer.extend_and_update()
diff --git a/tests/algorithms/test_mce_irl.py b/tests/algorithms/test_mce_irl.py
index 73a70786d..32909a1e0 100644
--- a/tests/algorithms/test_mce_irl.py
+++ b/tests/algorithms/test_mce_irl.py
@@ -340,7 +340,11 @@ def test_mce_irl_demo_formats(random_state_fixed):
                 hid_sizes=[],
             )
             mce_irl = MCEIRL(
-                demo, mdp, reward_net, linf_eps=1e-3, random_state=random_state
+                demo,
+                mdp,
+                reward_net,
+                linf_eps=1e-3,
+                random_state=random_state,
             )
             assert np.allclose(mce_irl.demo_state_om.sum(), mdp.horizon + 1)
             final_counts[kind] = mce_irl.train(max_iter=5)
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 6e488187f..b7c5f3ebf 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -59,7 +59,8 @@ def agent(venv):
 def random_fragmenter(random_state_fixed):
     random_state = random_state_fixed
     return preference_comparisons.RandomFragmenter(
-        random_state=random.Random(util.make_seeds(random_state)), warning_threshold=0
+        random_state=random.Random(util.make_seeds(random_state)),
+        warning_threshold=0,
     )
 
 
@@ -100,7 +101,10 @@ def test_mismatched_spaces(venv, agent, random_state_fixed):
         match="spaces do not match",
     ):
         preference_comparisons.AgentTrainer(
-            agent, bad_reward_net, venv, random_state=random_state
+            agent,
+            bad_reward_net,
+            venv,
+            random_state=random_state,
         )
 
 
@@ -305,7 +309,11 @@ def test_init_raises_error_when_trying_use_improperly_wrapped_ensemble(
 
 
 def test_discount_rate_no_crash(
-    agent_trainer, venv, random_fragmenter, custom_logger, random_state_fixed
+    agent_trainer,
+    venv,
+    random_fragmenter,
+    custom_logger,
+    random_state_fixed,
 ):
     random_state = random_state_fixed
     # also use a non-zero noise probability to check that doesn't cause errors
@@ -337,11 +345,14 @@ def test_discount_rate_no_crash(
 
 
 def test_synthetic_gatherer_deterministic(
-    agent_trainer, random_fragmenter, random_state_fixed
+    agent_trainer,
+    random_fragmenter,
+    random_state_fixed,
 ):
     random_state = random_state_fixed
     gatherer = preference_comparisons.SyntheticGatherer(
-        temperature=0, random_state=random_state
+        temperature=0,
+        random_state=random_state,
     )
     trajectories = agent_trainer.sample(10)
     fragments = random_fragmenter(trajectories, fragment_length=2, num_pairs=2)
@@ -421,7 +432,10 @@ def test_preference_dataset_queue(agent_trainer, random_fragmenter, random_state
 
 
 def test_store_and_load_preference_dataset(
-    agent_trainer, random_fragmenter, tmp_path, random_state_fixed
+    agent_trainer,
+    random_fragmenter,
+    tmp_path,
+    random_state_fixed,
 ):
     random_state = random_state_fixed
     dataset = preference_comparisons.PreferenceDataset()
diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index db845d3dd..794448a09 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -83,6 +83,7 @@ def test_complete_trajectories(policy_type, random_state_fixed) -> None:
 
     Args:
         policy_type: Kind of policy to use when generating trajectories.
+        random_state_fixed: Random state to use.
     """
     random_state = random_state_fixed
     min_episodes = 13
@@ -142,10 +143,13 @@ def test_unbiased_trajectories(
         min_episodes: The minimum number of episodes to sample.
         expected_counts: Mapping from episode length to expected number of episodes
             of that length (omit if 0 episodes of that length expected).
+        random_state_fixed: Random state to use.
     """
     random_state = random_state_fixed
     trajectories = _sample_fixed_length_trajectories(
-        episode_lengths, min_episodes, random_state
+        episode_lengths,
+        min_episodes,
+        random_state,
     )
     assert len(trajectories) == sum(expected_counts.values())
     traj_lens = np.array([len(traj) for traj in trajectories])
@@ -197,7 +201,10 @@ def test_rollout_stats(random_state_fixed):
 
     policy = serialize.load_policy("zero", "UNUSED", venv)
     trajs = rollout.generate_trajectories(
-        policy, venv, rollout.make_min_episodes(10), random_state=random_state
+        policy,
+        venv,
+        rollout.make_min_episodes(10),
+        random_state=random_state,
     )
     s = rollout.rollout_stats(trajs)
 
@@ -220,7 +227,10 @@ def test_unwrap_traj(random_state_fixed):
 
     policy = serialize.load_policy("zero", "UNUSED", venv)
     trajs = rollout.generate_trajectories(
-        policy, venv, rollout.make_min_episodes(10), random_state=random_state
+        policy,
+        venv,
+        rollout.make_min_episodes(10),
+        random_state=random_state,
     )
     trajs_unwrapped = [rollout.unwrap_traj(t) for t in trajs]
     trajs_unwrapped_twice = [rollout.unwrap_traj(t) for t in trajs_unwrapped]
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index 36b71bea6..163a02bb7 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -29,11 +29,17 @@ def test_actions_valid(env_name, policy_type, random_state_fixed):
     """Test output actions of our custom policies always lie in action space."""
     random_state = random_state_fixed
     venv = util.make_vec_env(
-        env_name, n_envs=1, parallel=False, random_state=random_state
+        env_name,
+        n_envs=1,
+        parallel=False,
+        random_state=random_state,
     )
     policy = serialize.load_policy(policy_type, "foobar", venv)
     transitions = rollout.generate_transitions(
-        policy, venv, n_timesteps=100, random_state=random_state
+        policy,
+        venv,
+        n_timesteps=100,
+        random_state=random_state,
     )
 
     for a in transitions.acts:
@@ -48,7 +54,9 @@ def test_actions_valid(env_name, policy_type, random_state_fixed):
     ],
 )
 def test_save_stable_model_errors_and_warnings(
-    tmpdir, policy_env_name_pair, random_state_fixed
+    tmpdir,
+    policy_env_name_pair,
+    random_state_fixed,
 ):
     """Check errors and warnings in `save_stable_model()`."""
     random_state = random_state_fixed
@@ -76,7 +84,10 @@ def test_save_stable_model_errors_and_warnings(
 def _test_serialize_identity(env_name, model_cfg, tmpdir, random_state):
     """Test output actions of deserialized policy are same as original."""
     venv = util.make_vec_env(
-        env_name, n_envs=1, parallel=False, random_state=random_state
+        env_name,
+        n_envs=1,
+        parallel=False,
+        random_state=random_state,
     )
 
     model_name, model_cls_name = model_cfg
@@ -126,7 +137,10 @@ def test_serialize_identity(env_name, model_cfg, tmpdir, random_state_fixed):
 @pytest.mark.parametrize("env_name", [SIMPLE_CONTINUOUS_ENV])
 @pytest.mark.parametrize("model_cfg", CONTINUOUS_ONLY_CONFIGS)
 def test_serialize_identity_continuous_only(
-    env_name, model_cfg, tmpdir, random_state_fixed
+    env_name,
+    model_cfg,
+    tmpdir,
+    random_state_fixed,
 ):
     """Test serialize identity for continuous_only algorithms."""
     _test_serialize_identity(env_name, model_cfg, tmpdir, random_state_fixed)
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index 7e96b40f4..de27269db 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -127,7 +127,10 @@ def _sample(space, n):
 
 def _make_env_and_save_reward_net(env_name, reward_type, tmpdir, random_state):
     venv = util.make_vec_env(
-        env_name, n_envs=1, parallel=False, random_state=random_state
+        env_name,
+        n_envs=1,
+        parallel=False,
+        random_state=random_state,
     )
     save_path = os.path.join(tmpdir, "norm_reward.pt")
 
@@ -162,10 +165,16 @@ def test_reward_valid(env_name, reward_type, tmpdir, random_state_fixed):
     random_state = random_state_fixed
     # TODO(juan) the line below is not being used?
     venv = util.make_vec_env(
-        env_name, n_envs=1, parallel=False, random_state=random_state
+        env_name,
+        n_envs=1,
+        parallel=False,
+        random_state=random_state,
     )
     venv, tmppath = _make_env_and_save_reward_net(
-        env_name, reward_type, tmpdir, random_state
+        env_name,
+        reward_type,
+        tmpdir,
+        random_state,
     )
 
     TRAJECTORY_LEN = 10
@@ -184,7 +193,10 @@ def test_reward_valid(env_name, reward_type, tmpdir, random_state_fixed):
 def test_strip_wrappers_basic(random_state_fixed):
     random_state = random_state_fixed
     venv = util.make_vec_env(
-        "FrozenLake-v1", n_envs=1, parallel=False, random_state=random_state
+        "FrozenLake-v1",
+        n_envs=1,
+        parallel=False,
+        random_state=random_state,
     )
     net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     net = reward_nets.NormalizedRewardNet(net, networks.RunningNorm)
@@ -201,7 +213,10 @@ def test_strip_wrappers_basic(random_state_fixed):
 def test_strip_wrappers_complex(random_state_fixed):
     random_state = random_state_fixed
     venv = util.make_vec_env(
-        "FrozenLake-v1", n_envs=1, parallel=False, random_state=random_state
+        "FrozenLake-v1",
+        n_envs=1,
+        parallel=False,
+        random_state=random_state,
     )
     net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     net = reward_nets.ShapedRewardNet(net, _potential, discount_factor=0.99)
@@ -303,7 +318,10 @@ def test_serialize_identity(
     logging.info(f"Testing {net_cls}")
 
     venv = util.make_vec_env(
-        env_name, n_envs=1, parallel=False, random_state=random_state
+        env_name,
+        n_envs=1,
+        parallel=False,
+        random_state=random_state,
     )
     original = net_cls(venv.observation_space, venv.action_space)
     if normalize_rewards:
@@ -318,7 +336,10 @@ def test_serialize_identity(
     assert original.action_space == loaded.action_space
 
     transitions = rollout.generate_transitions(
-        random, venv, n_timesteps=100, random_state=random_state
+        random,
+        venv,
+        n_timesteps=100,
+        random_state=random_state,
     )
 
     if isinstance(original, reward_nets.NormalizedRewardNet):
@@ -661,7 +682,10 @@ def test_training_regression(normalize_input_layer, random_state_fixed):
     random = base.RandomPolicy(venv.observation_space, venv.action_space)
     for _ in range(2):
         transitions = rollout.generate_transitions(
-            random, venv, n_timesteps=100, random_state=random_state
+            random,
+            venv,
+            n_timesteps=100,
+            random_state=random_state,
         )
         trans_args = (
             transitions.obs,
diff --git a/tests/scripts/test_scripts.py b/tests/scripts/test_scripts.py
index 9c257fcfb..07c66571f 100644
--- a/tests/scripts/test_scripts.py
+++ b/tests/scripts/test_scripts.py
@@ -682,7 +682,10 @@ def test_preference_comparisons_transfer_learning(
 def test_train_rl_double_normalization(tmpdir: str, random_state_fixed):
     random_state = random_state_fixed
     venv = util.make_vec_env(
-        "CartPole-v1", n_envs=1, parallel=False, random_state=random_state
+        "CartPole-v1",
+        n_envs=1,
+        parallel=False,
+        random_state=random_state,
     )
     basic_reward_net = reward_nets.BasicRewardNet(
         venv.observation_space,

From 32898c36f8ba12aa80175be77087aa24f5842e71 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 01:55:53 +0200
Subject: [PATCH 079/187] Remove automatically generated `_api` docs files too
 on `make clean`

---
 Makefile | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Makefile b/Makefile
index b6d0e79c1..01cc08ef0 100644
--- a/Makefile
+++ b/Makefile
@@ -25,6 +25,7 @@ docscheck:
 cleandocs:
 	pushd docs/ \
 	&& make clean \
+	&& rm -rf _api/ \
 	&& popd
 
 cleantests:

From 76eb8bb44f5fce70d029aaa5d427537fca17b8b6 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 12:03:38 +0200
Subject: [PATCH 080/187] Fix docstrings.

---
 src/imitation/algorithms/bc.py                     | 1 +
 src/imitation/algorithms/dagger.py                 | 1 +
 src/imitation/algorithms/preference_comparisons.py | 3 +++
 src/imitation/data/rollout.py                      | 1 +
 src/imitation/util/util.py                         | 1 -
 tests/algorithms/test_preference_comparisons.py    | 3 +++
 tests/data/test_rollout.py                         | 6 ++++++
 7 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index 41e95931d..1cd9bf23b 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -293,6 +293,7 @@ def __init__(
         Args:
             observation_space: the observation space of the environment.
             action_space: the action space of the environment.
+            random_state: the random state to use for the random number generator.
             policy: a Stable Baselines3 policy; if unspecified,
                 defaults to `FeedForward32Policy`.
             demonstrations: Demonstrations from an expert (optional). Transitions
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index 0c271a69c..8df29d63c 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -547,6 +547,7 @@ def __init__(
             scratch_dir: Directory to use to store intermediate training
                 information (e.g. for resuming training).
             expert_policy: The expert policy used to generate synthetic demonstrations.
+            random_state: Random state to use for the random number generator.
             expert_trajs: Optional starting dataset that is inserted into the round 0
                 dataset.
             dagger_trainer_kwargs: Other keyword arguments passed to the
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index c14eee9f7..c83e84e10 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -850,6 +850,9 @@ def __init__(
                 This threshold is therefore in logspace. The default value
                 of 50 means that probabilities below 2e-22 are rounded up to 2e-22.
             custom_logger: Where to log to; if None (default), creates a new logger.
+
+        Raises:
+            ValueError: if `sample` is true and no random state is provided.
         """
         super().__init__(custom_logger=custom_logger)
         self.temperature = temperature
diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index f8485c018..392bf1237 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -543,6 +543,7 @@ def generate_transitions(
             - None, in which case actions will be sampled randomly
         venv: The vectorized environments to interact with.
         n_timesteps: The minimum number of timesteps to sample.
+        random_state: The random state to use for sampling trajectories.
         truncate: If True, then drop any additional samples to ensure that exactly
             `n_timesteps` samples are returned.
         **kwargs: Passed-through to generate_trajectories.
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 916a65b83..ac8b44e7a 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -79,7 +79,6 @@ def make_vec_env(
         env_name: The Env's string id in Gym.
         random_state: The random state to use to seed the environment.
         n_envs: The number of duplicate environments.
-        seed: The environment seed.
         parallel: If True, uses SubprocVecEnv; otherwise, DummyVecEnv.
         log_dir: If specified, saves Monitor output to this directory.
         max_episode_steps: If specified, wraps each env in a TimeLimit wrapper
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index b7c5f3ebf..fb8fd8503 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -649,6 +649,9 @@ def test_agent_trainer_sample_image_observations(random_state_fixed):
     SB3 algorithms may rearrange the channel dimension in environments with image
     observations, but `sample()` should return observations matching the original
     environment.
+
+    Args:
+        random_state_fixed: Random state (with a fixed seed).
     """
     random_state = random_state_fixed
     venv = DummyVecEnv([lambda: FakeImageEnv()])
diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index 794448a09..0a96c8ccb 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -192,6 +192,9 @@ def test_rollout_stats(random_state_fixed):
     """Applying `ObsRewIncrementWrapper` halves the reward mean.
 
     `rollout_stats` should reflect this.
+
+    Args:
+        random_state_fixed: Random state to use (with fixed seed).
     """
     random_state = random_state_fixed
     env = gym.make("CartPole-v1")
@@ -218,6 +221,9 @@ def test_unwrap_traj(random_state_fixed):
     """Check that unwrap_traj reverses `ObsRewIncrementWrapper`.
 
     Also check that unwrapping twice is a no-op.
+
+    Args:
+        random_state_fixed: Random state to use (with fixed seed).
     """
     random_state = random_state_fixed
     env = gym.make("CartPole-v1")

From 0a1b06a6e2741e461ecb7ab4e9de783455b91b46 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 12:21:40 +0200
Subject: [PATCH 081/187] Fix issue with next(iter(iterable))

---
 src/imitation/algorithms/base.py    |  7 ++----
 src/imitation/algorithms/density.py |  4 +--
 src/imitation/algorithms/mce_irl.py |  6 +----
 src/imitation/util/util.py          | 38 +++++++++++++++++++++--------
 4 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index 03dd600bb..20ad8c8b5 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -9,7 +9,7 @@
 from stable_baselines3.common import policies
 
 from imitation.data import rollout, types
-from imitation.util import logger as imit_logger
+from imitation.util import logger as imit_logger, util
 
 
 class BaseImitationAlgorithm(abc.ABC):
@@ -242,10 +242,7 @@ def make_data_loader(
         raise ValueError(f"batch_size={batch_size} must be positive.")
 
     if isinstance(transitions, Iterable):
-        try:
-            first_item = next(iter(transitions))
-        except StopIteration:
-            first_item = None
+        first_item, transitions = util.get_first_iter_element(transitions)
         if isinstance(first_item, types.Trajectory):
             transitions = rollout.flatten_trajectories(list(transitions))
 
diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index 492658331..aa8691573 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -16,7 +16,7 @@
 from imitation.algorithms import base
 from imitation.data import rollout, types, wrappers
 from imitation.rewards import reward_wrapper
-from imitation.util import logger as imit_logger
+from imitation.util import logger as imit_logger, util
 
 
 class DensityType(enum.Enum):
@@ -154,7 +154,7 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
         self.transitions = {}
 
         if isinstance(demonstrations, Iterable):
-            first_item = next(iter(demonstrations))
+            first_item, demonstrations = util.get_first_iter_element(demonstrations)
             if isinstance(first_item, types.Trajectory):
                 # Demonstrations are trajectories.
                 # We have timestep information.
diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index b63ed8280..273655b03 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -381,11 +381,7 @@ def set_demonstrations(self, demonstrations: MCEDemonstrations) -> None:
         # Demonstrations are either trajectories or transitions;
         # we must compute occupancy measure from this.
         if isinstance(demonstrations, Iterable):
-            # TODO(juan) this is wrong; if this is a container (list-like)
-            #  object then a new fresh iterator will be generated every time,
-            #  but for a general iterator calling next() exhaust the
-            #  iterator.
-            first_item = next(iter(demonstrations))
+            first_item, demonstrations = util.get_first_iter_element(demonstrations)
             if isinstance(first_item, types.Trajectory):
                 self._set_demo_from_trajectories(demonstrations)
                 return
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index ac8b44e7a..5fd09507f 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -16,7 +16,7 @@
     Sequence,
     TypeVar,
     Union,
-    overload,
+    overload, Tuple,
 )
 
 import gym
@@ -215,15 +215,7 @@ def endless_iter(iterable: Iterable[T]) -> Iterator[T]:
     Raises:
         ValueError: `iterable` is empty -- the first call it to returns no elements.
     """
-    # TODO(juan) this is wrong; if the iterable is not a container then the first
-    #  element will be wasted if it's not stored. The sensible solution is to
-    #  restrict the type of `iterable` to `Sequence` (this same issue is present
-    #  in a few other places in the codebase).
-    try:
-        next(iter(iterable))
-    except StopIteration:
-        raise ValueError(f"iterable {iterable} had no elements to iterate over.")
-
+    _, iterable = get_first_iter_element(iterable)
     return itertools.chain.from_iterable(itertools.repeat(iterable))
 
 
@@ -275,3 +267,29 @@ def tensor_iter_norm(
     # = sum(x**ord for x in tensor for tensor in tensor_iter)**(1/ord)
     # = th.norm(concatenated tensors)
     return th.norm(norm_tensor, p=ord)
+
+
+def get_first_iter_element(
+        iterable: Iterable[T]
+) -> Tuple[T, Iterable[T]]:
+    """
+    Gets the first element of the iterable, and returns a new iterable that adds the
+    first element back using itertools.chain.
+
+    If the iterable is a tuple or list, this is equivalent to (iterable[0], iterable).
+
+    Args:
+        iterable: The iterable to get the first element of.
+
+    Returns:
+        A tuple containing the first element of the iterable, and a fresh iterable
+        with all the elements.
+    """
+    if isinstance(iterable, (list, tuple)):
+        return iterable[0], iterable
+    else:
+        try:
+            first_element = next(iter(iterable))
+            return first_element, itertools.chain([first_element], iterable)
+        except StopIteration:
+            raise ValueError(f"iterable {iterable} had no elements to iterate over.")

From 58eed54e647420e52d3a4efceb6b9f3a2edf192b Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 12:22:17 +0200
Subject: [PATCH 082/187] Formatting

---
 src/imitation/algorithms/base.py    | 3 ++-
 src/imitation/algorithms/density.py | 3 ++-
 src/imitation/util/util.py          | 7 +++----
 3 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index 20ad8c8b5..27371bf04 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -9,7 +9,8 @@
 from stable_baselines3.common import policies
 
 from imitation.data import rollout, types
-from imitation.util import logger as imit_logger, util
+from imitation.util import logger as imit_logger
+from imitation.util import util
 
 
 class BaseImitationAlgorithm(abc.ABC):
diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index aa8691573..6d0f1775a 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -16,7 +16,8 @@
 from imitation.algorithms import base
 from imitation.data import rollout, types, wrappers
 from imitation.rewards import reward_wrapper
-from imitation.util import logger as imit_logger, util
+from imitation.util import logger as imit_logger
+from imitation.util import util
 
 
 class DensityType(enum.Enum):
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 5fd09507f..823bf263b 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -14,9 +14,10 @@
     Mapping,
     Optional,
     Sequence,
+    Tuple,
     TypeVar,
     Union,
-    overload, Tuple,
+    overload,
 )
 
 import gym
@@ -269,9 +270,7 @@ def tensor_iter_norm(
     return th.norm(norm_tensor, p=ord)
 
 
-def get_first_iter_element(
-        iterable: Iterable[T]
-) -> Tuple[T, Iterable[T]]:
+def get_first_iter_element(iterable: Iterable[T]) -> Tuple[T, Iterable[T]]:
     """
     Gets the first element of the iterable, and returns a new iterable that adds the
     first element back using itertools.chain.

From 8c1655e0ad55a4ab0fefb7772da823c1cdc6ae61 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 12:25:08 +0200
Subject: [PATCH 083/187] Remove whitespace

---
 tests/util/test_wb_logger.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/util/test_wb_logger.py b/tests/util/test_wb_logger.py
index 9bd4a65a5..443c22ef8 100644
--- a/tests/util/test_wb_logger.py
+++ b/tests/util/test_wb_logger.py
@@ -89,8 +89,6 @@ def finish(self):
 
 # we ignore the type below as one should technically not access the
 # __init__ method directly but only by creating an instance.
-
-
 @mock.patch.object(wandb, "__init__", mock_wandb.__init__)  # type: ignore
 @mock.patch.object(wandb, "init", mock_wandb.init)
 @mock.patch.object(wandb, "log", mock_wandb.log)

From 59cff9d509cb93d8ab5b644dc037b5bb863fd03d Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 12:27:54 +0200
Subject: [PATCH 084/187] Add TODO message to remove type ignore later

---
 tests/policies/test_replay_buffer_wrapper.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/policies/test_replay_buffer_wrapper.py b/tests/policies/test_replay_buffer_wrapper.py
index 801b7ea8a..85365c0d5 100644
--- a/tests/policies/test_replay_buffer_wrapper.py
+++ b/tests/policies/test_replay_buffer_wrapper.py
@@ -38,8 +38,10 @@ def make_algo_with_wrapped_buffer(
         policy_kwargs=dict(),
         env=venv,
         seed=42,
-        # we ignore the type below because sb3 has a bug (forgot to put Type[...])
-        # https://github.com/DLR-RM/stable-baselines3/issues/1039
+        # TODO(juan) we ignore the type below due to
+        #  https://github.com/DLR-RM/stable-baselines3/issues/1039
+        #  PR fixing this has been merged to master,
+        #  remove the type ignore in the next sb3 release.
         replay_buffer_class=ReplayBufferRewardWrapper,  # type: ignore
         replay_buffer_kwargs=dict(
             replay_buffer_class=replay_buffer_class,

From 74b90da8b612feb1e1df05a51701964f72783fbb Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 12:30:01 +0200
Subject: [PATCH 085/187] Remove unnecessary assertion.

---
 tests/data/test_types.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index ba3c4db10..7d9330fa4 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -27,7 +27,6 @@
 
 
 def _check_1d_shape(fn: Callable[[np.ndarray], Any], length: int, expected_msg: str):
-    assert isinstance(length, int)
     for shape in [(), (length, 1), (length, 2), (length - 1,), (length + 1,)]:
         with pytest.raises(ValueError, match=expected_msg):
             fn(np.zeros(shape))

From 54eb44c3dfd104c29fec8946582e03d462b6bdb2 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 13:27:19 +0200
Subject: [PATCH 086/187] Fixed types in density.py set_demonstrations

---
 src/imitation/algorithms/base.py    |  5 ++++
 src/imitation/algorithms/density.py | 31 ++++++++++++++---------
 src/imitation/util/util.py          | 38 ++++++++++++++++++++++++++---
 3 files changed, 60 insertions(+), 14 deletions(-)

diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index 27371bf04..dacaabac6 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -114,6 +114,11 @@ def __setstate__(self, state):
 
 TransitionMapping = Mapping[str, Union[np.ndarray, th.Tensor]]
 TransitionKind = TypeVar("TransitionKind", bound=types.TransitionsMinimal)
+# TODO(juan) AnyTransitions is used in many places but TransitionKind
+#  is just a type variable. If what we want is to allow for subclasses of
+#  types.TransitionsMinimal, we should use that type directly, as that is already
+#  built into the type system. Most places that use AnyTransitions aren't actually
+#  treating the type as a generic.
 AnyTransitions = Union[
     Iterable[types.Trajectory],
     Iterable[TransitionMapping],
diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index 6d0f1775a..00d70f10d 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -6,7 +6,8 @@
 
 import enum
 import itertools
-from typing import Dict, Iterable, Optional
+from collections.abc import Mapping
+from typing import Dict, Iterable, Optional, cast
 
 import numpy as np
 from gym.spaces.utils import flatten
@@ -157,26 +158,34 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
         if isinstance(demonstrations, Iterable):
             first_item, demonstrations = util.get_first_iter_element(demonstrations)
             if isinstance(first_item, types.Trajectory):
-                # Demonstrations are trajectories.
-                # We have timestep information.
+                # we assume that all elements are also types.Trajectory.
+                # (this means we have timestamp information)
+                # It's not perfectly type safe, but it allows for the flexibility of
+                # using iterables, which is useful for large data structures.
+                demonstrations = cast(Iterable[types.Trajectory], demonstrations)
+
                 for traj in demonstrations:
                     for i, (obs, act, next_obs) in enumerate(
                         zip(traj.obs[:-1], traj.acts, traj.obs[1:]),
                     ):
                         flat_trans = self._preprocess_transition(obs, act, next_obs)
                         self.transitions.setdefault(i, []).append(flat_trans)
-            else:
-                # Demonstrations are a Torch DataLoader or other Mapping iterable
+            elif isinstance(first_item, Mapping):
+                # analogous to cast above.
+                demonstrations = cast(Iterable[base.TransitionMapping], demonstrations)
+
                 for batch in demonstrations:
                     self._set_demo_from_batch(
-                        batch["obs"],
-                        batch["acts"],
-                        batch.get("next_obs"),
+                        util.safe_to_numpy(batch["obs"], warn=True),
+                        util.safe_to_numpy(batch["acts"], warn=True),
+                        util.safe_to_numpy(batch.get("next_obs"), warn=True),
                     )
+            else:
+                raise TypeError(
+                    f"Unsupported demonstration type {type(demonstrations)}"
+                )
         elif isinstance(demonstrations, types.TransitionsMinimal):
-            next_obs_b = (
-                demonstrations.next_obs if hasattr(demonstrations, "next_obs") else None
-            )
+            next_obs_b = getattr(demonstrations, "next_obs", None)
             self._set_demo_from_batch(
                 demonstrations.obs,
                 demonstrations.acts,
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 823bf263b..d5a5809e3 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -5,6 +5,7 @@
 import itertools
 import os
 import uuid
+import warnings
 from typing import (
     Any,
     Callable,
@@ -240,6 +241,38 @@ def safe_to_tensor(numpy_array: np.ndarray, **kwargs) -> th.Tensor:
     return th.as_tensor(numpy_array, **kwargs)
 
 
+@overload
+def safe_to_numpy(obj: Union[np.ndarray, th.Tensor], warn: bool = False) -> np.ndarray:
+    ...
+
+
+@overload
+def safe_to_numpy(obj: None, warn: bool = False) -> None:
+    ...
+
+
+def safe_to_numpy(obj: Optional[Union[np.ndarray, th.Tensor]], warn=False):
+    """Convert torch tensor to numpy.
+
+    If the object is already a numpy array, return it as is.
+    If the object is none, returns none.
+
+    Args:
+        obj: torch tensor object to convert to numpy array
+
+    Returns: object converted to numpy array
+
+    """
+    if obj is None:
+        return None
+    elif isinstance(obj, np.ndarray):
+        return obj
+    else:
+        if warn:
+            warnings.warn(f"Converted tensor {obj} to numpy array.")
+        return obj.detach().cpu().numpy()
+
+
 def tensor_iter_norm(
     tensor_iter: Iterable[th.Tensor],
     ord: Union[int, float] = 2,  # noqa: A002
@@ -271,10 +304,9 @@ def tensor_iter_norm(
 
 
 def get_first_iter_element(iterable: Iterable[T]) -> Tuple[T, Iterable[T]]:
-    """
-    Gets the first element of the iterable, and returns a new iterable that adds the
-    first element back using itertools.chain.
+    """Get first element of an iterable and a new fresh iterable.
 
+    The fresh iterable has the first element added back using itertools.chain.
     If the iterable is a tuple or list, this is equivalent to (iterable[0], iterable).
 
     Args:

From ced18b253703d81ef452c2518260965d627b3149 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 9 Sep 2022 13:44:30 +0200
Subject: [PATCH 087/187] Added type ignore to pytype bug

---
 src/imitation/algorithms/density.py |  2 +-
 src/imitation/util/util.py          | 11 +++++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index 00d70f10d..34dcff2c7 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -182,7 +182,7 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
                     )
             else:
                 raise TypeError(
-                    f"Unsupported demonstration type {type(demonstrations)}"
+                    f"Unsupported demonstration type {type(demonstrations)}",
                 )
         elif isinstance(demonstrations, types.TransitionsMinimal):
             next_obs_b = getattr(demonstrations, "next_obs", None)
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index d5a5809e3..9caacf53f 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -251,7 +251,10 @@ def safe_to_numpy(obj: None, warn: bool = False) -> None:
     ...
 
 
-def safe_to_numpy(obj: Optional[Union[np.ndarray, th.Tensor]], warn=False):
+def safe_to_numpy(
+    obj: Optional[Union[np.ndarray, th.Tensor]],
+    warn=False,
+) -> Optional[np.ndarray]:
     """Convert torch tensor to numpy.
 
     If the object is already a numpy array, return it as is.
@@ -259,12 +262,16 @@ def safe_to_numpy(obj: Optional[Union[np.ndarray, th.Tensor]], warn=False):
 
     Args:
         obj: torch tensor object to convert to numpy array
+        warn: if True, warn if the object is not already a numpy array. Useful for
+            warning the user of a potential performance hit if a torch tensor is
+            not the expected input type.
 
     Returns: object converted to numpy array
 
     """
     if obj is None:
-        return None
+        # We ignore the type due to https://github.com/google/pytype/issues/445
+        return None  # type: ignore[bad-return-type]
     elif isinstance(obj, np.ndarray):
         return obj
     else:

From cbc7062f0e9faa3d92945f5956ceeaeb0da9ba3d Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Fri, 9 Sep 2022 17:07:25 -0700
Subject: [PATCH 088/187] Fix_get_first_iter_element and add tests

---
 src/imitation/util/util.py | 45 +++++++++++++++++++++++++++-----------
 tests/util/test_util.py    | 24 ++++++++++++++++++++
 2 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 9caacf53f..d40e9c92f 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -209,14 +209,18 @@ def endless_iter(iterable: Iterable[T]) -> Iterator[T]:
     0
 
     Args:
-        iterable: The object to endlessly iterate over.
+        iterable: The non-iterator iterable object to endlessly iterate over.
 
     Returns:
         An iterator that repeats the elements in `iterable` forever.
 
     Raises:
-        ValueError: `iterable` is empty -- the first call it to returns no elements.
+        ValueError: if iterable is an iterator -- that will be exhausted, so
+            cannot be iterated over endlessly.
     """
+    if iter(iterable) == iterable:
+        raise ValueError("endless_iter needs a non-iterator Iterable.")
+
     _, iterable = get_first_iter_element(iterable)
     return itertools.chain.from_iterable(itertools.repeat(iterable))
 
@@ -266,8 +270,8 @@ def safe_to_numpy(
             warning the user of a potential performance hit if a torch tensor is
             not the expected input type.
 
-    Returns: object converted to numpy array
-
+    Returns:
+        Object converted to numpy array
     """
     if obj is None:
         # We ignore the type due to https://github.com/google/pytype/issues/445
@@ -313,8 +317,9 @@ def tensor_iter_norm(
 def get_first_iter_element(iterable: Iterable[T]) -> Tuple[T, Iterable[T]]:
     """Get first element of an iterable and a new fresh iterable.
 
-    The fresh iterable has the first element added back using itertools.chain.
-    If the iterable is a tuple or list, this is equivalent to (iterable[0], iterable).
+    The fresh iterable has the first element added back using ``itertools.chain``.
+    If the iterable is not an iterator, this is equivalent to
+    ``(next(iter(iterable)), iterable)``.
 
     Args:
         iterable: The iterable to get the first element of.
@@ -322,12 +327,26 @@ def get_first_iter_element(iterable: Iterable[T]) -> Tuple[T, Iterable[T]]:
     Returns:
         A tuple containing the first element of the iterable, and a fresh iterable
         with all the elements.
+
+    Raises:
+        ValueError: `iterable` is empty -- the first call it to returns no elements.
     """
-    if isinstance(iterable, (list, tuple)):
-        return iterable[0], iterable
+    iterator = iter(iterable)
+    try:
+        first_element = next(iterator)
+    except StopIteration:
+        raise ValueError(f"iterable {iterable} had no elements to iterate over.")
+
+    if iterator == iterable:
+        # `iterable` was an iterator. Getting `first_element` will have removed it
+        # from `iterator`, so we need to add a fresh iterable with `first_element`
+        # added back in.
+        return_iterable = itertools.chain([first_element], iterator)
     else:
-        try:
-            first_element = next(iter(iterable))
-            return first_element, itertools.chain([first_element], iterable)
-        except StopIteration:
-            raise ValueError(f"iterable {iterable} had no elements to iterate over.")
+        # `iterable` was not an iterator; we can just return `iterable`.
+        # `iter(iterable)` will give a fresh iterator containing the first element.
+        # It's preferable to return `iterable` without modification so that users
+        # can generate new iterators from it as needed.
+        return_iterable = iterable
+
+    return first_element, return_iterable
diff --git a/tests/util/test_util.py b/tests/util/test_util.py
index d7cd18cbe..113fafbf2 100644
--- a/tests/util/test_util.py
+++ b/tests/util/test_util.py
@@ -25,6 +25,30 @@ def test_endless_iter_error():
     x = []
     with pytest.raises(ValueError, match="no elements"):
         util.endless_iter(x)
+    with pytest.raises(ValueError, match="needs a non-iterator Iterable"):
+        generator = (x for x in range(5))
+        util.endless_iter(generator)
+
+
+@given(
+    st.lists(
+        st.integers(),
+        min_size=1,
+    ),
+)
+def test_get_first_iter_element(input_seq):
+    with pytest.raises(ValueError, match="iterable.* had no elements"):
+        util.get_first_iter_element([])
+
+    first_element, new_iterable = util.get_first_iter_element(input_seq)
+    assert first_element == input_seq[0]
+    assert input_seq is new_iterable
+
+    an_iterator = (x for x in input_seq)
+    first_element, new_iterable = util.get_first_iter_element(input_seq)
+    assert first_element == input_seq[0]
+    assert list(an_iterator) == input_seq
+    assert list(an_iterator) == []
 
 
 @given(

From d5187629ec2544971d8153b418c69e8ae813c8a0 Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Fri, 9 Sep 2022 18:15:37 -0700
Subject: [PATCH 089/187] Bugfix in BC and tests -- masked as previously
 iterator ran out too early!

---
 src/imitation/algorithms/bc.py | 11 +++++++----
 tests/algorithms/test_bc.py    |  4 +++-
 tests/util/test_util.py        |  2 ++
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index 1cd9bf23b..72fa22317 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -58,15 +58,18 @@ def __iter__(self) -> Iterator[algo_base.TransitionMapping]:
         def batch_iterator():
 
             # Note: the islice here ensures we do not exceed self.n_epochs
+            num_batches = 0
             for epoch_num in itertools.islice(itertools.count(), self.n_epochs):
-                num_batches_in_epoch = 0
-                for num_batches_in_epoch, batch in enumerate(self.batch_loader):
+                got_data = False
+                for batch in self.batch_loader:
                     yield batch
+                    got_data = True
+                    num_batches += 1
 
-                if num_batches_in_epoch == 0:
+                if not got_data:
                     raise AssertionError(
                         f"Data loader returned no data after "
-                        f"{num_batches_in_epoch} batches, during epoch "
+                        f"{num_batches} batches, during epoch "
                         f"{epoch_num} -- did it reset correctly?",
                     )
                 if self.on_epoch_end is not None:
diff --git a/tests/algorithms/test_bc.py b/tests/algorithms/test_bc.py
index 87f6c18ed..8b55d9413 100644
--- a/tests/algorithms/test_bc.py
+++ b/tests/algorithms/test_bc.py
@@ -190,7 +190,9 @@ def test_bc_data_loader_empty_iter_error(
 
     bad_data_loader = _DataLoaderFailsOnNthIter(
         dummy_yield_value=dummy_yield_value,
-        no_yield_after_iter=no_yield_after_iter,
+        # add 1 as BC uses up an iteration from getting the first element
+        # for type checking
+        no_yield_after_iter=no_yield_after_iter + 1,
     )
     trainer = bc.BC(
         observation_space=cartpole_venv.observation_space,
diff --git a/tests/util/test_util.py b/tests/util/test_util.py
index 113fafbf2..36e03215d 100644
--- a/tests/util/test_util.py
+++ b/tests/util/test_util.py
@@ -19,6 +19,8 @@ def test_endless_iter():
     assert next(it) == 0
     assert next(it) == 1
     assert next(it) == 0
+    assert next(it) == 1
+    assert next(it) == 0
 
 
 def test_endless_iter_error():

From 4804f69b85456358ff7b18be298a559f20cd611f Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 16:31:02 +0200
Subject: [PATCH 090/187] Remove makefile for now

---
 Makefile | 41 -----------------------------------------
 1 file changed, 41 deletions(-)
 delete mode 100644 Makefile

diff --git a/Makefile b/Makefile
deleted file mode 100644
index 01cc08ef0..000000000
--- a/Makefile
+++ /dev/null
@@ -1,41 +0,0 @@
-SRC_FILES=src/ tests/ experiments/ examples/ docs/conf.py setup.py
-
-typecheck:
-	mypy ${SRC_FILES}
-	pytype -j auto ${SRC_FILES}
-
-shellcheck:
-	find . -path ./venv -prune -o -name '*.sh' -print0 | xargs -0 shellcheck
-
-format:
-	black ${SRC_FILES}
-	isort ${SRC_FILES}
-
-formatcheck:
-	flake8 --darglint-ignore-regex '.*' ${SRC_FILES}
-	black --check --diff ${SRC_FILES}
-	codespell -I .codespell.skip --skip='*.pyc,tests/testdata/*,*.ipynb,*.csv' ${SRC_FILES}
-
-docscheck:
-	pushd docs/ \
-	&& make clean \
-	&& make html \
-	&& popd
-
-cleandocs:
-	pushd docs/ \
-	&& make clean \
-	&& rm -rf _api/ \
-	&& popd
-
-cleantests:
-	rm -rf output/
-	rm -rf quickstart/
-
-cleancache:
-	rm -rf .mypy_cache/
-	rm -rf .pytest_cache/
-	rm -rf .pytype/
-	rm -rf .hypothesis/
-
-clean: cleandocs cleancache cleantests
\ No newline at end of file

From dcf80a88cae76b9f32a7b5195f413d3dfb01dceb Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 17:02:05 +0200
Subject: [PATCH 091/187] Added link to SB3 issue for future reference.

---
 tests/policies/test_policies.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index 163a02bb7..a408010a1 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -190,6 +190,7 @@ def test_normalize_features_extractor(obs_space: gym.Space) -> None:
         obs = th.as_tensor([obs_space.sample()])
         # TODO(juan) the cast below is because preprocess_obs has too general a type.
         #  this should be replaced with an overload or a generic.
+        #  https://github.com/DLR-RM/stable-baselines3/issues/1065
         obs = cast(th.Tensor, preprocessing.preprocess_obs(obs, obs_space))
         assert isinstance(obs, th.Tensor)
         flattened_obs = obs.flatten(1, -1)

From cee0c46efea8ea31d26628ffa7be3728f1d441c6 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 17:13:44 +0200
Subject: [PATCH 092/187] Fix types of train_imitation Only return
 "expert_stats" if all trajectories have reward.

---
 src/imitation/scripts/train_imitation.py | 24 +++++++-----------------
 1 file changed, 7 insertions(+), 17 deletions(-)

diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index f18ee6d23..b3f1e0f7d 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -174,23 +174,13 @@ def train_imitation(
 
         imit_stats = train.eval_policy(imit_policy, venv)
 
-    # TODO(juan): I'm not super happy with this solution for the type system.
-    #  is model._all_demos always Sequence[TrajectoryWithRew]? We can change
-    #  the type in the class definition. Same goes for demonstrations.load_expert_trajs.
-    #  using assert doesn't work because we'd have to loop over all the trajectories
-    #  and check that each one is a TrajectoryWithRew, which seems inefficient
-    #  just for adding type annotations.
-    trajectories: Sequence[types.TrajectoryWithRew]
-    if use_dagger:
-        trajectories = cast(Sequence[types.TrajectoryWithRew], model._all_demos)
-    else:
-        assert expert_trajs is not None
-        trajectories = cast(Sequence[types.TrajectoryWithRew], expert_trajs)
-
-    return {
-        "imit_stats": imit_stats,
-        "expert_stats": rollout.rollout_stats(trajectories),
-    }
+    stats = {"imit_stats": imit_stats}
+    trajectories = model._all_demos if use_dagger else expert_trajs
+    assert trajectories is not None
+    if all(isinstance(t, types.TrajectoryWithRew) for t in trajectories):
+        expert_stats = rollout.rollout_stats(cast(Sequence[types.TrajectoryWithRew], trajectories))
+        stats["expert_stats"] = expert_stats
+    return stats
 
 
 @train_imitation_ex.command

From 25aa3084be127caaafe1921a84388a499eaa5930 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 17:16:02 +0200
Subject: [PATCH 093/187] Modify assert in test_bc to reflect correct type

---
 tests/algorithms/test_bc.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/algorithms/test_bc.py b/tests/algorithms/test_bc.py
index 8b55d9413..7b168de1c 100644
--- a/tests/algorithms/test_bc.py
+++ b/tests/algorithms/test_bc.py
@@ -106,7 +106,7 @@ def test_bc(trainer: bc.BC, cartpole_venv):
         15,
         return_episode_rewards=True,
     )
-    assert isinstance(novice_rewards, (list, tuple))
+    assert isinstance(novice_rewards, list)
 
     trainer.train(
         n_epochs=1,
@@ -120,7 +120,7 @@ def test_bc(trainer: bc.BC, cartpole_venv):
         15,
         return_episode_rewards=True,
     )
-    assert isinstance(rewards_after_training, (list, tuple))
+    assert isinstance(rewards_after_training, list)
     assert reward_improvement.is_significant_reward_improvement(
         novice_rewards,
         rewards_after_training,

From fb7caff16c880593252cb10b1fefcbfe526c3854 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 17:19:29 +0200
Subject: [PATCH 094/187] Add ci/clean_notebooks.py to CI checks

---
 .circleci/config.yml | 4 ++--
 ci/code_checks.sh    | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 4b8f1a4fb..f199d03c1 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -28,14 +28,14 @@ executors:
     resource_class: xlarge
     environment:
       # If you change these, also change ci/code_checks.sh
-      SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py
+      SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/clean_notebooks.py
       NUM_CPUS: 8
   type:
     <<: *defaults
     resource_class: medium
     environment:
       # If you change these, also change ci/code_checks.sh
-      SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py
+      SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/clean_notebooks.py
       NUM_CPUS: 2
 
 commands:
diff --git a/ci/code_checks.sh b/ci/code_checks.sh
index 771ed2cee..fbb5e1a17 100755
--- a/ci/code_checks.sh
+++ b/ci/code_checks.sh
@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 
 # If you change these, also change .circleci/config.yml.
-SRC_FILES=(src/ tests/ experiments/ examples/ docs/conf.py setup.py)
+SRC_FILES=(src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/clean_notebooks.py)
 
 set -x  # echo commands
 set -e  # quit immediately on error

From b23b42d1dcaa8c5fae89c39c8aed4ae12fbaa25f Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 17:37:21 +0200
Subject: [PATCH 095/187] Improve clean_notebooks.py by allowing checking only
 mode.

---
 ci/clean_notebooks.py                    | 75 +++++++++++++++++++-----
 src/imitation/scripts/train_imitation.py |  4 +-
 2 files changed, 62 insertions(+), 17 deletions(-)
 mode change 100644 => 100755 ci/clean_notebooks.py

diff --git a/ci/clean_notebooks.py b/ci/clean_notebooks.py
old mode 100644
new mode 100755
index 505739bd7..51eb1a21f
--- a/ci/clean_notebooks.py
+++ b/ci/clean_notebooks.py
@@ -1,10 +1,31 @@
-import json
+#!/usr/bin/env python
+"""Clean all notebooks in the repository."""
 import pathlib
 
 import nbformat
 
 
-def clean_notebook(file: pathlib.Path):
+class UncleanNotebookError(Exception):
+    """Raised when a notebook is unclean."""
+
+
+def clean_notebook(file: pathlib.Path, check_only=False) -> None:
+    """Clean an ipynb notebook.
+
+    "Cleaning" means removing all output and metadata, as well as any other unnecessary
+    or vendor-dependent information or fields, so that it can be committed to the
+    repository, and so that artificial diffs are not introduced when the notebook is
+    executed.
+
+    Args:
+        file: Path to the notebook to clean.
+        check_only: If True, only check if the notebook is clean, and raise an
+            exception if it is not. If False, clean the notebook in-place.
+
+    Raises:
+        UncleanNotebookError: If `check_only` is True and the notebook is not clean.
+            Message contains brief description of the reason for the failure.
+    """
     # Read the notebook
     with open(file) as f:
         nb = nbformat.read(f, as_version=4)
@@ -13,21 +34,43 @@ def clean_notebook(file: pathlib.Path):
     # also reset the execution count
     # if the cell has no code, remove it
     for cell in nb.cells:
-        if 'outputs' in cell:
-            cell['outputs'] = []
-        if 'metadata' in cell:
-            cell['metadata'] = {}
-        if 'execution_count' in cell:
-            cell['execution_count'] = None
-        if cell['cell_type'] == 'code' and not cell['source']:
+        if "outputs" in cell and cell["outputs"]:
+            if check_only:
+                raise UncleanNotebookError(f"Notebook {file} has outputs")
+            cell["outputs"] = []
+        if "metadata" in cell and cell["metadata"]:
+            if check_only:
+                raise UncleanNotebookError(f"Notebook {file} has metadata")
+            cell["metadata"] = {}
+        if "execution_count" in cell and cell["execution_count"]:
+            if check_only:
+                raise UncleanNotebookError(f"Notebook {file} has execution count")
+            cell["execution_count"] = None
+        if cell["cell_type"] == "code" and not cell["source"]:
+            if check_only:
+                raise UncleanNotebookError(f"Notebook {file} has empty code cell")
             nb.cells.remove(cell)
 
-    # Write the notebook
-    with open(file, 'w') as f:
-        nbformat.write(nb, f)
+    if not check_only:
+        # Write the notebook
+        with open(file, "w") as f:
+            nbformat.write(nb, f)
+
 
+if __name__ == "__main__":
+    # if the argument --check has been passed, check if the notebooks are clean
+    # otherwise, clean them in-place
+    import argparse
+    import traceback
 
-if __name__ == '__main__':
-    for file in pathlib.Path.cwd().glob('**/*.ipynb'):
-        print(f'Cleaning {file}')
-        clean_notebook(file)
\ No newline at end of file
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--check", action="store_true")
+    args = parser.parse_args()
+    check_only = args.check
+    for file in pathlib.Path.cwd().glob("**/*.ipynb"):
+        print(f"Cleaning {file}" if not check_only else f"Checking {file}")
+        try:
+            clean_notebook(file, check_only=check_only)
+        except UncleanNotebookError:
+            traceback.print_exc()
+            exit(1)
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index b3f1e0f7d..a580a7ba3 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -178,7 +178,9 @@ def train_imitation(
     trajectories = model._all_demos if use_dagger else expert_trajs
     assert trajectories is not None
     if all(isinstance(t, types.TrajectoryWithRew) for t in trajectories):
-        expert_stats = rollout.rollout_stats(cast(Sequence[types.TrajectoryWithRew], trajectories))
+        expert_stats = rollout.rollout_stats(
+            cast(Sequence[types.TrajectoryWithRew], trajectories),
+        )
         stats["expert_stats"] = expert_stats
     return stats
 

From 56761d6deab473e7f9a4767bce5fdc67ea60e9c6 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 17:40:52 +0200
Subject: [PATCH 096/187] Add ipynb notebook checks to CI

---
 .circleci/config.yml | 4 ++++
 ci/code_checks.sh    | 1 +
 2 files changed, 5 insertions(+)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index f199d03c1..63095668b 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -223,6 +223,10 @@ jobs:
           # since they'll just get checked in a separate shellcheck invocation.
           exclude: SC1091
 
+      - run:
+          name: ipynb-check
+          command: ./ci/clean_notebooks.py --check
+
       - run:
           name: flake8
           command: flake8 --version && flake8 -j "${NUM_CPUS}" ${SRC_FILES}
diff --git a/ci/code_checks.sh b/ci/code_checks.sh
index fbb5e1a17..bcea3e6fa 100755
--- a/ci/code_checks.sh
+++ b/ci/code_checks.sh
@@ -7,6 +7,7 @@ set -x  # echo commands
 set -e  # quit immediately on error
 
 echo "Source format checking"
+./ci/clean_notebooks.py --check
 flake8 --darglint-ignore-regex '.*' "${SRC_FILES[@]}"
 black --check --diff "${SRC_FILES[@]}"
 codespell -I .codespell.skip --skip='*.pyc,tests/testdata/*,*.ipynb,*.csv' "${SRC_FILES[@]}"

From 7e87d4034471db3a031c6b2ea9011b7ada312661 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 17:49:02 +0200
Subject: [PATCH 097/187] Add support for explicit files for notebook cleaning

---
 .circleci/config.yml  |  2 +-
 ci/clean_notebooks.py | 25 ++++++++++++++++++++++++-
 ci/code_checks.sh     |  2 +-
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 63095668b..4911e2660 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -225,7 +225,7 @@ jobs:
 
       - run:
           name: ipynb-check
-          command: ./ci/clean_notebooks.py --check
+          command: ./ci/clean_notebooks.py --check ${SRC_FILES}
 
       - run:
           name: flake8
diff --git a/ci/clean_notebooks.py b/ci/clean_notebooks.py
index 51eb1a21f..ea9d270c6 100755
--- a/ci/clean_notebooks.py
+++ b/ci/clean_notebooks.py
@@ -64,10 +64,33 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
     import traceback
 
     parser = argparse.ArgumentParser()
+    # capture files and paths to clean
+    parser.add_argument(
+        "files",
+        nargs="+",
+        type=pathlib.Path,
+        help="List of files or paths to clean",
+    )
     parser.add_argument("--check", action="store_true")
     args = parser.parse_args()
     check_only = args.check
-    for file in pathlib.Path.cwd().glob("**/*.ipynb"):
+    # build list of files to scan from list of paths and files
+    files = []
+    if len(args.files) == 0:
+        parser.print_help()
+        exit(1)
+    for file in args.files:
+        if file.is_dir():
+            files.extend(file.glob("*.ipynb"))
+        else:
+            if file.suffix == ".ipynb":
+                files.append(file)
+            else:
+                print(f"Skipping {file} (not a notebook)")
+    if not files:
+        print("No notebooks found")
+        exit(1)
+    for file in files:
         print(f"Cleaning {file}" if not check_only else f"Checking {file}")
         try:
             clean_notebook(file, check_only=check_only)
diff --git a/ci/code_checks.sh b/ci/code_checks.sh
index bcea3e6fa..eaa9d651e 100755
--- a/ci/code_checks.sh
+++ b/ci/code_checks.sh
@@ -7,7 +7,7 @@ set -x  # echo commands
 set -e  # quit immediately on error
 
 echo "Source format checking"
-./ci/clean_notebooks.py --check
+./ci/clean_notebooks.py --check "${SRC_FILES[@]}"
 flake8 --darglint-ignore-regex '.*' "${SRC_FILES[@]}"
 black --check --diff "${SRC_FILES[@]}"
 codespell -I .codespell.skip --skip='*.pyc,tests/testdata/*,*.ipynb,*.csv' "${SRC_FILES[@]}"

From e1e09fdee5f5a18c976e0448357bae817a170272 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 17:49:19 +0200
Subject: [PATCH 098/187] Clean notebooks

---
 examples/1_train_bc.ipynb      | 2 +-
 examples/2_train_dagger.ipynb  | 8 ++------
 examples/3_train_gail.ipynb    | 2 +-
 examples/4_train_airl.ipynb    | 2 +-
 examples/6_train_mce.ipynb     | 2 +-
 examples/7_train_density.ipynb | 2 +-
 6 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/examples/1_train_bc.ipynb b/examples/1_train_bc.ipynb
index 7497c1dfc..63c4e4efc 100644
--- a/examples/1_train_bc.ipynb
+++ b/examples/1_train_bc.ipynb
@@ -199,4 +199,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/examples/2_train_dagger.ipynb b/examples/2_train_dagger.ipynb
index 6bab7af54..23ed2e59b 100644
--- a/examples/2_train_dagger.ipynb
+++ b/examples/2_train_dagger.ipynb
@@ -53,11 +53,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "import tempfile\n",
@@ -134,4 +130,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/examples/3_train_gail.ipynb b/examples/3_train_gail.ipynb
index 15eebedb0..9f3659e9f 100644
--- a/examples/3_train_gail.ipynb
+++ b/examples/3_train_gail.ipynb
@@ -179,4 +179,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/examples/4_train_airl.ipynb b/examples/4_train_airl.ipynb
index 11987f716..ef0794d48 100644
--- a/examples/4_train_airl.ipynb
+++ b/examples/4_train_airl.ipynb
@@ -176,4 +176,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/examples/6_train_mce.ipynb b/examples/6_train_mce.ipynb
index ff693ee7a..6f46e76a7 100644
--- a/examples/6_train_mce.ipynb
+++ b/examples/6_train_mce.ipynb
@@ -243,4 +243,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/examples/7_train_density.ipynb b/examples/7_train_density.ipynb
index 00070c50a..44ecced41 100644
--- a/examples/7_train_density.ipynb
+++ b/examples/7_train_density.ipynb
@@ -137,4 +137,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
\ No newline at end of file
+}

From 9bd0ff0ddfc75d3f41076f0832c39e7e79afb44f Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 18:02:22 +0200
Subject: [PATCH 099/187] Small improvements in util.py

---
 src/imitation/util/util.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index d40e9c92f..72fa6ff63 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -280,7 +280,8 @@ def safe_to_numpy(
         return obj
     else:
         if warn:
-            warnings.warn(f"Converted tensor {obj} to numpy array.")
+            warnings.warn(f"Converted tensor to numpy array, might affect performance. "
+                          f"Make sure this is the intended behavior.")
         return obj.detach().cpu().numpy()
 
 
@@ -337,6 +338,7 @@ def get_first_iter_element(iterable: Iterable[T]) -> Tuple[T, Iterable[T]]:
     except StopIteration:
         raise ValueError(f"iterable {iterable} had no elements to iterate over.")
 
+    return_iterable: Iterable[T]
     if iterator == iterable:
         # `iterable` was an iterator. Getting `first_element` will have removed it
         # from `iterator`, so we need to add a fresh iterable with `first_element`

From 42d1e978e52f2faad6316889ab9f02b3dc965224 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 18:24:16 +0200
Subject: [PATCH 100/187] Replace TransitionKind with TransitionsMinimal

---
 src/imitation/algorithms/base.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index dacaabac6..eb8f2602e 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -114,15 +114,10 @@ def __setstate__(self, state):
 
 TransitionMapping = Mapping[str, Union[np.ndarray, th.Tensor]]
 TransitionKind = TypeVar("TransitionKind", bound=types.TransitionsMinimal)
-# TODO(juan) AnyTransitions is used in many places but TransitionKind
-#  is just a type variable. If what we want is to allow for subclasses of
-#  types.TransitionsMinimal, we should use that type directly, as that is already
-#  built into the type system. Most places that use AnyTransitions aren't actually
-#  treating the type as a generic.
 AnyTransitions = Union[
     Iterable[types.Trajectory],
     Iterable[TransitionMapping],
-    TransitionKind,
+    types.TransitionsMinimal,
 ]
 
 

From d0e58aa1f364ed25d9b12e004f8299370422186a Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 18:25:52 +0200
Subject: [PATCH 101/187] Delete unused statement in test

---
 tests/rewards/test_reward_nets.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index de27269db..b6bb46cf7 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -163,13 +163,6 @@ def _make_env_and_save_reward_net(env_name, reward_type, tmpdir, random_state):
 def test_reward_valid(env_name, reward_type, tmpdir, random_state_fixed):
     """Test output of reward function is appropriate shape and type."""
     random_state = random_state_fixed
-    # TODO(juan) the line below is not being used?
-    venv = util.make_vec_env(
-        env_name,
-        n_envs=1,
-        parallel=False,
-        random_state=random_state,
-    )
     venv, tmppath = _make_env_and_save_reward_net(
         env_name,
         reward_type,

From 4117a8f141f0d0649451cf8fe1ba201c072a483c Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 18:26:19 +0200
Subject: [PATCH 102/187] Update src/imitation/util/util.py

Co-authored-by: Adam Gleave <adam@gleave.me>
---
 src/imitation/util/util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index d40e9c92f..e330f2a2f 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -257,7 +257,7 @@ def safe_to_numpy(obj: None, warn: bool = False) -> None:
 
 def safe_to_numpy(
     obj: Optional[Union[np.ndarray, th.Tensor]],
-    warn=False,
+    warn: bool = False,
 ) -> Optional[np.ndarray]:
     """Convert torch tensor to numpy.
 

From 908925a62e2de01209680d0020094b14e22f4d3c Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 18:26:33 +0200
Subject: [PATCH 103/187] Update src/imitation/util/util.py

Co-authored-by: Adam Gleave <adam@gleave.me>
---
 src/imitation/util/util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index e330f2a2f..3947af1f9 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -166,7 +166,7 @@ def make_seeds(random_state: np.random.RandomState, n: int) -> List[int]:
 def make_seeds(
     random_state: np.random.RandomState,
     n: Optional[int] = None,
-) -> Union[List[int], int]:
+) -> Union[Sequence[int], int]:
     """Generate n random seeds from a random state.
 
     Args:

From 16c440e6c3b0587089cc7fc4650cad42c036910e Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 18:27:55 +0200
Subject: [PATCH 104/187] Make type ignore specific to pytype

---
 src/imitation/util/util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 5f82985a5..3b8459325 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -275,7 +275,7 @@ def safe_to_numpy(
     """
     if obj is None:
         # We ignore the type due to https://github.com/google/pytype/issues/445
-        return None  # type: ignore[bad-return-type]
+        return None  # pytype: disable=bad-return-type
     elif isinstance(obj, np.ndarray):
         return obj
     else:

From 2c18c7e17e6cbc005ada2019ea9b31f1ac1f0593 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 18:29:21 +0200
Subject: [PATCH 105/187] Linting

---
 src/imitation/util/util.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 3b8459325..fdaee8c5e 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -280,8 +280,10 @@ def safe_to_numpy(
         return obj
     else:
         if warn:
-            warnings.warn(f"Converted tensor to numpy array, might affect performance. "
-                          f"Make sure this is the intended behavior.")
+            warnings.warn(
+                "Converted tensor to numpy array, might affect performance. "
+                "Make sure this is the intended behavior.",
+            )
         return obj.detach().cpu().numpy()
 
 

From 831bfe472c7299983edab2957d37a033d8f318e7 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 20:36:02 +0200
Subject: [PATCH 106/187] Migrate from RandomState (deprecated) to Generator

---
 examples/1_train_bc.ipynb                     |   6 +-
 examples/2_train_dagger.ipynb                 |   4 +-
 examples/3_train_gail.ipynb                   |   4 +-
 examples/4_train_airl.ipynb                   |   4 +-
 examples/5_train_preference_comparisons.ipynb |  11 +-
 examples/6_train_mce.ipynb                    |  10 +-
 examples/7_train_density.ipynb                |   6 +-
 examples/quickstart.py                        |   6 +-
 src/imitation/algorithms/bc.py                |  12 +--
 src/imitation/algorithms/dagger.py            |  24 ++---
 src/imitation/algorithms/density.py           |   8 +-
 src/imitation/algorithms/mce_irl.py           |  16 +--
 .../algorithms/preference_comparisons.py      |  82 +++++++-------
 src/imitation/data/rollout.py                 |  18 ++--
 src/imitation/envs/examples/model_envs.py     |  38 +++----
 src/imitation/envs/resettable_env.py          |   8 +-
 src/imitation/policies/exploration_wrapper.py |  10 +-
 src/imitation/scripts/common/common.py        |  15 ++-
 src/imitation/scripts/common/rl.py            |  12 +--
 src/imitation/scripts/common/seeding.py       |  26 -----
 src/imitation/scripts/common/train.py         |   8 +-
 src/imitation/scripts/config/eval_policy.py   |   6 +-
 .../config/train_preference_comparisons.py    |   3 +-
 src/imitation/scripts/config/train_rl.py      |   3 +-
 src/imitation/scripts/eval_policy.py          |   6 +-
 src/imitation/scripts/train_imitation.py      |   8 +-
 .../scripts/train_preference_comparisons.py   |  15 ++-
 src/imitation/scripts/train_rl.py             |   6 +-
 src/imitation/util/util.py                    |  16 +--
 tests/algorithms/test_adversarial.py          |  48 ++++-----
 tests/algorithms/test_bc.py                   |  17 ++-
 tests/algorithms/test_dagger.py               |  51 ++++-----
 tests/algorithms/test_density_baselines.py    |  12 +--
 tests/algorithms/test_mce_irl.py              |  24 ++---
 .../algorithms/test_preference_comparisons.py | 102 ++++++++----------
 tests/conftest.py                             |  22 ++--
 tests/data/test_rollout.py                    |  54 +++++-----
 tests/policies/test_exploration_wrapper.py    |  24 ++---
 tests/policies/test_policies.py               |  28 +++--
 tests/policies/test_replay_buffer_wrapper.py  |  16 ++-
 tests/rewards/test_reward_nets.py             |  36 +++----
 tests/rewards/test_reward_wrapper.py          |   9 +-
 tests/scripts/test_scripts.py                 |   5 +-
 43 files changed, 386 insertions(+), 453 deletions(-)
 delete mode 100644 src/imitation/scripts/common/seeding.py

diff --git a/examples/1_train_bc.ipynb b/examples/1_train_bc.ipynb
index 63c4e4efc..5de3ffd82 100644
--- a/examples/1_train_bc.ipynb
+++ b/examples/1_train_bc.ipynb
@@ -85,12 +85,12 @@
     "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "import numpy as np\n",
     "\n",
-    "random_state = np.random.RandomState()\n",
+    "rng = np.random.default_rng()\n",
     "rollouts = rollout.rollout(\n",
     "    expert,\n",
     "    DummyVecEnv([lambda: RolloutInfoWrapper(env)]),\n",
     "    rollout.make_sample_until(min_timesteps=None, min_episodes=50),\n",
-    "    random_state=random_state,\n",
+    "    rng=rng,\n",
     ")\n",
     "transitions = rollout.flatten_trajectories(rollouts)"
    ]
@@ -135,7 +135,7 @@
     "    observation_space=env.observation_space,\n",
     "    action_space=env.action_space,\n",
     "    demonstrations=transitions,\n",
-    "    random_state=random_state,\n",
+    "    rng=rng,\n",
     ")"
    ]
   },
diff --git a/examples/2_train_dagger.ipynb b/examples/2_train_dagger.ipynb
index 23ed2e59b..bd2077659 100644
--- a/examples/2_train_dagger.ipynb
+++ b/examples/2_train_dagger.ipynb
@@ -70,7 +70,7 @@
     "bc_trainer = bc.BC(\n",
     "    observation_space=env.observation_space,\n",
     "    action_space=env.action_space,\n",
-    "    random_state=np.random.RandomState(),\n",
+    "    rng=np.random.default_rng(),\n",
     ")\n",
     "\n",
     "with tempfile.TemporaryDirectory(prefix=\"dagger_example_\") as tmpdir:\n",
@@ -80,7 +80,7 @@
     "        scratch_dir=tmpdir,\n",
     "        expert_policy=expert,\n",
     "        bc_trainer=bc_trainer,\n",
-    "        random_state=np.random.RandomState(),\n",
+    "        rng=np.random.default_rng(),\n",
     "    )\n",
     "\n",
     "    dagger_trainer.train(2000)"
diff --git a/examples/3_train_gail.ipynb b/examples/3_train_gail.ipynb
index 9f3659e9f..abadd6b17 100644
--- a/examples/3_train_gail.ipynb
+++ b/examples/3_train_gail.ipynb
@@ -61,12 +61,12 @@
     "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "import numpy as np\n",
     "\n",
-    "random_state = np.random.RandomState()\n",
+    "rng = np.random.default_rng()\n",
     "rollouts = rollout.rollout(\n",
     "    expert,\n",
     "    DummyVecEnv([lambda: RolloutInfoWrapper(gym.make(\"seals/CartPole-v0\"))] * 5),\n",
     "    rollout.make_sample_until(min_timesteps=None, min_episodes=60),\n",
-    "    random_state=random_state,\n",
+    "    rng=rng,\n",
     ")"
    ]
   },
diff --git a/examples/4_train_airl.ipynb b/examples/4_train_airl.ipynb
index ef0794d48..96191894b 100644
--- a/examples/4_train_airl.ipynb
+++ b/examples/4_train_airl.ipynb
@@ -58,12 +58,12 @@
     "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "import numpy as np\n",
     "\n",
-    "random_state = np.random.RandomState()\n",
+    "rng = np.random.default_rng()\n",
     "rollouts = rollout.rollout(\n",
     "    expert,\n",
     "    DummyVecEnv([lambda: RolloutInfoWrapper(gym.make(\"seals/CartPole-v0\"))] * 5),\n",
     "    rollout.make_sample_until(min_timesteps=None, min_episodes=60),\n",
-    "    random_state=random_state,\n",
+    "    rng=rng,\n",
     ")"
    ]
   },
diff --git a/examples/5_train_preference_comparisons.ipynb b/examples/5_train_preference_comparisons.ipynb
index 511b29bad..2a537c6bd 100644
--- a/examples/5_train_preference_comparisons.ipynb
+++ b/examples/5_train_preference_comparisons.ipynb
@@ -32,18 +32,19 @@
     "from stable_baselines3 import PPO\n",
     "import numpy as np\n",
     "\n",
-    "random_state = np.random.RandomState(0)\n",
+    "rng = np.random.default_rng(0)\n",
     "\n",
-    "venv = make_vec_env(\"Pendulum-v1\", random_state=random_state)\n",
+    "venv = make_vec_env(\"Pendulum-v1\", rng=rng)\n",
     "\n",
     "reward_net = BasicRewardNet(\n",
     "    venv.observation_space, venv.action_space, normalize_input_layer=RunningNorm\n",
     ")\n",
     "\n",
     "fragmenter = preference_comparisons.RandomFragmenter(\n",
-    "    warning_threshold=0, random_state=random.Random(0)\n",
+    "    warning_threshold=0,\n",
+    "    rng=rng,\n",
     ")\n",
-    "gatherer = preference_comparisons.SyntheticGatherer(random_state=random_state)\n",
+    "gatherer = preference_comparisons.SyntheticGatherer(rng=rng)\n",
     "preference_model = preference_comparisons.PreferenceModel(reward_net)\n",
     "reward_trainer = preference_comparisons.BasicRewardTrainer(\n",
     "    model=reward_net,\n",
@@ -71,7 +72,7 @@
     "    reward_fn=reward_net,\n",
     "    venv=venv,\n",
     "    exploration_frac=0.0,\n",
-    "    random_state=random_state,\n",
+    "    rng=rng,\n",
     ")\n",
     "\n",
     "pref_comparisons = preference_comparisons.PreferenceComparisons(\n",
diff --git a/examples/6_train_mce.ipynb b/examples/6_train_mce.ipynb
index 6f46e76a7..662bff7a2 100644
--- a/examples/6_train_mce.ipynb
+++ b/examples/6_train_mce.ipynb
@@ -62,19 +62,19 @@
     "\n",
     "_, om = mce_occupancy_measures(env_single, pi=pi)\n",
     "\n",
-    "random_state = np.random.RandomState()\n",
+    "rng = np.random.default_rng()\n",
     "expert = TabularPolicy(\n",
     "    state_space=env_single.pomdp_state_space,\n",
     "    action_space=env_single.action_space,\n",
     "    pi=pi,\n",
-    "    random_state=random_state,\n",
+    "    rng=rng,\n",
     ")\n",
     "\n",
     "expert_trajs = rollout.generate_trajectories(\n",
     "    policy=expert,\n",
     "    venv=state_venv,\n",
     "    sample_until=rollout.make_min_timesteps(5000),\n",
-    "    random_state=random_state,\n",
+    "    rng=rng,\n",
     ")\n",
     "\n",
     "print(\"Expert stats: \", rollout.rollout_stats(expert_trajs))"
@@ -115,7 +115,7 @@
     "        reward_net,\n",
     "        log_interval=250,\n",
     "        optimizer_kwargs=dict(lr=lr),\n",
-    "        random_state=random_state,\n",
+    "        rng=rng,\n",
     "    )\n",
     "    occ_measure = mce_irl.train(**kwargs)\n",
     "\n",
@@ -123,7 +123,7 @@
     "        policy=mce_irl.policy,\n",
     "        venv=state_venv,\n",
     "        sample_until=rollout.make_min_timesteps(5000),\n",
-    "        random_state=random_state,\n",
+    "        rng=rng,\n",
     "    )\n",
     "    print(\"Imitation stats: \", rollout.rollout_stats(imitation_trajs))\n",
     "\n",
diff --git a/examples/7_train_density.ipynb b/examples/7_train_density.ipynb
index 44ecced41..f16b68b2b 100644
--- a/examples/7_train_density.ipynb
+++ b/examples/7_train_density.ipynb
@@ -59,16 +59,16 @@
     "from stable_baselines3 import PPO\n",
     "import numpy as np\n",
     "\n",
-    "random_state = np.random.RandomState()\n",
+    "rng = np.random.default_rng()\n",
     "env_name = \"Pendulum-v1\"\n",
-    "env = util.make_vec_env(env_name, random_state=random_state, n_envs=N_VEC)\n",
+    "env = util.make_vec_env(env_name, rng=rng, n_envs=N_VEC)\n",
     "rollouts = types.load(\"../tests/testdata/expert_models/pendulum_0/rollouts/final.pkl\")\n",
     "\n",
     "\n",
     "imitation_trainer = PPO(ActorCriticPolicy, env, learning_rate=3e-4, n_steps=2048)\n",
     "density_trainer = db.DensityAlgorithm(\n",
     "    venv=env,\n",
-    "    random_state=random_state,\n",
+    "    rng=rng,\n",
     "    demonstrations=rollouts,\n",
     "    rl_algo=imitation_trainer,\n",
     "    density_type=db.DensityType.STATE_ACTION_DENSITY,\n",
diff --git a/examples/quickstart.py b/examples/quickstart.py
index b255f971e..ae054186c 100644
--- a/examples/quickstart.py
+++ b/examples/quickstart.py
@@ -15,7 +15,7 @@
 from imitation.data.wrappers import RolloutInfoWrapper
 
 env = gym.make("CartPole-v1")
-random_state = np.random.RandomState(0)
+rng = np.random.default_rng(0)
 
 
 def train_expert():
@@ -42,7 +42,7 @@ def sample_expert_transitions():
         expert,
         DummyVecEnv([lambda: RolloutInfoWrapper(env)]),
         rollout.make_sample_until(min_timesteps=None, min_episodes=50),
-        random_state=random_state,
+        rng=rng,
     )
     return rollout.flatten_trajectories(rollouts)
 
@@ -52,7 +52,7 @@ def sample_expert_transitions():
     observation_space=env.observation_space,
     action_space=env.action_space,
     demonstrations=transitions,
-    random_state=random_state,
+    rng=rng,
 )
 
 reward, _ = evaluate_policy(bc_trainer.policy, env, n_eval_episodes=3, render=True)
diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index 72fa22317..676679161 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -191,14 +191,14 @@ class RolloutStatsComputer:
     def __call__(
         self,
         policy: policies.ActorCriticPolicy,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
     ) -> Mapping[str, float]:
         if self.venv is not None and self.n_episodes > 0:
             trajs = rollout.generate_trajectories(
                 policy,
                 self.venv,
                 rollout.make_min_episodes(self.n_episodes),
-                random_state=random_state,
+                rng=rng,
             )
             return rollout.rollout_stats(trajs)
         else:
@@ -280,7 +280,7 @@ def __init__(
         *,
         observation_space: gym.Space,
         action_space: gym.Space,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
         policy: Optional[policies.ActorCriticPolicy] = None,
         demonstrations: Optional[algo_base.AnyTransitions] = None,
         batch_size: int = 32,
@@ -296,7 +296,7 @@ def __init__(
         Args:
             observation_space: the observation space of the environment.
             action_space: the action space of the environment.
-            random_state: the random state to use for the random number generator.
+            rng: the random state to use for the random number generator.
             policy: a Stable Baselines3 policy; if unspecified,
                 defaults to `FeedForward32Policy`.
             demonstrations: Demonstrations from an expert (optional). Transitions
@@ -327,7 +327,7 @@ def __init__(
         self.action_space = action_space
         self.observation_space = observation_space
 
-        self.random_state = random_state
+        self.rng = rng
 
         if policy is None:
             policy = policy_base.FeedForward32Policy(
@@ -453,7 +453,7 @@ def _on_epoch_end(epoch_number: int):
             loss = self.trainer(batch)
 
             if batch_num % log_interval == 0:
-                rollout_stats = compute_rollout_stats(self.policy, self.random_state)
+                rollout_stats = compute_rollout_stats(self.policy, self.rng)
 
                 self._bc_logger.log_batch(
                     batch_num,
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index 8df29d63c..deea8b7a0 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -152,7 +152,7 @@ def __init__(
         get_robot_acts: Callable[[np.ndarray], np.ndarray],
         beta: float,
         save_dir: types.AnyPath,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
     ):
         """Builds InteractiveTrajectoryCollector.
 
@@ -165,7 +165,7 @@ def __init__(
                 robot action. The choice of robot or human action is independently
                 randomized for each individual `Env` at every timestep.
             save_dir: directory to save collected trajectories in.
-            random_state: random state for random number generation.
+            rng: random state for random number generation.
         """
         super().__init__(venv)
         self.get_robot_acts = get_robot_acts
@@ -177,7 +177,7 @@ def __init__(
         self._done_before = True
         self._is_reset = False
         self._last_user_actions = None
-        self.rng = random_state
+        self.rng = rng
 
     def seed(self, seed=Optional[int]) -> List[Union[None, int]]:
         """Set the seed for the DAgger random number generator and wrapped VecEnv.
@@ -192,7 +192,7 @@ def seed(self, seed=Optional[int]) -> List[Union[None, int]]:
             A list containing the seeds for each individual env. Note that all list
             elements may be None, if the env does not return anything when seeded.
         """
-        self.rng = np.random.RandomState(seed=seed)
+        self.rng = np.random.default_rng(seed=seed)
         return self.venv.seed(seed)
 
     def reset(self) -> np.ndarray:
@@ -308,7 +308,7 @@ def __init__(
         *,
         venv: vec_env.VecEnv,
         scratch_dir: types.AnyPath,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
         beta_schedule: Optional[Callable[[int], float]] = None,
         bc_trainer: bc.BC,
         custom_logger: Optional[logger.HierarchicalLogger] = None,
@@ -319,7 +319,7 @@ def __init__(
             venv: Vectorized training environment.
             scratch_dir: Directory to use to store intermediate training
                 information (e.g. for resuming training).
-            random_state: random state for random number generation.
+            rng: random state for random number generation.
             beta_schedule: Provides a value of `beta` (the probability of taking
                 expert action in any given state) at each round of training. If
                 `None`, then `linear_beta_schedule` will be used instead.
@@ -336,7 +336,7 @@ def __init__(
         self.round_num = 0
         self._last_loaded_round = -1
         self._all_demos = []
-        self.random_state = random_state
+        self.rng = rng
 
         utils.check_for_correct_spaces(
             self.venv,
@@ -478,7 +478,7 @@ def create_trajectory_collector(self) -> InteractiveTrajectoryCollector:
             get_robot_acts=lambda acts: self.bc_trainer.policy.predict(acts)[0],
             beta=beta,
             save_dir=save_dir,
-            random_state=self.random_state,
+            rng=self.rng,
         )
         return collector
 
@@ -533,7 +533,7 @@ def __init__(
         venv: vec_env.VecEnv,
         scratch_dir: types.AnyPath,
         expert_policy: policies.BasePolicy,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
         expert_trajs: Optional[Sequence[types.Trajectory]] = None,
         **dagger_trainer_kwargs,
     ):
@@ -547,7 +547,7 @@ def __init__(
             scratch_dir: Directory to use to store intermediate training
                 information (e.g. for resuming training).
             expert_policy: The expert policy used to generate synthetic demonstrations.
-            random_state: Random state to use for the random number generator.
+            rng: Random state to use for the random number generator.
             expert_trajs: Optional starting dataset that is inserted into the round 0
                 dataset.
             dagger_trainer_kwargs: Other keyword arguments passed to the
@@ -560,7 +560,7 @@ def __init__(
         super().__init__(
             venv=venv,
             scratch_dir=scratch_dir,
-            random_state=random_state,
+            rng=rng,
             **dagger_trainer_kwargs,
         )
         self.expert_policy = expert_policy
@@ -644,7 +644,7 @@ def train(
                 venv=collector,
                 sample_until=sample_until,
                 deterministic_policy=True,
-                random_state=collector.rng,
+                rng=collector.rng,
             )
 
             for traj in trajectories:
diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index 34dcff2c7..dd35e2b30 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -61,7 +61,7 @@ def __init__(
         *,
         demonstrations: Optional[base.AnyTransitions],
         venv: vec_env.VecEnv,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
         density_type: DensityType = DensityType.STATE_ACTION_DENSITY,
         kernel: str = "gaussian",
         kernel_bandwidth: float = 0.5,
@@ -85,7 +85,7 @@ def __init__(
                 any environment interaction to fit the reward model, but we use this
                 to extract the observation and action space, and to train the RL
                 algorithm `rl_algo` (if specified).
-            random_state: random state for sampling from demonstrations.
+            rng: random state for sampling from demonstrations.
             rl_algo: An RL algorithm to train on the resulting reward model (optional).
             is_stationary: if True, share same density models for all timesteps;
                 if False, use a different density model for each timestep.
@@ -122,7 +122,7 @@ def __init__(
         self.standardise = standardise_inputs
         self._scaler = None
         self._density_models = dict()
-        self.random_state = random_state
+        self.rng = rng
 
         self.rl_algo = rl_algo
         self.buffering_wrapper = wrappers.BufferingWrapper(self.venv)
@@ -357,7 +357,7 @@ def test_policy(self, *, n_trajectories: int = 10, true_reward: bool = True):
             self.rl_algo,
             self.venv if true_reward else self.venv_wrapped,
             sample_until=rollout.make_min_episodes(n_trajectories),
-            random_state=self.random_state,
+            rng=self.rng,
         )
         # We collect `trajs` above so disregard return value from `pop_trajectories`,
         # but still call it to clear out any saved trajectories.
diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index 273655b03..4b85717b0 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -142,14 +142,14 @@ class TabularPolicy(policies.BasePolicy):
     """A tabular policy. Cannot be trained -- prediction only."""
 
     pi: np.ndarray
-    rng: np.random.RandomState
+    rng: np.random.Generator
 
     def __init__(
         self,
         state_space: gym.Space,
         action_space: gym.Space,
         pi: np.ndarray,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
     ):
         """Builds TabularPolicy.
 
@@ -158,14 +158,14 @@ def __init__(
             action_space: The action space of the environment.
             pi: A tabular policy. Three-dimensional array, where pi[t,s,a]
                 is the probability of taking action a at state s at timestep t.
-            random_state: Random state, used for sampling when `predict` is called with
+            rng: Random state, used for sampling when `predict` is called with
                 `deterministic=False`.
         """
         assert isinstance(state_space, gym.spaces.Discrete), "state not tabular"
         assert isinstance(action_space, gym.spaces.Discrete), "action not tabular"
         # What we call state space here is observation space in SB3 nomenclature.
         super().__init__(observation_space=state_space, action_space=action_space)
-        self.rng = random_state
+        self.rng = rng
         self.set_pi(pi)
 
     def set_pi(self, pi: np.ndarray) -> None:
@@ -259,7 +259,7 @@ def __init__(
         demonstrations: Optional[MCEDemonstrations],
         env: resettable_env.TabularModelEnv,
         reward_net: reward_nets.RewardNet,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
         optimizer_cls: Type[th.optim.Optimizer] = th.optim.Adam,
         optimizer_kwargs: Optional[Mapping[str, Any]] = None,
         discount: float = 1.0,
@@ -279,7 +279,7 @@ def __init__(
                 The demonstrations must have observations one-hot coded unless
                 demonstrations is a state-occupancy measure.
             env: a tabular MDP.
-            random_state: random state used for sampling from policy.
+            rng: random state used for sampling from policy.
             reward_net: a neural network that computes rewards for the supplied
                 observations.
             optimizer_cls: optimizer to use for supervised training.
@@ -313,7 +313,7 @@ def __init__(
         self.linf_eps = linf_eps
         self.grad_l2_eps = grad_l2_eps
         self.log_interval = log_interval
-        self.rng = random_state
+        self.rng = rng
 
         # Initialize policy to be uniform random. We don't use this for MCE IRL
         # training, but it gives us something to return at all times with `policy`
@@ -324,7 +324,7 @@ def __init__(
             state_space=self.env.pomdp_state_space,
             action_space=self.env.action_space,
             pi=uniform_pi,
-            random_state=self.rng,
+            rng=self.rng,
         )
 
     def _set_demo_from_trajectories(self, trajs: Iterable[types.Trajectory]) -> None:
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index c83e84e10..36ebcab92 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -98,19 +98,19 @@ class TrajectoryDataset(TrajectoryGenerator):
     def __init__(
         self,
         trajectories: Sequence[TrajectoryWithRew],
-        random_state: random.Random,
+        rng: np.random.Generator,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Creates a dataset loaded from `path`.
 
         Args:
             trajectories: the dataset of rollouts.
-            random_state: Random state for RNG used for shuffling dataset.
+            rng: Random state for RNG used for shuffling dataset.
             custom_logger: Where to log to; if None (default), creates a new logger.
         """
         super().__init__(custom_logger=custom_logger)
         self._trajectories = trajectories
-        self.rng = random_state
+        self.rng = rng
 
     def sample(self, steps: int) -> Sequence[TrajectoryWithRew]:
         # make a copy before shuffling
@@ -127,7 +127,7 @@ def __init__(
         algorithm: base_class.BaseAlgorithm,
         reward_fn: Union[reward_function.RewardFn, reward_nets.RewardNet],
         venv: vec_env.VecEnv,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
         exploration_frac: float = 0.0,
         switch_prob: float = 0.5,
         random_prob: float = 0.5,
@@ -140,7 +140,7 @@ def __init__(
             reward_fn: either a RewardFn or a RewardNet instance that will supply
                 the rewards used for training the agent.
             venv: vectorized environment to train in.
-            random_state: random state used for exploration and for sampling.
+            rng: random state used for exploration and for sampling.
             exploration_frac: fraction of the trajectories that will be generated
                 partially randomly rather than only by the agent when sampling.
             switch_prob: the probability of switching the current policy at each
@@ -162,7 +162,7 @@ def __init__(
             reward_fn = reward_fn.predict_processed
         self.reward_fn = reward_fn
         self.exploration_frac = exploration_frac
-        self.random_state = random_state
+        self.rng = rng
 
         # The BufferingWrapper records all trajectories, so we can return
         # them after training. This should come first (before the wrapper that
@@ -200,7 +200,7 @@ def __init__(
             venv=algo_venv,
             random_prob=random_prob,
             switch_prob=switch_prob,
-            random_state=self.random_state,
+            rng=self.rng,
         )
 
     def train(self, steps: int, **kwargs) -> None:
@@ -270,7 +270,7 @@ def sample(self, steps: int) -> Sequence[types.TrajectoryWithRew]:
                 # deterministic. If self.algorithm is stochastic, then policy_callable
                 # will also be stochastic.
                 deterministic_policy=False,
-                random_state=self.random_state,
+                rng=self.rng,
             )
             additional_trajs, _ = self.buffering_wrapper.pop_finished_trajectories()
             agent_trajs = list(agent_trajs) + list(additional_trajs)
@@ -294,7 +294,7 @@ def sample(self, steps: int) -> Sequence[types.TrajectoryWithRew]:
                 # buffering_wrapper collects rollouts from a non-deterministic policy,
                 # so we do that here as well for consistency.
                 deterministic_policy=False,
-                random_state=self.random_state,
+                rng=self.rng,
             )
             exploration_trajs, _ = self.buffering_wrapper.pop_finished_trajectories()
             exploration_trajs = _get_trajectories(exploration_trajs, exploration_steps)
@@ -577,22 +577,21 @@ class RandomFragmenter(Fragmenter):
 
     def __init__(
         self,
-        random_state: random.Random,
+        rng: np.random.Generator,
         warning_threshold: int = 10,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initialize the fragmenter.
 
         Args:
-            random_state: the random state for the random number generator
+            rng: the random state for the random number generator
             warning_threshold: give a warning if the number of available
                 transitions is less than this many times the number of
                 required samples. Set to 0 to disable this warning.
             custom_logger: Where to log to; if None (default), creates a new logger.
         """
         super().__init__(custom_logger)
-        assert isinstance(random_state, random.Random)
-        self.rng = random_state
+        self.rng = rng
         self.warning_threshold = warning_threshold
 
     def __call__(
@@ -644,9 +643,9 @@ def __call__(
 
         # we need two fragments for each comparison
         for _ in range(2 * num_pairs):
-            traj = self.rng.choices(trajectories, weights, k=1)[0]
+            traj = self.rng.choice(trajectories, p=np.array(weights) / sum(weights))
             n = len(traj)
-            start = self.rng.randint(0, n - fragment_length)
+            start = self.rng.integers(0, n - fragment_length, endpoint=True)
             end = start + fragment_length
             terminal = (end == n) and traj.terminal
             fragment = TrajectoryWithRew(
@@ -782,20 +781,20 @@ class PreferenceGatherer(abc.ABC):
 
     def __init__(
         self,
-        random_state: Optional[np.random.RandomState] = None,
+        rng: Optional[np.random.Generator] = None,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Initializes the preference gatherer.
 
         Args:
-            random_state: random state for random number generation, if applicable.
+            rng: random state for random number generation, if applicable.
             custom_logger: Where to log to; if None (default), creates a new logger.
         """
         # The random seed isn't used here, but it's useful to have this
         # as an argument nevertheless because that means we can always
         # pass in a seed in training scripts (without worrying about whether
         # the PreferenceGatherer we use needs one).
-        del random_state
+        del rng
         self.logger = custom_logger or imit_logger.configure()
 
     @abc.abstractmethod
@@ -825,7 +824,7 @@ def __init__(
         temperature: float = 1,
         discount_factor: float = 1,
         sample: bool = True,
-        random_state: Optional[np.random.RandomState] = None,
+        rng: Optional[np.random.Generator] = None,
         threshold: float = 50,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
@@ -842,7 +841,7 @@ def __init__(
                 a Bernoulli distribution (or 0.5 in the case of ties with zero
                 temperature). If False, then the underlying Bernoulli probabilities
                 are returned instead.
-            random_state: random state for the internal RNG
+            rng: random state for the internal RNG
                 (only used if temperature > 0 and sample)
             threshold: preferences are sampled from a softmax of returns.
                 To avoid overflows, we clip differences in returns that are
@@ -858,11 +857,11 @@ def __init__(
         self.temperature = temperature
         self.discount_factor = discount_factor
         self.sample = sample
-        self.rng = random_state
+        self.rng = rng
         self.threshold = threshold
 
         if self.sample and self.rng is None:
-            raise ValueError("If sample is True, then random_state must be provided.")
+            raise ValueError("If sample is True, then rng must be provided.")
 
     def __call__(self, fragment_pairs: Sequence[TrajectoryWithRewPair]) -> np.ndarray:
         """Computes probability fragment 1 is preferred over fragment 2."""
@@ -1204,7 +1203,7 @@ def __init__(
         self,
         model: reward_nets.RewardEnsemble,
         loss: RewardLoss,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
         batch_size: int = 32,
         epochs: int = 1,
         lr: float = 1e-3,
@@ -1216,7 +1215,7 @@ def __init__(
         Args:
             model: the RewardNet instance to be trained
             loss: the loss to use
-            random_state: random state for the internal RNG used in bagging
+            rng: random state for the internal RNG used in bagging
             batch_size: number of fragment pairs per batch
             epochs: number of epochs in each training iteration (can be adjusted
                 on the fly by specifying an `epoch_multiplier` in `self.train()`
@@ -1244,7 +1243,7 @@ def __init__(
             weight_decay,
             custom_logger,
         )
-        self.rng = random_state
+        self.rng = rng
 
     def _training_inner_loop(
         self,
@@ -1294,7 +1293,7 @@ def get_base_model(
 def _make_reward_trainer(
     reward_model: reward_nets.RewardNet,
     loss: RewardLoss,
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
     reward_trainer_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> RewardTrainer:
     """Construct the correct type of reward trainer for this reward function."""
@@ -1316,7 +1315,7 @@ def _make_reward_trainer(
             return EnsembleTrainer(
                 base_model,
                 loss,
-                random_state=random_state,
+                rng=rng,
                 **reward_trainer_kwargs,
             )
         else:
@@ -1357,7 +1356,7 @@ def __init__(
         initial_epoch_multiplier: float = 200.0,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
         allow_variable_horizon: bool = False,
-        random_state: Optional[np.random.RandomState] = None,
+        rng: Optional[np.random.Generator] = None,
         query_schedule: Union[str, type_aliases.Schedule] = "hyperbolic",
     ):
         """Initialize the preference comparison trainer.
@@ -1412,7 +1411,7 @@ def __init__(
                 condition, and can seriously confound evaluation. Read
                 https://imitation.readthedocs.io/en/latest/guide/variable_horizon.html
                 before overriding this.
-            random_state: random state to use for initializing subcomponents such as
+            rng: random state to use for initializing subcomponents such as
                 fragmenter.
                 Only used when default components are used; if you instantiate your
                 own fragmenter, preference gatherer, etc., you are responsible for
@@ -1440,15 +1439,19 @@ def __init__(
         self._iteration = 0
 
         self.model = reward_model
-        self.random_state = random_state
+        self.rng = rng
 
-        if self.random_state is None and None in (preference_gatherer, fragmenter):
+        if self.rng is None and None in (
+            preference_gatherer,
+            fragmenter,
+            reward_trainer,
+        ):
             raise ValueError(
                 "If you don't provide a random state, you must provide your own "
-                "seeded fragmenter and preference gatherer. You can initialize"
-                "a random state with `np.random.RandomState(seed)`.",
+                "seeded fragmenter, preference gatherer, and reward_trainer. "
+                "You can initialize a random state with `np.random.default_rng(seed)`.",
             )
-        elif self.random_state is not None and None not in (
+        elif self.rng is not None and None not in (
             preference_gatherer,
             fragmenter,
         ):
@@ -1458,12 +1461,13 @@ def __init__(
             )
 
         if reward_trainer is None:
+            assert self.rng is not None
             preference_model = PreferenceModel(reward_model)
             loss = CrossEntropyRewardLoss(preference_model)
             self.reward_trainer = _make_reward_trainer(
                 reward_model,
                 loss,
-                random_state=self.random_state,
+                rng=self.rng,
             )
         else:
             self.reward_trainer = reward_trainer
@@ -1477,19 +1481,19 @@ def __init__(
         if fragmenter:
             self.fragmenter = fragmenter
         else:
-            assert self.random_state is not None
+            assert self.rng is not None
             self.fragmenter = RandomFragmenter(
                 custom_logger=self.logger,
-                random_state=random.Random(util.make_seeds(self.random_state)),
+                rng=self.rng,
             )
         self.fragmenter.logger = self.logger
         if preference_gatherer:
             self.preference_gatherer = preference_gatherer
         else:
-            assert self.random_state is not None
+            assert self.rng is not None
             self.preference_gatherer = SyntheticGatherer(
                 custom_logger=self.logger,
-                random_state=self.random_state,
+                rng=self.rng,
             )
 
         self.preference_gatherer.logger = self.logger
diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index 392bf1237..739d3e52e 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -323,7 +323,7 @@ def generate_trajectories(
     policy: AnyPolicy,
     venv: VecEnv,
     sample_until: GenTrajTerminationFn,
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
     *,
     deterministic_policy: bool = False,
 ) -> Sequence[types.TrajectoryWithRew]:
@@ -342,7 +342,7 @@ def generate_trajectories(
         deterministic_policy: If True, asks policy to deterministically return
             action. Note the trajectories might still be non-deterministic if the
             environment has non-determinism!
-        random_state: used for shuffling trajectories.
+        rng: used for shuffling trajectories.
 
     Returns:
         Sequence of trajectories, satisfying `sample_until`. Additional trajectories
@@ -405,7 +405,7 @@ def generate_trajectories(
     # `trajectories` sooner. Shuffle to avoid bias in order. This is important
     # when callees end up truncating the number of trajectories or transitions.
     # It is also cheap, since we're just shuffling pointers.
-    random_state.shuffle(trajectories)  # type: ignore[arg-type]
+    rng.shuffle(trajectories)  # type: ignore[arg-type]
 
     # Sanity checks.
     for trajectory in trajectories:
@@ -528,7 +528,7 @@ def generate_transitions(
     policy: AnyPolicy,
     venv: VecEnv,
     n_timesteps: int,
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
     *,
     truncate: bool = True,
     **kwargs,
@@ -543,7 +543,7 @@ def generate_transitions(
             - None, in which case actions will be sampled randomly
         venv: The vectorized environments to interact with.
         n_timesteps: The minimum number of timesteps to sample.
-        random_state: The random state to use for sampling trajectories.
+        rng: The random state to use for sampling trajectories.
         truncate: If True, then drop any additional samples to ensure that exactly
             `n_timesteps` samples are returned.
         **kwargs: Passed-through to generate_trajectories.
@@ -557,7 +557,7 @@ def generate_transitions(
         policy,
         venv,
         sample_until=make_min_timesteps(n_timesteps),
-        random_state=random_state,
+        rng=rng,
         **kwargs,
     )
     transitions = flatten_trajectories_with_rew(traj)
@@ -572,7 +572,7 @@ def rollout(
     policy: AnyPolicy,
     venv: VecEnv,
     sample_until: GenTrajTerminationFn,
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
     *,
     unwrap: bool = True,
     exclude_infos: bool = True,
@@ -591,7 +591,7 @@ def rollout(
             3) None, in which case actions will be sampled randomly.
         venv: The vectorized environments.
         sample_until: End condition for rollout sampling.
-        random_state: Random state to use for sampling.
+        rng: Random state to use for sampling.
         unwrap: If True, then save original observations and rewards (instead of
             potentially wrapped observations and rewards) by calling
             `unwrap_traj()`.
@@ -610,7 +610,7 @@ def rollout(
         policy,
         venv,
         sample_until,
-        random_state=random_state,
+        rng=rng,
         **kwargs,
     )
     if unwrap:
diff --git a/src/imitation/envs/examples/model_envs.py b/src/imitation/envs/examples/model_envs.py
index e5687b192..edc677fa3 100644
--- a/src/imitation/envs/examples/model_envs.py
+++ b/src/imitation/envs/examples/model_envs.py
@@ -12,7 +12,7 @@ def make_random_trans_mat(
     n_states,
     n_actions,
     max_branch_factor,
-    rand_state=np.random,
+    rng: np.random.Generator,
 ) -> np.ndarray:
     """Make a 'random' transition matrix.
 
@@ -28,7 +28,7 @@ def make_random_trans_mat(
         n_actions: Number of actions.
         max_branch_factor: Maximum number of states that can be reached from
             each state-action pair.
-        rand_state: NumPy random state.
+        rng: NumPy random state.
 
     Returns:
         The transition matrix `mat`, where `mat[s,a,next_s]` gives the probability
@@ -39,26 +39,26 @@ def make_random_trans_mat(
         for action in range(n_actions):
             # uniformly sample a number of successors in [1,max_branch_factor]
             # for this action
-            succs = rand_state.randint(1, max_branch_factor + 1)
-            next_states = rand_state.choice(n_states, size=(succs,), replace=False)
+            succs = rng.integers(1, max_branch_factor + 1)
+            next_states = rng.choice(n_states, size=(succs,), replace=False)
             # generate random vec in probability simplex
-            next_vec = rand_state.dirichlet(np.ones((succs,)))
+            next_vec = rng.dirichlet(np.ones((succs,)))
             next_vec = next_vec / np.sum(next_vec)
             out_mat[start_state, action, next_states] = next_vec
     return out_mat
 
 
-def make_random_state_dist(
+def make_rng_dist(
     n_avail: int,
     n_states: int,
-    rand_state: np.random.RandomState = np.random,
+    rng: np.random.Generator,
 ) -> np.ndarray:
     """Make a random initial state distribution over n_states.
 
     Args:
         n_avail: Number of states available to transition into.
         n_states: Total number of states.
-        rand_state: NumPy random state.
+        rng: NumPy random state.
 
     Returns:
         An initial state distribution that is zero at all but a uniformly random
@@ -71,8 +71,8 @@ def make_random_state_dist(
     """  # noqa: DAR402
     assert 0 < n_avail <= n_states
     init_dist = np.zeros((n_states,))
-    next_states = rand_state.choice(n_states, size=(n_avail,), replace=False)
-    avail_state_dist = rand_state.dirichlet(np.ones((n_avail,)))
+    next_states = rng.choice(n_states, size=(n_avail,), replace=False)
+    avail_state_dist = rng.dirichlet(np.ones((n_avail,)))
     init_dist[next_states] = avail_state_dist
     assert np.sum(init_dist > 0) == n_avail
     init_dist = init_dist / np.sum(init_dist)
@@ -83,7 +83,7 @@ def make_obs_mat(
     n_states: int,
     is_random: bool,
     obs_dim: Optional[int],
-    rand_state: np.random.RandomState = np.random,
+    rng: np.random.Generator,
 ) -> np.ndarray:
     """Makes an observation matrix with a single observation for each state.
 
@@ -94,7 +94,7 @@ def make_obs_mat(
                     If `False`, are unique one-hot vectors for each state.
         obs_dim (int or NoneType): Must be `None` if `is_random == False`.
                  Otherwise, this must be set to the size of the random vectors.
-        rand_state (np.random.RandomState): Random number generator.
+        rng (np.random.Generator): Random number generator.
 
     Returns:
         A matrix of shape `(n_states, obs_dim if is_random else n_states)`.
@@ -102,7 +102,7 @@ def make_obs_mat(
     if not is_random:
         assert obs_dim is None
     if is_random:
-        obs_mat = rand_state.normal(0, 2, (n_states, obs_dim))
+        obs_mat = rng.normal(0, 2, (n_states, obs_dim))
     else:
         obs_mat = np.identity(n_states)
     assert (
@@ -145,7 +145,7 @@ def __init__(
         super().__init__()
         # this generator is ONLY for constructing the MDP, not for controlling
         # random outcomes during rollouts
-        rand_gen = np.random.RandomState(generator_seed)
+        rng = np.random.default_rng(generator_seed)
         if random_obs:
             if obs_dim is None:
                 obs_dim = n_states
@@ -155,21 +155,21 @@ def __init__(
             n_states=n_states,
             is_random=random_obs,
             obs_dim=obs_dim,
-            rand_state=rand_gen,
+            rng=rng,
         )
         self._transition_matrix = make_random_trans_mat(
             n_states=n_states,
             n_actions=n_actions,
             max_branch_factor=branch_factor,
-            rand_state=rand_gen,
+            rng=rng,
         )
-        self._initial_state_dist = make_random_state_dist(
+        self._initial_state_dist = make_rng_dist(
             n_avail=branch_factor,
             n_states=n_states,
-            rand_state=rand_gen,
+            rng=rng,
         )
         self._horizon = horizon
-        self._reward_weights = rand_gen.randn(self._observation_matrix.shape[-1])
+        self._reward_weights = rng.standard_normal(self._observation_matrix.shape[-1])
         self._reward_matrix = self._observation_matrix @ self._reward_weights
         assert self._reward_matrix.shape == (self.n_states,)
 
diff --git a/src/imitation/envs/resettable_env.py b/src/imitation/envs/resettable_env.py
index 0e8780d13..877003bca 100644
--- a/src/imitation/envs/resettable_env.py
+++ b/src/imitation/envs/resettable_env.py
@@ -27,7 +27,7 @@ def __init__(self):
         self._action_space = None
         self.cur_state = None
         self._n_actions_taken = None
-        self.rand_state: Optional[np.random.RandomState] = None
+        self.rng: Optional[np.random.Generator] = None
         self.seed()
 
     @abc.abstractmethod
@@ -113,7 +113,7 @@ def seed(self, seed=None):
             # Gym API wants list of seeds to be returned for some reason, so
             # generate a seed explicitly in this case
             seed = np.random.randint(0, 1 << 31)
-        self.rand_state = np.random.RandomState(seed)
+        self.rng = np.random.default_rng(seed)
         return [seed]
 
     def reset(self):
@@ -177,12 +177,12 @@ def action_space(self) -> gym.Space:
         return self._action_space
 
     def initial_state(self):
-        return self.rand_state.choice(self.n_states, p=self.initial_state_dist)
+        return self.rng.choice(self.n_states, p=self.initial_state_dist)
 
     def transition(self, state, action):
         out_dist = self.transition_matrix[state, action]
         choice_states = np.arange(self.n_states)
-        return int(self.rand_state.choice(choice_states, p=out_dist, size=()))
+        return int(self.rng.choice(choice_states, p=out_dist, size=()))
 
     def reward(self, state, action, new_state):
         reward = self.reward_matrix[state]
diff --git a/src/imitation/policies/exploration_wrapper.py b/src/imitation/policies/exploration_wrapper.py
index 65855dce2..2a7da8060 100644
--- a/src/imitation/policies/exploration_wrapper.py
+++ b/src/imitation/policies/exploration_wrapper.py
@@ -24,7 +24,7 @@ def __init__(
         venv: vec_env.VecEnv,
         random_prob: float,
         switch_prob: float,
-        random_state: np.random.RandomState,
+        rng: np.random.Generator,
     ):
         """Initializes the ExplorationWrapper.
 
@@ -33,7 +33,7 @@ def __init__(
             venv: The environment to use (needed for sampling random actions).
             random_prob: The probability of picking the random policy when switching.
             switch_prob: The probability of switching away from the current policy.
-            random_state: The random state to use for seeding the environment and for
+            rng: The random state to use for seeding the environment and for
                 switching policies.
         """
         self.wrapped_policy = policy_callable
@@ -41,7 +41,7 @@ def __init__(
         self.switch_prob = switch_prob
         self.venv = venv
 
-        self.rng = random_state
+        self.rng = rng
         seed = util.make_seeds(self.rng)
         self.venv.action_space.seed(seed)
 
@@ -55,13 +55,13 @@ def _random_policy(self, obs: np.ndarray) -> np.ndarray:
 
     def _switch(self) -> None:
         """Pick a new policy at random."""
-        if self.rng.rand() < self.random_prob:
+        if self.rng.random() < self.random_prob:
             self.current_policy = self._random_policy
         else:
             self.current_policy = self.wrapped_policy
 
     def __call__(self, obs: np.ndarray) -> np.ndarray:
         acts = self.current_policy(obs)
-        if self.rng.rand() < self.switch_prob:
+        if self.rng.random() < self.switch_prob:
             self._switch()
         return acts
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index bc1bcdd5e..d50edbf1e 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -5,17 +5,18 @@
 import os
 from typing import Any, Generator, Mapping, Sequence, Tuple, Union
 
+import numpy as np
 import sacred
 from stable_baselines3.common import vec_env
 
-from imitation.scripts.common import seeding, wb
+from imitation.scripts.common import wb
 from imitation.util import logger as imit_logger
 from imitation.util import sacred as sacred_util
 from imitation.util import util
 
 common_ingredient = sacred.Ingredient(
     "common",
-    ingredients=[wb.wandb_ingredient, seeding.seeding_ingredient],
+    ingredients=[wb.wandb_ingredient],
 )
 logger = logging.getLogger(__name__)
 
@@ -77,6 +78,12 @@ def fast():
     locals()  # quieten flake8
 
 
+@common_ingredient.capture
+def make_rng(_seed) -> np.random.Generator:
+    """Creates a `np.random.Generator` with the given seed."""
+    return np.random.default_rng(_seed)
+
+
 @common_ingredient.capture
 def make_log_dir(
     _run,
@@ -158,11 +165,11 @@ def make_venv(
     Yields:
         The constructed vector environment.
     """
-    random_state = seeding.make_random_state()
+    rng = make_rng()
     try:
         venv = util.make_vec_env(
             env_name,
-            random_state=random_state,
+            rng=rng,
             n_envs=num_vec,
             parallel=parallel,
             max_episode_steps=max_episode_steps,
diff --git a/src/imitation/scripts/common/rl.py b/src/imitation/scripts/common/rl.py
index ce687732b..d000171d8 100644
--- a/src/imitation/scripts/common/rl.py
+++ b/src/imitation/scripts/common/rl.py
@@ -17,12 +17,12 @@
 from imitation.policies import serialize
 from imitation.policies.replay_buffer_wrapper import ReplayBufferRewardWrapper
 from imitation.rewards.reward_function import RewardFn
-from imitation.scripts.common import seeding
+from imitation.scripts.common import common
 from imitation.scripts.common.train import train_ingredient
 
 rl_ingredient = sacred.Ingredient(
     "rl",
-    ingredients=[train_ingredient, seeding.seeding_ingredient],
+    ingredients=[train_ingredient, common.common_ingredient],
 )
 logger = logging.getLogger(__name__)
 
@@ -106,6 +106,7 @@ def make_rl_algo(
     batch_size: int,
     rl_kwargs: Mapping[str, Any],
     train: Mapping[str, Any],
+    _seed,
     relabel_reward_fn: Optional[RewardFn] = None,
 ) -> base_class.BaseAlgorithm:
     """Instantiates a Stable Baselines3 RL algorithm.
@@ -127,7 +128,6 @@ def make_rl_algo(
         ValueError: `gen_batch_size` not divisible by `venv.num_envs`.
         TypeError: `rl_cls` is neither `OnPolicyAlgorithm` nor `OffPolicyAlgorithm`.
     """
-    seed = seeding.get_seed()
     if batch_size % venv.num_envs != 0:
         raise ValueError(
             f"num_envs={venv.num_envs} must evenly divide batch_size={batch_size}.",
@@ -161,7 +161,7 @@ def make_rl_algo(
         # https://github.com/DLR-RM/stable-baselines3/blob/30772aa9f53a4cf61571ee90046cdc454c1b11d7/sb3/common/off_policy_algorithm.py#L145
         policy_kwargs=dict(train["policy_kwargs"]),
         env=venv,
-        seed=seed,
+        seed=_seed,
         **rl_kwargs,
     )
     logger.info(f"RL algorithm: {type(rl_algo)}")
@@ -171,13 +171,13 @@ def make_rl_algo(
 
 @rl_ingredient.capture
 def load_rl_algo_from_path(
+    _seed,
     agent_path: str,
     venv: vec_env.VecEnv,
     rl_cls: Type[base_class.BaseAlgorithm],
     rl_kwargs: Mapping[str, Any],
     relabel_reward_fn: Optional[RewardFn] = None,
 ) -> base_class.BaseAlgorithm:
-    seed = seeding.get_seed()
     rl_kwargs = dict(rl_kwargs)
     if issubclass(rl_cls, off_policy_algorithm.OffPolicyAlgorithm):
         rl_kwargs = _maybe_add_relabel_buffer(
@@ -188,7 +188,7 @@ def load_rl_algo_from_path(
         cls=rl_cls,
         path=agent_path,
         venv=venv,
-        seed=seed,
+        seed=_seed,
         **rl_kwargs,
     )
     logger.info(f"Warm starting agent from '{agent_path}'")
diff --git a/src/imitation/scripts/common/seeding.py b/src/imitation/scripts/common/seeding.py
deleted file mode 100644
index bcbbb1957..000000000
--- a/src/imitation/scripts/common/seeding.py
+++ /dev/null
@@ -1,26 +0,0 @@
-"""Single source of truth for seeding random number generators in experiments."""
-
-from typing import Optional
-
-import numpy as np
-import sacred
-
-seeding_ingredient = sacred.Ingredient("seeding")
-
-
-@seeding_ingredient.config
-def config():
-    seed = 0
-    locals()  # quieten flake8
-
-
-@seeding_ingredient.capture
-def get_seed(seed: Optional[int]) -> Optional[int]:
-    """Returns the seed used by the seeding ingredient."""
-    return seed
-
-
-@seeding_ingredient.capture
-def make_random_state(seed: Optional[int]) -> np.random.RandomState:
-    """Creates a `np.random.RandomState` with the given seed."""
-    return np.random.RandomState(seed)
diff --git a/src/imitation/scripts/common/train.py b/src/imitation/scripts/common/train.py
index 384c930c9..5f291edf9 100644
--- a/src/imitation/scripts/common/train.py
+++ b/src/imitation/scripts/common/train.py
@@ -9,9 +9,9 @@
 import imitation.util.networks
 from imitation.data import rollout
 from imitation.policies import base
-from imitation.scripts.common import seeding
+from imitation.scripts.common import common
 
-train_ingredient = sacred.Ingredient("train", ingredients=[seeding.seeding_ingredient])
+train_ingredient = sacred.Ingredient("train", ingredients=[common.common_ingredient])
 logger = logging.getLogger(__name__)
 
 
@@ -90,13 +90,13 @@ def eval_policy(
         "monitor_return" key). "expert_stats" gives the return value of
         `rollout_stats()` on the expert demonstrations loaded from `rollout_path`.
     """
-    random_state = seeding.make_random_state()
+    rng = common.make_rng()
     sample_until_eval = rollout.make_min_episodes(n_episodes_eval)
     trajs = rollout.generate_trajectories(
         rl_algo,
         venv,
         sample_until=sample_until_eval,
-        random_state=random_state,
+        rng=rng,
     )
     return rollout.rollout_stats(trajs)
 
diff --git a/src/imitation/scripts/config/eval_policy.py b/src/imitation/scripts/config/eval_policy.py
index 4cd378a7a..2570eba75 100644
--- a/src/imitation/scripts/config/eval_policy.py
+++ b/src/imitation/scripts/config/eval_policy.py
@@ -2,11 +2,13 @@
 
 import sacred
 
-from imitation.scripts.common import common, seeding
+from imitation.scripts.common import common
 
 eval_policy_ex = sacred.Experiment(
     "eval_policy",
-    ingredients=[common.common_ingredient, seeding.seeding_ingredient],
+    ingredients=[
+        common.common_ingredient,
+    ],
 )
 
 
diff --git a/src/imitation/scripts/config/train_preference_comparisons.py b/src/imitation/scripts/config/train_preference_comparisons.py
index 64ee5bb07..ba4e9483c 100644
--- a/src/imitation/scripts/config/train_preference_comparisons.py
+++ b/src/imitation/scripts/config/train_preference_comparisons.py
@@ -3,7 +3,7 @@
 import sacred
 
 from imitation.algorithms import preference_comparisons
-from imitation.scripts.common import common, reward, rl, seeding, train
+from imitation.scripts.common import common, reward, rl, train
 
 train_preference_comparisons_ex = sacred.Experiment(
     "train_preference_comparisons",
@@ -12,7 +12,6 @@
         reward.reward_ingredient,
         rl.rl_ingredient,
         train.train_ingredient,
-        seeding.seeding_ingredient,
     ],
 )
 
diff --git a/src/imitation/scripts/config/train_rl.py b/src/imitation/scripts/config/train_rl.py
index e71d7b14f..b9ede3165 100644
--- a/src/imitation/scripts/config/train_rl.py
+++ b/src/imitation/scripts/config/train_rl.py
@@ -2,7 +2,7 @@
 
 import sacred
 
-from imitation.scripts.common import common, rl, seeding, train
+from imitation.scripts.common import common, rl, train
 
 train_rl_ex = sacred.Experiment(
     "train_rl",
@@ -10,7 +10,6 @@
         common.common_ingredient,
         train.train_ingredient,
         rl.rl_ingredient,
-        seeding.seeding_ingredient,
     ],
 )
 
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 5a2b377c1..e58f91c3f 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -14,7 +14,7 @@
 from imitation.policies import serialize
 from imitation.rewards import reward_wrapper
 from imitation.rewards.serialize import load_reward
-from imitation.scripts.common import common, seeding
+from imitation.scripts.common import common
 from imitation.scripts.config.eval_policy import eval_policy_ex
 from imitation.util import video_wrapper
 
@@ -90,7 +90,7 @@ def eval_policy(
     Returns:
         Return value of `imitation.util.rollout.rollout_stats()`.
     """
-    random_state = seeding.make_random_state()
+    rng = common.make_rng()
     log_dir = common.make_log_dir()
     sample_until = rollout.make_sample_until(eval_n_timesteps, eval_n_episodes)
     post_wrappers = [video_wrapper_factory(log_dir, **video_kwargs)] if videos else None
@@ -111,7 +111,7 @@ def eval_policy(
             policy,
             venv,
             sample_until,
-            random_state=random_state,
+            rng=rng,
         )
 
     if rollout_save_path:
diff --git a/src/imitation/scripts/train_imitation.py b/src/imitation/scripts/train_imitation.py
index a580a7ba3..7e55363e5 100644
--- a/src/imitation/scripts/train_imitation.py
+++ b/src/imitation/scripts/train_imitation.py
@@ -12,7 +12,7 @@
 from imitation.algorithms.dagger import SimpleDAggerTrainer
 from imitation.data import rollout, types
 from imitation.policies import serialize
-from imitation.scripts.common import common, demonstrations, seeding, train
+from imitation.scripts.common import common, demonstrations, train
 from imitation.scripts.config.train_imitation import train_imitation_ex
 
 logger = logging.getLogger(__name__)
@@ -123,7 +123,7 @@ def train_imitation(
     Returns:
         Statistics for rollouts from the trained policy and demonstration data.
     """
-    random_state = seeding.make_random_state()
+    rng = common.make_rng()
     custom_logger, log_dir = common.setup_logging()
 
     with common.make_venv() as venv:
@@ -139,7 +139,7 @@ def train_imitation(
             policy=imit_policy,
             demonstrations=expert_trajs,
             custom_logger=custom_logger,
-            random_state=random_state,
+            rng=rng,
             **bc_kwargs,
         )
         bc_train_kwargs = dict(log_rollouts_venv=venv, **bc_train_kwargs)
@@ -158,7 +158,7 @@ def train_imitation(
                 expert_policy=expert_policy,
                 custom_logger=custom_logger,
                 bc_trainer=bc_trainer,
-                random_state=random_state,
+                rng=rng,
             )
             model.train(
                 total_timesteps=int(dagger["total_timesteps"]),
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index 8bbbab0cf..eecf18ac2 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -18,7 +18,7 @@
 from imitation.policies import serialize
 from imitation.scripts.common import common, reward
 from imitation.scripts.common import rl as rl_common
-from imitation.scripts.common import seeding, train
+from imitation.scripts.common import train
 from imitation.scripts.config.train_preference_comparisons import (
     train_preference_comparisons_ex,
 )
@@ -147,8 +147,7 @@ def train_preference_comparisons(
         ValueError: Inconsistency between config and deserialized policy normalization.
     """
     custom_logger, log_dir = common.setup_logging()
-    random_state = seeding.make_random_state()
-    seed = seeding.get_seed()
+    rng = common.make_rng()
 
     with common.make_venv() as venv:
         reward_net = reward.make_reward_net(venv)
@@ -173,7 +172,7 @@ def train_preference_comparisons(
                 reward_fn=reward_net,
                 venv=venv,
                 exploration_frac=exploration_frac,
-                random_state=random_state,
+                rng=rng,
                 custom_logger=custom_logger,
                 **trajectory_generator_kwargs,
             )
@@ -187,7 +186,7 @@ def train_preference_comparisons(
                 )
             trajectory_generator = preference_comparisons.TrajectoryDataset(
                 trajectories=types.load_with_rewards(trajectory_path),
-                random_state=random.Random(seed),
+                rng=rng,
                 custom_logger=custom_logger,
                 **trajectory_generator_kwargs,
             )
@@ -195,7 +194,7 @@ def train_preference_comparisons(
         fragmenter: preference_comparisons.Fragmenter = (
             preference_comparisons.RandomFragmenter(
                 **fragmenter_kwargs,
-                random_state=random.Random(seed),
+                rng=rng,
                 custom_logger=custom_logger,
             )
         )
@@ -213,7 +212,7 @@ def train_preference_comparisons(
             )
         gatherer = gatherer_cls(
             **gatherer_kwargs,
-            random_state=random_state,
+            rng=rng,
             custom_logger=custom_logger,
         )
 
@@ -224,7 +223,7 @@ def train_preference_comparisons(
         reward_trainer = preference_comparisons._make_reward_trainer(
             reward_net,
             loss,
-            random_state,
+            rng,
             reward_trainer_kwargs,
         )
 
diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
index 997dd1600..722ef5eb7 100644
--- a/src/imitation/scripts/train_rl.py
+++ b/src/imitation/scripts/train_rl.py
@@ -22,7 +22,7 @@
 from imitation.policies import serialize
 from imitation.rewards.reward_wrapper import RewardVecEnvWrapper
 from imitation.rewards.serialize import load_reward
-from imitation.scripts.common import common, rl, seeding, train
+from imitation.scripts.common import common, rl, train
 from imitation.scripts.config.train_rl import train_rl_ex
 
 
@@ -87,7 +87,7 @@ def train_rl(
     Returns:
         The return value of `rollout_stats()` using the final policy.
     """
-    random_state = seeding.make_random_state()
+    rng = common.make_rng()
     custom_logger, log_dir = common.setup_logging()
     rollout_dir = osp.join(log_dir, "rollouts")
     policy_dir = osp.join(log_dir, "policies")
@@ -147,7 +147,7 @@ def train_rl(
             )
             types.save(
                 save_path,
-                rollout.rollout(rl_algo, venv, sample_until, random_state),
+                rollout.rollout(rl_algo, venv, sample_until, rng),
             )
         if policy_save_final:
             output_dir = os.path.join(policy_dir, "final")
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index fdaee8c5e..7884ed5a5 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -67,7 +67,7 @@ def make_unique_timestamp() -> str:
 
 def make_vec_env(
     env_name: str,
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
     n_envs: int = 8,
     parallel: bool = False,
     log_dir: Optional[str] = None,
@@ -79,7 +79,7 @@ def make_vec_env(
 
     Args:
         env_name: The Env's string id in Gym.
-        random_state: The random state to use to seed the environment.
+        rng: The random state to use to seed the environment.
         n_envs: The number of duplicate environments.
         parallel: If True, uses SubprocVecEnv; otherwise, DummyVecEnv.
         log_dir: If specified, saves Monitor output to this directory.
@@ -140,7 +140,7 @@ def make_env(i: int, this_seed: int) -> gym.Env:
 
         return env
 
-    env_seeds = make_seeds(random_state, n_envs)
+    env_seeds = make_seeds(rng, n_envs)
     env_fns: List[Callable[[], gym.Env]] = [
         functools.partial(make_env, i, s) for i, s in enumerate(env_seeds)
     ]
@@ -153,30 +153,30 @@ def make_env(i: int, this_seed: int) -> gym.Env:
 
 @overload
 def make_seeds(
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
 ) -> int:
     ...
 
 
 @overload
-def make_seeds(random_state: np.random.RandomState, n: int) -> List[int]:
+def make_seeds(rng: np.random.Generator, n: int) -> List[int]:
     ...
 
 
 def make_seeds(
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
     n: Optional[int] = None,
 ) -> Union[Sequence[int], int]:
     """Generate n random seeds from a random state.
 
     Args:
-        random_state: The random state to use to generate seeds.
+        rng: The random state to use to generate seeds.
         n: The number of seeds to generate.
 
     Returns:
         A list of n random seeds.
     """
-    seeds: List[int] = random_state.randint(0, (1 << 31) - 1, (n or 1,)).tolist()
+    seeds: List[int] = rng.integers(0, (1 << 31) - 1, (n or 1,)).tolist()
     if n is None:
         return seeds[0]
     else:
diff --git a/tests/algorithms/test_adversarial.py b/tests/algorithms/test_adversarial.py
index 81d37a30f..a61ca8d34 100644
--- a/tests/algorithms/test_adversarial.py
+++ b/tests/algorithms/test_adversarial.py
@@ -58,7 +58,7 @@ def make_trainer(
     algorithm_kwargs: Mapping[str, Any],
     tmpdir: str,
     expert_transitions: types.Transitions,
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
     expert_batch_size: int = 1,
     env_name: str = "seals/CartPole-v0",
     num_envs: int = 1,
@@ -81,7 +81,7 @@ def make_trainer(
         env_name,
         n_envs=num_envs,
         parallel=parallel,
-        random_state=random_state,
+        rng=rng,
     )
     model_cls = algorithm_kwargs["model_class"]
     gen_algo = model_cls(algorithm_kwargs["policy_class"], venv)
@@ -107,13 +107,12 @@ def make_trainer(
         venv.close()
 
 
-def test_airl_fail_fast(custom_logger, tmpdir, random_state_fixed):
-    random_state = random_state_fixed
+def test_airl_fail_fast(custom_logger, tmpdir, rng):
     venv = util.make_vec_env(
         "seals/CartPole-v0",
         n_envs=1,
         parallel=False,
-        random_state=random_state,
+        rng=rng,
     )
 
     gen_algo = stable_baselines3.DQN(stable_baselines3.dqn.MlpPolicy, venv)
@@ -121,7 +120,7 @@ def test_airl_fail_fast(custom_logger, tmpdir, random_state_fixed):
         gen_algo,
         venv,
         n_timesteps=20,
-        random_state=random_state,
+        rng=rng,
     )
     reward_net = reward_nets.BasicShapedRewardNet(
         observation_space=venv.observation_space,
@@ -141,13 +140,12 @@ def test_airl_fail_fast(custom_logger, tmpdir, random_state_fixed):
 
 
 @pytest.fixture(params=ALGORITHM_KWARGS.values(), ids=list(ALGORITHM_KWARGS.keys()))
-def trainer(request, tmpdir, expert_transitions, random_state_fixed):
-    random_state = random_state_fixed
+def trainer(request, tmpdir, expert_transitions, rng):
     with make_trainer(
         request.param,
         tmpdir,
         expert_transitions,
-        random_state,
+        rng,
     ) as trainer:
         yield trainer
 
@@ -196,14 +194,13 @@ def trainer_parametrized(
     _expert_batch_size,
     tmpdir,
     expert_transitions,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     with make_trainer(
         _algorithm_kwargs,
         tmpdir,
         expert_transitions,
-        random_state=random_state,
+        rng=rng,
         parallel=_parallel,
         convert_dataset=_convert_dataset,
         expert_batch_size=_expert_batch_size,
@@ -214,15 +211,14 @@ def trainer_parametrized(
 def test_train_disc_step_no_crash(
     trainer_parametrized,
     _expert_batch_size,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     transitions = rollout.generate_transitions(
         trainer_parametrized.gen_algo,
         trainer_parametrized.venv,
         n_timesteps=_expert_batch_size,
         truncate=True,
-        random_state=random_state,
+        rng=rng,
     )
     trainer_parametrized.train_disc(
         gen_samples=types.dataclass_quick_asdict(transitions),
@@ -243,15 +239,14 @@ def trainer_batch_sizes(
     _expert_batch_size,
     tmpdir,
     expert_transitions,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     with make_trainer(
         _algorithm_kwargs,
         tmpdir,
         expert_transitions,
         expert_batch_size=_expert_batch_size,
-        random_state=random_state,
+        rng=rng,
     ) as trainer:
         yield trainer
 
@@ -261,10 +256,9 @@ def test_train_disc_improve_D(
     tmpdir,
     expert_transitions,
     _expert_batch_size,
-    random_state_fixed,
+    rng,
     n_steps=3,
 ):
-    random_state = random_state_fixed
     expert_samples = expert_transitions[:_expert_batch_size]
     expert_samples = types.dataclass_quick_asdict(expert_samples)
     gen_samples = rollout.generate_transitions(
@@ -272,7 +266,7 @@ def test_train_disc_improve_D(
         trainer_batch_sizes.venv_train,
         n_timesteps=_expert_batch_size,
         truncate=True,
-        random_state=random_state,
+        rng=rng,
     )
     gen_samples = types.dataclass_quick_asdict(gen_samples)
     init_stats = final_stats = None
@@ -298,16 +292,15 @@ def trainer_diverse_env(
     _env_name,
     tmpdir,
     expert_transitions,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     if _algorithm_kwargs["model_class"] == stable_baselines3.DQN:
         pytest.skip("DQN does not support all environments.")
     with make_trainer(
         _algorithm_kwargs,
         tmpdir,
         expert_transitions,
-        random_state=random_state,
+        rng=rng,
         env_name=_env_name,
     ) as trainer:
         yield trainer
@@ -317,7 +310,7 @@ def trainer_diverse_env(
 def test_logits_expert_is_high_log_policy_act_prob(
     trainer_diverse_env: common.AdversarialTrainer,
     n_timesteps: int,
-    random_state_fixed,
+    rng,
 ):
     """Smoke test calling `logits_expert_is_high` on `AdversarialTrainer`.
 
@@ -327,14 +320,13 @@ def test_logits_expert_is_high_log_policy_act_prob(
     Args:
         trainer_diverse_env: The trainer to test.
         n_timesteps: The number of timesteps of rollouts to collect.
-        random_state_fixed: The random state to use.
+        rng: The random state to use.
     """
-    random_state = random_state_fixed
     trans = rollout.generate_transitions(
         policy=None,
         venv=trainer_diverse_env.venv,
         n_timesteps=n_timesteps,
-        random_state=random_state,
+        rng=rng,
     )
 
     obs, acts, next_obs, dones = trainer_diverse_env.reward_train.preprocess(
diff --git a/tests/algorithms/test_bc.py b/tests/algorithms/test_bc.py
index 7b168de1c..fce89e63c 100644
--- a/tests/algorithms/test_bc.py
+++ b/tests/algorithms/test_bc.py
@@ -47,9 +47,8 @@ def trainer(
     expert_data_type,
     custom_logger,
     cartpole_expert_trajectories,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     trans = rollout.flatten_trajectories(cartpole_expert_trajectories)
     if expert_data_type == "data_loader":
         expert_data = th_data.DataLoader(
@@ -72,12 +71,11 @@ def trainer(
         batch_size=batch_size,
         demonstrations=expert_data,
         custom_logger=custom_logger,
-        random_state=random_state,
+        rng=rng,
     )
 
 
-def test_weight_decay_init_error(cartpole_venv, custom_logger, random_state_fixed):
-    random_state = random_state_fixed
+def test_weight_decay_init_error(cartpole_venv, custom_logger, rng):
     with pytest.raises(ValueError, match=".*weight_decay.*"):
         bc.BC(
             observation_space=cartpole_venv.observation_space,
@@ -85,7 +83,7 @@ def test_weight_decay_init_error(cartpole_venv, custom_logger, random_state_fixe
             demonstrations=None,
             optimizer_kwargs=dict(weight_decay=1e-4),
             custom_logger=custom_logger,
-            random_state=random_state,
+            rng=rng,
         )
 
 
@@ -170,7 +168,7 @@ def test_bc_data_loader_empty_iter_error(
     no_yield_after_iter: bool,
     custom_logger: logger.HierarchicalLogger,
     cartpole_expert_trajectories,
-    random_state_fixed,
+    rng,
 ) -> None:
     """Check that we error out if the DataLoader suddenly stops yielding any batches.
 
@@ -181,10 +179,9 @@ def test_bc_data_loader_empty_iter_error(
         no_yield_after_iter: Data loader stops yielding after this many calls.
         custom_logger: Where to log to.
         cartpole_expert_trajectories: The expert trajectories to use.
-        random_state_fixed: Random state to use.
+        rng: Random state to use.
     """
     batch_size = 32
-    random_state = random_state_fixed
     trans = rollout.flatten_trajectories(cartpole_expert_trajectories)
     dummy_yield_value = dataclasses.asdict(trans[:batch_size])
 
@@ -199,7 +196,7 @@ def test_bc_data_loader_empty_iter_error(
         action_space=cartpole_venv.action_space,
         batch_size=batch_size,
         custom_logger=custom_logger,
-        random_state=random_state,
+        rng=rng,
     )
     trainer.set_demonstrations(bad_data_loader)
     with pytest.raises(AssertionError, match=".*no data.*"):
diff --git a/tests/algorithms/test_dagger.py b/tests/algorithms/test_dagger.py
index 908c3a1ef..683c34f8f 100644
--- a/tests/algorithms/test_dagger.py
+++ b/tests/algorithms/test_dagger.py
@@ -41,8 +41,7 @@ def test_beta_schedule():
         assert np.allclose(three_step_sched(i), (3 - i) / 3 if i <= 2 else 0)
 
 
-def test_traj_collector_seed(tmpdir, pendulum_venv, random_state_fixed):
-    random_state = random_state_fixed
+def test_traj_collector_seed(tmpdir, pendulum_venv, rng):
     collector = dagger.InteractiveTrajectoryCollector(
         venv=pendulum_venv,
         get_robot_acts=lambda o: [
@@ -50,7 +49,7 @@ def test_traj_collector_seed(tmpdir, pendulum_venv, random_state_fixed):
         ],
         beta=0.5,
         save_dir=tmpdir,
-        random_state=random_state,
+        rng=rng,
     )
     seeds1 = collector.seed(42)
     obs1 = collector.reset()
@@ -61,8 +60,7 @@ def test_traj_collector_seed(tmpdir, pendulum_venv, random_state_fixed):
     np.testing.assert_array_equal(obs1, obs2)
 
 
-def test_traj_collector(tmpdir, pendulum_venv, random_state_fixed):
-    random_state = random_state_fixed
+def test_traj_collector(tmpdir, pendulum_venv, rng):
     robot_calls = 0
     num_episodes = 0
 
@@ -76,7 +74,7 @@ def get_random_acts(obs):
         get_robot_acts=get_random_acts,
         beta=0.5,
         save_dir=tmpdir,
-        random_state=random_state,
+        rng=rng,
     )
     collector.reset()
     zero_acts = np.zeros(
@@ -114,7 +112,7 @@ def _build_dagger_trainer(
     expert_policy,
     pendulum_expert_rollouts: List[TrajectoryWithRew],
     custom_logger,
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
 ):
     del expert_policy
     if pendulum_expert_rollouts is not None:
@@ -127,7 +125,7 @@ def _build_dagger_trainer(
         action_space=venv.action_space,
         optimizer_kwargs=dict(lr=1e-3),
         custom_logger=custom_logger,
-        random_state=random_state,
+        rng=rng,
     )
     return dagger.DAggerTrainer(
         venv=venv,
@@ -135,7 +133,7 @@ def _build_dagger_trainer(
         beta_schedule=beta_schedule,
         bc_trainer=bc_trainer,
         custom_logger=custom_logger,
-        random_state=random_state,
+        rng=rng,
     )
 
 
@@ -146,14 +144,14 @@ def _build_simple_dagger_trainer(
     expert_policy,
     pendulum_expert_rollouts: Optional[List[TrajectoryWithRew]],
     custom_logger,
-    random_state,
+    rng,
 ):
     bc_trainer = bc.BC(
         observation_space=venv.observation_space,
         action_space=venv.action_space,
         optimizer_kwargs=dict(lr=1e-3),
         custom_logger=custom_logger,
-        random_state=random_state,
+        rng=rng,
     )
     return dagger.SimpleDAggerTrainer(
         venv=venv,
@@ -163,7 +161,7 @@ def _build_simple_dagger_trainer(
         expert_policy=expert_policy,
         expert_trajs=pendulum_expert_rollouts,
         custom_logger=custom_logger,
-        random_state=random_state,
+        rng=rng,
     )
 
 
@@ -181,7 +179,7 @@ def init_trainer_fn(
     pendulum_expert_policy,
     maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
     custom_logger,
-    random_state_fixed,
+    rng,
 ):
     # Provide a trainer initialization fixture in addition `trainer` fixture below
     # for tests that want to initialize multiple DAggerTrainer.
@@ -193,7 +191,7 @@ def init_trainer_fn(
         pendulum_expert_policy,
         maybe_pendulum_expert_trajectories,
         custom_logger,
-        random_state_fixed,
+        rng,
     )
 
 
@@ -210,7 +208,7 @@ def simple_dagger_trainer(
     pendulum_expert_policy,
     maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
     custom_logger,
-    random_state_fixed,
+    rng,
 ):
     return _build_simple_dagger_trainer(
         tmpdir,
@@ -219,7 +217,7 @@ def simple_dagger_trainer(
         pendulum_expert_policy,
         maybe_pendulum_expert_trajectories,
         custom_logger,
-        random_state_fixed,
+        rng,
     )
 
 
@@ -252,16 +250,14 @@ def test_trainer_needs_demos_exception_error(
         trainer.extend_and_update(dict(n_epochs=1))
 
 
-def test_trainer_train_arguments(trainer, pendulum_expert_policy, random_state_fixed):
-    random_state = random_state_fixed
-
+def test_trainer_train_arguments(trainer, pendulum_expert_policy, rng):
     def add_samples():
         collector = trainer.create_trajectory_collector()
         rollout.generate_trajectories(
             pendulum_expert_policy,
             collector,
             sample_until=rollout.make_min_timesteps(40),
-            random_state=random_state,
+            rng=rng,
         )
 
     # Lower default number of epochs for the no-arguments call that follows.
@@ -389,7 +385,7 @@ def test_simple_dagger_space_mismatch_error(
     pendulum_expert_policy,
     maybe_pendulum_expert_trajectories: Optional[List[TrajectoryWithRew]],
     custom_logger,
-    random_state_fixed,
+    rng,
 ):
     class MismatchedSpace(gym.spaces.Space):
         """Dummy space that is not equal to any other space."""
@@ -407,27 +403,26 @@ class MismatchedSpace(gym.spaces.Space):
                     pendulum_expert_policy,
                     maybe_pendulum_expert_trajectories,
                     custom_logger,
-                    random_state_fixed,
+                    rng,
                 )
 
 
-def test_dagger_not_enough_transitions_error(tmpdir, custom_logger, random_state_fixed):
-    random_state = random_state_fixed
-    venv = util.make_vec_env("CartPole-v0", random_state=random_state)
+def test_dagger_not_enough_transitions_error(tmpdir, custom_logger, rng):
+    venv = util.make_vec_env("CartPole-v0", rng=rng)
     # Initialize with large batch size to ensure error down the line.
     bc_trainer = bc.BC(
         observation_space=venv.observation_space,
         action_space=venv.action_space,
         batch_size=100_000,
         custom_logger=custom_logger,
-        random_state=random_state,
+        rng=rng,
     )
     trainer = dagger.DAggerTrainer(
         venv=venv,
         scratch_dir=tmpdir,
         bc_trainer=bc_trainer,
         custom_logger=custom_logger,
-        random_state=random_state,
+        rng=rng,
     )
     collector = trainer.create_trajectory_collector()
     policy = base.RandomPolicy(venv.observation_space, venv.action_space)
@@ -435,7 +430,7 @@ def test_dagger_not_enough_transitions_error(tmpdir, custom_logger, random_state
         policy,
         collector,
         rollout.make_min_episodes(1),
-        random_state=random_state,
+        rng=rng,
     )
     with pytest.raises(ValueError, match="Not enough transitions.*"):
         trainer.extend_and_update()
diff --git a/tests/algorithms/test_density_baselines.py b/tests/algorithms/test_density_baselines.py
index 4a10ca057..c66ad3b36 100644
--- a/tests/algorithms/test_density_baselines.py
+++ b/tests/algorithms/test_density_baselines.py
@@ -44,9 +44,8 @@ def test_density_reward(
     is_stationary,
     pendulum_venv,
     pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     # use only a subset of trajectories
     expert_trajectories_all = pendulum_expert_trajectories[:8]
     n_experts = len(expert_trajectories_all)
@@ -59,7 +58,7 @@ def test_density_reward(
         is_stationary=is_stationary,
         kernel_bandwidth=0.2,
         standardise_inputs=True,
-        random_state=random_state,
+        rng=rng,
     )
     reward_fn.train()
 
@@ -74,7 +73,7 @@ def test_density_reward(
         random_policy,
         pendulum_venv,
         sample_until=sample_until,
-        random_state=random_state,
+        rng=rng,
     )
     expert_trajectories_test = expert_trajectories_all[n_experts // 2 :]
     random_returns = score_trajectories(random_trajectories, reward_fn)
@@ -89,9 +88,8 @@ def test_density_reward(
 def test_density_trainer_smoke(
     pendulum_venv,
     pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     # tests whether density trainer runs, not whether it's good
     # (it's actually really poor)
     rollouts = pendulum_expert_trajectories[:2]
@@ -100,7 +98,7 @@ def test_density_trainer_smoke(
         demonstrations=rollouts,
         venv=pendulum_venv,
         rl_algo=rl_algo,
-        random_state=random_state,
+        rng=rng,
     )
     density_trainer.train()
     density_trainer.train_policy(n_timesteps=2)
diff --git a/tests/algorithms/test_mce_irl.py b/tests/algorithms/test_mce_irl.py
index 32909a1e0..91b1a6294 100644
--- a/tests/algorithms/test_mce_irl.py
+++ b/tests/algorithms/test_mce_irl.py
@@ -233,19 +233,18 @@ def test_policy_om_reasonable_mdp(discount: float):
     assert np.allclose(Dt[0], mdp.initial_state_dist)
 
 
-def test_tabular_policy():
+def test_tabular_policy(rng):
     """Tests tabular policy prediction, especially timestep calculation and masking."""
     state_space = gym.spaces.Discrete(2)
     action_space = gym.spaces.Discrete(2)
     pi = np.stack(
         [np.eye(2), 1 - np.eye(2)],
     )
-    random_state = np.random.RandomState(42)
     tabular = TabularPolicy(
         state_space=state_space,
         action_space=action_space,
         pi=pi,
-        random_state=random_state,
+        rng=rng,
     )
 
     states = np.array([0, 1, 1, 0, 1])
@@ -269,8 +268,7 @@ def test_tabular_policy():
     np.testing.assert_equal(timesteps[0], 2 - mask.astype(int))
 
 
-def test_tabular_policy_randomness(random_state_fixed):
-    random_state = random_state_fixed
+def test_tabular_policy_randomness(rng):
     state_space = gym.spaces.Discrete(2)
     action_space = gym.spaces.Discrete(2)
     pi = np.array(
@@ -285,7 +283,7 @@ def test_tabular_policy_randomness(random_state_fixed):
         state_space=state_space,
         action_space=action_space,
         pi=pi,
-        random_state=random_state,
+        rng=rng,
     )
 
     actions, _ = tabular.predict(np.zeros((100,), dtype=int))
@@ -297,8 +295,7 @@ def test_tabular_policy_randomness(random_state_fixed):
     np.testing.assert_equal(actions, 0)
 
 
-def test_mce_irl_demo_formats(random_state_fixed):
-    random_state = random_state_fixed
+def test_mce_irl_demo_formats(rng):
     mdp = model_envs.RandomMDP(
         n_states=5,
         n_actions=3,
@@ -314,7 +311,7 @@ def test_mce_irl_demo_formats(random_state_fixed):
         policy=None,
         venv=state_venv,
         sample_until=rollout.make_min_timesteps(100),
-        random_state=random_state,
+        rng=rng,
     )
     demonstrations = {
         "trajs": trajs,
@@ -344,7 +341,7 @@ def test_mce_irl_demo_formats(random_state_fixed):
                 mdp,
                 reward_net,
                 linf_eps=1e-3,
-                random_state=random_state,
+                rng=rng,
             )
             assert np.allclose(mce_irl.demo_state_om.sum(), mdp.horizon + 1)
             final_counts[kind] = mce_irl.train(max_iter=5)
@@ -365,9 +362,8 @@ def test_mce_irl_demo_formats(random_state_fixed):
 def test_mce_irl_reasonable_mdp(
     model_kwargs: Mapping[str, Any],
     discount: float,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     with th.random.fork_rng():
         th.random.manual_seed(715298)
 
@@ -393,7 +389,7 @@ def test_mce_irl_reasonable_mdp(
             reward_net,
             linf_eps=1e-3,
             discount=discount,
-            random_state=random_state,
+            rng=rng,
         )
         final_counts = mce_irl.train()
 
@@ -407,7 +403,7 @@ def test_mce_irl_reasonable_mdp(
             mce_irl.policy,
             state_venv,
             sample_until=rollout.make_min_episodes(5),
-            random_state=random_state,
+            rng=rng,
         )
         stats = rollout.rollout_stats(trajs)
         if discount > 0.0:  # skip check when discount==0.0 (random policy)
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index fb8fd8503..8ab001b1c 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -22,12 +22,12 @@
 
 
 @pytest.fixture
-def venv(random_state_fixed):
-    random_state = random_state_fixed
+def venv(rng):
+    rng
     return util.make_vec_env(
         "seals/CartPole-v0",
         n_envs=1,
-        random_state=random_state,
+        rng=rng,
     )
 
 
@@ -56,18 +56,17 @@ def agent(venv):
 
 
 @pytest.fixture
-def random_fragmenter(random_state_fixed):
-    random_state = random_state_fixed
+def random_fragmenter(rng):
     return preference_comparisons.RandomFragmenter(
-        random_state=random.Random(util.make_seeds(random_state)),
+        rng=rng,
         warning_threshold=0,
     )
 
 
 @pytest.fixture
-def agent_trainer(agent, reward_net, venv, random_state_fixed):
-    random_state = random_state_fixed
-    return preference_comparisons.AgentTrainer(agent, reward_net, venv, random_state)
+def agent_trainer(agent, reward_net, venv, rng):
+    rng
+    return preference_comparisons.AgentTrainer(agent, reward_net, venv, rng)
 
 
 def _check_trajs_equal(
@@ -85,12 +84,11 @@ def _check_trajs_equal(
         assert traj1.terminal == traj2.terminal
 
 
-def test_mismatched_spaces(venv, agent, random_state_fixed):
-    random_state = random_state_fixed
+def test_mismatched_spaces(venv, agent, rng):
     other_venv = util.make_vec_env(
         "seals/MountainCar-v0",
         n_envs=1,
-        random_state=random_state,
+        rng=rng,
     )
     bad_reward_net = reward_nets.BasicRewardNet(
         other_venv.observation_space,
@@ -104,7 +102,7 @@ def test_mismatched_spaces(venv, agent, random_state_fixed):
             agent,
             bad_reward_net,
             venv,
-            random_state=random_state,
+            rng=rng,
         )
 
 
@@ -114,12 +112,12 @@ def test_trajectory_dataset_seeding(
 ):
     dataset1 = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        random_state=random.Random(0),
+        rng=np.random.default_rng(0),
     )
     sample1 = dataset1.sample(num_samples)
     dataset2 = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        random_state=random.Random(0),
+        rng=np.random.default_rng(0),
     )
     sample2 = dataset2.sample(num_samples)
 
@@ -127,7 +125,7 @@ def test_trajectory_dataset_seeding(
 
     dataset3 = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        random_state=random.Random(42),
+        rng=np.random.default_rng(42),
     )
     sample3 = dataset3.sample(num_samples)
     with pytest.raises(AssertionError):
@@ -139,10 +137,11 @@ def test_trajectory_dataset_seeding(
 def test_trajectory_dataset_len(
     cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
     num_steps: int,
+    rng,
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        random_state=random.Random(0),
+        rng=rng,
     )
     sample = dataset.sample(num_steps)
     lengths = [len(t) for t in sample]
@@ -153,10 +152,11 @@ def test_trajectory_dataset_len(
 
 def test_trajectory_dataset_too_long(
     cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+    rng,
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        random_state=random.Random(0),
+        rng=rng,
     )
     with pytest.raises(RuntimeError, match="Asked for.*but only.* available"):
         dataset.sample(100000)
@@ -164,11 +164,12 @@ def test_trajectory_dataset_too_long(
 
 def test_trajectory_dataset_shuffle(
     cartpole_expert_trajectories: Sequence[TrajectoryWithRew],
+    rng,
     num_steps: int = 400,
 ):
     dataset = preference_comparisons.TrajectoryDataset(
         cartpole_expert_trajectories,
-        random_state=random.Random(0),
+        rng,
     )
     sample = dataset.sample(num_steps)
     sample2 = dataset.sample(num_steps)
@@ -200,9 +201,8 @@ def test_trainer_no_crash(
     random_fragmenter,
     custom_logger,
     schedule,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     main_trainer = preference_comparisons.PreferenceComparisons(
         agent_trainer,
         reward_net,
@@ -212,7 +212,7 @@ def test_trainer_no_crash(
         fragmenter=random_fragmenter,
         custom_logger=custom_logger,
         query_schedule=schedule,
-        random_state=random_state,
+        rng=rng,
     )
     result = main_trainer.train(100, 10)
     # We don't expect good performance after training for 10 (!) timesteps,
@@ -221,8 +221,7 @@ def test_trainer_no_crash(
     assert 0.0 < result["reward_accuracy"] <= 1.0
 
 
-def test_reward_ensemble_trainer_raises_type_error(venv, random_state_fixed):
-    random_state = random_state_fixed
+def test_reward_ensemble_trainer_raises_type_error(venv, rng):
     reward_net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     preference_model = preference_comparisons.PreferenceModel(
         model=reward_net,
@@ -239,7 +238,7 @@ def test_reward_ensemble_trainer_raises_type_error(venv, random_state_fixed):
         preference_comparisons.EnsembleTrainer(
             reward_net,
             loss,
-            random_state=random_state,
+            rng=rng,
         )
 
 
@@ -248,9 +247,8 @@ def test_correct_reward_trainer_used_by_default(
     reward_net,
     random_fragmenter,
     custom_logger,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     main_trainer = preference_comparisons.PreferenceComparisons(
         agent_trainer,
         reward_net,
@@ -258,7 +256,7 @@ def test_correct_reward_trainer_used_by_default(
         transition_oversampling=2,
         fragment_length=2,
         fragmenter=random_fragmenter,
-        random_state=random_state,
+        rng=rng,
         custom_logger=custom_logger,
     )
 
@@ -280,9 +278,8 @@ def test_init_raises_error_when_trying_use_improperly_wrapped_ensemble(
     venv,
     random_fragmenter,
     custom_logger,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     reward_net = testing_reward_nets.make_ensemble(
         venv.observation_space,
         venv.action_space,
@@ -303,7 +300,7 @@ def test_init_raises_error_when_trying_use_improperly_wrapped_ensemble(
             transition_oversampling=2,
             fragment_length=2,
             fragmenter=random_fragmenter,
-            random_state=random_state,
+            rng=rng,
             custom_logger=custom_logger,
         )
 
@@ -313,9 +310,8 @@ def test_discount_rate_no_crash(
     venv,
     random_fragmenter,
     custom_logger,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     # also use a non-zero noise probability to check that doesn't cause errors
     reward_net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     preference_model = preference_comparisons.PreferenceModel(
@@ -337,7 +333,7 @@ def test_discount_rate_no_crash(
         transition_oversampling=2,
         fragment_length=2,
         fragmenter=random_fragmenter,
-        random_state=random_state,
+        rng=rng,
         reward_trainer=reward_trainer,
         custom_logger=custom_logger,
     )
@@ -347,12 +343,11 @@ def test_discount_rate_no_crash(
 def test_synthetic_gatherer_deterministic(
     agent_trainer,
     random_fragmenter,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     gatherer = preference_comparisons.SyntheticGatherer(
         temperature=0,
-        random_state=random_state,
+        rng=rng,
     )
     trajectories = agent_trainer.sample(10)
     fragments = random_fragmenter(trajectories, fragment_length=2, num_pairs=2)
@@ -387,7 +382,7 @@ def test_fragments_terminal(random_fragmenter):
 def test_fragments_too_short_error(agent_trainer):
     trajectories = agent_trainer.sample(2)
     random_fragmenter = preference_comparisons.RandomFragmenter(
-        random_state=random.Random(0),
+        rng=random.Random(0),
         warning_threshold=0,
     )
     with pytest.raises(
@@ -414,12 +409,11 @@ def test_preference_dataset_errors(agent_trainer, random_fragmenter):
         dataset.push(fragments, preferences)
 
 
-def test_preference_dataset_queue(agent_trainer, random_fragmenter, random_state_fixed):
-    random_state = random_state_fixed
+def test_preference_dataset_queue(agent_trainer, random_fragmenter, rng):
     dataset = preference_comparisons.PreferenceDataset(max_size=5)
     trajectories = agent_trainer.sample(10)
 
-    gatherer = preference_comparisons.SyntheticGatherer(random_state=random_state)
+    gatherer = preference_comparisons.SyntheticGatherer(rng=rng)
     for i in range(6):
         fragments = random_fragmenter(trajectories, fragment_length=2, num_pairs=1)
         preferences = gatherer(fragments)
@@ -435,13 +429,12 @@ def test_store_and_load_preference_dataset(
     agent_trainer,
     random_fragmenter,
     tmp_path,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     dataset = preference_comparisons.PreferenceDataset()
     trajectories = agent_trainer.sample(10)
     fragments = random_fragmenter(trajectories, fragment_length=2, num_pairs=2)
-    gatherer = preference_comparisons.SyntheticGatherer(random_state=random_state)
+    gatherer = preference_comparisons.SyntheticGatherer(rng=rng)
     preferences = gatherer(fragments)
     dataset.push(fragments, preferences)
 
@@ -463,14 +456,13 @@ def test_exploration_no_crash(
     venv,
     random_fragmenter,
     custom_logger,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     agent_trainer = preference_comparisons.AgentTrainer(
         agent,
         reward_net,
         venv,
-        random_state=random_state,
+        rng=rng,
         exploration_frac=0.5,
     )
     main_trainer = preference_comparisons.PreferenceComparisons(
@@ -480,7 +472,7 @@ def test_exploration_no_crash(
         transition_oversampling=2,
         fragment_length=5,
         fragmenter=random_fragmenter,
-        random_state=random_state,
+        rng=rng,
         custom_logger=custom_logger,
     )
     main_trainer.train(100, 10)
@@ -493,9 +485,8 @@ def test_active_fragmenter_discount_rate_no_crash(
     random_fragmenter,
     uncertainty_on,
     custom_logger,
-    random_state_fixed,
+    rng,
 ):
-    random_state = random_state_fixed
     # also use a non-zero noise probability to check that doesn't cause errors
     reward_net = reward_nets.RewardEnsemble(
         venv.observation_space,
@@ -531,7 +522,7 @@ def test_active_fragmenter_discount_rate_no_crash(
     reward_trainer = preference_comparisons.EnsembleTrainer(
         reward_net,
         loss,
-        random_state=random_state,
+        rng=rng,
     )
 
     main_trainer = preference_comparisons.PreferenceComparisons(
@@ -541,7 +532,7 @@ def test_active_fragmenter_discount_rate_no_crash(
         transition_oversampling=2,
         fragment_length=2,
         fragmenter=fragmenter,
-        random_state=random_state,
+        rng=rng,
         reward_trainer=reward_trainer,
         custom_logger=custom_logger,
     )
@@ -643,7 +634,7 @@ def test_agent_trainer_sample(venv, agent_trainer):
     )
 
 
-def test_agent_trainer_sample_image_observations(random_state_fixed):
+def test_agent_trainer_sample_image_observations(rng):
     """Test `AgentTrainer.sample()` in an image environment.
 
     SB3 algorithms may rearrange the channel dimension in environments with image
@@ -651,9 +642,8 @@ def test_agent_trainer_sample_image_observations(random_state_fixed):
     environment.
 
     Args:
-        random_state_fixed: Random state (with a fixed seed).
+        rng: Random state (with a fixed seed).
     """
-    random_state = random_state_fixed
     venv = DummyVecEnv([lambda: FakeImageEnv()])
     reward_net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     agent = stable_baselines3.PPO(
@@ -668,7 +658,7 @@ def test_agent_trainer_sample_image_observations(random_state_fixed):
         reward_net,
         venv,
         exploration_frac=0.5,
-        random_state=random_state,
+        rng=rng,
     )
     trajectories = agent_trainer.sample(2)
     assert len(trajectories) > 0
diff --git a/tests/conftest.py b/tests/conftest.py
index 6bd429e35..294da97ec 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -58,7 +58,7 @@ def load_or_rollout_trajectories(
     cache_path,
     policy,
     venv,
-    random_state,
+    rng,
 ) -> Sequence[TrajectoryWithRew]:
     os.makedirs(os.path.dirname(cache_path), exist_ok=True)
     with FileLock(cache_path + ".lock"):
@@ -73,7 +73,7 @@ def load_or_rollout_trajectories(
                 policy,
                 venv,
                 rollout.make_sample_until(min_timesteps=2000, min_episodes=57),
-                random_state=random_state,
+                rng=rng,
             )
             types.save(cache_path, rollouts)
             return rollouts
@@ -141,9 +141,8 @@ def cartpole_expert_trajectories(
     cartpole_expert_policy,
     cartpole_venv,
     pytestconfig,
-    random_state_fixed,
+    rng,
 ) -> Sequence[TrajectoryWithRew]:
-    random_state = random_state_fixed
     rollouts_path = str(
         pytestconfig.cache.makedir("experts") / CARTPOLE_ENV_NAME / "rollout.npz",
     )
@@ -151,7 +150,7 @@ def cartpole_expert_trajectories(
         rollouts_path,
         cartpole_expert_policy,
         cartpole_venv,
-        random_state,
+        rng,
     )
 
 
@@ -204,17 +203,16 @@ def pendulum_expert_trajectories(
     pendulum_expert_policy,
     pendulum_venv,
     pytestconfig,
-    random_state_fixed,
+    rng,
 ) -> Sequence[TrajectoryWithRew]:
     rollouts_path = str(
         pytestconfig.cache.makedir("experts") / PENDULUM_ENV_NAME / "rollout.npz",
     )
-    random_state = random_state_fixed
     return load_or_rollout_trajectories(
         rollouts_path,
         pendulum_expert_policy,
         pendulum_venv,
-        random_state=random_state,
+        rng=rng,
     )
 
 
@@ -239,10 +237,10 @@ def custom_logger(tmpdir: str) -> logger.HierarchicalLogger:
 
 
 @pytest.fixture()
-def random_state_fixed() -> np.random.RandomState:
-    return np.random.RandomState(0)
+def rng_fixed() -> np.random.Generator:
+    return np.random.default_rng(0)
 
 
 @pytest.fixture()
-def random_state() -> np.random.RandomState:
-    return np.random.RandomState()
+def rng() -> np.random.Generator:
+    return np.random.default_rng()
diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index 0a96c8ccb..522a6477d 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -38,7 +38,7 @@ def step(self, action):
 def _sample_fixed_length_trajectories(
     episode_lengths: Sequence[int],
     min_episodes: int,
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
     policy_type: str = "policy",
     **kwargs,
 ) -> Sequence[types.Trajectory]:
@@ -65,7 +65,7 @@ def policy(x):
         policy,
         venv,
         sample_until=sample_until,
-        random_state=random_state,
+        rng=rng,
         **kwargs,
     )
     return trajectories
@@ -75,7 +75,7 @@ def policy(x):
     "policy_type",
     ["policy", "callable", "random"],
 )
-def test_complete_trajectories(policy_type, random_state_fixed) -> None:
+def test_complete_trajectories(policy_type, rng) -> None:
     """Checks trajectories include the terminal observation.
 
     This is hidden by default by VecEnv's auto-reset; we add it back in using
@@ -83,9 +83,8 @@ def test_complete_trajectories(policy_type, random_state_fixed) -> None:
 
     Args:
         policy_type: Kind of policy to use when generating trajectories.
-        random_state_fixed: Random state to use.
+        rng: Random state to use.
     """
-    random_state = random_state_fixed
     min_episodes = 13
     max_acts = 5
     num_envs = 4
@@ -93,7 +92,7 @@ def test_complete_trajectories(policy_type, random_state_fixed) -> None:
         [max_acts] * num_envs,
         min_episodes,
         policy_type=policy_type,
-        random_state=random_state,
+        rng=rng,
     )
     assert len(trajectories) >= min_episodes
     expected_obs = np.array([[0]] * max_acts + [[1]])
@@ -123,7 +122,7 @@ def test_unbiased_trajectories(
     episode_lengths: Sequence[int],
     min_episodes: int,
     expected_counts: Mapping[int, int],
-    random_state_fixed,
+    rng,
 ) -> None:
     """Checks trajectories are sampled without bias towards shorter episodes.
 
@@ -143,13 +142,12 @@ def test_unbiased_trajectories(
         min_episodes: The minimum number of episodes to sample.
         expected_counts: Mapping from episode length to expected number of episodes
             of that length (omit if 0 episodes of that length expected).
-        random_state_fixed: Random state to use.
+        rng: Random state to use.
     """
-    random_state = random_state_fixed
     trajectories = _sample_fixed_length_trajectories(
         episode_lengths,
         min_episodes,
-        random_state,
+        rng,
     )
     assert len(trajectories) == sum(expected_counts.values())
     traj_lens = np.array([len(traj) for traj in trajectories])
@@ -166,12 +164,12 @@ def test_seed_trajectories():
     However, `TerminalSentinelEnv` is fixed-length deterministic, so there are no
     such confounders in this test.
     """
-    rng_a1 = np.random.RandomState(0)
-    rng_a2 = np.random.RandomState(0)
-    rng_b = np.random.RandomState(1)
-    traj_a1 = _sample_fixed_length_trajectories([3, 5], 2, random_state=rng_a1)
-    traj_a2 = _sample_fixed_length_trajectories([3, 5], 2, random_state=rng_a2)
-    traj_b = _sample_fixed_length_trajectories([3, 5], 2, random_state=rng_b)
+    rng_a1 = np.random.default_rng(0)
+    rng_a2 = np.random.default_rng(0)
+    rng_b = np.random.default_rng(1)
+    traj_a1 = _sample_fixed_length_trajectories([3, 5], 2, rng=rng_a1)
+    traj_a2 = _sample_fixed_length_trajectories([3, 5], 2, rng=rng_a2)
+    traj_b = _sample_fixed_length_trajectories([3, 5], 2, rng=rng_b)
     assert [len(traj) for traj in traj_a1] == [len(traj) for traj in traj_a2]
     assert [len(traj) for traj in traj_a1] != [len(traj) for traj in traj_b]
 
@@ -188,15 +186,14 @@ def step(self, action):
         return obs / 2, rew / 2, done, info
 
 
-def test_rollout_stats(random_state_fixed):
+def test_rollout_stats(rng):
     """Applying `ObsRewIncrementWrapper` halves the reward mean.
 
     `rollout_stats` should reflect this.
 
     Args:
-        random_state_fixed: Random state to use (with fixed seed).
+        rng: Random state to use (with fixed seed).
     """
-    random_state = random_state_fixed
     env = gym.make("CartPole-v1")
     env = monitor.Monitor(env, None)
     env = ObsRewHalveWrapper(env)
@@ -207,7 +204,7 @@ def test_rollout_stats(random_state_fixed):
         policy,
         venv,
         rollout.make_min_episodes(10),
-        random_state=random_state,
+        rng=rng,
     )
     s = rollout.rollout_stats(trajs)
 
@@ -217,15 +214,14 @@ def test_rollout_stats(random_state_fixed):
     np.testing.assert_allclose(s["return_max"], s["monitor_return_max"] / 2)
 
 
-def test_unwrap_traj(random_state_fixed):
+def test_unwrap_traj(rng):
     """Check that unwrap_traj reverses `ObsRewIncrementWrapper`.
 
     Also check that unwrapping twice is a no-op.
 
     Args:
-        random_state_fixed: Random state to use (with fixed seed).
+        rng: Random state to use (with fixed seed).
     """
-    random_state = random_state_fixed
     env = gym.make("CartPole-v1")
     env = wrappers.RolloutInfoWrapper(env)
     env = ObsRewHalveWrapper(env)
@@ -236,7 +232,7 @@ def test_unwrap_traj(random_state_fixed):
         policy,
         venv,
         rollout.make_min_episodes(10),
-        random_state=random_state,
+        rng=rng,
     )
     trajs_unwrapped = [rollout.unwrap_traj(t) for t in trajs]
     trajs_unwrapped_twice = [rollout.unwrap_traj(t) for t in trajs_unwrapped]
@@ -282,21 +278,19 @@ def test_compute_returns(gamma):
     assert abs(rollout.discounted_sum(rewards, gamma) - returns) < 1e-8
 
 
-def test_generate_trajectories_type_error(random_state_fixed):
-    random_state = random_state_fixed
+def test_generate_trajectories_type_error(rng):
     venv = vec_env.DummyVecEnv([functools.partial(TerminalSentinelEnv, 1)])
     sample_until = rollout.make_min_episodes(1)
     with pytest.raises(TypeError, match="Policy must be.*got <class 'str'> instead"):
         rollout.generate_trajectories(
             "strings_are_not_valid_policies",  # type: ignore
             venv,
-            random_state=random_state,
+            rng=rng,
             sample_until=sample_until,
         )
 
 
-def test_generate_trajectories_value_error(random_state_fixed):
-    random_state = random_state_fixed
+def test_generate_trajectories_value_error(rng):
     venv = vec_env.DummyVecEnv([functools.partial(TerminalSentinelEnv, 1)])
     sample_until = rollout.make_min_episodes(1)
 
@@ -305,6 +299,6 @@ def test_generate_trajectories_value_error(random_state_fixed):
             lambda obs: np.zeros(len(obs), dtype=int),
             venv,
             sample_until=sample_until,
-            random_state=random_state,
+            rng=rng,
             deterministic_policy=True,
         )
diff --git a/tests/policies/test_exploration_wrapper.py b/tests/policies/test_exploration_wrapper.py
index 5b76f6f29..675bd84b9 100644
--- a/tests/policies/test_exploration_wrapper.py
+++ b/tests/policies/test_exploration_wrapper.py
@@ -11,11 +11,11 @@ def constant_policy(obs):
     return np.zeros(len(obs), dtype=int)
 
 
-def make_wrapper(random_prob, switch_prob):
+def make_wrapper(random_prob, switch_prob, rng):
     venv = util.make_vec_env(
         "seals/CartPole-v0",
         n_envs=1,
-        random_state=np.random.RandomState(),
+        rng=rng,
     )
     return (
         exploration_wrapper.ExplorationWrapper(
@@ -23,13 +23,13 @@ def make_wrapper(random_prob, switch_prob):
             venv=venv,
             random_prob=random_prob,
             switch_prob=switch_prob,
-            random_state=np.random.RandomState(0),
+            rng=rng,
         ),
         venv,
     )
 
 
-def test_random_prob():
+def test_random_prob(rng):
     """Test that `random_prob` produces right behaviors of policy switching.
 
     The policy always makes an initial switch when ExplorationWrapper is applied.
@@ -43,19 +43,19 @@ def test_random_prob():
     Raises:
         ValueError: Unknown policy type to switch.
     """
-    wrapper, _ = make_wrapper(random_prob=0.0, switch_prob=0.5)
+    wrapper, _ = make_wrapper(random_prob=0.0, switch_prob=0.5, rng=rng)
     assert wrapper.current_policy == constant_policy
     for _ in range(100):
         wrapper._switch()
         assert wrapper.current_policy == constant_policy
 
-    wrapper, _ = make_wrapper(random_prob=1.0, switch_prob=0.5)
+    wrapper, _ = make_wrapper(random_prob=1.0, switch_prob=0.5, rng=rng)
     assert wrapper.current_policy == wrapper._random_policy
     for _ in range(100):
         wrapper._switch()
         assert wrapper.current_policy == wrapper._random_policy
 
-    wrapper, _ = make_wrapper(random_prob=0.5, switch_prob=0.5)
+    wrapper, _ = make_wrapper(random_prob=0.5, switch_prob=0.5, rng=rng)
     num_random = 0
     num_constant = 0
     for _ in range(1000):
@@ -71,7 +71,7 @@ def test_random_prob():
     assert num_constant > 450
 
 
-def test_switch_prob():
+def test_switch_prob(rng):
     """Test that `switch_prob` produces right behaviors of policy switching.
 
     The policy always makes an initial switch when ExplorationWrapper is applied.
@@ -82,7 +82,7 @@ def test_switch_prob():
     (2) `switch_prob=1.0`: The policy always switches and the distribution of
         policies is determined by `random_prob`.
     """
-    wrapper, venv = make_wrapper(random_prob=0.5, switch_prob=0.0)
+    wrapper, venv = make_wrapper(random_prob=0.5, switch_prob=0.0, rng=rng)
     policy = wrapper.current_policy
     np.random.seed(0)
     obs = np.random.rand(100, 2)
@@ -91,7 +91,7 @@ def test_switch_prob():
         assert wrapper.current_policy == policy
 
     def _always_switch(random_prob, num_steps, seed):
-        wrapper, _ = make_wrapper(random_prob=random_prob, switch_prob=1.0)
+        wrapper, _ = make_wrapper(random_prob=random_prob, switch_prob=1.0, rng=rng)
         np.random.seed(seed)
         num_random = 0
         num_constant = 0
@@ -117,10 +117,10 @@ def _always_switch(random_prob, num_steps, seed):
     assert num_constant == 1000
 
 
-def test_valid_output():
+def test_valid_output(rng):
     """Ensure that we test both the random and the wrapped policy at least once."""
     for random_prob in [0.0, 0.5, 1.0]:
-        wrapper, venv = make_wrapper(random_prob=random_prob, switch_prob=0.5)
+        wrapper, venv = make_wrapper(random_prob=random_prob, switch_prob=0.5, rng=rng)
         np.random.seed(0)
         obs = np.random.rand(100, 2)
         for action in wrapper(obs):
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index a408010a1..2693ed398 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -25,21 +25,20 @@
 
 @pytest.mark.parametrize("env_name", SIMPLE_ENVS)
 @pytest.mark.parametrize("policy_type", HARDCODED_TYPES)
-def test_actions_valid(env_name, policy_type, random_state_fixed):
+def test_actions_valid(env_name, policy_type, rng):
     """Test output actions of our custom policies always lie in action space."""
-    random_state = random_state_fixed
     venv = util.make_vec_env(
         env_name,
         n_envs=1,
         parallel=False,
-        random_state=random_state,
+        rng=rng,
     )
     policy = serialize.load_policy(policy_type, "foobar", venv)
     transitions = rollout.generate_transitions(
         policy,
         venv,
         n_timesteps=100,
-        random_state=random_state,
+        rng=rng,
     )
 
     for a in transitions.acts:
@@ -56,13 +55,12 @@ def test_actions_valid(env_name, policy_type, random_state_fixed):
 def test_save_stable_model_errors_and_warnings(
     tmpdir,
     policy_env_name_pair,
-    random_state_fixed,
+    rng,
 ):
     """Check errors and warnings in `save_stable_model()`."""
-    random_state = random_state_fixed
     policy, env_name = policy_env_name_pair
     tmpdir = pathlib.Path(tmpdir)
-    venv = util.make_vec_env(env_name, random_state=random_state)
+    venv = util.make_vec_env(env_name, rng=rng)
 
     # Trigger FileNotFoundError for no model.{zip,pkl}
     dir_a = tmpdir / "a"
@@ -81,13 +79,13 @@ def test_save_stable_model_errors_and_warnings(
         serialize.load_policy(policy, str(dir_nonexistent), venv)
 
 
-def _test_serialize_identity(env_name, model_cfg, tmpdir, random_state):
+def _test_serialize_identity(env_name, model_cfg, tmpdir, rng):
     """Test output actions of deserialized policy are same as original."""
     venv = util.make_vec_env(
         env_name,
         n_envs=1,
         parallel=False,
-        random_state=random_state,
+        rng=rng,
     )
 
     model_name, model_cls_name = model_cfg
@@ -103,7 +101,7 @@ def _test_serialize_identity(env_name, model_cfg, tmpdir, random_state):
         venv,
         n_timesteps=1000,
         deterministic_policy=True,
-        random_state=np.random.RandomState(0),
+        rng=np.random.default_rng(0),
     )
 
     serialize.save_stable_model(tmpdir, model)
@@ -115,7 +113,7 @@ def _test_serialize_identity(env_name, model_cfg, tmpdir, random_state):
         venv,
         n_timesteps=1000,
         deterministic_policy=True,
-        random_state=np.random.RandomState(0),
+        rng=np.random.default_rng(0),
     )
 
     assert np.allclose(orig_rollout.acts, new_rollout.acts)
@@ -129,9 +127,9 @@ def _test_serialize_identity(env_name, model_cfg, tmpdir, random_state):
 
 @pytest.mark.parametrize("env_name", SIMPLE_ENVS)
 @pytest.mark.parametrize("model_cfg", NORMAL_CONFIGS)
-def test_serialize_identity(env_name, model_cfg, tmpdir, random_state_fixed):
+def test_serialize_identity(env_name, model_cfg, tmpdir, rng):
     """Test output actions of deserialized policy are same as original."""
-    _test_serialize_identity(env_name, model_cfg, tmpdir, random_state_fixed)
+    _test_serialize_identity(env_name, model_cfg, tmpdir, rng)
 
 
 @pytest.mark.parametrize("env_name", [SIMPLE_CONTINUOUS_ENV])
@@ -140,10 +138,10 @@ def test_serialize_identity_continuous_only(
     env_name,
     model_cfg,
     tmpdir,
-    random_state_fixed,
+    rng,
 ):
     """Test serialize identity for continuous_only algorithms."""
-    _test_serialize_identity(env_name, model_cfg, tmpdir, random_state_fixed)
+    _test_serialize_identity(env_name, model_cfg, tmpdir, rng)
 
 
 class ZeroModule(nn.Module):
diff --git a/tests/policies/test_replay_buffer_wrapper.py b/tests/policies/test_replay_buffer_wrapper.py
index 85365c0d5..3c26e254c 100644
--- a/tests/policies/test_replay_buffer_wrapper.py
+++ b/tests/policies/test_replay_buffer_wrapper.py
@@ -29,10 +29,10 @@ def make_algo_with_wrapped_buffer(
     rl_cls: Type[off_policy_algorithm.OffPolicyAlgorithm],
     policy_cls: Type[BasePolicy],
     replay_buffer_class: Type[buffers.ReplayBuffer],
-    random_state: np.random.RandomState,
+    rng: np.random.Generator,
     buffer_size: int = 100,
 ) -> off_policy_algorithm.OffPolicyAlgorithm:
-    venv = util.make_vec_env("Pendulum-v1", n_envs=1, random_state=random_state)
+    venv = util.make_vec_env("Pendulum-v1", n_envs=1, rng=rng)
     rl_algo = rl_cls(
         policy=policy_cls,
         policy_kwargs=dict(),
@@ -52,8 +52,7 @@ def make_algo_with_wrapped_buffer(
     return rl_algo
 
 
-def test_invalid_args(random_state_fixed):
-    random_state = random_state_fixed
+def test_invalid_args(rng):
     with pytest.raises(
         TypeError,
         match=r".*unexpected keyword argument 'replay_buffer_class'.*",
@@ -64,7 +63,7 @@ def test_invalid_args(random_state_fixed):
             rl_cls=sb3.PPO,  # type: ignore
             policy_cls=policies.ActorCriticPolicy,
             replay_buffer_class=buffers.ReplayBuffer,
-            random_state=random_state,
+            rng=rng,
         )
 
     with pytest.raises(AssertionError, match=r".*only ReplayBuffer is supported.*"):
@@ -72,12 +71,11 @@ def test_invalid_args(random_state_fixed):
             rl_cls=sb3.SAC,
             policy_cls=sb3.sac.policies.SACPolicy,
             replay_buffer_class=buffers.DictReplayBuffer,
-            random_state=random_state,
+            rng=rng,
         )
 
 
-def test_wrapper_class(tmpdir, random_state_fixed):
-    random_state = random_state_fixed
+def test_wrapper_class(tmpdir, rng):
     buffer_size = 15
     total_timesteps = 20
 
@@ -86,7 +84,7 @@ def test_wrapper_class(tmpdir, random_state_fixed):
         policy_cls=sb3.sac.policies.SACPolicy,
         replay_buffer_class=buffers.ReplayBuffer,
         buffer_size=buffer_size,
-        random_state=random_state,
+        rng=rng,
     )
 
     rl_algo.learn(total_timesteps=total_timesteps)
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index b6bb46cf7..4db5ceefc 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -125,12 +125,12 @@ def _sample(space, n):
     return np.array([space.sample() for _ in range(n)])
 
 
-def _make_env_and_save_reward_net(env_name, reward_type, tmpdir, random_state):
+def _make_env_and_save_reward_net(env_name, reward_type, tmpdir, rng):
     venv = util.make_vec_env(
         env_name,
         n_envs=1,
         parallel=False,
-        random_state=random_state,
+        rng=rng,
     )
     save_path = os.path.join(tmpdir, "norm_reward.pt")
 
@@ -160,14 +160,13 @@ def _make_env_and_save_reward_net(env_name, reward_type, tmpdir, random_state):
 
 @pytest.mark.parametrize("env_name", ENVS)
 @pytest.mark.parametrize("reward_type", DESERIALIZATION_TYPES)
-def test_reward_valid(env_name, reward_type, tmpdir, random_state_fixed):
+def test_reward_valid(env_name, reward_type, tmpdir, rng):
     """Test output of reward function is appropriate shape and type."""
-    random_state = random_state_fixed
     venv, tmppath = _make_env_and_save_reward_net(
         env_name,
         reward_type,
         tmpdir,
-        random_state,
+        rng,
     )
 
     TRAJECTORY_LEN = 10
@@ -183,13 +182,12 @@ def test_reward_valid(env_name, reward_type, tmpdir, random_state_fixed):
     assert isinstance(pred_reward[0], numbers.Number)
 
 
-def test_strip_wrappers_basic(random_state_fixed):
-    random_state = random_state_fixed
+def test_strip_wrappers_basic(rng):
     venv = util.make_vec_env(
         "FrozenLake-v1",
         n_envs=1,
         parallel=False,
-        random_state=random_state,
+        rng=rng,
     )
     net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     net = reward_nets.NormalizedRewardNet(net, networks.RunningNorm)
@@ -203,13 +201,12 @@ def test_strip_wrappers_basic(random_state_fixed):
     assert isinstance(net, reward_nets.BasicRewardNet)
 
 
-def test_strip_wrappers_complex(random_state_fixed):
-    random_state = random_state_fixed
+def test_strip_wrappers_complex(rng):
     venv = util.make_vec_env(
         "FrozenLake-v1",
         n_envs=1,
         parallel=False,
-        random_state=random_state,
+        rng=rng,
     )
     net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     net = reward_nets.ShapedRewardNet(net, _potential, discount_factor=0.99)
@@ -284,13 +281,12 @@ def forward(*args):
 
 
 @pytest.mark.parametrize("env_name", ENVS)
-def test_cant_load_unnorm_as_norm(env_name, tmpdir, random_state_fixed):
-    random_state = random_state_fixed
+def test_cant_load_unnorm_as_norm(env_name, tmpdir, rng):
     venv, tmppath = _make_env_and_save_reward_net(
         env_name,
         "RewardNet_unnormalized",
         tmpdir,
-        random_state=random_state,
+        rng=rng,
     )
     with pytest.raises(TypeError):
         serialize.load_reward("RewardNet_normalized", tmppath, venv)
@@ -304,17 +300,16 @@ def test_serialize_identity(
     net_cls,
     normalize_rewards,
     tmpdir,
-    random_state_fixed,
+    rng,
 ):
     """Does output of deserialized reward network match that of original?"""
-    random_state = random_state_fixed
     logging.info(f"Testing {net_cls}")
 
     venv = util.make_vec_env(
         env_name,
         n_envs=1,
         parallel=False,
-        random_state=random_state,
+        rng=rng,
     )
     original = net_cls(venv.observation_space, venv.action_space)
     if normalize_rewards:
@@ -332,7 +327,7 @@ def test_serialize_identity(
         random,
         venv,
         n_timesteps=100,
-        random_state=random_state,
+        rng=rng,
     )
 
     if isinstance(original, reward_nets.NormalizedRewardNet):
@@ -653,9 +648,8 @@ def forward(self):
 
 
 @pytest.mark.parametrize("normalize_input_layer", [None, networks.RunningNorm])
-def test_training_regression(normalize_input_layer, random_state_fixed):
+def test_training_regression(normalize_input_layer, rng):
     """Test reward_net normalization by training a regression model."""
-    random_state = random_state_fixed
     venv = DummyVecEnv([lambda: gym.make("CartPole-v0")] * 2)
     reward_net = reward_nets.BasicRewardNet(
         venv.observation_space,
@@ -678,7 +672,7 @@ def test_training_regression(normalize_input_layer, random_state_fixed):
             random,
             venv,
             n_timesteps=100,
-            random_state=random_state,
+            rng=rng,
         )
         trans_args = (
             transitions.obs,
diff --git a/tests/rewards/test_reward_wrapper.py b/tests/rewards/test_reward_wrapper.py
index fee403880..893bff59a 100644
--- a/tests/rewards/test_reward_wrapper.py
+++ b/tests/rewards/test_reward_wrapper.py
@@ -17,21 +17,20 @@ def __call__(self, obs, act, next_obs, steps=None):
         return (np.arange(len(obs)) + 1).astype("float32")
 
 
-def test_reward_overwrite(random_state_fixed):
+def test_reward_overwrite(rng):
     """Test that reward wrapper actually overwrites base rewards."""
-    random_state = random_state_fixed
     env_name = "Pendulum-v1"
     num_envs = 3
-    env = util.make_vec_env(env_name, random_state=random_state, n_envs=num_envs)
+    env = util.make_vec_env(env_name, rng=rng, n_envs=num_envs)
     reward_fn = FunkyReward()
     wrapped_env = reward_wrapper.RewardVecEnvWrapper(env, reward_fn)
     policy = RandomPolicy(env.observation_space, env.action_space)
     sample_until = rollout.make_min_episodes(10)
     default_stats = rollout.rollout_stats(
-        rollout.generate_trajectories(policy, env, sample_until, random_state),
+        rollout.generate_trajectories(policy, env, sample_until, rng),
     )
     wrapped_stats = rollout.rollout_stats(
-        rollout.generate_trajectories(policy, wrapped_env, sample_until, random_state),
+        rollout.generate_trajectories(policy, wrapped_env, sample_until, rng),
     )
     # Pendulum-v1 always has negative rewards
     assert default_stats["return_max"] < 0
diff --git a/tests/scripts/test_scripts.py b/tests/scripts/test_scripts.py
index 07c66571f..9f05b0dd5 100644
--- a/tests/scripts/test_scripts.py
+++ b/tests/scripts/test_scripts.py
@@ -679,13 +679,12 @@ def test_preference_comparisons_transfer_learning(
     _check_rollout_stats(run.result)
 
 
-def test_train_rl_double_normalization(tmpdir: str, random_state_fixed):
-    random_state = random_state_fixed
+def test_train_rl_double_normalization(tmpdir: str, rng):
     venv = util.make_vec_env(
         "CartPole-v1",
         n_envs=1,
         parallel=False,
-        random_state=random_state,
+        rng=rng,
     )
     basic_reward_net = reward_nets.BasicRewardNet(
         venv.observation_space,

From 3fa312bf8feb7fa46d5c685dfd508926d838364b Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 20:37:17 +0200
Subject: [PATCH 107/187] Add backticks to error message

---
 src/imitation/algorithms/preference_comparisons.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 36ebcab92..e73380c8a 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -861,7 +861,7 @@ def __init__(
         self.threshold = threshold
 
         if self.sample and self.rng is None:
-            raise ValueError("If sample is True, then rng must be provided.")
+            raise ValueError("If `sample` is True, then `rng` must be provided.")
 
     def __call__(self, fragment_pairs: Sequence[TrajectoryWithRewPair]) -> np.ndarray:
         """Computes probability fragment 1 is preferred over fragment 2."""

From d00b4d67cda6f76bd81d0b8fa51fe12354750567 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 20:47:17 +0200
Subject: [PATCH 108/187] Create "AnyNorm" alias

---
 src/imitation/util/networks.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index 50fb5350f..60d90f43d 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -3,10 +3,11 @@
 import collections
 import contextlib
 import functools
-from typing import Iterable, Optional, OrderedDict, Type
+from typing import Iterable, Optional, OrderedDict, Type, Union
 
 import torch as th
 from torch import nn
+from torch.nn.modules import normalization, batchnorm
 
 
 @contextlib.contextmanager
@@ -93,6 +94,15 @@ def update_stats(self, batch: th.Tensor) -> None:
         """Update `self.running_mean`, `self.running_var` and `self.count`."""
 
 
+AnyNorm = Union[
+    normalization.LayerNorm,
+    normalization.GroupNorm,
+    normalization.LocalResponseNorm,
+    batchnorm._BatchNorm,
+    BaseNorm,
+]
+
+
 class RunningNorm(BaseNorm):
     """Normalizes input to mean 0 and standard deviation 1 using a running average.
 
@@ -205,7 +215,7 @@ def build_mlp(
     dropout_prob: float = 0.0,
     squeeze_output: bool = False,
     flatten_input: bool = False,
-    normalize_input_layer: Optional[Type[BaseNorm]] = None,
+    normalize_input_layer: Optional[Type[AnyNorm]] = None,
 ) -> nn.Module:
     """Constructs a Torch MLP.
 

From c1d7d1215b6951557526c1bb421aeec8606af5e9 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 20:48:52 +0200
Subject: [PATCH 109/187] Small fix

---
 src/imitation/util/networks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index 60d90f43d..ac4033b85 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -96,7 +96,7 @@ def update_stats(self, batch: th.Tensor) -> None:
 
 AnyNorm = Union[
     normalization.LayerNorm,
-    normalization.GroupNorm,
+    # normalization.GroupNorm,  # does not work as it requires kwargs
     normalization.LocalResponseNorm,
     batchnorm._BatchNorm,
     BaseNorm,

From e456bb7449fb898a5a7ce87a8b83b1e35b914cca Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 21:21:38 +0200
Subject: [PATCH 110/187] Add additional checks to shapes in
 _set_demo_from_batch

---
 src/imitation/algorithms/density.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index dd35e2b30..bcd347682 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -143,8 +143,13 @@ def _set_demo_from_batch(
                 "STATE_STATE_DENSITY requires next_obs_b "
                 "to be provided, but it was None",
             )
+
+        assert act_b.shape[1:] == self.venv.action_space.shape
+        assert obs_b.shape[1:] == self.venv.observation_space.shape
+        assert len(act_b) == len(obs_b)
         if next_obs_b is not None:
             assert next_obs_b.shape[1:] == self.venv.observation_space.shape
+            assert len(next_obs_b) == len(obs_b)
 
         next_obs_b_iterator = next_obs_b or itertools.repeat(None)
         for obs, act, next_obs in zip(obs_b, act_b, next_obs_b_iterator):

From 52b71cdeaee7352301588ad3ba1ae6805ccd33cd Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 21:23:43 +0200
Subject: [PATCH 111/187] Fix RolloutStatsComputer type

---
 src/imitation/algorithms/bc.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index 676679161..61f2c28e2 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -180,7 +180,7 @@ class RolloutStatsComputer:
         n_episodes: The number of episodes to base the statistics on.
     """
 
-    venv: vec_env.VecEnv
+    venv: Optional[vec_env.VecEnv]
     n_episodes: int
 
     # TODO(shwang): Maybe instead use a callback that can be shared between
@@ -413,9 +413,6 @@ def train(
             self._bc_logger.reset_tensorboard_steps()
         self._bc_logger.log_epoch(0)
 
-        # TODO(juan) docstrings above say that this can be none and that no rollouts
-        #  are generated. However initializing this requires passing a non-None
-        #  venv.
         compute_rollout_stats = RolloutStatsComputer(
             log_rollouts_venv,
             log_rollouts_n_episodes,

From a23fd29daf9e7d724e3b1952a905270a2e35f687 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 21:27:16 +0200
Subject: [PATCH 112/187] Improved logging/messages in clean_notebooks.py

---
 ci/clean_notebooks.py                               | 13 +++++++++++--
 src/imitation/algorithms/preference_comparisons.py  |  1 -
 .../scripts/train_preference_comparisons.py         |  1 -
 src/imitation/util/networks.py                      |  2 +-
 4 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/ci/clean_notebooks.py b/ci/clean_notebooks.py
index ea9d270c6..fb1d326a0 100755
--- a/ci/clean_notebooks.py
+++ b/ci/clean_notebooks.py
@@ -30,6 +30,11 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
     with open(file) as f:
         nb = nbformat.read(f, as_version=4)
 
+    was_dirty = False
+
+    if check_only:
+        print(f"Checking {file}")
+
     # Remove the output and metadata from each cell
     # also reset the execution count
     # if the cell has no code, remove it
@@ -38,23 +43,28 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
             if check_only:
                 raise UncleanNotebookError(f"Notebook {file} has outputs")
             cell["outputs"] = []
+            was_dirty = True
         if "metadata" in cell and cell["metadata"]:
             if check_only:
                 raise UncleanNotebookError(f"Notebook {file} has metadata")
             cell["metadata"] = {}
+            was_dirty = True
         if "execution_count" in cell and cell["execution_count"]:
             if check_only:
                 raise UncleanNotebookError(f"Notebook {file} has execution count")
             cell["execution_count"] = None
+            was_dirty = True
         if cell["cell_type"] == "code" and not cell["source"]:
             if check_only:
                 raise UncleanNotebookError(f"Notebook {file} has empty code cell")
             nb.cells.remove(cell)
+            was_dirty = True
 
-    if not check_only:
+    if not check_only and was_dirty:
         # Write the notebook
         with open(file, "w") as f:
             nbformat.write(nb, f)
+        print(f"Cleaned {file}")
 
 
 if __name__ == "__main__":
@@ -91,7 +101,6 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
         print("No notebooks found")
         exit(1)
     for file in files:
-        print(f"Cleaning {file}" if not check_only else f"Checking {file}")
         try:
             clean_notebook(file, check_only=check_only)
         except UncleanNotebookError:
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index e73380c8a..554f8d9f9 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -6,7 +6,6 @@
 import abc
 import math
 import pickle
-import random
 from typing import (
     Any,
     Callable,
diff --git a/src/imitation/scripts/train_preference_comparisons.py b/src/imitation/scripts/train_preference_comparisons.py
index eecf18ac2..a36ea65c3 100644
--- a/src/imitation/scripts/train_preference_comparisons.py
+++ b/src/imitation/scripts/train_preference_comparisons.py
@@ -6,7 +6,6 @@
 
 import functools
 import os
-import random
 from typing import Any, Mapping, Optional, Type, Union
 
 import torch as th
diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index ac4033b85..382ef1afb 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -7,7 +7,7 @@
 
 import torch as th
 from torch import nn
-from torch.nn.modules import normalization, batchnorm
+from torch.nn.modules import batchnorm, normalization
 
 
 @contextlib.contextmanager

From 8989c4bfdf277daafa93aa0db6d561690f03c7d4 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 23:46:52 +0200
Subject: [PATCH 113/187] Fix issues resulting from merge

---
 examples/3_train_gail.ipynb                   |  5 +-
 examples/4_train_airl.ipynb                   |  9 ++-
 ...rain_preference_comparisons_with_cnn.ipynb | 12 ++--
 examples/7_train_density.ipynb                | 10 ++-
 setup.py                                      |  2 +-
 .../algorithms/preference_comparisons.py      |  4 +-
 src/imitation/util/util.py                    |  1 +
 .../algorithms/test_preference_comparisons.py | 72 +++++++++++++++++++
 tests/policies/test_policies.py               |  2 +-
 tests/rewards/test_reward_nets.py             |  8 ++-
 10 files changed, 104 insertions(+), 21 deletions(-)

diff --git a/examples/3_train_gail.ipynb b/examples/3_train_gail.ipynb
index 03c4ee68e..206f73fcc 100644
--- a/examples/3_train_gail.ipynb
+++ b/examples/3_train_gail.ipynb
@@ -69,6 +69,7 @@
     "        \"seals/CartPole-v0\",\n",
     "        n_envs=5,\n",
     "        post_wrappers=[lambda env, _: RolloutInfoWrapper(env)],\n",
+    "        rng=rng,\n",
     "    ),\n",
     "    rollout.make_sample_until(min_timesteps=None, min_episodes=60),\n",
     "    rng=rng,\n",
@@ -101,7 +102,7 @@
     "import gym\n",
     "\n",
     "\n",
-    "venv = make_vec_env(\"seals/CartPole-v0\", n_envs=8)\n",
+    "venv = make_vec_env(\"seals/CartPole-v0\", n_envs=8, rng=rng)\n",
     "learner = PPO(\n",
     "    env=venv,\n",
     "    policy=MlpPolicy,\n",
@@ -185,4 +186,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/examples/4_train_airl.ipynb b/examples/4_train_airl.ipynb
index 528b8bfa2..293af98fa 100644
--- a/examples/4_train_airl.ipynb
+++ b/examples/4_train_airl.ipynb
@@ -66,6 +66,7 @@
     "        \"seals/CartPole-v0\",\n",
     "        n_envs=5,\n",
     "        post_wrappers=[lambda env, _: RolloutInfoWrapper(env)],\n",
+    "        rng=rng,\n",
     "    ),\n",
     "    rollout.make_sample_until(min_timesteps=None, min_episodes=60),\n",
     "    rng=rng,\n",
@@ -93,13 +94,12 @@
     "from imitation.util.util import make_vec_env\n",
     "from stable_baselines3 import PPO\n",
     "from stable_baselines3.common.evaluation import evaluate_policy\n",
-    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "\n",
     "import gym\n",
     "import seals\n",
     "\n",
     "\n",
-    "venv = make_vec_env(\"seals/CartPole-v0\", n_envs=8)\n",
+    "venv = make_vec_env(\"seals/CartPole-v0\", n_envs=8, rng=rng)\n",
     "learner = PPO(\n",
     "    env=venv,\n",
     "    policy=MlpPolicy,\n",
@@ -176,9 +176,8 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.10"
-  },
-  "orig_nbformat": 4
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/examples/5a_train_preference_comparisons_with_cnn.ipynb b/examples/5a_train_preference_comparisons_with_cnn.ipynb
index 61f1e8889..85f0e9064 100644
--- a/examples/5a_train_preference_comparisons_with_cnn.ipynb
+++ b/examples/5a_train_preference_comparisons_with_cnn.ipynb
@@ -28,6 +28,7 @@
     "import torch as th\n",
     "import gym\n",
     "from gym.wrappers import TimeLimit\n",
+    "import numpy as np\n",
     "\n",
     "from seals.util import AutoResetWrapper\n",
     "\n",
@@ -43,6 +44,8 @@
     "\n",
     "device = th.device(\"cuda\" if th.cuda.is_available() else \"cpu\")\n",
     "\n",
+    "rng = np.random.default_rng()\n",
+    "\n",
     "# Here we ensure that our environment has constant-length episodes by resetting\n",
     "# it when done, and running until 100 timesteps have elapsed.\n",
     "# For real training, you will want a much longer time limit.\n",
@@ -63,8 +66,8 @@
     "    venv.action_space,\n",
     ").to(device)\n",
     "\n",
-    "fragmenter = preference_comparisons.RandomFragmenter(warning_threshold=0, seed=0)\n",
-    "gatherer = preference_comparisons.SyntheticGatherer(seed=0)\n",
+    "fragmenter = preference_comparisons.RandomFragmenter(warning_threshold=0, rng=rng)\n",
+    "gatherer = preference_comparisons.SyntheticGatherer(rng=rng)\n",
     "preference_model = preference_comparisons.PreferenceModel(reward_net)\n",
     "reward_trainer = preference_comparisons.BasicRewardTrainer(\n",
     "    model=reward_net,\n",
@@ -88,7 +91,7 @@
     "    reward_fn=reward_net,\n",
     "    venv=venv,\n",
     "    exploration_frac=0.0,\n",
-    "    seed=0,\n",
+    "    rng=rng,\n",
     ")\n",
     "\n",
     "pref_comparisons = preference_comparisons.PreferenceComparisons(\n",
@@ -102,7 +105,6 @@
     "    transition_oversampling=1,\n",
     "    initial_comparison_frac=0.1,\n",
     "    allow_variable_horizon=False,\n",
-    "    seed=0,\n",
     "    initial_epoch_multiplier=1,\n",
     ")"
    ]
@@ -231,4 +233,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}
\ No newline at end of file
diff --git a/examples/7_train_density.ipynb b/examples/7_train_density.ipynb
index 8a86b34f5..c062fb75d 100644
--- a/examples/7_train_density.ipynb
+++ b/examples/7_train_density.ipynb
@@ -62,6 +62,7 @@
     "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "from imitation.data.wrappers import RolloutInfoWrapper\n",
     "import gym\n",
+    "import numpy as np\n",
     "\n",
     "\n",
     "rng = np.random.default_rng()\n",
@@ -73,10 +74,13 @@
     "    [lambda: RolloutInfoWrapper(gym.make(env_name)) for _ in range(N_VEC)]\n",
     ")\n",
     "rollouts = rollout.rollout(\n",
-    "    expert, rollout_env, rollout.make_sample_until(min_timesteps=2000, min_episodes=57)\n",
+    "    expert,\n",
+    "    rollout_env,\n",
+    "    rollout.make_sample_until(min_timesteps=2000, min_episodes=57),\n",
+    "    rng=rng,\n",
     ")\n",
     "\n",
-    "env = util.make_vec_env(env_name, N_VEC)\n",
+    "env = util.make_vec_env(env_name, n_envs=N_VEC, rng=rng)\n",
     "\n",
     "\n",
     "imitation_trainer = PPO(ActorCriticPolicy, env, learning_rate=3e-4, n_steps=2048)\n",
@@ -153,4 +157,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 402dfa9cc..7c69699a0 100644
--- a/setup.py
+++ b/setup.py
@@ -217,7 +217,7 @@ def get_local_version(version: "ScmVersion", time_format="%Y%m%d") -> str:
         # recommended packages for development
         "dev": [
             "autopep8",
-            "awscli",
+            # "awscli",
             "ntfy[slack]",
             "ipdb",
             "isort~=5.0",
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 554f8d9f9..14351f131 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -1455,8 +1455,8 @@ def __init__(
             fragmenter,
         ):
             raise ValueError(
-                "If you provide your own fragmenter and preference gatherer, you "
-                "don't need to provide a random state.",
+                "If you provide your own fragmenter, preference gatherer, "
+                "and reward trainer, you don't need to provide a random state.",
             )
 
         if reward_trainer is None:
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 7884ed5a5..6e3ac15b8 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -67,6 +67,7 @@ def make_unique_timestamp() -> str:
 
 def make_vec_env(
     env_name: str,
+    *,
     rng: np.random.Generator,
     n_envs: int = 8,
     parallel: bool = False,
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 8ab001b1c..71a4e017a 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -191,6 +191,78 @@ def test_transitions_left_in_buffer(agent_trainer):
         agent_trainer.train(steps=1)
 
 
+@pytest.mark.parametrize(
+    "schedule",
+    ["constant", "hyperbolic", "inverse_quadratic", lambda t: 1 / (1 + t**3)],
+)
+def test_preference_comparisons_raises(
+    agent_trainer,
+    reward_net,
+    random_fragmenter,
+    preference_model,
+    custom_logger,
+    schedule,
+    rng,
+):
+    loss = preference_comparisons.CrossEntropyRewardLoss(preference_model)
+    reward_trainer = preference_comparisons.BasicRewardTrainer(
+        reward_net,
+        loss,
+    )
+    gatherer = preference_comparisons.SyntheticGatherer(rng=rng)
+    # no rng, must provide fragmenter, preference gatherer, reward trainer
+    no_rng_msg = (
+        ".*don't provide.*random state.*provide.*fragmenter"
+        ".*preference gatherer.*reward_trainer.*"
+    )
+
+    def build_preference_comparsions(gatherer, reward_trainer, fragmenter, rng):
+        preference_comparisons.PreferenceComparisons(
+            agent_trainer,
+            reward_net,
+            num_iterations=2,
+            transition_oversampling=2,
+            reward_trainer=reward_trainer,
+            preference_gatherer=gatherer,
+            fragmenter=fragmenter,
+            custom_logger=custom_logger,
+            query_schedule=schedule,
+            rng=rng,
+        )
+
+    with pytest.raises(ValueError, match=no_rng_msg):
+        build_preference_comparsions(gatherer, None, None, rng=None)
+
+    with pytest.raises(ValueError, match=no_rng_msg):
+        build_preference_comparsions(None, reward_trainer, None, rng=None)
+
+    with pytest.raises(ValueError, match=no_rng_msg):
+        build_preference_comparsions(None, None, random_fragmenter, rng=None)
+
+    # This should not raise
+    build_preference_comparsions(gatherer, reward_trainer, random_fragmenter, rng=None)
+
+    # if providing fragmenter, preference gatherer, reward trainer, does not need rng.
+    with_rng_msg = (
+        "provide.*fragmenter.*preference gatherer.*reward trainer"
+        ".*don't need.*random state.*"
+    )
+
+    with pytest.raises(ValueError, match=with_rng_msg):
+        build_preference_comparsions(
+            gatherer,
+            reward_trainer,
+            random_fragmenter,
+            rng=rng,
+        )
+
+    # This should not raise
+    build_preference_comparsions(None, None, None, rng=rng)
+    build_preference_comparsions(gatherer, None, None, rng=rng)
+    build_preference_comparsions(None, reward_trainer, None, rng=rng)
+    build_preference_comparsions(None, None, random_fragmenter, rng=rng)
+
+
 @pytest.mark.parametrize(
     "schedule",
     ["constant", "hyperbolic", "inverse_quadratic", lambda t: 1 / (1 + t**3)],
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index 38a8e6bf3..c7953a006 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -79,7 +79,7 @@ def test_save_stable_model_errors_and_warnings(
         serialize.load_policy(policy, venv, path=str(dir_nonexistent))
 
 
-def _test_serialize_identity(env_name, model_cfg, tmpdir):
+def _test_serialize_identity(env_name, model_cfg, tmpdir, rng):
     """Test output actions of deserialized policy are same as original."""
     venv = util.make_vec_env(
         env_name,
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index d7ab6c1b5..6bfb13958 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -360,8 +360,8 @@ def test_strip_wrappers_complex(rng):
     assert isinstance(net, reward_nets.BasicRewardNet)
 
 
-def test_strip_wrappers_image_complex():
-    venv = util.make_vec_env("Asteroids-v4", n_envs=1, parallel=False)
+def test_strip_wrappers_image_complex(rng):
+    venv = util.make_vec_env("Asteroids-v4", n_envs=1, parallel=False, rng=rng)
     net = reward_nets.CnnRewardNet(venv.observation_space, venv.action_space)
     net = reward_nets.ShapedRewardNet(net, _potential, discount_factor=0.99)
     net = reward_nets.NormalizedRewardNet(net, networks.RunningNorm)
@@ -538,6 +538,7 @@ def test_serialize_identity(
     net_kwargs,
     normalize_rewards,
     tmpdir,
+    rng,
 ):
     """Does output of deserialized reward MLP match that of original?"""
     _serialize_deserialize_identity(
@@ -546,6 +547,7 @@ def test_serialize_identity(
         net_kwargs,
         normalize_rewards,
         tmpdir,
+        rng,
     )
 
 
@@ -559,6 +561,7 @@ def test_serialize_identity_images(
     net_kwargs,
     normalize_rewards,
     tmpdir,
+    rng,
 ):
     """Does output of deserialized reward CNN match that of original?"""
     _serialize_deserialize_identity(
@@ -567,6 +570,7 @@ def test_serialize_identity_images(
         net_kwargs,
         normalize_rewards,
         tmpdir,
+        rng,
     )
 
 

From 2e54951aed9b182d28be7976bbb2bfcc0b2f408e Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 23:50:34 +0200
Subject: [PATCH 114/187] Bug fix

---
 src/imitation/scripts/common/demonstrations.py | 2 ++
 src/imitation/util/networks.py                 | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/imitation/scripts/common/demonstrations.py b/src/imitation/scripts/common/demonstrations.py
index 5e8437374..cc82c2ab0 100644
--- a/src/imitation/scripts/common/demonstrations.py
+++ b/src/imitation/scripts/common/demonstrations.py
@@ -43,6 +43,7 @@ def get_expert_trajectories(
 @demonstrations_ingredient.capture
 def generate_expert_trajs(
     n_expert_demos: Optional[int],
+    rng,
 ) -> Optional[Sequence[types.Trajectory]]:
     """Generates expert demonstrations.
 
@@ -67,6 +68,7 @@ def generate_expert_trajs(
             expert.get_expert_policy(rollout_env),
             rollout_env,
             rollout.make_sample_until(min_episodes=n_expert_demos),
+            rng=rng,
         )
 
 
diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index dd362d08c..1c7c993ba 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -215,7 +215,7 @@ def build_mlp(
     dropout_prob: float = 0.0,
     squeeze_output: bool = False,
     flatten_input: bool = False,
-    normalize_input_layer: Optional[Type[AnyNorm]] = None,
+    normalize_input_layer: Optional[Type[nn.Module]] = None,
 ) -> nn.Module:
     """Constructs a Torch MLP.
 

From 1d001d33af13cf5f321178ac42ba07d47cefdf2e Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 13 Sep 2022 23:53:57 +0200
Subject: [PATCH 115/187] Bug fix (wasn't really fixed before)

---
 src/imitation/scripts/common/demonstrations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/scripts/common/demonstrations.py b/src/imitation/scripts/common/demonstrations.py
index cc82c2ab0..2cbe23a95 100644
--- a/src/imitation/scripts/common/demonstrations.py
+++ b/src/imitation/scripts/common/demonstrations.py
@@ -43,7 +43,6 @@ def get_expert_trajectories(
 @demonstrations_ingredient.capture
 def generate_expert_trajs(
     n_expert_demos: Optional[int],
-    rng,
 ) -> Optional[Sequence[types.Trajectory]]:
     """Generates expert demonstrations.
 
@@ -57,6 +56,7 @@ def generate_expert_trajs(
     Raises:
         ValueError: If n_expert_demos is None.
     """
+    rng = common.make_rng()
     if n_expert_demos is None:
         raise ValueError("n_expert_demos must be specified when rollout_path is None")
 

From 71b29d7ec3233e921594049a61545c9410c114f0 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 14 Sep 2022 00:09:23 +0200
Subject: [PATCH 116/187] Fixed docs example of BC

---
 docs/algorithms/bc.rst                     | 5 ++++-
 tests/policies/test_exploration_wrapper.py | 6 ++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/docs/algorithms/bc.rst b/docs/algorithms/bc.rst
index a60abc064..daf308272 100644
--- a/docs/algorithms/bc.rst
+++ b/docs/algorithms/bc.rst
@@ -18,7 +18,8 @@ Example
 Detailed example notebook: `1_train_bc.ipynb <https://github.com/HumanCompatibleAI/imitation/blob/master/examples/1_train_bc.ipynb>`_
 
 .. testcode::
-    
+
+    import numpy as np
     import gym
     from stable_baselines3 import PPO
     from stable_baselines3.common.evaluation import evaluate_policy
@@ -29,6 +30,7 @@ Detailed example notebook: `1_train_bc.ipynb <https://github.com/HumanCompatible
     from imitation.data import rollout
     from imitation.data.wrappers import RolloutInfoWrapper
 
+    rng = np.random.default_rng(0)
     env = gym.make("CartPole-v1")
     expert = PPO(policy=MlpPolicy, env=env)
     expert.learn(1000)
@@ -37,6 +39,7 @@ Detailed example notebook: `1_train_bc.ipynb <https://github.com/HumanCompatible
         expert,
         DummyVecEnv([lambda: RolloutInfoWrapper(env)]),
         rollout.make_sample_until(min_timesteps=None, min_episodes=50),
+        rng=rng,
     )
     transitions = rollout.flatten_trajectories(rollouts)
 
diff --git a/tests/policies/test_exploration_wrapper.py b/tests/policies/test_exploration_wrapper.py
index 675bd84b9..48b7b2469 100644
--- a/tests/policies/test_exploration_wrapper.py
+++ b/tests/policies/test_exploration_wrapper.py
@@ -40,6 +40,9 @@ def test_random_prob(rng):
     (2) `random_prob=1.0`: Initial and following policies are always random policies.
     (3) `random_prob=0.5`: Around half-half for constant and random policies.
 
+    Args:
+        rng (np.random.Generator): random number generator.
+
     Raises:
         ValueError: Unknown policy type to switch.
     """
@@ -81,6 +84,9 @@ def test_switch_prob(rng):
     (1) `switch_prob=0.0`: The policy never switches after initial switch.
     (2) `switch_prob=1.0`: The policy always switches and the distribution of
         policies is determined by `random_prob`.
+
+    Args:
+        rng (np.random.Generator): random number generator.
     """
     wrapper, venv = make_wrapper(random_prob=0.5, switch_prob=0.0, rng=rng)
     policy = wrapper.current_policy

From d1d12d59dbbc20019afe953815f0939e8bf9a9ea Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 14 Sep 2022 17:23:59 +0200
Subject: [PATCH 117/187] Fix bugs resulting from merge

---
 docs/algorithms/preference_comparisons.rst              | 5 +++++
 examples/5_train_preference_comparisons.ipynb           | 3 ++-
 examples/5a_train_preference_comparisons_with_cnn.ipynb | 1 +
 src/imitation/algorithms/preference_comparisons.py      | 9 +++++++--
 tests/algorithms/test_preference_comparisons.py         | 8 ++++++++
 5 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/docs/algorithms/preference_comparisons.rst b/docs/algorithms/preference_comparisons.rst
index fc1667c24..09cd9e0fe 100644
--- a/docs/algorithms/preference_comparisons.rst
+++ b/docs/algorithms/preference_comparisons.rst
@@ -22,6 +22,8 @@ Detailed example notebook: `5_train_preference_comparisons.ipynb <https://github
 
 .. testcode::
 
+    import numpy as np
+
     from stable_baselines3 import PPO
     from stable_baselines3.common.evaluation import evaluate_policy
     from stable_baselines3.ppo import MlpPolicy
@@ -33,6 +35,8 @@ Detailed example notebook: `5_train_preference_comparisons.ipynb <https://github
     from imitation.util.networks import RunningNorm
     from imitation.util.util import make_vec_env
 
+    rng = np.random.default_rng(0)
+
     venv = make_vec_env("Pendulum-v1")
 
     reward_net = BasicRewardNet(
@@ -46,6 +50,7 @@ Detailed example notebook: `5_train_preference_comparisons.ipynb <https://github
         model=reward_net,
         loss=preference_comparisons.CrossEntropyRewardLoss(preference_model),
         epochs=3,
+        rng=rng,
     )
 
     agent = PPO(
diff --git a/examples/5_train_preference_comparisons.ipynb b/examples/5_train_preference_comparisons.ipynb
index 2a537c6bd..158446939 100644
--- a/examples/5_train_preference_comparisons.ipynb
+++ b/examples/5_train_preference_comparisons.ipynb
@@ -50,6 +50,7 @@
     "    model=reward_net,\n",
     "    loss=preference_comparisons.CrossEntropyRewardLoss(preference_model),\n",
     "    epochs=3,\n",
+    "    rng=rng,\n",
     ")\n",
     "\n",
     "agent = PPO(\n",
@@ -201,4 +202,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/examples/5a_train_preference_comparisons_with_cnn.ipynb b/examples/5a_train_preference_comparisons_with_cnn.ipynb
index 85f0e9064..21937d04e 100644
--- a/examples/5a_train_preference_comparisons_with_cnn.ipynb
+++ b/examples/5a_train_preference_comparisons_with_cnn.ipynb
@@ -73,6 +73,7 @@
     "    model=reward_net,\n",
     "    loss=preference_comparisons.CrossEntropyRewardLoss(preference_model),\n",
     "    epochs=3,\n",
+    "    rng=rng,\n",
     ")\n",
     "\n",
     "agent = PPO(\n",
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 80d7e8045..587cc35fd 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -1189,7 +1189,7 @@ def _train(self, dataset: PreferenceDataset, epoch_multiplier: float = 1.0) -> N
                 dataset,
                 lengths=[train_length, val_length],
                 # we convert the numpy generator to the pytorch generator.
-                generator=th.Generator().manual_seed(util.make_seeds(self.rng))
+                generator=th.Generator().manual_seed(util.make_seeds(self.rng)),
             )
             dataloader = self._make_data_loader(train_dataset)
             val_dataloader = self._make_data_loader(val_dataset)
@@ -1395,7 +1395,12 @@ def _make_reward_trainer(
                 f" by AddSTDRewardWrapper but found {type(reward_model).__name__}.",
             )
     else:
-        return BasicRewardTrainer(reward_model, loss=loss, **reward_trainer_kwargs)
+        return BasicRewardTrainer(
+            reward_model,
+            loss=loss,
+            rng=rng,
+            **reward_trainer_kwargs,
+        )
 
 
 QUERY_SCHEDULES: Dict[str, type_aliases.Schedule] = {
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index f243447f1..73e831d86 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -208,6 +208,7 @@ def test_preference_comparisons_raises(
     reward_trainer = preference_comparisons.BasicRewardTrainer(
         reward_net,
         loss,
+        rng=rng,
     )
     gatherer = preference_comparisons.SyntheticGatherer(rng=rng)
     # no rng, must provide fragmenter, preference gatherer, reward trainer
@@ -396,6 +397,7 @@ def test_discount_rate_no_crash(
     reward_trainer = preference_comparisons.BasicRewardTrainer(
         reward_net,
         loss,
+        rng=rng,
     )
 
     main_trainer = preference_comparisons.PreferenceComparisons(
@@ -626,6 +628,7 @@ def test_reward_trainer_regularization_no_crash(
     custom_logger,
     preference_model,
     interval_param_scaler,
+    rng,
 ):
     reward_net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     loss = preference_comparisons.CrossEntropyRewardLoss(preference_model)
@@ -641,6 +644,7 @@ def test_reward_trainer_regularization_no_crash(
         loss,
         regularizer_factory=regularizer_factory,
         custom_logger=custom_logger,
+        rng=rng,
     )
 
     main_trainer = preference_comparisons.PreferenceComparisons(
@@ -652,6 +656,7 @@ def test_reward_trainer_regularization_no_crash(
         fragmenter=random_fragmenter,
         reward_trainer=reward_trainer,
         custom_logger=custom_logger,
+        rng=rng,
     )
     main_trainer.train(50, 50)
 
@@ -663,6 +668,7 @@ def test_reward_trainer_regularization_raises(
     custom_logger,
     preference_model,
     interval_param_scaler,
+    rng,
 ):
     reward_net = reward_nets.BasicRewardNet(venv.observation_space, venv.action_space)
     loss = preference_comparisons.CrossEntropyRewardLoss(preference_model)
@@ -678,6 +684,7 @@ def test_reward_trainer_regularization_raises(
         loss,
         regularizer_factory=regularizer_factory,
         custom_logger=custom_logger,
+        rng=rng,
     )
 
     main_trainer = preference_comparisons.PreferenceComparisons(
@@ -689,6 +696,7 @@ def test_reward_trainer_regularization_raises(
         fragmenter=random_fragmenter,
         reward_trainer=reward_trainer,
         custom_logger=custom_logger,
+        rng=rng,
     )
     with pytest.raises(
         ValueError,

From d8d355f63cdbc3cf96ede1321a6bd36d6ba5d394 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 14 Sep 2022 17:28:44 +0200
Subject: [PATCH 118/187] Fix docs (dagger.rst) caught by sphinx CI

---
 docs/algorithms/dagger.rst | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/docs/algorithms/dagger.rst b/docs/algorithms/dagger.rst
index cf71d4b1d..fd4f7820f 100644
--- a/docs/algorithms/dagger.rst
+++ b/docs/algorithms/dagger.rst
@@ -24,7 +24,7 @@ Detailed example notebook: `2_train_dagger.ipynb <https://github.com/HumanCompat
 .. testcode::
 
     import tempfile
-
+    import numpy as np
     import gym
     from stable_baselines3 import PPO
     from stable_baselines3.common.evaluation import evaluate_policy
@@ -34,6 +34,7 @@ Detailed example notebook: `2_train_dagger.ipynb <https://github.com/HumanCompat
     from imitation.algorithms import bc
     from imitation.algorithms.dagger import SimpleDAggerTrainer
 
+    rng = np.random.default_rng(0)
     env = gym.make("CartPole-v1")
     expert = PPO(policy=MlpPolicy, env=env)
     expert.learn(1000)
@@ -42,6 +43,7 @@ Detailed example notebook: `2_train_dagger.ipynb <https://github.com/HumanCompat
     bc_trainer = bc.BC(
         observation_space=env.observation_space,
         action_space=env.action_space,
+        rng=rng,
     )
     with tempfile.TemporaryDirectory(prefix="dagger_example_") as tmpdir:
         print(tmpdir)

From 7ba6dc606ff68748a0be8c113be71a935f9053d9 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 14 Sep 2022 17:42:08 +0200
Subject: [PATCH 119/187] Add mypy to CI

---
 .circleci/config.yml | 4 ++++
 ci/code_checks.sh    | 1 +
 2 files changed, 5 insertions(+)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index c23f36f41..66f77cf1c 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -257,6 +257,10 @@ jobs:
           name: pytype
           command: pytype --version && pytype -j "${NUM_CPUS}" ${SRC_FILES[@]}
 
+      - run:
+          name: mypy
+          command: mypy --version && mypy ${SRC_FILES[@]}
+
   unit-test-linux:
     executor: unit-test-linux
     steps:
diff --git a/ci/code_checks.sh b/ci/code_checks.sh
index eaa9d651e..94577deb2 100755
--- a/ci/code_checks.sh
+++ b/ci/code_checks.sh
@@ -23,6 +23,7 @@ fi
 if [ "${skipexpensive:-}" != "true" ]; then
   echo "Type checking"
   pytype -j auto "${SRC_FILES[@]}"
+  mypy --show-error-codes "${SRC_FILES[@]}"
 
   echo "Building docs (validates docstrings)"
   pushd docs/

From 024ca75cdc067709cdf2b89cfc1932ff6a5ebb8c Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 14 Sep 2022 18:36:16 +0200
Subject: [PATCH 120/187] Continue fixing miscellaneous type errors

---
 src/imitation/algorithms/bc.py            |  7 +++--
 src/imitation/algorithms/dagger.py        | 35 ++++++++++++++++-------
 src/imitation/algorithms/density.py       |  4 +++
 src/imitation/envs/examples/model_envs.py |  4 +--
 src/imitation/policies/serialize.py       | 10 +++----
 src/imitation/rewards/reward_nets.py      |  6 ++--
 src/imitation/util/networks.py            |  2 +-
 src/imitation/util/util.py                | 15 ++++++----
 tests/conftest.py                         |  4 ++-
 tests/data/test_types.py                  |  1 +
 tests/test_regularization.py              |  3 +-
 11 files changed, 59 insertions(+), 32 deletions(-)

diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index 1dc494894..3514270b5 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -27,7 +27,7 @@
 from imitation.algorithms import base as algo_base
 from imitation.data import rollout, types
 from imitation.policies import base as policy_base
-from imitation.util import logger as imit_logger
+from imitation.util import logger as imit_logger, util
 
 
 @dataclasses.dataclass(frozen=True)
@@ -113,6 +113,8 @@ def __call__(
             A BCTrainingMetrics object with the loss and all the components it
             consists of.
         """
+        obs = util.safe_to_tensor(obs)
+        acts = util.safe_to_tensor(acts)
         _, log_prob, entropy = policy.evaluate_actions(obs, acts)
         prob_true_act = th.exp(log_prob).mean()
         log_prob = log_prob.mean()
@@ -426,6 +428,7 @@ def _on_epoch_end(epoch_number: int):
             if on_epoch_end is not None:
                 on_epoch_end()
 
+        assert self._demo_data_loader is not None
         demonstration_batches = BatchIteratorWithEpochEndCallback(
             self._demo_data_loader,
             n_epochs,
@@ -466,4 +469,4 @@ def save_policy(self, policy_path: types.AnyPath) -> None:
         Args:
             policy_path: path to save policy to.
         """
-        th.save(self.policy, policy_path)
+        th.save(self.policy, types.path_to_str(policy_path))
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index deea8b7a0..24bff78fa 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -21,7 +21,7 @@
 
 from imitation.algorithms import base, bc
 from imitation.data import rollout, types
-from imitation.util import logger, util
+from imitation.util import logger as imit_logger, util
 
 
 class BetaSchedule(abc.ABC):
@@ -69,7 +69,7 @@ def __call__(self, round_num: int) -> float:
 def reconstruct_trainer(
     scratch_dir: types.AnyPath,
     venv: vec_env.VecEnv,
-    custom_logger: Optional[logger.HierarchicalLogger] = None,
+    custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     device: Union[th.device, str] = "auto",
 ) -> "DAggerTrainer":
     """Reconstruct trainer from the latest snapshot in some working directory.
@@ -88,8 +88,8 @@ def reconstruct_trainer(
     Returns:
         A deserialized `DAggerTrainer`.
     """
-    custom_logger = custom_logger or logger.configure()
-    checkpoint_path = pathlib.Path(scratch_dir, "checkpoint-latest.pt")
+    custom_logger = custom_logger or imit_logger.configure()
+    checkpoint_path = pathlib.Path(types.path_to_str(scratch_dir), "checkpoint-latest.pt")
     trainer = th.load(checkpoint_path, map_location=utils.get_device(device))
     trainer.venv = venv
     trainer._logger = custom_logger
@@ -105,14 +105,14 @@ def _save_dagger_demo(
     #   however that NPZ save here is likely more space efficient than
     #   pickle from types.save(), and types.save only accepts
     #   TrajectoryWithRew right now (subclass of Trajectory).
-    save_dir = pathlib.Path(save_dir)
+    save_dir_obj = pathlib.Path(types.path_to_str(save_dir))
     assert isinstance(trajectory, types.Trajectory)
     actual_prefix = f"{prefix}-" if prefix else ""
     timestamp = util.make_unique_timestamp()
     filename = f"{actual_prefix}dagger-demo-{timestamp}.npz"
 
-    save_dir.mkdir(parents=True, exist_ok=True)
-    npz_path = pathlib.Path(save_dir, filename)
+    save_dir_obj.mkdir(parents=True, exist_ok=True)
+    npz_path = save_dir_obj / filename
     np.savez_compressed(npz_path, **dataclasses.asdict(trajectory))
     logging.info(f"Saved demo at '{npz_path}'")
 
@@ -146,6 +146,10 @@ class InteractiveTrajectoryCollector(vec_env.VecEnvWrapper):
     of every episode.
     """
 
+    traj_accum: Optional[rollout.TrajectoryAccumulator]
+    _last_obs: Optional[np.ndarray]
+    _last_user_actions: Optional[np.ndarray]
+
     def __init__(
         self,
         venv: vec_env.VecEnv,
@@ -203,6 +207,7 @@ def reset(self) -> np.ndarray:
         """
         self.traj_accum = rollout.TrajectoryAccumulator()
         obs = self.venv.reset()
+        assert isinstance(obs, np.ndarray)
         for i, ob in enumerate(obs):
             self.traj_accum.add_step({"obs": ob}, key=i)
         self._last_obs = obs
@@ -230,6 +235,7 @@ def step_async(self, actions: np.ndarray) -> None:
                 and executed instead via `self.get_robot_act`.
         """
         assert self._is_reset, "call .reset() before .step()"
+        assert self._last_obs is not None
 
         # Replace each given action with a robot action 100*(1-beta)% of the time.
         actual_acts = np.array(actions)
@@ -250,6 +256,9 @@ def step_wait(self) -> VecEnvStepReturn:
             Observation, reward, dones (is terminal?) and info dict.
         """
         next_obs, rews, dones, infos = self.venv.step_wait()
+        assert isinstance(next_obs, np.ndarray)
+        assert self.traj_accum is not None
+        assert self._last_user_actions is not None
         self._last_obs = next_obs
         fresh_demos = self.traj_accum.add_steps_and_auto_finish(
             obs=next_obs,
@@ -311,7 +320,7 @@ def __init__(
         rng: np.random.Generator,
         beta_schedule: Optional[Callable[[int], float]] = None,
         bc_trainer: bc.BC,
-        custom_logger: Optional[logger.HierarchicalLogger] = None,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
     ):
         """Builds DAggerTrainer.
 
@@ -331,7 +340,7 @@ def __init__(
         if beta_schedule is None:
             beta_schedule = LinearBetaSchedule(15)
         self.beta_schedule = beta_schedule
-        self.scratch_dir = pathlib.Path(scratch_dir)
+        self.scratch_dir = pathlib.Path(types.path_to_str(scratch_dir))
         self.venv = venv
         self.round_num = 0
         self._last_loaded_round = -1
@@ -353,8 +362,12 @@ def __getstate__(self):
         del d["_logger"]
         return d
 
-    @base.BaseImitationAlgorithm.logger.setter
-    def logger(self, value: logger.HierarchicalLogger) -> None:
+    @property
+    def logger(self) -> imit_logger.HierarchicalLogger:
+        """Returns logger for this object."""
+        return super().logger
+    @logger.setter
+    def logger(self, value: imit_logger.HierarchicalLogger) -> None:
         # DAgger and inner-BC logger should stay in sync
         self._logger = value
         self.bc_trainer.logger = value
diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index bcd347682..1b95aea4f 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -304,10 +304,12 @@ def __call__(
         assert len(state) == len(action) and len(state) == len(next_state)
         for idx, (obs, act, next_obs) in enumerate(zip(state, action, next_state)):
             flat_trans = self._preprocess_transition(obs, act, next_obs)
+            assert self._scaler is not None
             scaled_padded_trans = self._scaler.transform(flat_trans[np.newaxis])
             if self.is_stationary:
                 rew = self._density_models[None].score(scaled_padded_trans)
             else:
+                assert steps is not None
                 time = steps[idx]
                 if time >= len(self._density_models):
                     # Can't do anything sensible here yet. Correct solution is to use
@@ -334,6 +336,7 @@ def train_policy(self, n_timesteps: int = int(1e6), **kwargs):
                 method of the imitation RL model. Refer to Stable Baselines docs for
                 details.
         """
+        assert self.rl_algo is not None
         self.rl_algo.set_env(self.venv_wrapped)
         self.rl_algo.learn(
             n_timesteps,
@@ -373,4 +376,5 @@ def test_policy(self, *, n_trajectories: int = 10, true_reward: bool = True):
 
     @property
     def policy(self) -> base_class.BasePolicy:
+        assert self.rl_algo is not None
         return self.rl_algo.policy
diff --git a/src/imitation/envs/examples/model_envs.py b/src/imitation/envs/examples/model_envs.py
index edc677fa3..fa4747764 100644
--- a/src/imitation/envs/examples/model_envs.py
+++ b/src/imitation/envs/examples/model_envs.py
@@ -99,11 +99,11 @@ def make_obs_mat(
     Returns:
         A matrix of shape `(n_states, obs_dim if is_random else n_states)`.
     """
-    if not is_random:
-        assert obs_dim is None
     if is_random:
+        assert obs_dim is not None
         obs_mat = rng.normal(0, 2, (n_states, obs_dim))
     else:
+        assert obs_dim is None
         obs_mat = np.identity(n_states)
     assert (
         obs_mat.ndim == 2 and obs_mat.shape[:1] == (n_states,) and obs_mat.shape[1] > 0
diff --git a/src/imitation/policies/serialize.py b/src/imitation/policies/serialize.py
index 5759b5f48..cfd68414d 100644
--- a/src/imitation/policies/serialize.py
+++ b/src/imitation/policies/serialize.py
@@ -48,17 +48,17 @@ def load_stable_baselines_model(
         The deserialized RL algorithm.
     """
     logging.info(f"Loading Stable Baselines policy for '{cls}' from '{path}'")
-    path = pathlib.Path(path)
+    path_obj = pathlib.Path(path)
 
-    if path.is_dir():
-        path = path / "model.zip"
-        if not path.exists():
+    if path_obj.is_dir():
+        path_obj = path_obj / "model.zip"
+        if not path_obj.exists():
             raise FileNotFoundError(
                 f"Expected '{path}' to be a directory containing a 'model.zip' file.",
             )
 
     # SOMEDAY(adam): added 2022-01, can probably remove this check in 2023
-    vec_normalize_path = path.parent / "vec_normalize.pkl"
+    vec_normalize_path = path_obj.parent / "vec_normalize.pkl"
     if vec_normalize_path.exists():
         raise FileExistsError(
             "Outdated policy format: we do not support restoring normalization "
diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index 442028a69..1abbba9ee 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -1,7 +1,7 @@
 """Constructs deep network reward models."""
 
 import abc
-from typing import Callable, Iterable, Optional, Sequence, Tuple, Type, cast
+from typing import Callable, Iterable, Optional, Sequence, Tuple, Type, cast, Dict, Any
 
 import gym
 import numpy as np
@@ -427,7 +427,7 @@ def __init__(
         if self.use_done:
             combined_size += 1
 
-        full_build_mlp_kwargs = {
+        full_build_mlp_kwargs: Dict[str, Any] = {
             "hid_sizes": (32, 32),
             **kwargs,
             # we do not want the values below to be overridden
@@ -529,7 +529,7 @@ def __init__(
         if self.use_done:
             output_size *= 2
 
-        full_build_cnn_kwargs = {
+        full_build_cnn_kwargs: Dict[str, Any] = {
             "hid_channels": (32, 32),
             **kwargs,
             # we do not want the values below to be overridden
diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index 1c7c993ba..45435ce0b 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -319,7 +319,7 @@ def build_cnn(
     Raises:
         ValueError: if squeeze_output was supplied with out_size!=1.
     """
-    layers = collections.OrderedDict()
+    layers: OrderedDict[str, nn.Module] = collections.OrderedDict()
 
     if name is None:
         prefix = ""
diff --git a/src/imitation/util/util.py b/src/imitation/util/util.py
index 6e3ac15b8..c5900fcfe 100644
--- a/src/imitation/util/util.py
+++ b/src/imitation/util/util.py
@@ -226,7 +226,7 @@ def endless_iter(iterable: Iterable[T]) -> Iterator[T]:
     return itertools.chain.from_iterable(itertools.repeat(iterable))
 
 
-def safe_to_tensor(numpy_array: np.ndarray, **kwargs) -> th.Tensor:
+def safe_to_tensor(array: Union[np.ndarray, th.Tensor], **kwargs) -> th.Tensor:
     """Converts a NumPy array to a PyTorch tensor.
 
     The data is copied in the case where the array is non-writable. Unfortunately if
@@ -234,16 +234,19 @@ def safe_to_tensor(numpy_array: np.ndarray, **kwargs) -> th.Tensor:
     undefined behavior if you try to write to the tensor.
 
     Args:
-        numpy_array: The numpy array to convert to a PyTorch tensor.
+        array: The numpy array to convert to a PyTorch tensor.
         kwargs: Additional keyword arguments to pass to `th.as_tensor`.
 
     Returns:
-        A PyTorch tensor with the same content as `numpy_array`.
+        A PyTorch tensor with the same content as `array`.
     """
-    if not numpy_array.flags.writeable:
-        numpy_array = numpy_array.copy()
+    if isinstance(array, th.Tensor):
+        return array
 
-    return th.as_tensor(numpy_array, **kwargs)
+    if not array.flags.writeable:
+        array = array.copy()
+
+    return th.as_tensor(array, **kwargs)
 
 
 @overload
diff --git a/tests/conftest.py b/tests/conftest.py
index 8f7c74d5b..765da98f3 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -97,12 +97,14 @@ def pendulum_venv() -> VecEnv:
 
 @pytest.fixture
 def pendulum_expert_policy() -> BasePolicy:
-    return PPO.load(
+    policy = PPO.load(
         load_from_hub(
             "HumanCompatibleAI/ppo-Pendulum-v1",
             "ppo-Pendulum-v1.zip",
         ),
     ).policy
+    assert policy is not None
+    return policy
 
 
 @pytest.fixture
diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index 7d9330fa4..dbfacf3b2 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -272,6 +272,7 @@ def test_invalid_trajectories(
                 ValueError,
                 match=r"infos when present must be present for each action.*",
             ):
+                assert traj.infos is not None
                 dataclasses.replace(traj, infos=traj.infos[:-1])
             with pytest.raises(
                 ValueError,
diff --git a/tests/test_regularization.py b/tests/test_regularization.py
index 7c09a1ccf..07f7163ae 100644
--- a/tests/test_regularization.py
+++ b/tests/test_regularization.py
@@ -1,6 +1,7 @@
 """Tests for `imitation.regularization.*`."""
 import itertools
 import tempfile
+from typing import Union
 
 import numpy as np
 import pytest
@@ -326,7 +327,7 @@ class SimpleLossRegularizer(regularizers.LossRegularizer):
     It multiplies the total loss by lambda_+1.
     """
 
-    def _loss_penalty(self, loss: th.Tensor) -> th.Tensor:
+    def _loss_penalty(self, loss: regularizers.Scalar) -> regularizers.Scalar:
         return loss * self.lambda_  # this multiplies the total loss by lambda_+1.
 
 

From fca86247e5e87dadc703e57d45c8d25cd8bbdcf4 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 14 Sep 2022 18:40:37 +0200
Subject: [PATCH 121/187] Linting

---
 src/imitation/algorithms/bc.py       | 3 ++-
 src/imitation/algorithms/dagger.py   | 9 +++++++--
 src/imitation/rewards/reward_nets.py | 2 +-
 tests/test_regularization.py         | 1 -
 4 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index 3514270b5..0c387d1f2 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -27,7 +27,8 @@
 from imitation.algorithms import base as algo_base
 from imitation.data import rollout, types
 from imitation.policies import base as policy_base
-from imitation.util import logger as imit_logger, util
+from imitation.util import logger as imit_logger
+from imitation.util import util
 
 
 @dataclasses.dataclass(frozen=True)
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index 24bff78fa..9a3b086f8 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -21,7 +21,8 @@
 
 from imitation.algorithms import base, bc
 from imitation.data import rollout, types
-from imitation.util import logger as imit_logger, util
+from imitation.util import logger as imit_logger
+from imitation.util import util
 
 
 class BetaSchedule(abc.ABC):
@@ -89,7 +90,10 @@ def reconstruct_trainer(
         A deserialized `DAggerTrainer`.
     """
     custom_logger = custom_logger or imit_logger.configure()
-    checkpoint_path = pathlib.Path(types.path_to_str(scratch_dir), "checkpoint-latest.pt")
+    checkpoint_path = pathlib.Path(
+        types.path_to_str(scratch_dir),
+        "checkpoint-latest.pt",
+    )
     trainer = th.load(checkpoint_path, map_location=utils.get_device(device))
     trainer.venv = venv
     trainer._logger = custom_logger
@@ -366,6 +370,7 @@ def __getstate__(self):
     def logger(self) -> imit_logger.HierarchicalLogger:
         """Returns logger for this object."""
         return super().logger
+
     @logger.setter
     def logger(self, value: imit_logger.HierarchicalLogger) -> None:
         # DAgger and inner-BC logger should stay in sync
diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index 1abbba9ee..b712f5af1 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -1,7 +1,7 @@
 """Constructs deep network reward models."""
 
 import abc
-from typing import Callable, Iterable, Optional, Sequence, Tuple, Type, cast, Dict, Any
+from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple, Type, cast
 
 import gym
 import numpy as np
diff --git a/tests/test_regularization.py b/tests/test_regularization.py
index 07f7163ae..1e19f8610 100644
--- a/tests/test_regularization.py
+++ b/tests/test_regularization.py
@@ -1,7 +1,6 @@
 """Tests for `imitation.regularization.*`."""
 import itertools
 import tempfile
-from typing import Union
 
 import numpy as np
 import pytest

From 4923f59d42c849fe4ca02ba3e4d546fd249875c4 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 23 Sep 2022 15:58:33 +0100
Subject: [PATCH 122/187] Fix issue with normalize_input_layer type

---
 src/imitation/util/networks.py | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index 45435ce0b..98c3c65bb 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -94,15 +94,6 @@ def update_stats(self, batch: th.Tensor) -> None:
         """Update `self.running_mean`, `self.running_var` and `self.count`."""
 
 
-AnyNorm = Union[
-    normalization.LayerNorm,
-    # normalization.GroupNorm,  # does not work as it requires kwargs
-    normalization.LocalResponseNorm,
-    batchnorm._BatchNorm,
-    BaseNorm,
-]
-
-
 class RunningNorm(BaseNorm):
     """Normalizes input to mean 0 and standard deviation 1 using a running average.
 
@@ -256,7 +247,14 @@ def build_mlp(
 
     # Normalize input layer
     if normalize_input_layer:
-        layers[f"{prefix}normalize_input"] = normalize_input_layer(in_size)
+        try:
+            layer_instance = normalize_input_layer(in_size)  # type: ignore[call-arg]
+        except TypeError as exc:
+            raise ValueError(
+                f"normalize_input_layer={normalize_input_layer} is not a valid "
+                "normalization layer type accepting only one argument (in_size)."
+            ) from exc
+        layers[f"{prefix}normalize_input"] = layer_instance
 
     # Hidden layers
     prev_size = in_size

From 44bb66c6459182f89545fd8641d849ae84dbd51a Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 23 Sep 2022 16:08:44 +0100
Subject: [PATCH 123/187] Add support for checking presence of generic type
 ignores

---
 ci/check_typeignore.py | 67 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)
 create mode 100755 ci/check_typeignore.py

diff --git a/ci/check_typeignore.py b/ci/check_typeignore.py
new file mode 100755
index 000000000..d43d762a1
--- /dev/null
+++ b/ci/check_typeignore.py
@@ -0,0 +1,67 @@
+#!/usr/bin/env python
+
+"""This script checks that no files in our source code have a "#type: ignore" comment without explicitly indicating the
+reason for the ignore. This is to ensure that we don't accidentally ignore errors that we should be fixing."""
+
+
+import os
+import re
+import pathlib
+import sys
+from typing import List
+
+import click
+
+# Regex to match a "# type: ignore" comment not followed by a reason.
+TYPE_IGNORE_COMMENT = re.compile(r"#\s*type:\s*ignore\s*(?![^\[]*\[)")
+
+# Regex to match a "# type: ignore[<reason>]" comment.
+TYPE_IGNORE_REASON_COMMENT = re.compile(r"#\s*type:\s*ignore\[(?P<reason>.*)\]")
+
+
+def check_file(file: pathlib.Path):
+    """Checks that the given file has no "# type: ignore" comments without a reason."""
+    with open(file, "r") as f:
+        for i, line in enumerate(f):
+            if TYPE_IGNORE_COMMENT.search(line):
+                raise ValueError(f"{file}:{i+1}: Found a '# type: ignore' comment without a reason.")
+
+            if search := TYPE_IGNORE_REASON_COMMENT.search(line):
+                reason = search.group("reason")
+                if reason == "":
+                    raise ValueError(f"{file}:{i+1}: Found a '# type: ignore[]' comment without a reason.")
+
+
+def check_files(files: List[pathlib.Path]):
+    """Checks that the given files have no "# type: ignore" comments without a reason."""
+    for file in files:
+        check_file(file)
+
+
+def get_files_to_check(root_dir: pathlib.Path) -> List[pathlib.Path]:
+    """Returns a list of files that should be checked for "# type: ignore" comments."""
+    # Get the list of files that should be checked.
+    files = []
+    for root, _, filenames in os.walk(root_dir):
+        for filename in filenames:
+            if filename.endswith(".py"):
+                files.append(pathlib.Path(root) / filename)
+
+    return files
+
+
+@click.command()
+@click.option("--root-dir", type=click.Path(exists=True, file_okay=False, dir_okay=True), default="src")
+def main(root_dir: str):
+    """Checks that no files in our source code have a "#type: ignore" comment without explicitly indicating the
+    reason for the ignore. This is to ensure that we don't accidentally ignore errors that we should be fixing."""
+    files = get_files_to_check(pathlib.Path(root_dir))
+    try:
+        check_files(files)
+    except ValueError as e:
+        print(e)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file

From 849d70a7d9e8d5c94e4d30269887e5a0bbb1d9f4 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 23 Sep 2022 20:02:01 +0100
Subject: [PATCH 124/187] Allow subdirectories in notebook clean

---
 ci/clean_notebooks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ci/clean_notebooks.py b/ci/clean_notebooks.py
index fb1d326a0..f8f59176f 100755
--- a/ci/clean_notebooks.py
+++ b/ci/clean_notebooks.py
@@ -91,7 +91,7 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
         exit(1)
     for file in args.files:
         if file.is_dir():
-            files.extend(file.glob("*.ipynb"))
+            files.extend(file.glob("**/*.ipynb"))
         else:
             if file.suffix == ".ipynb":
                 files.append(file)

From 2639fd190a81b22c25e328383d11cca011613646 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 23 Sep 2022 20:53:53 +0100
Subject: [PATCH 125/187] Add full typing support for TransitionsMinimal as a
 sequence

---
 src/imitation/data/types.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index a64c10264..e466fc596 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -5,7 +5,7 @@
 import os
 import pathlib
 import warnings
-from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, TypeVar, Union, cast
+from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, TypeVar, Union, cast, Iterator, overload
 
 import numpy as np
 import torch as th
@@ -169,8 +169,11 @@ def transitions_collate_fn(
     return result
 
 
+TransitionsMinimalSelf = TypeVar("TransitionsMinimalSelf", bound="TransitionsMinimal")
+
+
 @dataclasses.dataclass(frozen=True)
-class TransitionsMinimal(th_data.Dataset):
+class TransitionsMinimal(th_data.Dataset, Sequence[Mapping[str, np.ndarray]]):
     """A Torch-compatible `Dataset` of obs-act transitions.
 
     This class and its subclasses are usually instantiated via
@@ -200,7 +203,7 @@ class TransitionsMinimal(th_data.Dataset):
     infos: np.ndarray
     """Array of info dicts. Shape: (batch_size,)."""
 
-    def __len__(self):
+    def __len__(self) -> int:
         """Returns number of transitions. Always positive."""
         return len(self.obs)
 
@@ -239,6 +242,14 @@ def __post_init__(self):
     # def __getitem__(self, key: int) -> Mapping[str, np.ndarray]:
     #     pass  # pragma: no cover
 
+    @overload
+    def __getitem__(self, key: int) -> Mapping[str, np.ndarray]:
+        pass
+
+    @overload
+    def __getitem__(self: TransitionsMinimalSelf, key: slice) -> TransitionsMinimalSelf:
+        pass
+
     def __getitem__(self, key):
         """See TransitionsMinimal docstring for indexing and slicing semantics."""
         d = dataclass_quick_asdict(self)

From cca71d0c0ca9aad2782e04611c6bced86ca3223b Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Fri, 23 Sep 2022 20:54:02 +0100
Subject: [PATCH 126/187] Fix types for density.py

---
 src/imitation/algorithms/density.py | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index 1b95aea4f..806def71e 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -7,7 +7,7 @@
 import enum
 import itertools
 from collections.abc import Mapping
-from typing import Dict, Iterable, Optional, cast
+from typing import Dict, Iterable, Optional, cast, List
 
 import numpy as np
 from gym.spaces.utils import flatten
@@ -137,7 +137,7 @@ def _set_demo_from_batch(
         obs_b: np.ndarray,
         act_b: np.ndarray,
         next_obs_b: Optional[np.ndarray],
-    ) -> None:
+    ) -> Dict[Optional[int], List[np.ndarray]]:
         if next_obs_b is None and self.density_type == DensityType.STATE_STATE_DENSITY:
             raise ValueError(
                 "STATE_STATE_DENSITY requires next_obs_b "
@@ -151,14 +151,16 @@ def _set_demo_from_batch(
             assert next_obs_b.shape[1:] == self.venv.observation_space.shape
             assert len(next_obs_b) == len(obs_b)
 
+        transitions: Dict[Optional[int], List[np.ndarray]] = {}
         next_obs_b_iterator = next_obs_b or itertools.repeat(None)
         for obs, act, next_obs in zip(obs_b, act_b, next_obs_b_iterator):
             flat_trans = self._preprocess_transition(obs, act, next_obs)
-            self.transitions.setdefault(None, []).append(flat_trans)
+            transitions.setdefault(None, []).append(flat_trans)
+        return transitions
 
     def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
         """Sets the demonstration data."""
-        self.transitions = {}
+        transitions: Dict[Optional[int], List[np.ndarray]] = {}
 
         if isinstance(demonstrations, Iterable):
             first_item, demonstrations = util.get_first_iter_element(demonstrations)
@@ -174,17 +176,17 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
                         zip(traj.obs[:-1], traj.acts, traj.obs[1:]),
                     ):
                         flat_trans = self._preprocess_transition(obs, act, next_obs)
-                        self.transitions.setdefault(i, []).append(flat_trans)
+                        transitions.setdefault(i, []).append(flat_trans)
             elif isinstance(first_item, Mapping):
                 # analogous to cast above.
                 demonstrations = cast(Iterable[base.TransitionMapping], demonstrations)
 
                 for batch in demonstrations:
-                    self._set_demo_from_batch(
+                    transitions.update(self._set_demo_from_batch(
                         util.safe_to_numpy(batch["obs"], warn=True),
                         util.safe_to_numpy(batch["acts"], warn=True),
                         util.safe_to_numpy(batch.get("next_obs"), warn=True),
-                    )
+                    ))
             else:
                 raise TypeError(
                     f"Unsupported demonstration type {type(demonstrations)}",
@@ -199,7 +201,7 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
         else:
             raise TypeError(f"Unsupported demonstration type {type(demonstrations)}")
 
-        self.transitions = {k: np.stack(v, axis=0) for k, v in self.transitions.items()}
+        self.transitions = {k: np.stack(v, axis=0) for k, v in transitions.items()}
 
         if not self.is_stationary and None in self.transitions:
             raise ValueError(
@@ -213,7 +215,7 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
     def train(self):
         """Fits the density model to demonstration data `self.transitions`."""
         # if requested, we'll scale demonstration transitions so that they have
-        # zero mean and unit variance (i.e all components are equally important)
+        # zero mean and unit variance (i.e. all components are equally important)
         self._scaler = preprocessing.StandardScaler(
             with_mean=self.standardise,
             with_std=self.standardise,
@@ -377,4 +379,5 @@ def test_policy(self, *, n_trajectories: int = 10, true_reward: bool = True):
     @property
     def policy(self) -> base_class.BasePolicy:
         assert self.rl_algo is not None
+        assert self.rl_algo.policy is not None
         return self.rl_algo.policy

From b888f130fe1b5583f9c2639848a247042d8028e1 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 26 Sep 2022 11:41:38 +0100
Subject: [PATCH 127/187] Misc fixes

---
 src/imitation/algorithms/preference_comparisons.py | 10 ++++++----
 src/imitation/util/networks.py                     |  2 +-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 307df9841..1a5c160fe 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -1211,7 +1211,6 @@ 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}")
@@ -1324,7 +1323,6 @@ def __init__(
             rng=rng,
             regularizer_factory=regularizer_factory,
         )
-        self.rng = rng
         self.member_trainers = []
         for member_pref_model in self._preference_model.member_pref_models:
             reward_trainer = BasicRewardTrainer(
@@ -1335,11 +1333,15 @@ def __init__(
                 lr=lr,
                 custom_logger=self.logger,
                 regularizer_factory=regularizer_factory,
+                rng=self.rng,
             )
             self.member_trainers.append(reward_trainer)
-        self.rng = rng
 
-    @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:
diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index 98c3c65bb..bd6592b11 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -194,7 +194,7 @@ def update_stats(self, batch: th.Tensor) -> None:
         self.running_var += learning_rate * delta_var
 
         self.count += b_size
-        self.num_batches += 1
+        self.num_batches += 1  # type: ignore[misc]
 
 
 def build_mlp(

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 128/187] 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 129/187] 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 130/187] 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 131/187] 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 e53dbcc6278f619c12de2a8b7abfb0dde12c4f1e Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 12:08:15 +0100
Subject: [PATCH 132/187] Bug fixes

---
 src/imitation/algorithms/density.py                | 14 ++++++++------
 src/imitation/algorithms/preference_comparisons.py |  7 ++++---
 src/imitation/data/types.py                        | 14 +++++++++++++-
 tests/algorithms/test_preference_comparisons.py    |  4 ++--
 4 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index 806def71e..7a08938d4 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -7,7 +7,7 @@
 import enum
 import itertools
 from collections.abc import Mapping
-from typing import Dict, Iterable, Optional, cast, List
+from typing import Dict, Iterable, List, Optional, cast
 
 import numpy as np
 from gym.spaces.utils import flatten
@@ -182,11 +182,13 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
                 demonstrations = cast(Iterable[base.TransitionMapping], demonstrations)
 
                 for batch in demonstrations:
-                    transitions.update(self._set_demo_from_batch(
-                        util.safe_to_numpy(batch["obs"], warn=True),
-                        util.safe_to_numpy(batch["acts"], warn=True),
-                        util.safe_to_numpy(batch.get("next_obs"), warn=True),
-                    ))
+                    transitions.update(
+                        self._set_demo_from_batch(
+                            util.safe_to_numpy(batch["obs"], warn=True),
+                            util.safe_to_numpy(batch["acts"], warn=True),
+                            util.safe_to_numpy(batch.get("next_obs"), warn=True),
+                        )
+                    )
             else:
                 raise TypeError(
                     f"Unsupported demonstration type {type(demonstrations)}",
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index b7a1b4db4..0553e4a33 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -1302,7 +1302,7 @@ def __init__(
         Raises:
             TypeError: if model is not a RewardEnsemble.
         """
-        if not preference_model.is_ensemble:
+        if preference_model.ensemble_model is None:
             raise TypeError(
                 "PreferenceModel of a RewardEnsemble expected by EnsembleTrainer.",
             )
@@ -1347,7 +1347,8 @@ def _train(self, dataset: PreferenceDataset, epoch_multiplier: float = 1.0) -> N
             dataset,
             replacement=True,
             num_samples=len(dataset),
-            generator=self.rng,
+            # we convert the numpy generator to the pytorch generator.
+            generator=th.Generator().manual_seed(util.make_seeds(self.rng)),
         )
         for member_idx in range(len(self.member_trainers)):
             # sampler gives new indexes on every call
@@ -1390,7 +1391,7 @@ def _make_reward_trainer(
     if reward_trainer_kwargs is None:
         reward_trainer_kwargs = {}
 
-    if preference_model.is_ensemble:
+    if preference_model.ensemble_model is not None:
         return EnsembleTrainer(
             preference_model,
             loss,
diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index e466fc596..3481714ad 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -5,7 +5,19 @@
 import os
 import pathlib
 import warnings
-from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, TypeVar, Union, cast, Iterator, overload
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 
 import numpy as np
 import torch as th
diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 1566e14d8..aaccb8d29 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -204,9 +204,9 @@ def test_preference_comparisons_raises(
     schedule,
     rng,
 ):
-    loss = preference_comparisons.CrossEntropyRewardLoss(preference_model)
+    loss = preference_comparisons.CrossEntropyRewardLoss()
     reward_trainer = preference_comparisons.BasicRewardTrainer(
-        reward_net,
+        preference_model,
         loss,
         rng=rng,
     )

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 133/187] 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 3c58cb09c61c3ea41c0052c1c5ea85b3858ab25c Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 12:51:23 +0100
Subject: [PATCH 134/187] Fix bug in serialize.py

---
 src/imitation/policies/serialize.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/policies/serialize.py b/src/imitation/policies/serialize.py
index cfd68414d..03c939831 100644
--- a/src/imitation/policies/serialize.py
+++ b/src/imitation/policies/serialize.py
@@ -68,7 +68,7 @@ def load_stable_baselines_model(
     # TODO(juan) remove the type ignore when this SB3 PR gets merged
     #  and released:
     #  https://github.com/DLR-RM/stable-baselines3/pull/1043
-    return cls.load(path, env=venv, **kwargs)  # type: ignore[return-value]
+    return cls.load(path_obj, env=venv, **kwargs)  # type: ignore[return-value]
 
 
 def _load_stable_baselines_from_file(

From 765aa3fc1ac500a578d9475f28f52ba4f35e3d92 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 14:49:32 +0100
Subject: [PATCH 135/187] Fixed codecheck by pointing notebook checks to docs

---
 .circleci/config.yml                | 2 +-
 ci/code_checks.sh                   | 2 +-
 docs/tutorials/1_train_bc.ipynb     | 2 +-
 src/imitation/algorithms/density.py | 2 +-
 src/imitation/data/types.py         | 1 -
 src/imitation/util/networks.py      | 3 +--
 6 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 90c4c7f2d..c8cb04aec 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -225,7 +225,7 @@ jobs:
 
       - run:
           name: ipynb-check
-          command: ./ci/clean_notebooks.py --check ${SRC_FILES}
+          command: ./ci/clean_notebooks.py --check ./docs/tutorials
 
       - run:
           name: flake8
diff --git a/ci/code_checks.sh b/ci/code_checks.sh
index 94577deb2..b39b61198 100755
--- a/ci/code_checks.sh
+++ b/ci/code_checks.sh
@@ -7,7 +7,7 @@ set -x  # echo commands
 set -e  # quit immediately on error
 
 echo "Source format checking"
-./ci/clean_notebooks.py --check "${SRC_FILES[@]}"
+./ci/clean_notebooks.py --check ./docs/tutorials/
 flake8 --darglint-ignore-regex '.*' "${SRC_FILES[@]}"
 black --check --diff "${SRC_FILES[@]}"
 codespell -I .codespell.skip --skip='*.pyc,tests/testdata/*,*.ipynb,*.csv' "${SRC_FILES[@]}"
diff --git a/docs/tutorials/1_train_bc.ipynb b/docs/tutorials/1_train_bc.ipynb
index 500f0a3db..5c4f16b86 100644
--- a/docs/tutorials/1_train_bc.ipynb
+++ b/docs/tutorials/1_train_bc.ipynb
@@ -200,4 +200,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index 7a08938d4..d327adf19 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -187,7 +187,7 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
                             util.safe_to_numpy(batch["obs"], warn=True),
                             util.safe_to_numpy(batch["acts"], warn=True),
                             util.safe_to_numpy(batch.get("next_obs"), warn=True),
-                        )
+                        ),
                     )
             else:
                 raise TypeError(
diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index 3481714ad..36f3ea2a1 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -8,7 +8,6 @@
 from typing import (
     Any,
     Dict,
-    Iterator,
     Mapping,
     Optional,
     Sequence,
diff --git a/src/imitation/util/networks.py b/src/imitation/util/networks.py
index bd6592b11..187fecc59 100644
--- a/src/imitation/util/networks.py
+++ b/src/imitation/util/networks.py
@@ -7,7 +7,6 @@
 
 import torch as th
 from torch import nn
-from torch.nn.modules import batchnorm, normalization
 
 
 @contextlib.contextmanager
@@ -252,7 +251,7 @@ def build_mlp(
         except TypeError as exc:
             raise ValueError(
                 f"normalize_input_layer={normalize_input_layer} is not a valid "
-                "normalization layer type accepting only one argument (in_size)."
+                "normalization layer type accepting only one argument (in_size).",
             ) from exc
         layers[f"{prefix}normalize_input"] = layer_instance
 

From 568446badc4ab494e515ded74671f7945916b1d0 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 14:54:08 +0100
Subject: [PATCH 136/187] Add rng to mce_irl.rst (doctest)

---
 docs/algorithms/mce_irl.rst | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/docs/algorithms/mce_irl.rst b/docs/algorithms/mce_irl.rst
index 0b2e8ce21..889b7b9c8 100644
--- a/docs/algorithms/mce_irl.rst
+++ b/docs/algorithms/mce_irl.rst
@@ -13,6 +13,8 @@ Detailed example notebook: :doc:`../tutorials/6_train_mce`
 
     from functools import partial
 
+    import numpy as np
+
     from stable_baselines3.common.vec_env import DummyVecEnv
 
     from imitation.algorithms.mce_irl import (
@@ -25,6 +27,8 @@ Detailed example notebook: :doc:`../tutorials/6_train_mce`
     from imitation.envs.examples.model_envs import CliffWorld
     from imitation.rewards import reward_nets
 
+    rng = np.random.RandomState(0)
+
     env_creator = partial(CliffWorld, height=4, horizon=8, width=7, use_xy_obs=True)
     env_single = env_creator()
 
@@ -51,6 +55,7 @@ Detailed example notebook: :doc:`../tutorials/6_train_mce`
         reward_net,
         log_interval=250,
         optimizer_kwargs={"lr": 0.01},
+        rng=rng,
     )
     occ_measure = mce_irl.train()
 

From 1aacc95540f417d9d339a67929fb6d0a89d931a1 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 15:01:54 +0100
Subject: [PATCH 137/187] Add rng to density.rst (doctest)

---
 docs/algorithms/density.rst | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/docs/algorithms/density.rst b/docs/algorithms/density.rst
index ed822bbbc..424cf7d4b 100644
--- a/docs/algorithms/density.rst
+++ b/docs/algorithms/density.rst
@@ -18,7 +18,9 @@ Detailed example notebook: :doc:`../tutorials/7_train_density`
     from imitation.data import types
     from imitation.util import util
 
-    env = util.make_vec_env("Pendulum-v1", 2)
+
+
+    env = util.make_vec_env("Pendulum-v1", rng=rng, n_envs=2)
     rollouts = types.load("../tests/testdata/expert_models/pendulum_0/rollouts/final.pkl")
 
     imitation_trainer = PPO(ActorCriticPolicy, env)

From 630f2e193605e0fe87c3a6fbc7f8bf74d8209ee5 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 15:35:53 +0100
Subject: [PATCH 138/187] Fix remaining rst files

---
 docs/algorithms/airl.rst                   |  7 ++++++-
 docs/algorithms/bc.rst                     |  1 +
 docs/algorithms/dagger.rst                 |  6 +++++-
 docs/algorithms/density.rst                |  4 +++-
 docs/algorithms/gail.rst                   |  7 ++++++-
 docs/algorithms/mce_irl.rst                |  3 ++-
 docs/algorithms/preference_comparisons.rst |  9 ++++-----
 src/imitation/util/logger.py               | 14 ++++++--------
 8 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/docs/algorithms/airl.rst b/docs/algorithms/airl.rst
index 13fbec8b2..fea05abed 100644
--- a/docs/algorithms/airl.rst
+++ b/docs/algorithms/airl.rst
@@ -21,6 +21,7 @@ Detailed example notebook: :doc:`../tutorials/4_train_airl`
 
 .. testcode::
 
+    import numpy as np
     import gym
     from stable_baselines3 import PPO
     from stable_baselines3.common.evaluation import evaluate_policy
@@ -34,6 +35,8 @@ Detailed example notebook: :doc:`../tutorials/4_train_airl`
     from imitation.util.networks import RunningNorm
     from imitation.util.util import make_vec_env
 
+    rng = np.random.default_rng(0)
+
     env = gym.make("seals/CartPole-v0")
     expert = PPO(policy=MlpPolicy, env=env)
     expert.learn(1000)
@@ -42,13 +45,15 @@ Detailed example notebook: :doc:`../tutorials/4_train_airl`
         expert,
         make_vec_env(
             "seals/CartPole-v0",
+            rng=rng,
             n_envs=5,
             post_wrappers=[lambda env, _: RolloutInfoWrapper(env)],
         ),
         rollout.make_sample_until(min_timesteps=None, min_episodes=60),
+        rng=rng,
     )
 
-    venv = make_vec_env("seals/CartPole-v0", n_envs=8)
+    venv = make_vec_env("seals/CartPole-v0", rng=rng, n_envs=8)
     learner = PPO(env=venv, policy=MlpPolicy)
     reward_net = BasicShapedRewardNet(
         venv.observation_space,
diff --git a/docs/algorithms/bc.rst b/docs/algorithms/bc.rst
index f2ffb5401..b7026894a 100644
--- a/docs/algorithms/bc.rst
+++ b/docs/algorithms/bc.rst
@@ -47,6 +47,7 @@ Detailed example notebook: :doc:`../tutorials/1_train_bc`
         observation_space=env.observation_space,
         action_space=env.action_space,
         demonstrations=transitions,
+        rng=rng,
     )
     bc_trainer.train(n_epochs=1)
     reward, _ = evaluate_policy(bc_trainer.policy, env, 10)
diff --git a/docs/algorithms/dagger.rst b/docs/algorithms/dagger.rst
index dbb4689a1..5be16dc96 100644
--- a/docs/algorithms/dagger.rst
+++ b/docs/algorithms/dagger.rst
@@ -48,7 +48,11 @@ Detailed example notebook: :doc:`../tutorials/2_train_dagger`
     with tempfile.TemporaryDirectory(prefix="dagger_example_") as tmpdir:
         print(tmpdir)
         dagger_trainer = SimpleDAggerTrainer(
-            venv=venv, scratch_dir=tmpdir, expert_policy=expert, bc_trainer=bc_trainer,
+            venv=venv,
+            scratch_dir=tmpdir,
+            expert_policy=expert,
+            bc_trainer=bc_trainer,
+            rng=rng,
         )
         dagger_trainer.train(2000)
 
diff --git a/docs/algorithms/density.rst b/docs/algorithms/density.rst
index 424cf7d4b..c06436644 100644
--- a/docs/algorithms/density.rst
+++ b/docs/algorithms/density.rst
@@ -10,6 +10,7 @@ Detailed example notebook: :doc:`../tutorials/7_train_density`
 .. testcode::
 
     import pprint
+    import numpy as np
 
     from stable_baselines3 import PPO
     from stable_baselines3.common.policies import ActorCriticPolicy
@@ -18,7 +19,7 @@ Detailed example notebook: :doc:`../tutorials/7_train_density`
     from imitation.data import types
     from imitation.util import util
 
-
+    rng = np.random.default_rng(0)
 
     env = util.make_vec_env("Pendulum-v1", rng=rng, n_envs=2)
     rollouts = types.load("../tests/testdata/expert_models/pendulum_0/rollouts/final.pkl")
@@ -28,6 +29,7 @@ Detailed example notebook: :doc:`../tutorials/7_train_density`
         venv=env,
         demonstrations=rollouts,
         rl_algo=imitation_trainer,
+        rng=rng,
     )
     density_trainer.train()
 
diff --git a/docs/algorithms/gail.rst b/docs/algorithms/gail.rst
index 584cb36f9..75cc55fe1 100644
--- a/docs/algorithms/gail.rst
+++ b/docs/algorithms/gail.rst
@@ -19,6 +19,7 @@ Detailed example notebook: :doc:`../tutorials/3_train_gail`
 
 .. testcode::
 
+    import numpy as np
     import gym
     from stable_baselines3 import PPO
     from stable_baselines3.common.evaluation import evaluate_policy
@@ -32,6 +33,8 @@ Detailed example notebook: :doc:`../tutorials/3_train_gail`
     from imitation.util.networks import RunningNorm
     from imitation.util.util import make_vec_env
 
+    rng = np.random.default_rng(0)
+
     env = gym.make("seals/CartPole-v0")
     expert = PPO(policy=MlpPolicy, env=env, n_steps=64)
     expert.learn(1000)
@@ -42,11 +45,13 @@ Detailed example notebook: :doc:`../tutorials/3_train_gail`
             "seals/CartPole-v0",
             n_envs=5,
             post_wrappers=[lambda env, _: RolloutInfoWrapper(env)],
+            rng=rng,
         ),
         rollout.make_sample_until(min_timesteps=None, min_episodes=60),
+        rng=rng,
     )
 
-    venv = make_vec_env("seals/CartPole-v0", n_envs=8)
+    venv = make_vec_env("seals/CartPole-v0", n_envs=8, rng=rng)
     learner = PPO(env=venv, policy=MlpPolicy)
     reward_net = BasicRewardNet(
         venv.observation_space,
diff --git a/docs/algorithms/mce_irl.rst b/docs/algorithms/mce_irl.rst
index 889b7b9c8..353640f6e 100644
--- a/docs/algorithms/mce_irl.rst
+++ b/docs/algorithms/mce_irl.rst
@@ -27,7 +27,7 @@ Detailed example notebook: :doc:`../tutorials/6_train_mce`
     from imitation.envs.examples.model_envs import CliffWorld
     from imitation.rewards import reward_nets
 
-    rng = np.random.RandomState(0)
+    rng = np.random.default_rng(0)
 
     env_creator = partial(CliffWorld, height=4, horizon=8, width=7, use_xy_obs=True)
     env_single = env_creator()
@@ -63,6 +63,7 @@ Detailed example notebook: :doc:`../tutorials/6_train_mce`
         policy=mce_irl.policy,
         venv=state_venv,
         sample_until=rollout.make_min_timesteps(5000),
+        rng=rng,
     )
     print("Imitation stats: ", rollout.rollout_stats(imitation_trajs))
 
diff --git a/docs/algorithms/preference_comparisons.rst b/docs/algorithms/preference_comparisons.rst
index 63bb85a4b..4f9a4ac1d 100644
--- a/docs/algorithms/preference_comparisons.rst
+++ b/docs/algorithms/preference_comparisons.rst
@@ -37,14 +37,14 @@ Detailed example notebook: :doc:`../tutorials/5_train_preference_comparisons`
 
     rng = np.random.default_rng(0)
 
-    venv = make_vec_env("Pendulum-v1")
+    venv = make_vec_env("Pendulum-v1", rng=rng)
 
     reward_net = BasicRewardNet(
         venv.observation_space, venv.action_space, normalize_input_layer=RunningNorm,
     )
 
-    fragmenter = preference_comparisons.RandomFragmenter(warning_threshold=0, seed=0)
-    gatherer = preference_comparisons.SyntheticGatherer(seed=0)
+    fragmenter = preference_comparisons.RandomFragmenter(warning_threshold=0, rng=rng)
+    gatherer = preference_comparisons.SyntheticGatherer(rng=rng)
     preference_model = preference_comparisons.PreferenceModel(reward_net)
     reward_trainer = preference_comparisons.BasicRewardTrainer(
         preference_model=preference_model,
@@ -68,7 +68,7 @@ Detailed example notebook: :doc:`../tutorials/5_train_preference_comparisons`
         reward_fn=reward_net,
         venv=venv,
         exploration_frac=0.0,
-        seed=0,
+        rng=rng,
     )
 
     pref_comparisons = preference_comparisons.PreferenceComparisons(
@@ -78,7 +78,6 @@ Detailed example notebook: :doc:`../tutorials/5_train_preference_comparisons`
         fragmenter=fragmenter,
         preference_gatherer=gatherer,
         reward_trainer=reward_trainer,
-        seed=0,
         initial_epoch_multiplier=1,
     )
     pref_comparisons.train(total_timesteps=5_000, total_comparisons=200)
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 244e64268..354b6e950 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -114,14 +114,12 @@ class HierarchicalLogger(sb_logger.Logger):
     -------------------
     | loss | 1        |
     -------------------
-    ---------------------------------
-    | mean/              |          |
-    |    dataset/entropy | 5.5      |
-    ---------------------------------
-    -----------------------------
-    | mean/          |          |
-    |    foo/bar/biz | 42       |
-    -----------------------------
+    -----------------------------------
+    | mean:dataset:entropy | 5.5      |
+    -----------------------------------
+    -------------------------------
+    | mean:foo:bar:biz | 42       |
+    -------------------------------
     <BLANKLINE>
     """
 

From 4f3fea54fbaf272830359effc6c59f28253c16c7 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 15:44:57 +0100
Subject: [PATCH 139/187] Increase sample size to reduce flakiness

---
 tests/algorithms/test_mce_irl.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/algorithms/test_mce_irl.py b/tests/algorithms/test_mce_irl.py
index 91b1a6294..50c022eee 100644
--- a/tests/algorithms/test_mce_irl.py
+++ b/tests/algorithms/test_mce_irl.py
@@ -286,9 +286,9 @@ def test_tabular_policy_randomness(rng):
         rng=rng,
     )
 
-    actions, _ = tabular.predict(np.zeros((100,), dtype=int))
+    actions, _ = tabular.predict(np.zeros((1000,), dtype=int))
     assert 0.45 <= np.mean(actions) <= 0.55
-    ones_obs = np.ones((100,), dtype=int)
+    ones_obs = np.ones((1000,), dtype=int)
     actions, _ = tabular.predict(ones_obs)
     assert 0.05 <= np.mean(actions) <= 0.15
     actions, _ = tabular.predict(ones_obs, deterministic=True)

From a2a459bab7e2bd0277f871ad9380d72263c155f5 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 16:19:09 +0100
Subject: [PATCH 140/187] Ignore files not passing mypy for now

---
 .circleci/config.yml              | 2 +-
 src/imitation/algorithms/base.py  | 3 ++-
 src/imitation/algorithms/bc.py    | 2 ++
 tests/rewards/test_reward_nets.py | 3 ++-
 tests/util/test_networks.py       | 2 +-
 5 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index c8cb04aec..d850cb445 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -263,7 +263,7 @@ jobs:
 
       - run:
           name: mypy
-          command: mypy --version && mypy ${SRC_FILES[@]}
+          command: mypy --version && mypy ${SRC_FILES[@]} --exclude 'imitation/algorithms/preference_comparisons.py$|imitation/rewards/reward_nets.py$|imitation/util/sacred.py$|imitation/algorithms/base.py$|imitation/scripts/train_preference_comparisons.py$|imitation/rewards/serialize.py$|imitation/scripts/common/train.py$|imitation/algorithms/mce_irl.py$|imitation/algorithms/density.py$|tests/algorithms/test_bc.py$'
 
   unit-test-linux:
     executor: unit-test-linux
diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index 85131302f..861927de2 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -1,7 +1,7 @@
 """Module of base classes and helper methods for imitation learning algorithms."""
 
 import abc
-from typing import Any, Generic, Iterable, Mapping, Optional, TypeVar, Union
+from typing import Any, Generic, Iterable, Mapping, Optional, TypeVar, Union, cast
 
 import numpy as np
 import torch as th
@@ -245,6 +245,7 @@ def make_data_loader(
     if isinstance(transitions, Iterable):
         first_item, transitions = util.get_first_iter_element(transitions)
         if isinstance(first_item, types.Trajectory):
+            transitions = cast(Iterable[types.Trajectory], transitions)
             transitions = rollout.flatten_trajectories(list(transitions))
 
     if isinstance(transitions, types.TransitionsMinimal):
diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index 0c387d1f2..cc9a20cb2 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -123,6 +123,8 @@ def __call__(
 
         l2_norms = [th.sum(th.square(w)) for w in policy.parameters()]
         l2_norm = sum(l2_norms) / 2  # divide by 2 to cancel with gradient of square
+        # sum of list defaults to float(0) if len == 0.
+        assert isinstance(l2_norm, th.Tensor)
 
         ent_loss = -self.ent_weight * entropy
         neglogp = -log_prob
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index 6bfb13958..e7c85ac4e 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -767,7 +767,8 @@ def test_predict_processed_wrappers_pass_on_kwargs(
     zero_reward_net: testing_reward_nets.MockRewardNet,
     numpy_transitions: NumpyTransitions,
 ):
-    zero_reward_net.predict_processed = mock.Mock(return_value=np.zeros((10,)))
+    zero_reward_net.predict_processed = \  # type: ignore[assignment]
+        mock.Mock(return_value=np.zeros((10,)))
     wrapped_reward_net = make_predict_processed_wrapper(
         zero_reward_net,
     )
diff --git a/tests/util/test_networks.py b/tests/util/test_networks.py
index 6e8bfc533..aad4f729b 100644
--- a/tests/util/test_networks.py
+++ b/tests/util/test_networks.py
@@ -71,7 +71,7 @@ def update_stats(self, batch: th.Tensor) -> None:
             self.running_var += learning_rate * S - delta**2
 
         self.count += b_size
-        self.num_batches += 1
+        self.num_batches += 1  # type: ignore[misc]
 
 
 @pytest.mark.parametrize("normalization_layer", NORMALIZATION_LAYERS)

From e5bad0b2941174adee05d6cc55d47f89b603ad28 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 16:24:42 +0100
Subject: [PATCH 141/187] Comment in wrong line

---
 tests/rewards/test_reward_nets.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index e7c85ac4e..df5fe935a 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -767,8 +767,8 @@ def test_predict_processed_wrappers_pass_on_kwargs(
     zero_reward_net: testing_reward_nets.MockRewardNet,
     numpy_transitions: NumpyTransitions,
 ):
-    zero_reward_net.predict_processed = \  # type: ignore[assignment]
-        mock.Mock(return_value=np.zeros((10,)))
+    zero_reward_net.predict_processed = \
+        mock.Mock(return_value=np.zeros((10,)))  # type: ignore[assignment]
     wrapped_reward_net = make_predict_processed_wrapper(
         zero_reward_net,
     )

From 1352bf7df86f4081703c5078e7c0e1a8f513059d Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 16:24:42 +0100
Subject: [PATCH 142/187] Comment in wrong line

---
 tests/rewards/test_reward_nets.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index e7c85ac4e..cca2a39b7 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -767,8 +767,9 @@ def test_predict_processed_wrappers_pass_on_kwargs(
     zero_reward_net: testing_reward_nets.MockRewardNet,
     numpy_transitions: NumpyTransitions,
 ):
-    zero_reward_net.predict_processed = \  # type: ignore[assignment]
+    zero_reward_net.predict_processed = (  # type: ignore[assignment]
         mock.Mock(return_value=np.zeros((10,)))
+    )
     wrapped_reward_net = make_predict_processed_wrapper(
         zero_reward_net,
     )

From 94e97053b655a7b1648f5eef871859685db606be Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 16:51:15 +0100
Subject: [PATCH 143/187] Move excluded files to argument

---
 .circleci/config.yml | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index d850cb445..92ebf0bb8 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -37,6 +37,16 @@ executors:
       # If you change these, also change ci/code_checks.sh
       SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/clean_notebooks.py
       NUM_CPUS: 2
+      EXCLUDE_MYPY: "(?x)(
+      src/imitation/algorithms/preference_comparisons.py$
+      | src/imitation/rewards/reward_nets.py$|imitation/util/sacred.py$
+      | src/imitation/algorithms/base.py$|imitation/scripts/train_preference_comparisons.py$
+      | src/imitation/rewards/serialize.py$|imitation/scripts/common/train.py$
+      | src/imitation/algorithms/mce_irl.py$
+      | src/imitation/algorithms/density.py$
+      | tests/algorithms/test_bc.py$
+      | examples/quickstart.py$
+      )"
 
 commands:
   dependencies-linux:
@@ -263,7 +273,7 @@ jobs:
 
       - run:
           name: mypy
-          command: mypy --version && mypy ${SRC_FILES[@]} --exclude 'imitation/algorithms/preference_comparisons.py$|imitation/rewards/reward_nets.py$|imitation/util/sacred.py$|imitation/algorithms/base.py$|imitation/scripts/train_preference_comparisons.py$|imitation/rewards/serialize.py$|imitation/scripts/common/train.py$|imitation/algorithms/mce_irl.py$|imitation/algorithms/density.py$|tests/algorithms/test_bc.py$'
+          command: mypy --version && mypy ${SRC_FILES[@]} --exclude ${EXCLUDE_MYPY}
 
   unit-test-linux:
     executor: unit-test-linux

From 617444eed502c3fea643be1b4836fbf5126d34fb Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 18:07:10 +0100
Subject: [PATCH 144/187] Add quotes to mypy arg call

---
 .circleci/config.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 92ebf0bb8..88571aece 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -273,7 +273,7 @@ jobs:
 
       - run:
           name: mypy
-          command: mypy --version && mypy ${SRC_FILES[@]} --exclude ${EXCLUDE_MYPY}
+          command: mypy --version && mypy ${SRC_FILES[@]} --exclude "${EXCLUDE_MYPY}"
 
   unit-test-linux:
     executor: unit-test-linux

From d1eebbb1259bb613625dbc147fab4a1dcb9d46eb Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 18:19:42 +0100
Subject: [PATCH 145/187] Fix CI mypy call

---
 .circleci/config.yml              | 1 -
 examples/quickstart.py            | 8 ++++++--
 tests/rewards/test_reward_nets.py | 4 ++--
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 88571aece..2a19a479a 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -45,7 +45,6 @@ executors:
       | src/imitation/algorithms/mce_irl.py$
       | src/imitation/algorithms/density.py$
       | tests/algorithms/test_bc.py$
-      | examples/quickstart.py$
       )"
 
 commands:
diff --git a/examples/quickstart.py b/examples/quickstart.py
index ae054186c..5867c753d 100644
--- a/examples/quickstart.py
+++ b/examples/quickstart.py
@@ -55,11 +55,15 @@ def sample_expert_transitions():
     rng=rng,
 )
 
-reward, _ = evaluate_policy(bc_trainer.policy, env, n_eval_episodes=3, render=True)
+reward, _ = evaluate_policy(
+    bc_trainer.policy, env, n_eval_episodes=3, render=True  # type: ignore
+)
 print(f"Reward before training: {reward}")
 
 print("Training a policy using Behavior Cloning")
 bc_trainer.train(n_epochs=1)
 
-reward, _ = evaluate_policy(bc_trainer.policy, env, n_eval_episodes=3, render=True)
+reward, _ = evaluate_policy(
+    bc_trainer.policy, env, n_eval_episodes=3, render=True  # type: ignore
+)
 print(f"Reward after training: {reward}")
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index cca2a39b7..88bf9a8c7 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -767,8 +767,8 @@ def test_predict_processed_wrappers_pass_on_kwargs(
     zero_reward_net: testing_reward_nets.MockRewardNet,
     numpy_transitions: NumpyTransitions,
 ):
-    zero_reward_net.predict_processed = (  # type: ignore[assignment]
-        mock.Mock(return_value=np.zeros((10,)))
+    zero_reward_net.predict_processed = mock.Mock(  # type: ignore[assignment]
+        return_value=np.zeros((10,))
     )
     wrapped_reward_net = make_predict_processed_wrapper(
         zero_reward_net,

From 5dd02aa07bf9f6e6d369675908a9be69d69ca95e Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 18:31:57 +0100
Subject: [PATCH 146/187] Fix CI yaml

---
 .circleci/config.yml              | 19 ++++++++++---------
 examples/quickstart.py            | 10 ++++++++--
 tests/rewards/test_reward_nets.py |  2 +-
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 2a19a479a..d4b11f2e7 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -37,15 +37,16 @@ executors:
       # If you change these, also change ci/code_checks.sh
       SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/clean_notebooks.py
       NUM_CPUS: 2
-      EXCLUDE_MYPY: "(?x)(
-      src/imitation/algorithms/preference_comparisons.py$
-      | src/imitation/rewards/reward_nets.py$|imitation/util/sacred.py$
-      | src/imitation/algorithms/base.py$|imitation/scripts/train_preference_comparisons.py$
-      | src/imitation/rewards/serialize.py$|imitation/scripts/common/train.py$
-      | src/imitation/algorithms/mce_irl.py$
-      | src/imitation/algorithms/density.py$
-      | tests/algorithms/test_bc.py$
-      )"
+      EXCLUDE_MYPY: | 
+        (?x)(
+          src/imitation/algorithms/preference_comparisons.py$
+          | src/imitation/rewards/reward_nets.py$|imitation/util/sacred.py$
+          | src/imitation/algorithms/base.py$|imitation/scripts/train_preference_comparisons.py$
+          | src/imitation/rewards/serialize.py$|imitation/scripts/common/train.py$
+          | src/imitation/algorithms/mce_irl.py$
+          | src/imitation/algorithms/density.py$
+          | tests/algorithms/test_bc.py$
+        )
 
 commands:
   dependencies-linux:
diff --git a/examples/quickstart.py b/examples/quickstart.py
index 5867c753d..30d0e22f7 100644
--- a/examples/quickstart.py
+++ b/examples/quickstart.py
@@ -56,7 +56,10 @@ def sample_expert_transitions():
 )
 
 reward, _ = evaluate_policy(
-    bc_trainer.policy, env, n_eval_episodes=3, render=True  # type: ignore
+    bc_trainer.policy,  # type: ignore
+    env,
+    n_eval_episodes=3,
+    render=True,
 )
 print(f"Reward before training: {reward}")
 
@@ -64,6 +67,9 @@ def sample_expert_transitions():
 bc_trainer.train(n_epochs=1)
 
 reward, _ = evaluate_policy(
-    bc_trainer.policy, env, n_eval_episodes=3, render=True  # type: ignore
+    bc_trainer.policy,  # type: ignore
+    env,
+    n_eval_episodes=3,
+    render=True,
 )
 print(f"Reward after training: {reward}")
diff --git a/tests/rewards/test_reward_nets.py b/tests/rewards/test_reward_nets.py
index 88bf9a8c7..c65332d0c 100644
--- a/tests/rewards/test_reward_nets.py
+++ b/tests/rewards/test_reward_nets.py
@@ -768,7 +768,7 @@ def test_predict_processed_wrappers_pass_on_kwargs(
     numpy_transitions: NumpyTransitions,
 ):
     zero_reward_net.predict_processed = mock.Mock(  # type: ignore[assignment]
-        return_value=np.zeros((10,))
+        return_value=np.zeros((10,)),
     )
     wrapped_reward_net = make_predict_processed_wrapper(
         zero_reward_net,

From db7dc2237dee021f95845127e44f73f49ff15891 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 27 Sep 2022 18:48:40 +0100
Subject: [PATCH 147/187] Break ignored files up into one line each

---
 .circleci/config.yml | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index d4b11f2e7..6d0942478 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -40,9 +40,12 @@ executors:
       EXCLUDE_MYPY: | 
         (?x)(
           src/imitation/algorithms/preference_comparisons.py$
-          | src/imitation/rewards/reward_nets.py$|imitation/util/sacred.py$
-          | src/imitation/algorithms/base.py$|imitation/scripts/train_preference_comparisons.py$
-          | src/imitation/rewards/serialize.py$|imitation/scripts/common/train.py$
+          | src/imitation/rewards/reward_nets.py$
+          | src/imitation/util/sacred.py$
+          | src/imitation/algorithms/base.py$
+          | src/imitation/scripts/train_preference_comparisons.py$
+          | src/imitation/rewards/serialize.py$
+          | src/imitation/scripts/common/train.py$
           | src/imitation/algorithms/mce_irl.py$
           | src/imitation/algorithms/density.py$
           | tests/algorithms/test_bc.py$

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 148/187] 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 149/187] 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 2cc744a203d4a12da1c271da26f719356d2c6a7b Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 28 Sep 2022 16:00:00 +0100
Subject: [PATCH 150/187] Small bug fixes

---
 src/imitation/algorithms/bc.py | 2 +-
 src/imitation/util/logger.py   | 2 +-
 tests/data/test_types.py       | 9 +++++----
 3 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index cc9a20cb2..d080e6e3c 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -472,4 +472,4 @@ def save_policy(self, policy_path: types.AnyPath) -> None:
         Args:
             policy_path: path to save policy to.
         """
-        th.save(self.policy, types.path_to_str(policy_path))
+        th.save(self.policy, types.parse_path(policy_path))
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index a2d1128dd..38172c200 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -63,7 +63,7 @@ def _build_output_formats(
         if f == "wandb":
             output_formats.append(WandbOutputFormat())
         else:
-            output_formats.append(sb_logger.make_output_format(f, str(folder)))
+            output_formats.append(make_output_format(f, str(folder)))
     return output_formats
 
 
diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index a02533355..f1c10ff55 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -216,12 +216,13 @@ def test_save_trajectories(
         else:
             chdir_context = contextlib.nullcontext()
             save_dir_str = tmpdir
-        save_dir = types.parse_path(save_dir_str)
-
-        trajs = [trajectory_rew if use_rewards else trajectory]
-        save_path = save_dir / "trajs"
 
         with chdir_context:
+
+            save_dir = types.parse_path(save_dir_str)
+            trajs = [trajectory_rew if use_rewards else trajectory]
+            save_path = save_dir / "trajs"
+
             if use_pickle:
                 # Pickle format
                 with open(save_path, "wb") as f:

From 69f05008d96bd48873a91a708652a0e6974fe412 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 28 Sep 2022 16:06:56 +0100
Subject: [PATCH 151/187] Small bug fixes

---
 src/imitation/util/sacred.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/util/sacred.py b/src/imitation/util/sacred.py
index 8a85d146e..609ada89e 100644
--- a/src/imitation/util/sacred.py
+++ b/src/imitation/util/sacred.py
@@ -87,7 +87,7 @@ def build_sacred_symlink(log_dir: types.AnyPath, run: sacred.run.Run) -> None:
         warnings.warn(RuntimeWarning("Couldn't find sacred directory."))
         return
     symlink_path = log_dir / "sacred"
-    target_path = sacred_dir.relative_to(log_dir)
+    target_path = pathlib.Path(os.path.relpath(sacred_dir, start=log_dir))
 
     # Path.symlink_to errors if the symlink already exists. In our case, we actually
     # want to override the symlink to point to the most recent Sacred dir. The

From d3bc455d2fb528aabc1c29e241b4566ba982e908 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 28 Sep 2022 16:07:21 +0100
Subject: [PATCH 152/187] Sort import

---
 src/imitation/util/logger.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 38172c200..2955a3242 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -3,8 +3,8 @@
 import contextlib
 import datetime
 import os
-import sys
 import pathlib
+import sys
 import tempfile
 from typing import Any, Dict, Generator, List, Optional, Sequence, Tuple, Union
 

From f494fc2b2bc9aae0768a0f8ef246cbf4d09a1b8e Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 29 Sep 2022 10:50:54 +0100
Subject: [PATCH 153/187] Linting

---
 src/imitation/data/types.py                     | 5 -----
 src/imitation/scripts/config/train_imitation.py | 1 -
 src/imitation/scripts/convert_trajs.py          | 3 ++-
 tests/util/test_logger.py                       | 6 ++----
 4 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index d7084e9f2..e1958dfdc 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -121,11 +121,6 @@ def parse_optional_path(
 
     Returns:
         A `pathlib.Path` object, or None if `path` is None.
-
-    Raises:
-        ValueError: If `allow_relative` is False and the path is not absolute.
-        ValueError: If `base_directory` is specified and `allow_relative` is
-            False.
     """
     if path is None:
         return None
diff --git a/src/imitation/scripts/config/train_imitation.py b/src/imitation/scripts/config/train_imitation.py
index 3d99098cc..c2466a936 100644
--- a/src/imitation/scripts/config/train_imitation.py
+++ b/src/imitation/scripts/config/train_imitation.py
@@ -3,7 +3,6 @@
 import sacred
 import torch as th
 
-from imitation.data import types
 from imitation.scripts.common import common
 from imitation.scripts.common import demonstrations as demos_common
 from imitation.scripts.common import expert, train
diff --git a/src/imitation/scripts/convert_trajs.py b/src/imitation/scripts/convert_trajs.py
index 2e3c7c016..85db4d3f9 100644
--- a/src/imitation/scripts/convert_trajs.py
+++ b/src/imitation/scripts/convert_trajs.py
@@ -20,7 +20,8 @@ def update_traj_file_in_place(path_str: str, /) -> None:
     The new data is saved as `Sequence[imitation.types.TrajectoryWithRew]`.
 
     Args:
-        path: Path to a pickle file containing `Sequence[imitation.types.Trajectory]`
+        path_str: Path to a pickle file containing
+            `Sequence[imitation.types.Trajectory]`
             or `Sequence[imitation.old_types.TrajectoryWithRew]`.
     """
     path = types.parse_path(path_str)
diff --git a/tests/util/test_logger.py b/tests/util/test_logger.py
index a4ea315af..ec131b574 100644
--- a/tests/util/test_logger.py
+++ b/tests/util/test_logger.py
@@ -309,10 +309,8 @@ 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), h.accumulate_means("foo"):
+        h.add_accumulate_prefix("bar")
 
 
 def test_cant_add_key_prefix_outside_accumulate_means(tmpdir):

From 9025b5da8aecf08f7d58d5ab354f4e0b61108d80 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 29 Sep 2022 12:15:30 +0100
Subject: [PATCH 154/187] Do not follow imports for ignored files

---
 .circleci/config.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 6d0942478..f9cccaf19 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -276,7 +276,7 @@ jobs:
 
       - run:
           name: mypy
-          command: mypy --version && mypy ${SRC_FILES[@]} --exclude "${EXCLUDE_MYPY}"
+          command: mypy --version && mypy ${SRC_FILES[@]} --exclude "${EXCLUDE_MYPY}" --follow-imports=silent --show-error-codes
 
   unit-test-linux:
     executor: unit-test-linux

From 453e0158ccee0dd5cf43d45ce556c2ac9c1854d0 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 29 Sep 2022 12:18:35 +0100
Subject: [PATCH 155/187] Fix tests for context managers

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

diff --git a/tests/util/test_logger.py b/tests/util/test_logger.py
index ec131b574..7f076df9c 100644
--- a/tests/util/test_logger.py
+++ b/tests/util/test_logger.py
@@ -309,11 +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")
+    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 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 156/187] 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 157/187] 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 158/187] 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 7293915160814556e32ca78d899304af47092418 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 29 Sep 2022 14:46:01 +0100
Subject: [PATCH 159/187] Upgrade Python version in Windows CI

---
 .circleci/config.yml | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 8af7eb5cc..9f90432c6 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -127,10 +127,16 @@ commands:
       - run:
           name: install python and binary dependencies
           command: |
-            choco install --side-by-side -y python --version 3.8
+            choco install --side-by-side -y python --version=3.8.10
             choco install -y ffmpeg
           shell: powershell.exe
 
+      - run:
+          name: upgrade pip
+          command: |
+              python -m pip install --upgrade pip
+          shell: powershell.exe
+
       - run:
           name: install virtualenv
           command: pip install virtualenv

From 9f50e3bc2a75689e975dfb6127196c47090a9392 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 16:12:40 +0100
Subject: [PATCH 160/187] Remove unused import

---
 src/imitation/algorithms/preference_comparisons.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index ea7d21d3e..91d0949dd 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -7,7 +7,6 @@
 import math
 import pickle
 import re
-import random
 from collections import defaultdict
 from typing import (
     Any,

From 2a4d54170e904377fac5421b7f24597c99a0697b Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 16:46:51 +0100
Subject: [PATCH 161/187] Remove unused fixture

---
 tests/conftest.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 765da98f3..63e88b6b8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -145,11 +145,6 @@ def custom_logger(tmpdir: str) -> logger.HierarchicalLogger:
     return logger.configure(tmpdir)
 
 
-@pytest.fixture()
-def rng_fixed() -> np.random.Generator:
-    return np.random.default_rng(0)
-
-
 @pytest.fixture()
 def rng() -> np.random.Generator:
     return np.random.default_rng()

From 4526e335788ce5c87ab6ee5b48d8a48b5d979c33 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 16:46:58 +0100
Subject: [PATCH 162/187] Add coveragerc file

---
 .coveragerc | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 .coveragerc

diff --git a/.coveragerc b/.coveragerc
new file mode 100644
index 000000000..a9abf0ed9
--- /dev/null
+++ b/.coveragerc
@@ -0,0 +1,4 @@
+[report]
+exclude_lines =
+    pragma: not covered
+    @overload
\ No newline at end of file

From d0920ceb38c28e4e1e2d792d3547ceb31ebfeb5a Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 16:56:50 +0100
Subject: [PATCH 163/187] Add utils test

---
 tests/util/test_util.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tests/util/test_util.py b/tests/util/test_util.py
index 36e03215d..a5790a444 100644
--- a/tests/util/test_util.py
+++ b/tests/util/test_util.py
@@ -93,6 +93,13 @@ def test_safe_to_tensor():
         assert not np.may_share_memory(numpy, torch)
 
 
+def test_safe_to_numpy():
+    tensor = th.tensor([1, 2, 3])
+    numpy = util.safe_to_numpy(tensor)
+    assert (numpy == tensor.numpy()).all()
+    assert util.safe_to_numpy(None) is None
+
+
 def test_tensor_iter_norm():
     # vector is [1,0,1,1,-5,-6]; its 2-norm is 8, and 1-norm is 14
     tensor_list = [

From 8b63ed9023363b62f5209deecf510dac9bcf88c4 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 17:13:30 +0100
Subject: [PATCH 164/187] Add tests and asserts

---
 src/imitation/scripts/common/reward.py |  6 +-----
 tests/util/test_networks.py            | 11 +++++++++++
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/src/imitation/scripts/common/reward.py b/src/imitation/scripts/common/reward.py
index a00498138..c40d3751f 100644
--- a/src/imitation/scripts/common/reward.py
+++ b/src/imitation/scripts/common/reward.py
@@ -156,11 +156,7 @@ def make_reward_net(
         )
 
         if add_std_alpha is not None:
-            if not isinstance(reward_net, reward_nets.RewardNetWithVariance):
-                raise ValueError(
-                    "add_std_alpha is only supported for "
-                    "reward nets with variance tracking.",
-                )
+            assert isinstance(reward_net, reward_nets.RewardNetWithVariance)
             reward_net = reward_nets.AddSTDRewardWrapper(
                 reward_net,
                 default_alpha=add_std_alpha,
diff --git a/tests/util/test_networks.py b/tests/util/test_networks.py
index aad4f729b..45c9f0590 100644
--- a/tests/util/test_networks.py
+++ b/tests/util/test_networks.py
@@ -243,6 +243,17 @@ def test_build_mlp_norm_training(init_kwargs) -> None:
         optimizer.step()
 
 
+def test_build_mlp_raises_on_invalid_normalize_input_layer() -> None:
+    """Test that `networks.build_mlp()` raises on invalid input layer."""
+    with pytest.raises(ValueError, match="normalize_input_layer.*not a valid normalization layer.*"):
+        networks.build_mlp(
+            in_size=1,
+            hid_sizes=[16, 16],
+            out_size=1,
+            normalize_input_layer=th.nn.Module,
+        )
+
+
 def test_input_validation_on_ema_norm():
     with pytest.raises(ValueError):
         networks.EMANorm(128, decay=1.1)

From cc2d75253ae87a8d6a3cd9e341dd9b4bf37a9b13 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 17:20:27 +0100
Subject: [PATCH 165/187] Add test to synthetic gatherer

---
 tests/algorithms/test_preference_comparisons.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index aaccb8d29..86620ade1 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -431,6 +431,17 @@ def test_synthetic_gatherer_deterministic(
     assert np.all(preferences1 == preferences2)
 
 
+def test_synthetic_gatherer_raises(
+        agent_trainer,
+        random_fragmenter,
+):
+    with pytest.raises(ValueError, match="If `sample` is True, then `rng` must be provided"):
+        preference_comparisons.SyntheticGatherer(
+            temperature=0,
+            sample=True,
+        )
+
+
 def test_fragments_terminal(random_fragmenter):
     trajectories = [
         types.TrajectoryWithRew(

From 77f672d782d44b525c50e9412eb8e20342125621 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 17:26:03 +0100
Subject: [PATCH 166/187] Add trajectory unwrap tests

---
 .../algorithms/test_preference_comparisons.py |  8 +++---
 tests/data/test_rollout.py                    | 27 +++++++++++++++++++
 tests/util/test_networks.py                   |  4 ++-
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 86620ade1..47b3013ec 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -432,10 +432,12 @@ def test_synthetic_gatherer_deterministic(
 
 
 def test_synthetic_gatherer_raises(
-        agent_trainer,
-        random_fragmenter,
+    agent_trainer,
+    random_fragmenter,
 ):
-    with pytest.raises(ValueError, match="If `sample` is True, then `rng` must be provided"):
+    with pytest.raises(
+        ValueError, match="If `sample` is True, then `rng` must be provided"
+    ):
         preference_comparisons.SyntheticGatherer(
             temperature=0,
             sample=True,
diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index b92b556dd..da3148acf 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -248,6 +248,33 @@ def test_unwrap_traj(rng):
         np.testing.assert_equal(t1.rews, t2.rews)
 
 
+def test_unwrap_traj_raises_no_infos():
+    """Check that unwrap_traj raises ValueError if no infos in trajectory."""
+    with pytest.raises(ValueError, match="Trajectory must have infos to unwrap"):
+        rollout.unwrap_traj(
+            types.TrajectoryWithRew(
+                acts=np.array(
+                    [
+                        0,
+                    ]
+                ),
+                obs=np.array(
+                    [
+                        0,
+                        0,
+                    ]
+                ),
+                terminal=False,
+                rews=np.array(
+                    [
+                        0.0,
+                    ]
+                ),
+                infos=None,
+            )
+        )
+
+
 def test_make_sample_until_errors():
     with pytest.raises(ValueError, match="At least one.*"):
         rollout.make_sample_until(min_timesteps=None, min_episodes=None)
diff --git a/tests/util/test_networks.py b/tests/util/test_networks.py
index 45c9f0590..ecfb9a56d 100644
--- a/tests/util/test_networks.py
+++ b/tests/util/test_networks.py
@@ -245,7 +245,9 @@ def test_build_mlp_norm_training(init_kwargs) -> None:
 
 def test_build_mlp_raises_on_invalid_normalize_input_layer() -> None:
     """Test that `networks.build_mlp()` raises on invalid input layer."""
-    with pytest.raises(ValueError, match="normalize_input_layer.*not a valid normalization layer.*"):
+    with pytest.raises(
+        ValueError, match="normalize_input_layer.*not a valid normalization layer.*"
+    ):
         networks.build_mlp(
             in_size=1,
             hid_sizes=[16, 16],

From 0edb9b601ee3cda0864e060398bfd944672c8dad Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 17:40:43 +0100
Subject: [PATCH 167/187] Formatting

---
 .../algorithms/test_preference_comparisons.py |  3 ++-
 tests/data/test_rollout.py                    | 24 ++++++-------------
 tests/util/test_networks.py                   |  3 ++-
 3 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/tests/algorithms/test_preference_comparisons.py b/tests/algorithms/test_preference_comparisons.py
index 47b3013ec..331ae6d1e 100644
--- a/tests/algorithms/test_preference_comparisons.py
+++ b/tests/algorithms/test_preference_comparisons.py
@@ -436,7 +436,8 @@ def test_synthetic_gatherer_raises(
     random_fragmenter,
 ):
     with pytest.raises(
-        ValueError, match="If `sample` is True, then `rng` must be provided"
+        ValueError,
+        match="If `sample` is True, then `rng` must be provided",
     ):
         preference_comparisons.SyntheticGatherer(
             temperature=0,
diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index da3148acf..ac1e3eba1 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -251,27 +251,17 @@ def test_unwrap_traj(rng):
 def test_unwrap_traj_raises_no_infos():
     """Check that unwrap_traj raises ValueError if no infos in trajectory."""
     with pytest.raises(ValueError, match="Trajectory must have infos to unwrap"):
+        acts = np.array([0])
+        obs = np.array([[0, 0]])
+        rews = np.array([0.0])
         rollout.unwrap_traj(
             types.TrajectoryWithRew(
-                acts=np.array(
-                    [
-                        0,
-                    ]
-                ),
-                obs=np.array(
-                    [
-                        0,
-                        0,
-                    ]
-                ),
+                acts=acts,
+                obs=obs,
                 terminal=False,
-                rews=np.array(
-                    [
-                        0.0,
-                    ]
-                ),
+                rews=rews,
                 infos=None,
-            )
+            ),
         )
 
 
diff --git a/tests/util/test_networks.py b/tests/util/test_networks.py
index ecfb9a56d..147e1185b 100644
--- a/tests/util/test_networks.py
+++ b/tests/util/test_networks.py
@@ -246,7 +246,8 @@ def test_build_mlp_norm_training(init_kwargs) -> None:
 def test_build_mlp_raises_on_invalid_normalize_input_layer() -> None:
     """Test that `networks.build_mlp()` raises on invalid input layer."""
     with pytest.raises(
-        ValueError, match="normalize_input_layer.*not a valid normalization layer.*"
+        ValueError,
+        match="normalize_input_layer.*not a valid normalization layer.*",
     ):
         networks.build_mlp(
             in_size=1,

From 92972f2a780c663e185ce66819255585150c0293 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 17:51:44 +0100
Subject: [PATCH 168/187] Remove bracket typo

---
 tests/data/test_rollout.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index ac1e3eba1..9c7acf226 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -252,7 +252,7 @@ def test_unwrap_traj_raises_no_infos():
     """Check that unwrap_traj raises ValueError if no infos in trajectory."""
     with pytest.raises(ValueError, match="Trajectory must have infos to unwrap"):
         acts = np.array([0])
-        obs = np.array([[0, 0]])
+        obs = np.array([0, 0])
         rews = np.array([0.0])
         rollout.unwrap_traj(
             types.TrajectoryWithRew(

From 09eb3092c0c40dd123c4c6be9f01854113247cdd Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 3 Oct 2022 21:44:53 +0100
Subject: [PATCH 169/187] Fix .coveragerc instruction

---
 .coveragerc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.coveragerc b/.coveragerc
index a9abf0ed9..d12534191 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -1,4 +1,4 @@
 [report]
 exclude_lines =
-    pragma: not covered
+    pragma: no cover
     @overload
\ No newline at end of file

From 8780d082e7aa06252dbe34b470ac69c56890e84d Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 5 Oct 2022 16:14:40 +0100
Subject: [PATCH 170/187] Improve density algo coverage and bug fixes

---
 src/imitation/algorithms/density.py        | 24 +++++-----
 tests/algorithms/test_density_baselines.py | 54 ++++++++++++++++++++++
 2 files changed, 67 insertions(+), 11 deletions(-)

diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index d327adf19..bf346303f 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -132,7 +132,7 @@ def __init__(
         )
         self.wrapper_callback = self.venv_wrapped.make_log_callback()
 
-    def _set_demo_from_batch(
+    def _get_demo_from_batch(
         self,
         obs_b: np.ndarray,
         act_b: np.ndarray,
@@ -152,7 +152,7 @@ def _set_demo_from_batch(
             assert len(next_obs_b) == len(obs_b)
 
         transitions: Dict[Optional[int], List[np.ndarray]] = {}
-        next_obs_b_iterator = next_obs_b or itertools.repeat(None)
+        next_obs_b_iterator = next_obs_b if next_obs_b is not None else itertools.repeat(None)
         for obs, act, next_obs in zip(obs_b, act_b, next_obs_b_iterator):
             flat_trans = self._preprocess_transition(obs, act, next_obs)
             transitions.setdefault(None, []).append(flat_trans)
@@ -162,7 +162,16 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
         """Sets the demonstration data."""
         transitions: Dict[Optional[int], List[np.ndarray]] = {}
 
-        if isinstance(demonstrations, Iterable):
+        if isinstance(demonstrations, types.TransitionsMinimal):
+            next_obs_b = getattr(demonstrations, "next_obs", None)
+            transitions.update(
+                self._get_demo_from_batch(
+                    demonstrations.obs,
+                    demonstrations.acts,
+                    next_obs_b,
+                ),
+            )
+        elif isinstance(demonstrations, Iterable):
             first_item, demonstrations = util.get_first_iter_element(demonstrations)
             if isinstance(first_item, types.Trajectory):
                 # we assume that all elements are also types.Trajectory.
@@ -183,7 +192,7 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
 
                 for batch in demonstrations:
                     transitions.update(
-                        self._set_demo_from_batch(
+                        self._get_demo_from_batch(
                             util.safe_to_numpy(batch["obs"], warn=True),
                             util.safe_to_numpy(batch["acts"], warn=True),
                             util.safe_to_numpy(batch.get("next_obs"), warn=True),
@@ -193,13 +202,6 @@ def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
                 raise TypeError(
                     f"Unsupported demonstration type {type(demonstrations)}",
                 )
-        elif isinstance(demonstrations, types.TransitionsMinimal):
-            next_obs_b = getattr(demonstrations, "next_obs", None)
-            self._set_demo_from_batch(
-                demonstrations.obs,
-                demonstrations.acts,
-                next_obs_b,
-            )
         else:
             raise TypeError(f"Unsupported demonstration type {type(demonstrations)}")
 
diff --git a/tests/algorithms/test_density_baselines.py b/tests/algorithms/test_density_baselines.py
index c66ad3b36..201f27e1c 100644
--- a/tests/algorithms/test_density_baselines.py
+++ b/tests/algorithms/test_density_baselines.py
@@ -1,6 +1,7 @@
 """Tests for `imitation.algorithms.density_baselines`."""
 
 from typing import Sequence
+from dataclasses import asdict
 
 import numpy as np
 import pytest
@@ -103,3 +104,56 @@ def test_density_trainer_smoke(
     density_trainer.train()
     density_trainer.train_policy(n_timesteps=2)
     density_trainer.test_policy(n_trajectories=2)
+
+
+def test_density_with_other_trajectory_types(
+        pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
+        pendulum_venv,
+        rng,
+):
+    rl_algo = stable_baselines3.PPO(policies.ActorCriticPolicy, pendulum_venv)
+    rollouts = pendulum_expert_trajectories[:2]
+    transitions = rollout.flatten_trajectories_with_rew(rollouts)
+    transitions_mappings = [asdict(transitions),]
+
+    minimal_transitions = types.TransitionsMinimal(
+        obs=transitions.obs,
+        acts=transitions.acts,
+        infos=transitions.infos,
+    )
+    DensityAlgorithm(
+        demonstrations=transitions_mappings,
+        venv=pendulum_venv,
+        rl_algo=rl_algo,
+        rng=rng,
+    )
+
+    DensityAlgorithm(
+        demonstrations=minimal_transitions,
+        venv=pendulum_venv,
+        rl_algo=rl_algo,
+        rng=rng,
+    )
+
+
+def test_density_trainer_raises(
+        pendulum_venv,
+        rng,
+):
+    rl_algo = stable_baselines3.PPO(policies.ActorCriticPolicy, pendulum_venv)
+    density_trainer = DensityAlgorithm(
+        venv=pendulum_venv,
+        rl_algo=rl_algo,
+        rng=rng,
+        demonstrations=None,
+        density_type=DensityType.STATE_STATE_DENSITY,
+    )
+    with pytest.raises(ValueError, match="STATE_STATE_DENSITY requires next_obs_b"):
+        density_trainer._get_demo_from_batch(
+            np.zeros((1, 3)),
+            np.zeros((1, 1)),
+            None,
+        )
+
+    with pytest.raises(TypeError, match="Unsupported demonstration type"):
+        density_trainer.set_demonstrations("foo")
\ No newline at end of file

From 9a66f0faafa96fa14d957b8a41631e9d3c0713c1 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 5 Oct 2022 16:24:21 +0100
Subject: [PATCH 171/187] Fix bug in test

---
 src/imitation/algorithms/density.py        |  4 +++-
 tests/algorithms/test_density_baselines.py | 18 ++++++++++--------
 tests/util/test_util.py                    |  2 +-
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
index bf346303f..c585b396a 100644
--- a/src/imitation/algorithms/density.py
+++ b/src/imitation/algorithms/density.py
@@ -152,7 +152,9 @@ def _get_demo_from_batch(
             assert len(next_obs_b) == len(obs_b)
 
         transitions: Dict[Optional[int], List[np.ndarray]] = {}
-        next_obs_b_iterator = next_obs_b if next_obs_b is not None else itertools.repeat(None)
+        next_obs_b_iterator = (
+            next_obs_b if next_obs_b is not None else itertools.repeat(None)
+        )
         for obs, act, next_obs in zip(obs_b, act_b, next_obs_b_iterator):
             flat_trans = self._preprocess_transition(obs, act, next_obs)
             transitions.setdefault(None, []).append(flat_trans)
diff --git a/tests/algorithms/test_density_baselines.py b/tests/algorithms/test_density_baselines.py
index 201f27e1c..ebddaa95b 100644
--- a/tests/algorithms/test_density_baselines.py
+++ b/tests/algorithms/test_density_baselines.py
@@ -1,7 +1,7 @@
 """Tests for `imitation.algorithms.density_baselines`."""
 
-from typing import Sequence
 from dataclasses import asdict
+from typing import Sequence
 
 import numpy as np
 import pytest
@@ -107,14 +107,16 @@ def test_density_trainer_smoke(
 
 
 def test_density_with_other_trajectory_types(
-        pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
-        pendulum_venv,
-        rng,
+    pendulum_expert_trajectories: Sequence[TrajectoryWithRew],
+    pendulum_venv,
+    rng,
 ):
     rl_algo = stable_baselines3.PPO(policies.ActorCriticPolicy, pendulum_venv)
     rollouts = pendulum_expert_trajectories[:2]
     transitions = rollout.flatten_trajectories_with_rew(rollouts)
-    transitions_mappings = [asdict(transitions),]
+    transitions_mappings = [
+        asdict(transitions),
+    ]
 
     minimal_transitions = types.TransitionsMinimal(
         obs=transitions.obs,
@@ -137,8 +139,8 @@ def test_density_with_other_trajectory_types(
 
 
 def test_density_trainer_raises(
-        pendulum_venv,
-        rng,
+    pendulum_venv,
+    rng,
 ):
     rl_algo = stable_baselines3.PPO(policies.ActorCriticPolicy, pendulum_venv)
     density_trainer = DensityAlgorithm(
@@ -156,4 +158,4 @@ def test_density_trainer_raises(
         )
 
     with pytest.raises(TypeError, match="Unsupported demonstration type"):
-        density_trainer.set_demonstrations("foo")
\ No newline at end of file
+        density_trainer.set_demonstrations("foo")
diff --git a/tests/util/test_util.py b/tests/util/test_util.py
index a5790a444..b7661473f 100644
--- a/tests/util/test_util.py
+++ b/tests/util/test_util.py
@@ -47,7 +47,7 @@ def test_get_first_iter_element(input_seq):
     assert input_seq is new_iterable
 
     an_iterator = (x for x in input_seq)
-    first_element, new_iterable = util.get_first_iter_element(input_seq)
+    first_element, new_iterable = util.get_first_iter_element(an_iterator)
     assert first_element == input_seq[0]
     assert list(an_iterator) == input_seq
     assert list(an_iterator) == []

From 3d0db4b31dec8c59ae7ae492004c52eacf647423 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 5 Oct 2022 16:26:38 +0100
Subject: [PATCH 172/187] Add pragma no cover updates

---
 .coveragerc                         | 4 +++-
 src/imitation/algorithms/mce_irl.py | 6 ++++--
 src/imitation/policies/base.py      | 2 +-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/.coveragerc b/.coveragerc
index d12534191..a312bb801 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -1,4 +1,6 @@
 [report]
 exclude_lines =
     pragma: no cover
-    @overload
\ No newline at end of file
+    @overload
+    @typing.overload
+    raise NotImplementedError
\ No newline at end of file
diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index 4b85717b0..5cd8cd223 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -176,14 +176,16 @@ def set_pi(self, pi: np.ndarray) -> None:
         self.pi = pi
 
     def _predict(self, observation: th.Tensor, deterministic: bool = False):
-        raise NotImplementedError("Should never be called as predict overridden.")
+        raise NotImplementedError(  # pragma: no cover
+            "Should never be called as predict overridden."
+        )
 
     def forward(  # type: ignore[override]
         self,
         observation: th.Tensor,
         deterministic: bool = False,
     ):
-        raise NotImplementedError("Should never be called.")
+        raise NotImplementedError("Should never be called.")  # pragma: no cover
 
     def predict(
         self,
diff --git a/src/imitation/policies/base.py b/src/imitation/policies/base.py
index f0ee6d588..ffa3a281d 100644
--- a/src/imitation/policies/base.py
+++ b/src/imitation/policies/base.py
@@ -40,7 +40,7 @@ def _choose_action(self, obs: np.ndarray) -> np.ndarray:
     def forward(self, *args):
         # technically BasePolicy is a Torch module, so this needs a forward()
         # method
-        raise NotImplementedError()
+        raise NotImplementedError()  # pragma: no cover
 
 
 class RandomPolicy(HardCodedPolicy):

From 6f9ff973c8a026f7f2015be0fbcdb49caf48c0c0 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Wed, 5 Oct 2022 16:29:43 +0100
Subject: [PATCH 173/187] Minor coverage tweaks

---
 src/imitation/algorithms/mce_irl.py        | 2 +-
 src/imitation/util/logger.py               | 2 +-
 tests/algorithms/test_density_baselines.py | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/imitation/algorithms/mce_irl.py b/src/imitation/algorithms/mce_irl.py
index 5cd8cd223..8cd53b8b4 100644
--- a/src/imitation/algorithms/mce_irl.py
+++ b/src/imitation/algorithms/mce_irl.py
@@ -177,7 +177,7 @@ def set_pi(self, pi: np.ndarray) -> None:
 
     def _predict(self, observation: th.Tensor, deterministic: bool = False):
         raise NotImplementedError(  # pragma: no cover
-            "Should never be called as predict overridden."
+            "Should never be called as predict overridden.",
         )
 
     def forward(  # type: ignore[override]
diff --git a/src/imitation/util/logger.py b/src/imitation/util/logger.py
index 810c8dd7e..77e54df8c 100644
--- a/src/imitation/util/logger.py
+++ b/src/imitation/util/logger.py
@@ -354,7 +354,7 @@ def __init__(self):
         """
         try:
             import wandb
-        except ModuleNotFoundError as e:
+        except ModuleNotFoundError as e:  # pragma: no cover
             raise ModuleNotFoundError(
                 "Trying to log data with `WandbOutputFormat` "
                 "but `wandb` not installed: try `pip install wandb`.",
diff --git a/tests/algorithms/test_density_baselines.py b/tests/algorithms/test_density_baselines.py
index ebddaa95b..8a6a72cb2 100644
--- a/tests/algorithms/test_density_baselines.py
+++ b/tests/algorithms/test_density_baselines.py
@@ -158,4 +158,4 @@ def test_density_trainer_raises(
         )
 
     with pytest.raises(TypeError, match="Unsupported demonstration type"):
-        density_trainer.set_demonstrations("foo")
+        density_trainer.set_demonstrations("foo")  # type: ignore

From f8d6dbdd5e5ba84b8aceb5eba49de795211b5be0 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Thu, 6 Oct 2022 13:18:36 +0100
Subject: [PATCH 174/187] Fix iterator test

---
 tests/util/test_util.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tests/util/test_util.py b/tests/util/test_util.py
index b7661473f..ce663d8e0 100644
--- a/tests/util/test_util.py
+++ b/tests/util/test_util.py
@@ -46,11 +46,16 @@ def test_get_first_iter_element(input_seq):
     assert first_element == input_seq[0]
     assert input_seq is new_iterable
 
-    an_iterator = (x for x in input_seq)
-    first_element, new_iterable = util.get_first_iter_element(an_iterator)
+    def generator_fn():
+        for x in input_seq:
+            yield x
+
+    generator = generator_fn()
+    assert generator == iter(generator)
+    first_element, new_iterable = util.get_first_iter_element(generator)
     assert first_element == input_seq[0]
-    assert list(an_iterator) == input_seq
-    assert list(an_iterator) == []
+    assert list(new_iterable) == input_seq
+    assert list(new_iterable) == []
 
 
 @given(

From 99f03fca1676f0f8c796e1e08f0632ce5d265b32 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Sat, 8 Oct 2022 19:11:02 +0100
Subject: [PATCH 175/187] Add test for parse_path

---
 src/imitation/data/types.py |  1 -
 tests/data/test_types.py    | 35 ++++++++++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index e1958dfdc..fe98704fd 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -87,7 +87,6 @@ def parse_path(
     else:
         parsed_path = pathlib.Path(str(path))
 
-    parsed_path = parsed_path.resolve()
     if parsed_path.is_absolute():
         return parsed_path
     else:
diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index f1c10ff55..dfa287e6d 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -218,7 +218,6 @@ def test_save_trajectories(
             save_dir_str = tmpdir
 
         with chdir_context:
-
             save_dir = types.parse_path(save_dir_str)
             trajs = [trajectory_rew if use_rewards else trajectory]
             save_path = save_dir / "trajs"
@@ -389,3 +388,37 @@ def test_zero_length_fails():
     empty = np.array([])
     with pytest.raises(ValueError, match=r"Degenerate trajectory.*"):
         types.Trajectory(obs=np.array([42]), acts=empty, infos=None, terminal=True)
+
+
+def test_parse_path():
+    # absolute paths
+    assert types.parse_path("/foo/bar") == pathlib.Path("/foo/bar")
+    assert types.parse_path(pathlib.Path("/foo/bar")) == pathlib.Path("/foo/bar")
+    assert types.parse_path(b"/foo/bar") == pathlib.Path("/foo/bar")
+
+    # relative paths. implicit conversion to cwd
+    assert types.parse_path("foo/bar") == pathlib.Path.cwd() / "foo/bar"
+    assert types.parse_path(pathlib.Path("foo/bar")) == pathlib.Path.cwd() / "foo/bar"
+    assert types.parse_path(b"foo/bar") == pathlib.Path.cwd() / "foo/bar"
+
+    # relative paths. conversion using custom base directory
+    base_dir = pathlib.Path("/foo/bar")
+    assert types.parse_path("baz", base_directory=base_dir) == base_dir / "baz"
+    assert (
+        types.parse_path(pathlib.Path("baz"), base_directory=base_dir)
+        == base_dir / "baz"
+    )
+    assert types.parse_path(b"baz", base_directory=base_dir) == base_dir / "baz"
+
+    # pass a relative path but disallowing relative paths. should raise error.
+    with pytest.raises(ValueError, match="Path .* is not absolute"):
+        types.parse_path("foo/bar", allow_relative=False)
+
+    # pass a base direectory but disallowing relative paths. should raise error.
+    with pytest.raises(
+        ValueError,
+        match="If `base_directory` is specified, then `allow_relative` must be True.",
+    ):
+        types.parse_path(
+            "foo/bar", base_directory=pathlib.Path("/foo/bar"), allow_relative=False
+        )

From 15d015768816ab05f9c048bfbf065b4c744b1685 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Sat, 8 Oct 2022 22:21:16 +0100
Subject: [PATCH 176/187] Updates on sacred util

---
 .circleci/config.yml             |  1 -
 src/imitation/scripts/analyze.py | 11 ++++----
 src/imitation/util/sacred.py     | 45 +++++++++++++++-----------------
 tests/data/test_types.py         |  4 ++-
 4 files changed, 29 insertions(+), 32 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 5a0823d1e..0bcb17cf7 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -41,7 +41,6 @@ executors:
         (?x)(
           src/imitation/algorithms/preference_comparisons.py$
           | src/imitation/rewards/reward_nets.py$
-          | src/imitation/util/sacred.py$
           | src/imitation/algorithms/base.py$
           | src/imitation/scripts/train_preference_comparisons.py$
           | src/imitation/rewards/serialize.py$
diff --git a/src/imitation/scripts/analyze.py b/src/imitation/scripts/analyze.py
index a3a2bf683..417799752 100644
--- a/src/imitation/scripts/analyze.py
+++ b/src/imitation/scripts/analyze.py
@@ -5,7 +5,6 @@
 import json
 import logging
 import os
-import os.path as osp
 import pathlib
 import tempfile
 import warnings
@@ -49,7 +48,8 @@ def _gather_sacred_dicts(
     # e.g. chain.from_iterable([["pathone", "pathtwo"], [], ["paththree"]]) =>
     # ("pathone", "pathtwo", "paththree")
     sacred_dirs = itertools.chain.from_iterable(
-        sacred_util.filter_subdirs(source_dir) for source_dir in source_dirs
+        sacred_util.filter_subdirs(types.parse_path(source_dir))
+        for source_dir in source_dirs
     )
     sacred_dicts_list = []
 
@@ -109,8 +109,7 @@ def gather_tb_directories() -> dict:
         # Expecting a path like "~/ray_results/{run_name}/sacred/1".
         # Want to search for all Tensorboard dirs inside
         # "~/ray_results/{run_name}".
-        sacred_dir = types.parse_path(sd.sacred_dir)
-        run_dir = sacred_dir.parent.parent
+        run_dir = sd.sacred_dir.parent.parent
         run_name = run_dir.name
 
         # log is what we use as subdirectory in new code.
@@ -118,8 +117,8 @@ def gather_tb_directories() -> dict:
         for basename in ["log", "rl", "tb", "sb_tb"]:
             tb_src_dirs = tuple(
                 sacred_util.filter_subdirs(
-                    str(run_dir),
-                    lambda path: osp.basename(path) == basename,
+                    run_dir,
+                    lambda path: path.name == basename,
                 ),
             )
             if tb_src_dirs:
diff --git a/src/imitation/util/sacred.py b/src/imitation/util/sacred.py
index 609ada89e..2c8ee421a 100644
--- a/src/imitation/util/sacred.py
+++ b/src/imitation/util/sacred.py
@@ -16,35 +16,31 @@
 class SacredDicts(NamedTuple):
     """Each dict `foo` is loaded from `f"{sacred_dir}/foo.json"`."""
 
-    sacred_dir: str
+    sacred_dir: pathlib.Path
     config: dict
     run: dict
 
     @classmethod
-    def load_from_dir(cls, sacred_dir: str):
-        args = []
-        for field in cls._fields:
-            if field == "sacred_dir":
-                args.append(sacred_dir)
-            else:
-                json_path = os.path.join(sacred_dir, f"{field}.json")
-                with open(json_path, "r") as f:
-                    args.append(json.load(f))
-        return cls(*args)
+    def load_from_dir(cls, sacred_dir: pathlib.Path):
+        return cls(
+            sacred_dir=sacred_dir,
+            config=json.loads((sacred_dir / "config.json").read_text()),
+            run=json.loads((sacred_dir / "run.json").read_text()),
+        )
 
 
-def dir_contains_sacred_jsons(dir_path: str) -> bool:
-    run_path = os.path.join(dir_path, "run.json")
-    config_path = os.path.join(dir_path, "config.json")
-    return os.path.isfile(run_path) and os.path.isfile(config_path)
+def dir_contains_sacred_jsons(dir_path: pathlib.Path) -> bool:
+    run_path = dir_path / "run.json"
+    config_path = dir_path / "config.json"
+    return run_path.is_file() and config_path.is_file()
 
 
 def filter_subdirs(
-    root_dir: str,
-    filter_fn: Callable[[str], bool] = dir_contains_sacred_jsons,
+    root_dir: pathlib.Path,
+    filter_fn: Callable[[pathlib.Path], bool] = dir_contains_sacred_jsons,
     *,
     nested_ok: bool = False,
-) -> Sequence[str]:
+) -> Sequence[pathlib.Path]:
     """Walks through a directory tree, returning paths to filtered subdirectories.
 
     Does not follow symlinks.
@@ -64,17 +60,18 @@ def filter_subdirs(
             paths is a subdirecotry of another.
     """
     filtered_dirs = set()
-    for root, _, _ in os.walk(root_dir, followlinks=False):
+    for root_str, _, _ in os.walk(root_dir, followlinks=False):
+        root = pathlib.Path(root_str)
         if filter_fn(root):
             filtered_dirs.add(root)
 
     if not nested_ok:
         for dirpath in filtered_dirs:
-            components = os.path.split(dirpath)
-            for i in range(1, len(components)):
-                prefix = os.path.join(*components[0:i])
-                if prefix in filtered_dirs:
-                    raise ValueError(f"Parent {prefix} to {dir} also a dir directory")
+            for other_dirpath in filtered_dirs:
+                if dirpath != other_dirpath and other_dirpath in dirpath.parents:
+                    raise ValueError(
+                        f"Found nested directories: {dirpath} and {other_dirpath}",
+                    )
     return list(filtered_dirs)
 
 
diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index dfa287e6d..089dcf136 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -420,5 +420,7 @@ def test_parse_path():
         match="If `base_directory` is specified, then `allow_relative` must be True.",
     ):
         types.parse_path(
-            "foo/bar", base_directory=pathlib.Path("/foo/bar"), allow_relative=False
+            "foo/bar",
+            base_directory=pathlib.Path("/foo/bar"),
+            allow_relative=False,
         )

From 26603687072ed95e1692efe189cfa649da774015 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 10 Oct 2022 12:02:59 +0100
Subject: [PATCH 177/187] Mark type ignore rule

---
 tests/data/test_rollout.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index 9c7acf226..db125908a 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -300,7 +300,7 @@ def test_generate_trajectories_type_error(rng):
     sample_until = rollout.make_min_episodes(1)
     with pytest.raises(TypeError, match="Policy must be.*got <class 'str'> instead"):
         rollout.generate_trajectories(
-            "strings_are_not_valid_policies",  # type: ignore
+            "strings_are_not_valid_policies",  # type: ignore[arg-type]
             venv,
             rng=rng,
             sample_until=sample_until,

From bb37a5527f5301e25157ce8a62788508443053b8 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 10 Oct 2022 12:35:37 +0100
Subject: [PATCH 178/187] Mark type ignore rule

---
 tests/policies/test_replay_buffer_wrapper.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/policies/test_replay_buffer_wrapper.py b/tests/policies/test_replay_buffer_wrapper.py
index 1fc38988a..40fc6eac5 100644
--- a/tests/policies/test_replay_buffer_wrapper.py
+++ b/tests/policies/test_replay_buffer_wrapper.py
@@ -56,7 +56,7 @@ def test_invalid_args(rng):
         # we ignore the type because we are intentionally
         # passing the wrong type for the test
         make_algo_with_wrapped_buffer(
-            rl_cls=sb3.PPO,  # type: ignore
+            rl_cls=sb3.PPO,  # type: ignore[arg-type]
             policy_cls=policies.ActorCriticPolicy,
             replay_buffer_class=buffers.ReplayBuffer,
             rng=rng,

From d8f76ea6c9e55dc2c62a132523a24d2907b66b33 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 10 Oct 2022 12:40:21 +0100
Subject: [PATCH 179/187] Miscellaneous bug fixes and improvements

---
 src/imitation/policies/serialize.py        | 2 +-
 src/imitation/scripts/common/common.py     | 8 ++------
 src/imitation/scripts/eval_policy.py       | 3 +--
 src/imitation/scripts/train_adversarial.py | 2 +-
 4 files changed, 5 insertions(+), 10 deletions(-)

diff --git a/src/imitation/policies/serialize.py b/src/imitation/policies/serialize.py
index 52d2a88ea..70fa2fc02 100644
--- a/src/imitation/policies/serialize.py
+++ b/src/imitation/policies/serialize.py
@@ -198,7 +198,7 @@ def save_stable_model(
     # information in future. (E.g. we used to save `VecNormalize` statistics here,
     # although that is no longer necessary.)
     output_dir.mkdir(parents=True, exist_ok=True)
-    model.save(str(output_dir / filename))
+    model.save(output_dir / filename)
     logging.info(f"Saved policy to {output_dir}")
 
 
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index 0516e46f6..5e486072c 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -52,12 +52,8 @@ def hook(config, command_name: str, logger):
     if config["common"]["log_dir"] is None:
         env_sanitized = config["common"]["env_name"].replace("/", "_")
         assert isinstance(env_sanitized, str)
-        config_log_root = types.parse_optional_path(config["common"]["log_root"])
-        log_root = (
-            config_log_root
-            if config_log_root
-            else types.parse_path("output")  # relative to cwd
-        )
+        config_log_root = config["common"]["log_root"] or "output"
+        log_root = types.parse_path(config_log_root)
         log_dir = log_root / command_name / env_sanitized / util.make_unique_timestamp()
         updates["log_dir"] = str(log_dir)
     return updates
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
index 53f136547..06eab4820 100644
--- a/src/imitation/scripts/eval_policy.py
+++ b/src/imitation/scripts/eval_policy.py
@@ -104,8 +104,7 @@ def eval_policy(
         )
 
     if rollout_save_path:
-
-        types.save(log_dir / rollout_save_path.lstrip("{log_dir}/"), trajs)
+        types.save(log_dir / rollout_save_path.replace("{log_dir}/", ""), trajs)
 
     return rollout.rollout_stats(trajs)
 
diff --git a/src/imitation/scripts/train_adversarial.py b/src/imitation/scripts/train_adversarial.py
index cfefc0cfc..b84aec720 100644
--- a/src/imitation/scripts/train_adversarial.py
+++ b/src/imitation/scripts/train_adversarial.py
@@ -141,7 +141,7 @@ def train_adversarial(
             venv=venv,
             demonstrations=expert_trajs,
             gen_algo=gen_algo,
-            log_dir=str(log_dir),
+            log_dir=log_dir,
             reward_net=reward_net,
             custom_logger=custom_logger,
             **algorithm_kwargs,

From 6f4cb4ee27bb5f934706901b844c337bbef7f7eb Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 10 Oct 2022 14:04:27 +0100
Subject: [PATCH 180/187] Reformat hanging line

---
 tests/scripts/test_scripts.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/scripts/test_scripts.py b/tests/scripts/test_scripts.py
index e3b3913af..a2db0ee0c 100644
--- a/tests/scripts/test_scripts.py
+++ b/tests/scripts/test_scripts.py
@@ -56,8 +56,10 @@
 TEST_DATA_PATH = types.parse_path("tests/testdata")
 
 if not TEST_DATA_PATH.exists():
-    raise RuntimeError("Folder with test data has not been found. Make sure you are "
-                       "running tests relative to the base imitation project folder.")
+    raise RuntimeError(
+        "Folder with test data has not been found. Make sure you are "
+        "running tests relative to the base imitation project folder.",
+    )
 
 CARTPOLE_TEST_DATA_PATH = TEST_DATA_PATH / "expert_models/cartpole_0/"
 CARTPOLE_TEST_ROLLOUT_PATH = CARTPOLE_TEST_DATA_PATH / "rollouts/final.pkl"

From 9f28bbf37d1fa17cf224eb5cd6b08de7366c4bba Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 10 Oct 2022 16:11:15 +0100
Subject: [PATCH 181/187] Ignore parse path checks for windows

---
 tests/data/test_types.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index 089dcf136..5183c7296 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -391,6 +391,11 @@ def test_zero_length_fails():
 
 
 def test_parse_path():
+    if os.name == "nt":  # pragma: no cover
+        pytest.skip(
+            "Windows uses path.WindowsPath instead when paths are resolved, which"
+            "cannot be compared directly to pathlib.Path objects."
+        )
     # absolute paths
     assert types.parse_path("/foo/bar") == pathlib.Path("/foo/bar")
     assert types.parse_path(pathlib.Path("/foo/bar")) == pathlib.Path("/foo/bar")

From 8dc3bcd3083e2c1895dda7b02561a86c477b3bf6 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 10 Oct 2022 16:31:25 +0100
Subject: [PATCH 182/187] Add trailing comma

---
 tests/data/test_types.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index 5183c7296..11f380fab 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -394,7 +394,7 @@ def test_parse_path():
     if os.name == "nt":  # pragma: no cover
         pytest.skip(
             "Windows uses path.WindowsPath instead when paths are resolved, which"
-            "cannot be compared directly to pathlib.Path objects."
+            "cannot be compared directly to pathlib.Path objects.",
         )
     # absolute paths
     assert types.parse_path("/foo/bar") == pathlib.Path("/foo/bar")

From 9ee38c8c3f4bf09bc786f295e1075088c810ad66 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 10 Oct 2022 17:30:35 +0100
Subject: [PATCH 183/187] Minor changes

---
 tests/data/test_types.py      | 4 ++++
 tests/scripts/test_scripts.py | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index 11f380fab..2a2efcac8 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -429,3 +429,7 @@ def test_parse_path():
             base_directory=pathlib.Path("/foo/bar"),
             allow_relative=False,
         )
+
+    # Parse optional path. Works the same way but passes None down the line.
+    assert types.parse_optional_path(None) is None
+    assert types.parse_optional_path("/foo/bar") == types.parse_path("/foo/bar")
\ No newline at end of file
diff --git a/tests/scripts/test_scripts.py b/tests/scripts/test_scripts.py
index a2db0ee0c..3b90628e0 100644
--- a/tests/scripts/test_scripts.py
+++ b/tests/scripts/test_scripts.py
@@ -55,7 +55,7 @@
 
 TEST_DATA_PATH = types.parse_path("tests/testdata")
 
-if not TEST_DATA_PATH.exists():
+if not TEST_DATA_PATH.exists():  # pragma: no cover
     raise RuntimeError(
         "Folder with test data has not been found. Make sure you are "
         "running tests relative to the base imitation project folder.",

From 750bcc99ffb5a9c13ff86e721d6e499cd4e1071e Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Mon, 10 Oct 2022 17:41:38 +0100
Subject: [PATCH 184/187] No newline end of file

---
 tests/data/test_types.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/data/test_types.py b/tests/data/test_types.py
index 2a2efcac8..f7223b93e 100644
--- a/tests/data/test_types.py
+++ b/tests/data/test_types.py
@@ -432,4 +432,4 @@ def test_parse_path():
 
     # Parse optional path. Works the same way but passes None down the line.
     assert types.parse_optional_path(None) is None
-    assert types.parse_optional_path("/foo/bar") == types.parse_path("/foo/bar")
\ No newline at end of file
+    assert types.parse_optional_path("/foo/bar") == types.parse_path("/foo/bar")

From b6bf8450b70dedec86ec3fc7da5d09caa0ce7bf6 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 11 Oct 2022 14:15:11 +0200
Subject: [PATCH 185/187] Update src/imitation/data/types.py

Co-authored-by: Adam Gleave <adam@gleave.me>
---
 src/imitation/data/types.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index 900ef5ed4..abf035ad4 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -113,7 +113,7 @@ def parse_optional_path(
     Args:
         path: The path to parse. Can be a string, bytes, or `os.PathLike`.
         allow_relative: If True, then relative paths are allowed as input, and
-            are resolved relative to the current working directory. If false,
+            are resolved relative to the current working directory. If False,
             an error is raised if the path is not absolute.
         base_directory: If specified, then relative paths are resolved relative
             to this directory, instead of the current working directory.

From 066ae6fb2f674fe5f2f69a031324eab4386a871f Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 11 Oct 2022 14:15:17 +0200
Subject: [PATCH 186/187] Update src/imitation/data/types.py

Co-authored-by: Adam Gleave <adam@gleave.me>
---
 src/imitation/data/types.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index abf035ad4..9b66260d1 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -59,7 +59,7 @@ def parse_path(
     Args:
         path: The path to parse. Can be a string, bytes, or `os.PathLike`.
         allow_relative: If True, then relative paths are allowed as input, and
-            are resolved relative to the current working directory. If false,
+            are resolved relative to the current working directory. If False,
             an error is raised if the path is not absolute.
         base_directory: If specified, then relative paths are resolved relative
             to this directory, instead of the current working directory.

From 0594039938f58b0c7f9fc23eda4a40d3b3ce3486 Mon Sep 17 00:00:00 2001
From: Juan Rocamonde <juancarlosrocamonde@gmail.com>
Date: Tue, 11 Oct 2022 13:28:37 +0100
Subject: [PATCH 187/187]  Include suggestions from Adam

---
 src/imitation/data/types.py            | 6 ++----
 src/imitation/scripts/common/common.py | 2 +-
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index 9b66260d1..70c3f03f3 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -493,12 +493,10 @@ def save(path: AnyPath, trajectories: Sequence[Trajectory]):
         trajectories: The trajectories to save.
 
     Raises:
-        ValueError: If the trajectories are not all of the same type, i.e. some are
+        ValueError: If not all trajectories have the same type, i.e. some are
             `Trajectory` and others are `TrajectoryWithRew`.
     """
-    if isinstance(path, bytes):
-        path = path.decode("utf-8")
-    p = pathlib.Path(path)
+    p = parse_path(path)
     p.parent.mkdir(parents=True, exist_ok=True)
     tmp_path = f"{p}.tmp"
 
diff --git a/src/imitation/scripts/common/common.py b/src/imitation/scripts/common/common.py
index 5e486072c..72d44f2f4 100644
--- a/src/imitation/scripts/common/common.py
+++ b/src/imitation/scripts/common/common.py
@@ -55,7 +55,7 @@ def hook(config, command_name: str, logger):
         config_log_root = config["common"]["log_root"] or "output"
         log_root = types.parse_path(config_log_root)
         log_dir = log_root / command_name / env_sanitized / util.make_unique_timestamp()
-        updates["log_dir"] = str(log_dir)
+        updates["log_dir"] = log_dir
     return updates