Repository: HumanCompatibleAI/imitation (language: python)

Issue:
Common reward model interface
Currently our `RewardNet` class is designed for `AIRL`. In particular, having separate `train` and `test` outputs is specific to AIRL, and does not make any sense for most others algorithms.

We currently have an alternative interface `tabular_irl.RewardModel`. This is intended for JAX -- but it'd be incompatible with `RewardNet` even after being ported to PyTorch. (And note that it's not actually specific to tabular methods -- although some of the concrete subclasses are.)

In the future we might add other algorithms like GCL, DRLHP, etc which also use rewards but do not have this train vs test distinction.

My proposal would be to make an abstract `RewardNet` or `RewardModel` class (either name seems fine) that is minimal and so should work across all algorithms. In particular, I imagine the `__init__` would take in an `observation_space` and `action_space`, there would be a single abstract method `reward` (v.s. `reward_train` and `reward_test` now) and a concrete implementation `predict` taking NumPy arrays. (You might also need some other things for PyTorch, e.g. `device`, that we have now.)

We can then have concrete instances for various different reward models. I would suggest having one that is an MLP reward model (with boolean flags indicating whether it uses observation, action, next observation and dones as input), one that is an MLP potential shaping only, a class that lets you take linear combinations of reward models, and then a shaped MLP can be defined in terms of these.

For AIRL you could then feed in an MLP reward model and an MLP potential shaping, it can internally sum them and have the train vs test distinction -- with no need for a separate class for this.

Note I've done something similar in https://github.com/HumanCompatibleAI/evaluating-rewards/blob/master/src/evaluating_rewards/rewards/base.py which has been working well (and I iterated on a couple of designs before getting to this stage). But that's all TF code, so cannot be used verbatim.


Discussion / hints:
#311 mostly fixes these, except that IRL still uses its own reward models. Should be very easy to change so that IRL uses the models in `rewards.reward_nets`. But I don't see a training script for IRL. Do people just use `mce_irl` directly or is there some other place we'd need to modify?
I don't think we have a training script for `tabular_irl`. Which is bad, we should, but I think that can probably be postponed -- it'd be a natural thing to do if/when we unify training different algorithms into a single script. So I think the only other place to update is `test_tabular`.

Fix patch:
diff --git a/ci/code_checks.sh b/ci/code_checks.sh
index 8400d5bac..eaec6c00c 100755
--- a/ci/code_checks.sh
+++ b/ci/code_checks.sh
@@ -16,12 +16,12 @@ if [ -x "`which circleci`" ]; then
 fi
 
 if [ "$skipexpensive" != "true" ]; then
+  echo "Type checking"
+  pytype -j auto ${SRC_FILES[@]}
+
   echo "Building docs (validates docstrings)"
   pushd docs/
   make clean
   make html
   popd
-
-  echo "Type checking"
-  pytype -j auto ${SRC_FILES[@]}
 fi
diff --git a/docs/conf.py b/docs/conf.py
index 2d69b909d..162070056 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -20,7 +20,7 @@
 import imitation  # pytype: disable=import-error
 
 project = "imitation"
-copyright = "2019, Center for Human-Compatible AI"  # noqa: A001
+copyright = "2019-2021, Center for Human-Compatible AI"  # noqa: A001
 author = "Center for Human-Compatible AI"
 
 # The full version, including alpha/beta/rc tags
@@ -33,7 +33,7 @@
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
-    "sphinxcontrib.napoleon",
+    "sphinx.ext.napoleon",
     "sphinx.ext.autodoc",
     "sphinx_autodoc_typehints",
     "sphinx.ext.autosummary",
diff --git a/docs/guide/gettingstarted.rst b/docs/guide/gettingstarted.rst
index 359dcf1d2..e6dd1f4cc 100644
--- a/docs/guide/gettingstarted.rst
+++ b/docs/guide/gettingstarted.rst
@@ -43,71 +43,5 @@ to run this.
 
 .. _example script: https://github.com/HumanCompatibleAI/imitation/blob/master/examples/quickstart.py
 
-.. code-block:: python
-
-    """Loads CartPole demonstrations and trains BC, GAIL, and AIRL models on that data.
-    """
-    import pathlib
-    import pickle
-    import tempfile
-
-    import seals  # noqa: F401
-    import stable_baselines3 as sb3
-
-    from imitation.algorithms import adversarial, bc
-    from imitation.data import rollout
-    from imitation.util import logger, util
-
-    # Load pickled test demonstrations.
-    with open("tests/testdata/expert_models/cartpole_0/rollouts/final.pkl", "rb") as f:
-        # This is a list of `imitation.data.types.Trajectory`, where
-        # every instance contains observations and actions for a single expert
-        # demonstration.
-        trajectories = pickle.load(f)
-
-    # Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
-    # This is a more general dataclass containing unordered
-    # (observation, actions, next_observation) transitions.
-    transitions = rollout.flatten_trajectories(trajectories)
-
-    venv = util.make_vec_env("seals/CartPole-v0", n_envs=2)
-
-    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
-    tempdir_path = pathlib.Path(tempdir.name)
-    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
-
-    # Train BC on expert data.
-    # BC also accepts as `expert_data` any PyTorch-style DataLoader that iterates over
-    # dictionaries containing observations and actions.
-    bc_logger = logger.configure(tempdir_path / "BC/")
-    bc_trainer = bc.BC(
-        venv.observation_space,
-        venv.action_space,
-        expert_data=transitions,
-        custom_logger=bc_logger,
-    )
-    bc_trainer.train(n_epochs=1)
-
-    # Train GAIL on expert data.
-    # GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
-    # iterates over dictionaries containing observations, actions, and next_observations.
-    gail_logger = logger.configure(tempdir_path / "GAIL/")
-    gail_trainer = adversarial.GAIL(
-        venv,
-        expert_data=transitions,
-        expert_batch_size=32,
-        gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024),
-        custom_logger=gail_logger,
-    )
-    gail_trainer.train(total_timesteps=2048)
-
-    # Train AIRL on expert data.
-    airl_logger = logger.configure(tempdir_path / "AIRL/")
-    airl_trainer = adversarial.AIRL(
-        venv,
-        expert_data=transitions,
-        expert_batch_size=32,
-        gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024),
-        custom_logger=airl_logger,
-    )
-    airl_trainer.train(total_timesteps=2048)
+.. literalinclude :: ../../examples/quickstart.py
+   :language: python
\ No newline at end of file
diff --git a/docs/modules/imitation.algorithms.rst b/docs/modules/imitation.algorithms.rst
index 3a2ab7f40..a4cfc99c6 100644
--- a/docs/modules/imitation.algorithms.rst
+++ b/docs/modules/imitation.algorithms.rst
@@ -19,10 +19,10 @@ imitation.algorithms.dagger module
 
 .. automodule:: imitation.algorithms.dagger
 
-imitation.algorithms.density\_baselines module
-----------------------------------------------
+imitation.algorithms.density module
+-----------------------------------
 
-.. automodule:: imitation.algorithms.density_baselines
+.. automodule:: imitation.algorithms.density
 
 imitation.algorithms.tabular\_irl module
 ----------------------------------------
diff --git a/examples/quickstart.py b/examples/quickstart.py
index c0443c84b..acfb2c0b9 100644
--- a/examples/quickstart.py
+++ b/examples/quickstart.py
@@ -31,25 +31,25 @@
 print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
 
 # Train BC on expert data.
-# BC also accepts as `expert_data` any PyTorch-style DataLoader that iterates over
+# BC also accepts as `demonstrations` any PyTorch-style DataLoader that iterates over
 # dictionaries containing observations and actions.
 bc_logger = logger.configure(tempdir_path / "BC/")
 bc_trainer = bc.BC(
-    venv.observation_space,
-    venv.action_space,
-    expert_data=transitions,
+    observation_space=venv.observation_space,
+    action_space=venv.action_space,
+    demonstrations=transitions,
     custom_logger=bc_logger,
 )
 bc_trainer.train(n_epochs=1)
 
 # Train GAIL on expert data.
-# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
+# GAIL, and AIRL also accept as `demonstrations` any Pytorch-style DataLoader that
 # iterates over dictionaries containing observations, actions, and next_observations.
 gail_logger = logger.configure(tempdir_path / "GAIL/")
 gail_trainer = adversarial.GAIL(
-    venv,
-    expert_data=transitions,
-    expert_batch_size=32,
+    venv=venv,
+    demonstrations=transitions,
+    demo_batch_size=32,
     gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024),
     custom_logger=gail_logger,
 )
@@ -58,9 +58,9 @@
 # Train AIRL on expert data.
 airl_logger = logger.configure(tempdir_path / "AIRL/")
 airl_trainer = adversarial.AIRL(
-    venv,
-    expert_data=transitions,
-    expert_batch_size=32,
+    venv=venv,
+    demonstrations=transitions,
+    demo_batch_size=32,
     gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024),
     custom_logger=airl_logger,
 )
diff --git a/src/imitation/algorithms/adversarial.py b/src/imitation/algorithms/adversarial.py
index 51d6aeb10..dc00d87f8 100644
--- a/src/imitation/algorithms/adversarial.py
+++ b/src/imitation/algorithms/adversarial.py
@@ -1,12 +1,13 @@
+"""Module for adversarial imitation learning algorithms, GAIL and AIRL."""
+
 import dataclasses
 import functools
 import logging
 import os
-from typing import Callable, Dict, Iterable, Mapping, Optional, Type, Union
+from typing import Callable, Dict, Mapping, Optional, Type
 
 import numpy as np
 import torch as th
-import torch.utils.data as th_data
 import torch.utils.tensorboard as thboard
 import tqdm
 from stable_baselines3.common import base_class, vec_env
@@ -18,7 +19,7 @@
 from imitation.util import logger, reward_wrapper, util
 
 
-class AdversarialTrainer(base.BaseImitationAlgorithm):
+class AdversarialTrainer(base.DemonstrationAlgorithm[types.Transitions]):
     """Base class for adversarial imitation learning algorithms like GAIL and AIRL."""
 
     venv: vec_env.VecEnv
@@ -40,13 +41,13 @@ class AdversarialTrainer(base.BaseImitationAlgorithm):
 
     def __init__(
         self,
+        *,
+        demonstrations: base.AnyTransitions,
+        demo_batch_size: int,
         venv: vec_env.VecEnv,
-        gen_algo: base_class.BaseAlgorithm,
         discrim_net: discrim_nets.DiscrimNet,
-        expert_data: Union[Iterable[Mapping], types.Transitions],
-        expert_batch_size: int,
+        gen_algo: base_class.BaseAlgorithm,
         n_disc_updates_per_round: int = 2,
-        *,
         log_dir: str = "output/",
         normalize_obs: bool = True,
         normalize_reward: bool = True,
@@ -63,30 +64,20 @@ def __init__(
         """Builds AdversarialTrainer.
 
         Args:
+            demonstrations: Demonstrations from an expert (optional). Transitions
+                expressed directly as a `types.TransitionsMinimal` object, a sequence
+                of trajectories, or an iterable of transition batches (mappings from
+                keywords to arrays containing observations, etc).
+            demo_batch_size: The number of samples in each batch of expert data. The
+                discriminator batch size is twice this number because each discriminator
+                batch contains a generator sample for every expert sample.
             venv: The vectorized environment to train in.
             gen_algo: The generator RL algorithm that is trained to maximize
                 discriminator confusion. Environment and logger will be set to
                 `venv` and `custom_logger`.
             discrim_net: The discriminator network. This will be moved to the same
                 device as `gen_algo`.
-            expert_data: Either a `torch.utils.data.DataLoader`-like object or an
-                instance of `Transitions` which is automatically converted into a
-                shuffled version of the former type.
-
-                If the argument passed is a `DataLoader`, then it must yield batches of
-                expert data via its `__iter__` method. Each batch is a dictionary whose
-                keys "obs", "acts", "next_obs", and "dones", correspond to Tensor or
-                NumPy array values each with batch dimension equal to
-                `expert_batch_size`. If any batch dimension doesn't equal
-                `expert_batch_size` then a `ValueError` is raised.
-
-                If the argument is a `Transitions` instance, then `len(expert_data)`
-                must be at least `expert_batch_size`.
-            expert_batch_size: The number of samples in each batch yielded from
-                the expert data loader. The discriminator batch size is twice this
-                number because each discriminator batch contains a generator sample for
-                every expert sample.
-            n_discrim_updates_per_round: The number of discriminator updates after each
+            n_disc_updates_per_round: The number of discriminator updates after each
                 round of generator updates in AdversarialTrainer.learn().
             log_dir: Directory to store TensorBoard logs, plots, etc. in.
             normalize_obs: Whether to normalize observations with `VecNormalize`.
@@ -98,10 +89,9 @@ def __init__(
                 or number of environments (for off-policy).
             gen_replay_buffer_capacity: The capacity of the
                 generator replay buffer (the number of obs-action-obs samples from
-                the generator that can be stored).
-
-                By default this is equal to `gen_train_timesteps`, meaning that we
-                sample only from the most recent batch of generator samples.
+                the generator that can be stored). By default this is equal to
+                `gen_train_timesteps`, meaning that we sample only from the most
+                recent batch of generator samples.
             custom_logger: Where to log to; if None (default), creates a new logger.
             init_tensorboard: If True, makes various discriminator
                 TensorBoard summaries.
@@ -120,7 +110,11 @@ def __init__(
                 https://imitation.readthedocs.io/en/latest/guide/variable_horizon.html
                 before overriding this.
         """
+        self.demo_batch_size = demo_batch_size
+        self._demo_data_loader = None
+        self._endless_expert_iterator = None
         super().__init__(
+            demonstrations=demonstrations,
             custom_logger=custom_logger,
             allow_variable_horizon=allow_variable_horizon,
         )
@@ -129,29 +123,6 @@ def __init__(
         self._disc_step = 0
         self.n_disc_updates_per_round = n_disc_updates_per_round
 
-        if expert_batch_size <= 0:
-            raise ValueError(f"expert_batch_size={expert_batch_size} must be positive.")
-
-        self.expert_batch_size = expert_batch_size
-        if isinstance(expert_data, types.Transitions):
-            if len(expert_data) < expert_batch_size:
-                raise ValueError(
-                    "Provided Transitions instance as `expert_data` argument but "
-                    "len(expert_data) < expert_batch_size. "
-                    f"({len(expert_data)} < {expert_batch_size}).",
-                )
-
-            self.expert_data_loader = th_data.DataLoader(
-                expert_data,
-                batch_size=expert_batch_size,
-                collate_fn=types.transitions_collate_fn,
-                shuffle=True,
-                drop_last=True,
-            )
-        else:
-            self.expert_data_loader = expert_data
-        self._endless_expert_iterator = util.endless_iter(self.expert_data_loader)
-
         self.debug_use_ground_truth = debug_use_ground_truth
         self.venv = venv
         self.gen_algo = gen_algo
@@ -213,6 +184,13 @@ def __init__(
             self.venv,
         )
 
+    def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
+        self._demo_data_loader = base.make_data_loader(
+            demonstrations,
+            self.demo_batch_size,
+        )
+        self._endless_expert_iterator = util.endless_iter(self._demo_data_loader)
+
     def _next_expert_batch(self) -> Mapping:
         return next(self._endless_expert_iterator)
 
@@ -229,17 +207,16 @@ def train_disc(
                 If provided, must contain keys corresponding to every field of the
                 `Transitions` dataclass except "infos". All corresponding values can be
                 either NumPy arrays or Tensors. Extra keys are ignored. Must contain
-                `self.expert_batch_size` samples.
-
-                If this argument is not provided, then `self.expert_batch_size` expert
-                samples from `self.expert_data_loader` are used by default.
+                `self.demo_batch_size` samples. If this argument is not provided, then
+                `self.demo_batch_size` expert samples from `self.demo_data_loader` are
+                used by default.
             gen_samples: Transition samples from the generator policy in same dictionary
                 form as `expert_samples`. If provided, must contain exactly
-                `self.expert_batch_size` samples. If not provided, then take
+                `self.demo_batch_size` samples. If not provided, then take
                 `len(expert_samples)` samples from the generator replay buffer.
 
         Returns:
-           dict: Statistics for discriminator (e.g. loss, accuracy).
+            Statistics for discriminator (e.g. loss, accuracy).
         """
         with self.logger.accumulate_means("disc"):
             # optionally write TB summaries for collected ops
@@ -375,17 +352,17 @@ def _make_disc_train_batch(
                 raise RuntimeError(
                     "No generator samples for training. " "Call `train_gen()` first.",
                 )
-            gen_samples = self._gen_replay_buffer.sample(self.expert_batch_size)
+            gen_samples = self._gen_replay_buffer.sample(self.demo_batch_size)
             gen_samples = types.dataclass_quick_asdict(gen_samples)
 
         n_gen = len(gen_samples["obs"])
         n_expert = len(expert_samples["obs"])
-        if not (n_gen == n_expert == self.expert_batch_size):
+        if not (n_gen == n_expert == self.demo_batch_size):
             raise ValueError(
-                "Need to have exactly self.expert_batch_size number of expert and "
+                "Need to have exactly self.demo_batch_size number of expert and "
                 "generator samples, each. "
                 f"(n_gen={n_gen} n_expert={n_expert} "
-                f"expert_batch_size={self.expert_batch_size})",
+                f"demo_batch_size={self.demo_batch_size})",
             )
 
         # Guarantee that Mapping arguments are in mutable form.
@@ -462,27 +439,38 @@ def _make_disc_train_batch(
 
 
 class GAIL(AdversarialTrainer):
+    """Generative Adversarial Imitation Learning (`GAIL`_).
+
+    .. _GAIL: https://arxiv.org/abs/1606.03476
+    """
+
     def __init__(
         self,
+        *,
+        demonstrations: base.AnyTransitions,
+        demo_batch_size: int,
         venv: vec_env.VecEnv,
-        expert_data: Union[Iterable[Mapping], types.Transitions],
-        expert_batch_size: int,
         gen_algo: base_class.BaseAlgorithm,
-        *,
-        # FIXME(sam) pass in discrim net directly; don't ask for kwargs indirectly
         discrim_kwargs: Optional[Mapping] = None,
         **kwargs,
     ):
         """Generative Adversarial Imitation Learning.
 
-        Most parameters are described in and passed to `AdversarialTrainer.__init__`.
-        Additional parameters that `GAIL` adds on top of its superclass initializer are
-        as follows:
-
         Args:
+            demonstrations: Demonstrations from an expert (optional). Transitions
+                expressed directly as a `types.TransitionsMinimal` object, a sequence
+                of trajectories, or an iterable of transition batches (mappings from
+                keywords to arrays containing observations, etc).
+            demo_batch_size: The number of samples in each batch of expert data. The
+                discriminator batch size is twice this number because each discriminator
+                batch contains a generator sample for every expert sample.
+            venv: The vectorized environment to train in.
+            gen_algo: The generator RL algorithm that is trained to maximize
+                discriminator confusion. Environment and logger will be set to
+                `venv` and `custom_logger`.
             discrim_kwargs: Optional keyword arguments to use while constructing the
                 DiscrimNetGAIL.
-
+            **kwargs: Passed through to `AdversarialTrainer.__init__`.
         """
         discrim_kwargs = discrim_kwargs or {}
         discrim = discrim_nets.DiscrimNetGAIL(
@@ -491,44 +479,60 @@ def __init__(
             **discrim_kwargs,
         )
         super().__init__(
-            venv,
-            gen_algo,
-            discrim,
-            expert_data,
-            expert_batch_size,
+            demonstrations=demonstrations,
+            demo_batch_size=demo_batch_size,
+            venv=venv,
+            discrim_net=discrim,
+            gen_algo=gen_algo,
             **kwargs,
         )
 
 
 class AIRL(AdversarialTrainer):
-    """The AIRL reward network, used by the imitation policy."""
+    """Adversarial Inverse Reinforcement Learning (`AIRL`_).
+
+    .. _AIRL: https://arxiv.org/abs/1710.11248
+    """
 
     def __init__(
         self,
+        *,
+        demonstrations: base.AnyTransitions,
+        demo_batch_size: int,
         venv: vec_env.VecEnv,
-        expert_data: Union[Iterable[Mapping], types.Transitions],
-        expert_batch_size: int,
         gen_algo: base_class.BaseAlgorithm,
-        *,
         # FIXME(sam): pass in reward net directly, not via _cls and _kwargs
         reward_net_cls: Type[reward_nets.RewardNet] = reward_nets.BasicShapedRewardNet,
         reward_net_kwargs: Optional[Mapping] = None,
         discrim_kwargs: Optional[Mapping] = None,
         **kwargs,
     ):
-        """Adversarial Inverse Reinforcement Learning.
-
-        Most parameters are described in and passed to `AdversarialTrainer.__init__`.
-        Additional parameters that `AIRL` adds on top of its superclass initializer are
-        as follows:
+        """Builds an AIRL trainer.
 
         Args:
+            demonstrations: Demonstrations from an expert (optional). Transitions
+                expressed directly as a `types.TransitionsMinimal` object, a sequence
+                of trajectories, or an iterable of transition batches (mappings from
+                keywords to arrays containing observations, etc).
+            demo_batch_size: The number of samples in each batch of expert data. The
+                discriminator batch size is twice this number because each discriminator
+                batch contains a generator sample for every expert sample.
+            venv: The vectorized environment to train in.
+            gen_algo: The generator RL algorithm that is trained to maximize
+                discriminator confusion. Environment and logger will be set to
+                `venv` and `custom_logger`.
             reward_net_cls: Reward network constructor. The reward network is part of
                 the AIRL discriminator.
             reward_net_kwargs: Optional keyword arguments to use while constructing
                 the reward network.
             discrim_kwargs: Optional keyword arguments to use while constructing the
                 DiscrimNetAIRL.
+            **kwargs: Passed through to `AdversarialTrainer.__init__`.
+
+        Raises:
+            TypeError: If `gen_algo.policy` does not have an `evaluate_actions`
+                attribute (present in `ActorCriticPolicy`), needed to compute
+                log-probability of actions.
         """
         # TODO(shwang): Maybe offer str=>RewardNet conversion like
         #  stable_baselines3 does with policy classes.
@@ -544,11 +548,11 @@ def __init__(
         discrim_kwargs = discrim_kwargs or {}
         discrim = discrim_nets.DiscrimNetAIRL(reward_network, **discrim_kwargs)
         super().__init__(
-            venv,
-            gen_algo,
-            discrim,
-            expert_data,
-            expert_batch_size,
+            demonstrations=demonstrations,
+            demo_batch_size=demo_batch_size,
+            venv=venv,
+            discrim_net=discrim,
+            gen_algo=gen_algo,
             **kwargs,
         )
 
diff --git a/src/imitation/algorithms/base.py b/src/imitation/algorithms/base.py
index 97388c6c5..5b31df610 100644
--- a/src/imitation/algorithms/base.py
+++ b/src/imitation/algorithms/base.py
@@ -1,7 +1,13 @@
+"""Module of base classes and helper methods for imitation learning algorithms."""
+
 import abc
-from typing import Iterable, Optional
+from typing import Any, Generic, Iterable, Mapping, Optional, TypeVar, Union
+
+import numpy as np
+import torch as th
+import torch.utils.data as th_data
 
-from imitation.data import types
+from imitation.data import rollout, types
 from imitation.util import logger as imit_logger
 
 
@@ -19,7 +25,8 @@ class BaseImitationAlgorithm(abc.ABC):
 
     def __init__(
         self,
-        custom_logger: Optional[imit_logger.HierarchicalLogger],
+        *,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
         allow_variable_horizon: bool = False,
     ):
         """Creates an imitation learning algorithm.
@@ -58,7 +65,7 @@ def logger(self, value: imit_logger.HierarchicalLogger) -> None:
         self._logger = value
 
     def _check_fixed_horizon(self, trajs: Iterable[types.Trajectory]) -> None:
-        """Check that `trajs` has fixed episode length and equal to prior calls.
+        """Checks that `trajs` has fixed episode length and equal to prior calls.
 
         If algorithm is safe to use with variable horizon episodes (e.g. behavioral
         cloning), then just don't call this method.
@@ -67,8 +74,8 @@ def _check_fixed_horizon(self, trajs: Iterable[types.Trajectory]) -> None:
             trajs: An iterable sequence of trajectories.
 
         Raises:
-            ValueError if the length of trajectories in trajs are different from one
-            another, or from trajectory lengths in previous calls to this method.
+            ValueError: The length of trajectories in trajs differs from one
+                another, or from trajectory lengths in previous calls to this method.
         """
         if self.allow_variable_horizon:  # skip check -- YOLO
             return
@@ -101,3 +108,155 @@ def __setstate__(self, state):
         self.__dict__.update(state)
         # callee should modify self.logger directly if they want to override this
         self.logger = state.get("_logger") or imit_logger.configure()
+
+
+TransitionMapping = Mapping[str, Union[np.ndarray, th.Tensor]]
+TransitionKind = TypeVar("TransitionKind", bound=types.TransitionsMinimal)
+AnyTransitions = Union[
+    Iterable[types.Trajectory],
+    Iterable[TransitionMapping],
+    TransitionKind,
+]
+
+
+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,
+    ):
+        """Creates an algorithm that learns from demonstrations.
+
+        Args:
+            demonstrations: Demonstrations from an expert (optional). Transitions
+                expressed directly as a `types.TransitionsMinimal` object, a sequence
+                of trajectories, or an iterable of transition batches (mappings from
+                keywords to arrays containing observations, etc).
+            custom_logger: Where to log to; if None (default), creates a new logger.
+            allow_variable_horizon: If False (default), algorithm will raise an
+                exception if it detects trajectories of different length during
+                training. If True, overrides this safety check. WARNING: variable
+                horizon episodes leak information about the reward via termination
+                condition, and can seriously confound evaluation. Read
+                https://imitation.readthedocs.io/en/latest/guide/variable_horizon.html
+                before overriding this.
+        """
+        super().__init__(
+            custom_logger=custom_logger,
+            allow_variable_horizon=allow_variable_horizon,
+        )
+
+        if demonstrations is not None:
+            self.set_demonstrations(demonstrations)
+
+    @abc.abstractmethod
+    def set_demonstrations(self, demonstrations: AnyTransitions) -> None:
+        """Sets the demonstration data.
+
+        Changing the demonstration data on-demand can be useful for
+        interactive algorithms like DAgger.
+
+        Args:
+             demonstrations: Either a Torch `DataLoader`, any other iterator that
+                yields dictionaries containing "obs" and "acts" Tensors or NumPy arrays,
+                `TransitionKind` instance, or a Sequence of Trajectory objects.
+        """
+
+
+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,
+    ):
+        """Builds _WrapedDataLoader.
+
+        Args:
+            data_loader: The data loader (batch iterable) to wrap.
+            expected_batch_size: The batch size to check for.
+        """
+        self.data_loader = data_loader
+        self.expected_batch_size = expected_batch_size
+
+    def __iter__(self):
+        """Iterator yielding data from `self.data_loader`, checking `self.expected_batch_size`.
+
+        Raises:
+            ValueError: `self.data_loader` returns a batch of size not equal to
+                `self.expected_batch_size`.
+        """
+        for batch in self.data_loader:
+            if len(batch["obs"]) != self.expected_batch_size:
+                raise ValueError(
+                    f"Expected batch size {self.expected_batch_size} "
+                    f"!= {len(batch['obs'])} = len(batch['obs'])",
+                )
+            if len(batch["acts"]) != self.expected_batch_size:
+                raise ValueError(
+                    f"Expected batch size {self.expected_batch_size} "
+                    f"!= {len(batch['acts'])} = len(batch['acts'])",
+                )
+            yield batch
+
+
+def make_data_loader(
+    transitions: AnyTransitions,
+    batch_size: int,
+    data_loader_kwargs: Optional[Mapping[str, Any]] = None,
+) -> Iterable[TransitionMapping]:
+    """Converts demonstration data to Torch data loader.
+
+    Args:
+        transitions: Transitions expressed directly as a `types.TransitionsMinimal`
+            object, a sequence of trajectories, or an iterable of transition
+            batches (mappings from keywords to arrays containing observations, etc).
+        batch_size: The size of the batch to create. Does not change the batch size
+            if `transitions` is already an iterable of transition batches.
+        data_loader_kwargs: Arguments to pass to `th_data.DataLoader`.
+
+    Returns:
+        An iterable of transition batches.
+
+    Raises:
+        ValueError: if `transitions` is an iterable over transition batches with batch
+            size not equal to `batch_size`; or if `transitions` is transitions or a
+            sequence of trajectories with total timesteps less than `batch_size`.
+        TypeError: if `transitions` is an unsupported type.
+    """
+    if batch_size <= 0:
+        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
+        if isinstance(first_item, types.Trajectory):
+            transitions = rollout.flatten_trajectories(list(transitions))
+
+    if isinstance(transitions, types.TransitionsMinimal):
+        if len(transitions) < batch_size:
+            raise ValueError(
+                f"Number of transitions in `demonstrations` {len(transitions)} "
+                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,
+        )
+    elif isinstance(transitions, Iterable):
+        return _WrappedDataLoader(transitions, batch_size)
+    else:
+        raise TypeError(f"`demonstrations` unexpected type {type(transitions)}")
diff --git a/src/imitation/algorithms/bc.py b/src/imitation/algorithms/bc.py
index d77ce48eb..c013b30dc 100644
--- a/src/imitation/algorithms/bc.py
+++ b/src/imitation/algorithms/bc.py
@@ -5,12 +5,11 @@
 """
 
 import contextlib
-from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Tuple, Type, Union
+from typing import Any, Callable, Iterable, Mapping, Optional, Tuple, Type, Union
 
 import gym
 import numpy as np
 import torch as th
-import torch.utils.data as th_data
 import tqdm.autonotebook as tqdm
 from stable_baselines3.common import policies, utils, vec_env
 
@@ -42,16 +41,15 @@ class ConstantLRSchedule:
     """A callable that returns a constant learning rate."""
 
     def __init__(self, lr: float = 1e-3):
-        """
+        """Builds ConstantLRSchedule.
+
         Args:
             lr: the constant learning rate that calls to this object will return.
         """
         self.lr = lr
 
     def __call__(self, _):
-        """
-        Returns the constant learning rate.
-        """
+        """Returns the constant learning rate."""
         return self.lr
 
 
@@ -69,18 +67,21 @@ def update(self, n):
 
 
 class EpochOrBatchIteratorWithProgress:
+    """Wraps DataLoader so that all BC batches can be processed in one for-loop.
+
+    Also uses `tqdm` to show progress in stdout.
+    """
+
     def __init__(
         self,
-        data_loader: Iterable[dict],
+        data_loader: Iterable[algo_base.TransitionMapping],
         n_epochs: Optional[int] = None,
         n_batches: Optional[int] = None,
         on_epoch_end: Optional[Callable[[], None]] = None,
         on_batch_end: Optional[Callable[[], None]] = None,
         progress_bar_visible: bool = True,
     ):
-        """Wraps DataLoader so that all BC batches can be processed in a one for-loop.
-
-        Also uses `tqdm` to show progress in stdout.
+        """Builds EpochOrBatchIteratorWithProgress.
 
         Args:
             data_loader: An iterable over data dicts, as used in `BC`.
@@ -93,6 +94,9 @@ def __init__(
             on_batch_end: A callback function without parameters to be called at the
                 end of every batch.
             progress_bar_visible: If True, then show a tqdm progress bar.
+
+        Raises:
+            ValueError: If neither or both of `n_epochs` and `n_batches` are non-None.
         """
         if n_epochs is not None and n_batches is None:
             self.use_epochs = True
@@ -110,9 +114,10 @@ def __init__(
         self.on_batch_end = on_batch_end
         self.progress_bar_visible = progress_bar_visible
 
-    def __iter__(self) -> Iterable[Tuple[dict, dict]]:
+    def __iter__(
+        self,
+    ) -> Iterable[Tuple[algo_base.TransitionMapping, Mapping[str, Any]]]:
         """Yields batches while updating tqdm display to display progress."""
-
         samples_so_far = 0
         epoch_num = 0
         batch_num = 0
@@ -172,53 +177,59 @@ def update_desc():
                         return
 
 
-class BC(algo_base.BaseImitationAlgorithm):
+class BC(algo_base.DemonstrationAlgorithm):
+    """Behavioral cloning (BC).
 
-    DEFAULT_BATCH_SIZE: int = 32
-    """Default batch size for DataLoader automatically constructed from Transitions.
-
-    See `set_expert_data_loader()`.
+    Recovers a policy via supervised learning from observation-action pairs.
     """
 
     # TODO(scottemmons): pass BasePolicy into BC directly (rather than passing its
     #  arguments)
     def __init__(
         self,
+        *,
         observation_space: gym.Space,
         action_space: gym.Space,
-        *,
         policy_class: Type[policies.BasePolicy] = policy_base.FeedForward32Policy,
         policy_kwargs: Optional[Mapping[str, Any]] = None,
-        expert_data: Union[Iterable[Mapping], types.TransitionsMinimal, None] = None,
+        demonstrations: Optional[algo_base.AnyTransitions] = None,
+        demo_batch_size: int = 32,
         optimizer_cls: Type[th.optim.Optimizer] = th.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
+        optimizer_kwargs: Optional[Mapping[str, Any]] = None,
         ent_weight: float = 1e-3,
         l2_weight: float = 0.0,
         device: Union[str, th.device] = "auto",
         custom_logger: Optional[logger.HierarchicalLogger] = None,
     ):
-        """Behavioral cloning (BC).
-
-        Recovers a policy via supervised learning on observation-action Tensor
-        pairs, sampled from a Torch DataLoader or any Iterator that ducktypes
-        `torch.utils.data.DataLoader`.
+        """Builds BC.
 
         Args:
             observation_space: the observation space of the environment.
             action_space: the action space of the environment.
             policy_class: used to instantiate imitation policy.
             policy_kwargs: keyword arguments passed to policy's constructor.
-            expert_data: If not None, then immediately call
-                  `self.set_expert_data_loader(expert_data)` during initialization.
+            demonstrations: Demonstrations from an expert (optional). Transitions
+                expressed directly as a `types.TransitionsMinimal` object, a sequence
+                of trajectories, or an iterable of transition batches (mappings from
+                keywords to arrays containing observations, etc).
+            demo_batch_size: The number of samples in each batch of expert data.
             optimizer_cls: optimiser to use for supervised training.
             optimizer_kwargs: keyword arguments, excluding learning rate and
-                  weight decay, for optimiser construction.
+                weight decay, for optimiser construction.
             ent_weight: scaling applied to the policy's entropy regularization.
             l2_weight: scaling applied to the policy's L2 regularization.
             device: name/identity of device to place policy on.
             custom_logger: Where to log to; if None (default), creates a new logger.
+
+        Raises:
+            ValueError: If `weight_decay` is specified in `optimizer_kwargs` (use the
+                parameter `l2_weight` instead.)
         """
-        super().__init__(custom_logger=custom_logger)
+        self.demo_batch_size = demo_batch_size
+        super().__init__(
+            demonstrations=demonstrations,
+            custom_logger=custom_logger,
+        )
 
         if optimizer_kwargs:
             if "weight_decay" in optimizer_kwargs:
@@ -248,48 +259,21 @@ def __init__(
             **optimizer_kwargs,
         )
 
-        self.expert_data_loader: Optional[Iterable[Mapping]] = None
         self.ent_weight = ent_weight
         self.l2_weight = l2_weight
 
-        if expert_data is not None:
-            self.set_expert_data_loader(expert_data)
-
-    def set_expert_data_loader(
-        self,
-        expert_data: Union[Iterable[Mapping], types.TransitionsMinimal],
-    ) -> None:
-        """Set the expert data loader, which yields batches of obs-act pairs.
-
-        Changing the expert data loader on-demand is useful for DAgger and other
-        interactive algorithms.
-
-        Args:
-             expert_data: Either a Torch `DataLoader`, any other iterator that
-                yields dictionaries containing "obs" and "acts" Tensors or Numpy arrays,
-                or a `TransitionsMinimal` instance.
-
-                If this is a `TransitionsMinimal` instance, then it is automatically
-                converted into a shuffled `DataLoader` with batch size
-                `BC.DEFAULT_BATCH_SIZE`.
-        """
-        if isinstance(expert_data, types.TransitionsMinimal):
-            self.expert_data_loader = th_data.DataLoader(
-                expert_data,
-                shuffle=True,
-                batch_size=BC.DEFAULT_BATCH_SIZE,
-                collate_fn=types.transitions_collate_fn,
-            )
-        else:
-            self.expert_data_loader = expert_data
+    def set_demonstrations(self, demonstrations: algo_base.AnyTransitions) -> None:
+        self._demo_data_loader = algo_base.make_data_loader(
+            demonstrations,
+            self.demo_batch_size,
+        )
 
     def _calculate_loss(
         self,
         obs: Union[th.Tensor, np.ndarray],
         acts: Union[th.Tensor, np.ndarray],
-    ) -> Tuple[th.Tensor, Dict[str, float]]:
-        """
-        Calculate the supervised learning loss used to train the behavioral clone.
+    ) -> Tuple[th.Tensor, Mapping[str, float]]:
+        """Calculate the supervised learning loss used to train the behavioral clone.
 
         Args:
             obs: The observations seen by the expert. If this is a Tensor, then
@@ -371,7 +355,7 @@ def train(
                 effect if `.train()` is being called for the first time.
         """
         it = EpochOrBatchIteratorWithProgress(
-            self.expert_data_loader,
+            self._demo_data_loader,
             n_epochs=n_epochs,
             n_batches=n_batches,
             on_epoch_end=on_epoch_end,
diff --git a/src/imitation/algorithms/dagger.py b/src/imitation/algorithms/dagger.py
index 2f36457c0..2d0dfc359 100644
--- a/src/imitation/algorithms/dagger.py
+++ b/src/imitation/algorithms/dagger.py
@@ -24,38 +24,43 @@
 
 
 class BetaSchedule(abc.ABC):
-    """
-    Determines the value of beta (% of time that demonstrator action is used) over the
-    progression of training rounds.
-    """
+    """Computes beta (% of time demonstration action used) from training round."""
 
     @abc.abstractmethod
     def __call__(self, round_num: int) -> float:
-        """Gives the value of beta for the current round.
+        """Computes the value of beta for the current round.
 
         Args:
-            round: the current round number. Rounds are assumed to be numbered 0, 1, 2,
-              etc.
+            round_num: the current round number. Rounds are assumed to be sequentially
+                numbered from 0.
 
         Returns:
-            beta: the fraction of the time to sample a demonstrator action. Robot
-              actions will be sampled the remainder of the time.
-        """
+            The fraction of the time to sample a demonstrator action. Robot
+                actions will be sampled the remainder of the time.
+        """  # noqa: DAR202
 
 
 class LinearBetaSchedule(BetaSchedule):
-    """
-    Linearly-decreasing schedule for beta (% of time that demonstrator action is used).
-    """
+    """Linearly-decreasing schedule for beta."""
 
     def __init__(self, rampdown_rounds: int):
-        """
+        """Builds LinearBetaSchedule.
+
         Args:
             rampdown_rounds: number of rounds over which to anneal beta.
         """
         self.rampdown_rounds = rampdown_rounds
 
     def __call__(self, round_num: int) -> float:
+        """Computes beta value.
+
+        Args:
+            round_num: the current round number.
+
+        Returns:
+            beta linearly decreasing from `1` to `0` between round `0` and
+            `self.rampdown_rounds`. After that, it is 0.
+        """
         assert round_num >= 0
         return min(1, max(0, (self.rampdown_rounds - round_num) / self.rampdown_rounds))
 
@@ -137,17 +142,17 @@ def __init__(
         beta: float,
         save_dir: types.AnyPath,
     ):
-        """Trajectory collector constructor.
+        """Builds InteractiveTrajectoryCollector.
 
         Args:
-          venv: vectorized environment to sample trajectories from.
-          get_robot_acts: get robot actions that can be substituted for
-              human actions. Takes a vector of observations as input & returns a
-              vector of actions.
-          beta: fraction of the time to use action given to .step() instead of
-              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.
+            venv: vectorized environment to sample trajectories from.
+            get_robot_acts: get robot actions that can be substituted for
+                human actions. Takes a vector of observations as input & returns a
+                vector of actions.
+            beta: fraction of the time to use action given to .step() instead of
+                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.
         """
         super().__init__(venv)
         self.get_robot_acts = get_robot_acts
@@ -169,10 +174,10 @@ def seed(self, seed=Optional[int]) -> List[Union[None, int]]:
 
         Args:
             seed: The random seed. May be None for completely random seeding.
+
         Returns:
-            Returns 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 being
-            seeded.
+            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)
         return self.venv.seed(seed)
@@ -281,28 +286,27 @@ class DAggerTrainer(base.BaseImitationAlgorithm):
 
     def __init__(
         self,
+        *,
         venv: vec_env.VecEnv,
         scratch_dir: types.AnyPath,
-        *,
         beta_schedule: Callable[[int], float] = None,
         batch_size: int = 32,
         bc_kwargs: Optional[dict] = None,
         custom_logger: Optional[logger.HierarchicalLogger] = None,
     ):
-        """DaggerTrainer constructor.
+        """Builds DAggerTrainer.
 
         Args:
             venv: Vectorized training environment.
             scratch_dir: Directory to use to store intermediate training
                 information (e.g. for resuming training).
-
-        Keyword Args:
             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.
             batch_size: Number of samples in each batch during BC training.
             bc_kwargs: Additional arguments for constructing the `BC` instance that
                 will be used to train the underlying policy.
+            custom_logger: Where to log to; if None (default), creates a new logger.
         """
         super().__init__(custom_logger=custom_logger)
 
@@ -318,8 +322,8 @@ def __init__(
         self._all_demos = []
 
         self.bc_trainer = bc.BC(
-            self.venv.observation_space,
-            self.venv.action_space,
+            observation_space=self.venv.observation_space,
+            action_space=self.venv.action_space,
             custom_logger=custom_logger,
             **self.bc_kwargs,
         )
@@ -386,7 +390,7 @@ def _try_load_demos(self) -> None:
                 shuffle=True,
                 collate_fn=types.transitions_collate_fn,
             )
-            self.bc_trainer.set_expert_data_loader(data_loader)
+            self.bc_trainer.set_demonstrations(data_loader)
             self._last_loaded_round = self.round_num
 
     def extend_and_update(self, bc_train_kwargs: Optional[Mapping] = None) -> int:
@@ -408,7 +412,7 @@ def extend_and_update(self, bc_train_kwargs: Optional[Mapping] = None) -> int:
                 keys are provided, then `n_epochs` is set to `self.DEFAULT_N_EPOCHS`.
 
         Returns:
-            round_num: new round number after advancing the round counter.
+            New round number after advancing the round counter.
         """
         if bc_train_kwargs is None:
             bc_train_kwargs = {}
@@ -434,10 +438,9 @@ def get_trajectory_collector(self) -> InteractiveTrajectoryCollector:
         """Create trajectory collector to extend current round's demonstration set.
 
         Returns:
-            collector: an `InteractiveTrajectoryCollector` configured with the
-                appropriate beta, appropriate imitator policy, etc. for the current
-                round. Refer to the documentation for
-                `InteractiveTrajectoryCollector` to see how to use this.
+            A collector configured with the appropriate beta, imitator policy, etc.
+            for the current round. Refer to the documentation for
+            `InteractiveTrajectoryCollector` to see how to use this.
         """
         save_dir = self._demo_dir_path_for_round()
         beta = self.beta_schedule(self.round_num)
@@ -496,17 +499,14 @@ class SimpleDAggerTrainer(DAggerTrainer):
 
     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,
     ):
-        """SimpleDAggerTrainer constructor.
-
-        This constructor also accepts keyword arguments for `DaggerTrainer`. See
-        the `DaggerTrainer.__init__` docstring for more details.
+        """Builds SimpleDAggerTrainer.
 
         Args:
             venv: Vectorized training environment. Note that when the robot
@@ -520,15 +520,19 @@ def __init__(
                 dataset.
             dagger_trainer_kwargs: Other keyword arguments passed to the
                 superclass initializer `DAggerTrainer.__init__`.
+
+        Raises:
+            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)
         self.expert_policy = expert_policy
         if expert_policy.observation_space != self.venv.observation_space:
             raise ValueError(
-                "Mismatched observation space between expert_policy and env",
+                "Mismatched observation space between expert_policy and venv",
             )
         if expert_policy.action_space != self.venv.action_space:
-            raise ValueError("Mismatched action space between expert_policy and env")
+            raise ValueError("Mismatched action space between expert_policy and venv")
 
         # TODO(shwang):
         #   Might welcome Transitions and DataLoaders as sources of expert data
diff --git a/src/imitation/algorithms/density.py b/src/imitation/algorithms/density.py
new file mode 100644
index 000000000..14bafc51e
--- /dev/null
+++ b/src/imitation/algorithms/density.py
@@ -0,0 +1,328 @@
+"""Density-based baselines for imitation learning.
+
+Each of these algorithms learns a density estimate on some aspect of the demonstrations,
+then rewards the agent for following that estimate.
+"""
+
+import enum
+import itertools
+from typing import Iterable, Mapping, Optional
+
+import numpy as np
+from gym.spaces.utils import flatten
+from sklearn import neighbors, preprocessing
+from stable_baselines3.common import base_class, vec_env
+
+from imitation.algorithms import base
+from imitation.data import rollout, types, wrappers
+from imitation.util import logger as imit_logger
+from imitation.util import reward_wrapper
+
+
+class DensityType(enum.Enum):
+    """Input type the density model should use."""
+
+    STATE_DENSITY = enum.auto()
+    """Density on state s."""
+
+    STATE_ACTION_DENSITY = enum.auto()
+    """Density on (s,a) pairs."""
+
+    STATE_STATE_DENSITY = enum.auto()
+    """Density on (s,s') pairs."""
+
+
+class DensityAlgorithm(base.DemonstrationAlgorithm):
+    """Learns a reward function based on density modeling.
+
+    Specifically, it constructs a non-parametric estimate of `p(s)`, `p(s,a)`, `p(s,s')`
+    and then computes a reward using the log of these probabilities.
+    """
+
+    transitions: Mapping[Optional[int], np.ndarray]
+
+    def __init__(
+        self,
+        *,
+        demonstrations: Optional[base.AnyTransitions],
+        venv: vec_env.VecEnv,
+        density_type: DensityType = DensityType.STATE_ACTION_DENSITY,
+        kernel: str = "gaussian",
+        kernel_bandwidth: float = 0.5,
+        rl_algo: Optional[base_class.BaseAlgorithm] = None,
+        is_stationary: bool = True,
+        standardise_inputs: bool = True,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+        allow_variable_horizon: bool = False,
+    ):
+        """Builds DensityAlgorithm.
+
+        Args:
+            demonstrations: expert demonstration trajectories.
+            density_type: type of density to train on: single state, state-action pairs,
+                or state-state pairs.
+            kernel: kernel to use for density estimation with `sklearn.KernelDensity`.
+            kernel_bandwidth: bandwidth of kernel. If `standardise_inputs` is
+                true and you are using a Gaussian kernel, then it probably makes sense
+                to set this somewhere between 0.1 and 1.
+            venv: The environment to learn a reward model in. We don't actually need
+                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).
+            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.
+                A non-stationary model is particularly likely to be useful when using
+                STATE_DENSITY, to encourage agent to imitate entire trajectories, not
+                just a few states that have high frequency in the demonstration dataset.
+                If non-stationary, demonstrations must be trajectories, not transitions
+                (which do not contain timesteps).
+            standardise_inputs: if True, then the inputs to the reward model
+                will be standardised to have zero mean and unit variance over the
+                demonstration trajectories. Otherwise, inputs will be passed to the
+                reward model with their ordinary scale.
+            custom_logger: Where to log to; if None (default), creates a new logger.
+            allow_variable_horizon: If False (default), algorithm will raise an
+                exception if it detects trajectories of different length during
+                training. If True, overrides this safety check. WARNING: variable
+                horizon episodes leak information about the reward via termination
+                condition, and can seriously confound evaluation. Read
+                https://imitation.readthedocs.io/en/latest/guide/variable_horizon.html
+                before overriding this.
+        """
+        self.is_stationary = is_stationary
+        self.density_type = density_type
+        self.venv = venv
+        self.transitions = {}
+        super().__init__(
+            demonstrations=demonstrations,
+            custom_logger=custom_logger,
+            allow_variable_horizon=allow_variable_horizon,
+        )
+
+        self.kernel = kernel
+        self.kernel_bandwidth = kernel_bandwidth
+        self.standardise = standardise_inputs
+        self._scaler = None
+        self._density_models = {}
+
+        self.rl_algo = rl_algo
+        self.buffering_wrapper = wrappers.BufferingWrapper(self.venv)
+        self.venv_wrapped = reward_wrapper.RewardVecEnvWrapper(
+            self.buffering_wrapper,
+            self,
+        )
+        self.wrapper_callback = self.venv_wrapped.make_log_callback()
+
+    def _set_demo_from_batch(
+        self,
+        obs_b: np.ndarray,
+        act_b: np.ndarray,
+        next_obs_b: Optional[np.ndarray],
+    ) -> None:
+        next_obs_b = next_obs_b or itertools.repeat(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)
+
+    def set_demonstrations(self, demonstrations: base.AnyTransitions) -> None:
+        """Sets the demonstration data."""
+        self.transitions = {}
+
+        if isinstance(demonstrations, Iterable):
+            first_item = next(iter(demonstrations))
+            if isinstance(first_item, types.Trajectory):
+                # Demonstrations are trajectories.
+                # We have timestep information.
+                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
+                for batch in demonstrations:
+                    self._set_demo_from_batch(
+                        batch["obs"],
+                        batch["acts"],
+                        batch.get("next_obs"),
+                    )
+        elif isinstance(demonstrations, types.TransitionsMinimal):
+            next_obs_b = (
+                demonstrations.next_obs if hasattr(demonstrations, "next_obs") else None
+            )
+            self._set_demo_from_batch(
+                demonstrations.obs,
+                demonstrations.acts,
+                next_obs_b,
+            )
+        else:
+            raise TypeError(f"Unsupported demonstration type {type(demonstrations)}")
+
+        self.transitions = {k: np.stack(v, axis=0) for k, v in self.transitions.items()}
+
+        if not self.is_stationary and None in self.transitions:
+            raise ValueError(
+                "Non-stationary model incompatible with non-trajectory demonstrations.",
+            )
+        if self.is_stationary:
+            self.transitions = {
+                None: np.concatenate(list(self.transitions.values()), axis=0),
+            }
+
+    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)
+        self._scaler = preprocessing.StandardScaler(
+            with_mean=self.standardise,
+            with_std=self.standardise,
+        )
+        flattened_dataset = np.concatenate(list(self.transitions.values()), axis=0)
+        self._scaler.fit(flattened_dataset)
+
+        # now fit density model
+        self._density_models = {
+            k: self._fit_density(self._scaler.transform(v))
+            for k, v in self.transitions.items()
+        }
+
+    def _fit_density(self, transitions: np.ndarray) -> neighbors.KernelDensity:
+        density_model = neighbors.KernelDensity(
+            kernel=self.kernel,
+            bandwidth=self.kernel_bandwidth,
+        )
+        density_model.fit(transitions)
+        return density_model
+
+    def _preprocess_transition(
+        self,
+        obs: np.ndarray,
+        act: np.ndarray,
+        next_obs: np.ndarray,
+    ) -> np.ndarray:
+        """Compute flattened transition on subset specified by `self.density_type`."""
+        if self.density_type == DensityType.STATE_DENSITY:
+            return flatten(self.venv.observation_space, obs)
+        elif self.density_type == DensityType.STATE_ACTION_DENSITY:
+            return np.concatenate(
+                [
+                    flatten(self.venv.observation_space, obs),
+                    flatten(self.venv.action_space, act),
+                ],
+            )
+        elif self.density_type == DensityType.STATE_STATE_DENSITY:
+            return np.concatenate(
+                [
+                    flatten(self.venv.observation_space, obs),
+                    flatten(self.venv.observation_space, next_obs),
+                ],
+            )
+        else:
+            raise ValueError(f"Unknown density type {self.density_type}")
+
+    def __call__(
+        self,
+        obs_b: np.ndarray,
+        act_b: np.ndarray,
+        next_obs_b: np.ndarray,
+        dones: np.ndarray,
+        steps: Optional[np.ndarray] = None,
+    ) -> np.ndarray:
+        r"""Compute reward from given (s,a,s') transition batch.
+
+        This handles *batches* of observations, since it's designed to work with
+        VecEnvs.
+
+        Args:
+            obs_b: current batch of observations.
+            act_b: batch of actions that agent took in response to those
+                observations.
+            next_obs_b: batch of observations encountered after the
+                agent took those actions.
+            dones: is it terminal state?
+            steps: What timestep is this from? Used if `self.is_stationary` is false,
+                otherwise ignored.
+
+        Returns:
+            Array of scalar rewards of the form `r_t(s,a,s') = \log \hat p_t(s,a,s')`
+            (one for each environment), where `\log \hat p` is the underlying density
+            model (and may be independent of s', a, or t, depending on options passed
+            to constructor).
+
+        Raises:
+            ValueError: Non-stationary model (`self.is_stationary` false) and `steps`
+                is not provided.
+        """
+        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?
+
+        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)):
+            flat_trans = self._preprocess_transition(obs, act, next_obs)
+            scaled_padded_trans = self._scaler.transform(flat_trans[np.newaxis])
+            if self.is_stationary:
+                rew = self._density_models[None].score(scaled_padded_trans)
+            else:
+                time = steps[idx]
+                if time >= len(self._density_models):
+                    # Can't do anything sensible here yet. Correct solution is to use
+                    # hierarchical model in which we first check whether state is
+                    # absorbing, then assign either constant score or a score based on
+                    # density.
+                    raise ValueError(
+                        f"Time {time} out of range (0, {len(self._density_models)}], "
+                        "and absorbing states not currently supported",
+                    )
+                else:
+                    time_model = self._density_models[time]
+                    rew = time_model.score(scaled_padded_trans)
+            rew_list.append(rew)
+        rew_array = np.asarray(rew_list, dtype="float32")
+        return rew_array
+
+    def train_policy(self, n_timesteps: int = int(1e6), **kwargs):
+        """Train the imitation policy for a given number of timesteps.
+
+        Args:
+            n_timesteps: number of timesteps to train the policy for.
+            kwargs (dict): extra arguments that will be passed to the `learn()`
+                method of the imitation RL model. Refer to Stable Baselines docs for
+                details.
+        """
+        self.rl_algo.set_env(self.venv_wrapped)
+        self.rl_algo.learn(
+            n_timesteps,
+            # ensure we can see total steps for all
+            # learn() calls, not just for this call
+            reset_num_timesteps=False,
+            callback=self.wrapper_callback,
+            **kwargs,
+        )
+        trajs = self.buffering_wrapper.pop_trajectories()
+        self._check_fixed_horizon(trajs)
+
+    def test_policy(self, *, n_trajectories: int = 10, true_reward: bool = True):
+        """Test current imitation policy on environment & give some rollout stats.
+
+        Args:
+            n_trajectories: number of rolled-out trajectories.
+            true_reward: should this use ground truth reward from underlying
+                environment (True), or imitation reward (False)?
+
+        Returns:
+            dict: rollout statistics collected by
+                `imitation.utils.rollout.rollout_stats()`.
+        """
+        trajs = rollout.generate_trajectories(
+            self.rl_algo,
+            self.venv if true_reward else self.venv_wrapped,
+            sample_until=rollout.make_min_episodes(n_trajectories),
+        )
+        self._check_fixed_horizon(trajs)
+        reward_stats = rollout.rollout_stats(trajs)
+        return reward_stats
diff --git a/src/imitation/algorithms/density_baselines.py b/src/imitation/algorithms/density_baselines.py
deleted file mode 100644
index 3d2dd06a8..000000000
--- a/src/imitation/algorithms/density_baselines.py
+++ /dev/null
@@ -1,299 +0,0 @@
-"""Density-based baselines for imitation learning.
-
-Each of these algorithms learns a density estimate on some aspect of the demonstrations,
-then rewards the agent for following that estimate.
-"""
-
-from typing import Sequence
-
-import numpy as np
-from gym.spaces.utils import flatten
-from sklearn.neighbors import KernelDensity
-from sklearn.preprocessing import StandardScaler
-from stable_baselines3.common import on_policy_algorithm, vec_env
-
-from imitation.data import rollout, types
-from imitation.util import reward_wrapper
-
-# Constants identifying different kinds of density we can use. Note that all
-# can be augmented to depend on the time step by passing `is_stationary = True`
-# to `DensityReward`.
-
-# Density on state s
-STATE_DENSITY = "state_density"
-# Density on (s,a) pairs
-STATE_ACTION_DENSITY = "state_action_density"
-# Density (s,s') pairs
-STATE_STATE_DENSITY = "state_state_density"
-
-
-class DensityReward:
-    def __init__(
-        self,
-        *,
-        trajectories,
-        density_type,
-        kernel,
-        kernel_bandwidth,
-        obs_space,
-        act_space,
-        is_stationary=True,
-        standardise_inputs=True,
-    ):
-        """Reward function based on a density estimate of trajectories.
-
-        Args:
-          trajectories ([dict]): expert demonstration trajectories generated by,
-            e.g., `imitation.util.rollout.generate_trajectories`.
-          obs_space (gym.Space): observation space for underlying environment.
-          act_space (gym.Space): action space for underlying environment.
-          density_type (str): type of density to train on. Valid choices are
-            defined as constants in this module, including `STATE_DENSITY` (density
-            on single states), `STATE_ACTION_DENSITY` (density on state-action
-            pairs), and `STATE_STATE_DENSITY` (density on state-state transitions).
-          is_stationary (bool): should this trainer use the same density model for
-            all time steps (True), or a different density model for each time step
-            (False)? `is_stationary=False` might be useful if using
-            `density_type=STATE_DENSITY`, since it encourages the agent to imitate
-            entire trajectories instead of simply occupying a few states that have
-            high frequency in the demonstration dataset.
-          standardise_inputs (bool): if True, then the inputs to the reward model
-            will be standardised to have zero mean and unit variance over the
-            demonstration trajectories. Otherwise, inputs will be passed to the
-            reward model with their ordinary scale.
-          kernel (str): kernel to use for density estimation. Current valid options
-            include `gaussian` and `exponential`.
-          kernel_bandwidth (float): bandwidth of kernel. If `standardise_inputs` is
-            true and you are using a Gaussian kernel, then it probably makes sense
-            to set this somewhere between 0.1 and 1.
-        """
-        self.density_type = density_type
-        self.is_stationary = is_stationary
-        self.kernel = kernel
-        self.kernel_bandwidth = kernel_bandwidth
-        self.standardise = standardise_inputs
-        self.obs_space = obs_space
-        self.act_space = act_space
-        self._fit_models(trajectories)
-
-    def _fit_models(self, trajectories):
-        flat_trajs = self._preprocess_trajectories(trajectories)
-
-        # if requested, we'll scale demonstration transitions so that they have
-        # zero mean and unit variance (i.e all components are equally important)
-        flattened_dataset = np.stack(sum(flat_trajs, []), axis=0)
-        self._scaler = StandardScaler(
-            with_mean=self.standardise,
-            with_std=self.standardise,
-        )
-        self._scaler.fit(flattened_dataset)
-
-        # now fit density model
-        # TODO: add absorbing state fix that I describe in __call__
-        if self.is_stationary:
-            # fit to all pairs, since density model is stationary
-            self._density_model = self._fit_single_density(
-                self._scaler.transform(flattened_dataset),
-            )
-        else:
-            # fit separately for samples at each time step
-            T = max(map(len, flat_trajs))
-            traj_groups = [[]] * T
-            for traj in flat_trajs:
-                for t, flat_trans in enumerate(traj):
-                    traj_groups[t].append(flat_trans)
-            traj_groups_scaled = [
-                self._scaler.transform(np.stack(step_transitions, axis=0))
-                for step_transitions in traj_groups
-            ]
-            self._density_models = [
-                self._fit_single_density(scaled_flat_trans)
-                for scaled_flat_trans in traj_groups_scaled
-            ]
-
-    def _fit_single_density(self, flat_transitions):
-        # This bandwidth was chosen to make sense with standardised inputs that
-        # have unit variance in each component. There might be a better way to
-        # choose it automatically.
-        density_model = KernelDensity(
-            kernel=self.kernel,
-            bandwidth=self.kernel_bandwidth,
-        )
-        density_model.fit(flat_transitions)
-        return density_model
-
-    def _preprocess_trajectories(self, trajectories: Sequence[types.Trajectory]):
-        """Preprocess a list of trajectories into atomic units that we can learn a
-        density function on. Depending on configuration, that could mean a sequence
-        state/state pairs, or state/action pairs, or single states, etc.
-
-        Args:
-          trajectories: trajectories to process.
-
-        Returns:
-          flat_trajectories (np.ndarray): a corresponding list of "flattened"
-            trajectories. Each element of a flattened trajectory is a
-            one-dimensional ndarray representing a (state,state) pair at a
-            particular time step, or (state,action) pair at that time step, or
-            whatever the class configuration calls for.
-        """
-        flat_trajectories = []
-        for traj in trajectories:
-            obs_vec = traj.obs
-            act_vec = traj.acts
-            assert len(obs_vec) == len(act_vec) + 1
-            flat_traj = []
-            for step_num in range(len(traj.acts)):
-                flat_trans = self._preprocess_transition(
-                    obs_vec[step_num],
-                    act_vec[step_num],
-                    obs_vec[step_num + 1],
-                )
-                flat_traj.append(flat_trans)
-            flat_trajectories.append(flat_traj)
-        return flat_trajectories
-
-    def _preprocess_transition(self, obs, act, next_obs):
-        if self.density_type == STATE_DENSITY:
-            return flatten(self.obs_space, obs)
-        elif self.density_type == STATE_ACTION_DENSITY:
-            return np.concatenate(
-                [flatten(self.obs_space, obs), flatten(self.act_space, act)],
-            )
-        elif self.density_type == STATE_STATE_DENSITY:
-            return np.concatenate(
-                [flatten(self.obs_space, obs), flatten(self.obs_space, next_obs)],
-            )
-        else:
-            raise ValueError(f"Unknown density type {self.density_type}")
-
-    def __call__(self, obs_b, act_b, next_obs_b, steps=None):
-        r"""Compute reward from given (s,a,s') transition batch.
-
-        This handles *batches* of observations, since it's designed to work with
-        VecEnvs.
-
-        Args:
-          obs_b (np.ndarray): current batch of observations.
-          act_b (np.ndarray): batch of actions that agent took in response to those
-            observations.
-          next_obs_b (np.ndarray): batch of observations encountered after the
-            agent took those actions.
-          steps (np.ndarray): number of time steps elapsed in each environment.
-
-        Returns:
-          rew_array (np.ndarray): array of scalar rewards of the form `r_t(s,a,s')
-            = \log \hat p_t(s,a,s')` (one for each environment), where `\log \hat
-            p` is the underlying density model (and may be independent of s', a, or
-            t, depending on options passed to constructor)."""
-        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)):
-            flat_trans = self._preprocess_transition(obs, act, next_obs)
-            scaled_padded_trans = self._scaler.transform(flat_trans[np.newaxis])
-            if self.is_stationary:
-                rew = self._density_model.score(scaled_padded_trans)
-            else:
-                time = steps[idx]
-                if time >= len(self._density_models):
-                    # Can't do anything sensible here yet. Correct solution is to use
-                    # hierarchical model in which we first check whether state is
-                    # absorbing, then assign either constant score or a score based on
-                    # density.
-                    raise Exception(
-                        f"Time {time} out of range (0, {len(self._density_models)}], "
-                        "and I haven't implemented absorbing states etc. yet",
-                    )
-                else:
-                    time_model = self._density_models[time]
-                    rew = time_model.score(scaled_padded_trans)
-            rew_list.append(rew)
-        rew_array = np.asarray(rew_list, dtype="float32")
-        return rew_array
-
-
-class DensityTrainer:
-    def __init__(
-        self,
-        venv: vec_env.VecEnv,
-        rollouts: Sequence[types.Trajectory],
-        imitation_trainer: on_policy_algorithm.OnPolicyAlgorithm,
-        *,
-        standardise_inputs: bool = True,
-        kernel: str = "gaussian",
-        kernel_bandwidth: float = 0.5,
-        density_type: str = STATE_ACTION_DENSITY,
-        is_stationary: bool = False,
-    ):
-        r"""Family of simple imitation learning baseline algorithms that apply RL to
-        maximise a rough density estimate of the demonstration trajectories.
-        Specifically, it constructs a non-parametric estimate of `p(s)`, `p(s,s')`,
-        `p_t(s,a)`, etc. (depending on options), then rewards the imitation learner
-        with `r_t(s,a,s')=\log p_t(s,a,s')` (or `\log p(s,s')`, or whatever the
-        user wants the model to condition on).
-
-        Args:
-            venv: environment to train on.
-            rollouts: list of expert trajectories to imitate.
-            imitation_trainer: RL algorithm & initial policy that will
-                be used to train the imitation learner.
-            kernel, kernel_bandwidth, density_type, is_stationary,
-                standardise_inputs: these are passed directly to `DensityReward`;
-                refer to documentation for that class."""
-        self.venv = venv
-        self.imitation_trainer = imitation_trainer
-        self.reward_fn = DensityReward(
-            trajectories=rollouts,
-            density_type=density_type,
-            obs_space=self.venv.observation_space,
-            act_space=self.venv.action_space,
-            is_stationary=is_stationary,
-            kernel=kernel,
-            kernel_bandwidth=kernel_bandwidth,
-            standardise_inputs=standardise_inputs,
-        )
-        self.wrapped_env = reward_wrapper.RewardVecEnvWrapper(self.venv, self.reward_fn)
-
-    def train_policy(self, n_timesteps=int(1e6), **kwargs):
-        """Train the imitation policy for a given number of timesteps.
-
-        Args:
-            n_timesteps (int): number of timesteps to train the policy for.
-            kwargs (dict): extra arguments that will be passed to the `learn()`
-                method of the imitation RL model. Refer to Stable Baselines docs for
-                details.
-        """
-        self.imitation_trainer.set_env(self.wrapped_env)
-        # FIXME: learn() is not meant to be called frequently; there are
-        # significant per-call overheads (see Adam's comment in adversarial.py)
-        # FIXME: the ep_reward_mean reported by SB is wrong; it comes from a
-        # Monitor() that is being (incorrectly) used to wrap the underlying
-        # environment.
-        self.imitation_trainer.learn(
-            n_timesteps,
-            # ensure we can see total steps for all
-            # learn() calls, not just for this call
-            reset_num_timesteps=False,
-            **kwargs,
-        )
-
-    def test_policy(self, *, n_trajectories=10, true_reward=True):
-        """Test current imitation policy on environment & give some rollout stats.
-
-        Args:
-            n_trajectories (int): number of rolled-out trajectories.
-            true_reward (bool): should this use ground truth reward from underlying
-                environment (True), or imitation reward (False)?
-
-        Returns:
-            dict: rollout statistics collected by
-                `imitation.utils.rollout.rollout_stats()`.
-        """
-        self.imitation_trainer.set_env(self.venv)
-        trajs = rollout.generate_trajectories(
-            self.imitation_trainer,
-            self.venv if true_reward else self.wrapped_env,
-            sample_until=rollout.make_min_episodes(n_trajectories),
-        )
-        reward_stats = rollout.rollout_stats(trajs)
-        return reward_stats
diff --git a/src/imitation/algorithms/preference_comparisons.py b/src/imitation/algorithms/preference_comparisons.py
index 9fc5fde19..ef4fdf98a 100644
--- a/src/imitation/algorithms/preference_comparisons.py
+++ b/src/imitation/algorithms/preference_comparisons.py
@@ -36,7 +36,7 @@ class TrajectoryGenerator(abc.ABC):
     """Object to log statistics and natural language messages to."""
 
     def __init__(self, custom_logger: Optional[imit_logger.HierarchicalLogger] = None):
-        """Initialize the trajectory generator
+        """Builds TrajectoryGenerator.
 
         Args:
             custom_logger: Where to log to; if None (default), creates a new logger.
@@ -54,7 +54,7 @@ def sample(self, steps: int) -> Sequence[TrajectoryWithRew]:
         Returns:
             A list of sampled trajectories with rewards (which should
             be the environment rewards, not ones from a reward model).
-        """
+        """  # noqa: DAR202
 
     def train(self, steps: int, **kwargs):
         """Train an agent if the trajectory generator uses one.
@@ -121,6 +121,9 @@ def __init__(
             reward_fn: either a RewardFn or a RewardNet instance that will supply
                 the rewards used for training the agent.
             custom_logger: Where to log to; if None (default), creates a new logger.
+
+        Raises:
+            ValueError: `algorithm` does not have an environment set.
         """
         self.algorithm = algorithm
         # NOTE: this has to come after setting self.algorithm because super().__init__
@@ -144,16 +147,16 @@ def __init__(
         )
         self.algorithm.set_env(self.venv)
 
-    def train(self, steps: int, **kwargs):
+    def train(self, steps: int, **kwargs) -> None:
         """Train the agent using the reward function specified during instantiation.
 
         Args:
             steps: number of environment timesteps to train for
             **kwargs: other keyword arguments to pass to BaseAlgorithm.train()
 
-        Returns:
-            a list of all trajectories that occurred during training, including their
-            original environment rewards (rather than the ones computed using reward_fn)
+        Raises:
+            RuntimeError: Transitions left in `self.buffering_wrapper`; call
+                `self.sample` first to clear them.
         """
         n_transitions = self.buffering_wrapper.n_transitions
         if n_transitions:
@@ -267,7 +270,7 @@ def __call__(
 
         Returns:
             a sequence of fragment pairs
-        """
+        """  # noqa: DAR202
 
 
 class RandomFragmenter(Fragmenter):
@@ -371,8 +374,10 @@ def __call__(
 
 
 class PreferenceGatherer(abc.ABC):
+    """Base class for gathering preference comparisons between trajectory fragments."""
+
     def __init__(self, custom_logger: Optional[imit_logger.HierarchicalLogger] = None):
-        """Initialize the preference gatherer.
+        """Initializes the preference gatherer.
 
         Args:
             custom_logger: Where to log to; if None (default), creates a new logger.
@@ -381,8 +386,7 @@ def __init__(self, custom_logger: Optional[imit_logger.HierarchicalLogger] = Non
 
     @abc.abstractmethod
     def __call__(self, fragment_pairs: Sequence[TrajectoryWithRewPair]) -> np.ndarray:
-        """Gathers the probabilities that fragment 1 is preferred
-        for a batch of fragments.
+        """Gathers the probabilities that fragment 1 is preferred in `fragment_pairs`.
 
         Args:
             fragment_pairs: sequence of pairs of trajectory fragments
@@ -396,7 +400,7 @@ def __call__(self, fragment_pairs: Sequence[TrajectoryWithRewPair]) -> np.ndarra
             Note that for human feedback, these probabilities are simply 0 or 1
             (or 0.5 in case of indifference), but synthetic models may yield other
             probabilities.
-        """
+        """  # noqa: DAR202
 
 
 class SyntheticGatherer(PreferenceGatherer):
@@ -470,8 +474,8 @@ def _reward_sums(self, fragment_pairs) -> Tuple[np.ndarray, np.ndarray]:
         rews1, rews2 = zip(
             *[
                 (
-                    rollout.compute_returns(f1.rews, self.discount_factor),
-                    rollout.compute_returns(f2.rews, self.discount_factor),
+                    rollout.discounted_sum(f1.rews, self.discount_factor),
+                    rollout.discounted_sum(f2.rews, self.discount_factor),
                 )
                 for f1, f2 in fragment_pairs
             ],
@@ -491,6 +495,7 @@ class PreferenceDataset(th.utils.data.Dataset):
     """
 
     def __init__(self):
+        """Builds an empty PreferenceDataset."""
         self.fragments1: List[TrajectoryWithRew] = []
         self.fragments2: List[TrajectoryWithRew] = []
         self.preferences = np.array([])
@@ -506,6 +511,10 @@ def push(
             fragments: list of pairs of trajectory fragments to add
             preferences: corresponding preference probabilities (probability
                 that fragment 1 is preferred over fragment 2)
+
+        Raises:
+            ValueError: `preferences` shape does not match `fragments` or
+                has non-float32 dtype.
         """
         fragments1, fragments2 = zip(*fragments)
         if preferences.shape != (len(fragments),):
@@ -820,6 +829,9 @@ def train(self, total_timesteps: int, total_comparisons: int) -> Mapping[str, An
         Returns:
             A dictionary with final metrics such as loss and accuracy
             of the reward model.
+
+        Raises:
+            ValueError: `total_comparisons < self.comparisons_per_iteration`.
         """
         iterations, extra_comparisons = divmod(
             total_comparisons,
diff --git a/src/imitation/algorithms/tabular_irl.py b/src/imitation/algorithms/tabular_irl.py
index 94e5023bd..794855b4d 100644
--- a/src/imitation/algorithms/tabular_irl.py
+++ b/src/imitation/algorithms/tabular_irl.py
@@ -6,53 +6,58 @@
     http://www.cs.cmu.edu/~bziebart/publications/thesis-bziebart.pdf
 """
 
-import logging
-from typing import Callable, Iterable, Optional, Tuple
+from typing import Iterable, Optional, Tuple, Union
 
 import numpy as np
 import scipy.special
 import torch as th
-import torch.nn as nn
-from torch.optim.optimizer import Optimizer
+from torch.optim import optimizer
 
-from imitation.envs.resettable_env import TabularModelEnv
-from imitation.util.util import tensor_iter_norm
+from imitation.algorithms import base
+from imitation.data import rollout, types
+from imitation.envs import resettable_env
+from imitation.rewards import reward_nets
+from imitation.util import logger as imit_logger
+from imitation.util import util
+
+# TODO(adam): support reward linear in state-action or arbitrary features?
 
 
 def mce_partition_fh(
-    env: TabularModelEnv,
+    env: resettable_env.TabularModelEnv,
     *,
-    R: Optional[np.ndarray] = None,
+    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, undiscounted MDP.
+    r"""Performs the soft Bellman backup for a finite-horizon MDP.
 
     Calculates V^{soft}, Q^{soft}, and pi using recurrences (9.1), (9.2), and
     (9.3) from Ziebart (2010).
 
     Args:
         env: a tabular, known-dynamics MDP.
-        R: a reward matrix. Defaults to env.reward_matrix.
+        reward: a reward matrix. Defaults to env.reward_matrix.
+        discount: discount rate.
 
     Returns:
         (V, Q, \pi) corresponding to the soft values, Q-values and MCE policy.
         V is a 2d array, indexed V[t,s]. Q is a 3d array, indexed Q[t,s,a].
         \pi is a 3d array, indexed \pi[t,s,a].
     """
-
     # shorthand
     horizon = env.horizon
     n_states = env.n_states
     n_actions = env.n_actions
     T = env.transition_matrix
-    if R is None:
-        R = env.reward_matrix
+    if reward is None:
+        reward = env.reward_matrix
 
     # Initialization
     # indexed as V[t,s]
     V = np.full((horizon, n_states), -np.inf)
     # indexed as Q[t,s,a]
     Q = np.zeros((horizon, n_states, n_actions))
-    broad_R = R[:, None]
+    broad_R = reward[:, None]
 
     # Base case: final timestep
     # final Q(s,a) is just reward
@@ -63,7 +68,7 @@ def mce_partition_fh(
     # Recursive case
     for t in reversed(range(horizon - 1)):
         next_values_s_a = T @ V[t + 1, :]
-        Q[t, :, :] = broad_R + next_values_s_a
+        Q[t, :, :] = broad_R + discount * next_values_s_a
         V[t, :] = scipy.special.logsumexp(Q[t, :, :], axis=1)
 
     pi = np.exp(Q - V[:, :, None])
@@ -72,10 +77,11 @@ def mce_partition_fh(
 
 
 def mce_occupancy_measures(
-    env: TabularModelEnv,
+    env: resettable_env.TabularModelEnv,
     *,
-    R: Optional[np.ndarray] = None,
+    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.
 
@@ -83,26 +89,26 @@ def mce_occupancy_measures(
 
     Args:
         env: a tabular MDP.
-        R: reward matrix. Defaults is env.reward_matrix.
+        reward: reward matrix. Defaults is env.reward_matrix.
         pi: policy to simulate. Defaults to soft-optimal policy w.r.t reward
             matrix.
+        discount: rate to discount the cumulative occupancy measure D.
 
     Returns:
-        Tuple of D (ndarray) and Dt (ndarray). D is an :math:`|S|`-dimensional
-        vector recording the expected number of times each state is visited.
-        Dt is a :math:`T*|S|`-dimensional vector recording the probability of
-        being in a given state at a given timestep.
+        Tuple of D (ndarray) and Dt (ndarray). D is an :math:`|S|`-dimensional vector
+        recording the expected discounted number of times each state is visited. Dt is
+        a :math:`T*|S|`-dimensional vector recording the probability of being in a given
+        state at a given timestep.
     """
-
     # shorthand
     horizon = env.horizon
     n_states = env.n_states
     n_actions = env.n_actions
     T = env.transition_matrix
-    if R is None:
-        R = env.reward_matrix
+    if reward is None:
+        reward = env.reward_matrix
     if pi is None:
-        _, _, pi = mce_partition_fh(env, R=R)
+        _, _, pi = mce_partition_fh(env, reward=reward)
 
     D = np.zeros((horizon, n_states))
     D[0, :] = env.initial_state_dist
@@ -111,7 +117,9 @@ def mce_occupancy_measures(
             E = D[t - 1] * pi[t - 1, :, a]
             D[t, :] += E @ T[:, a, :]
 
-    return D, D.sum(axis=0)
+    Dcum = rollout.discounted_sum(D, discount)
+    assert isinstance(Dcum, np.ndarray)
+    return D, Dcum
 
 
 def squeeze_r(r_output: th.Tensor) -> th.Tensor:
@@ -130,144 +138,180 @@ def squeeze_r(r_output: th.Tensor) -> th.Tensor:
     return r_output
 
 
-def mce_irl(
-    env: TabularModelEnv,
-    optimizer: Optimizer,
-    rmodel: nn.Module,
-    demo_state_om: np.ndarray,
-    linf_eps: float = 1e-3,
-    grad_l2_eps: float = 1e-4,
-    print_interval: Optional[int] = 100,
-) -> np.ndarray:
-    r"""Tabular MCE IRL.
-
-    Args:
-        env: a tabular MDP.
-        optimizer: an optimizer for `rmodel`.
-        rmodel: a neural network that computes rewards for the supplied
-            observations.
-        demo_state_om: matrix representing state occupancy measure for
-            demonstrator.
-        linf_eps: optimisation terminates if the $l_{\infty}$ distance between
-            the demonstrator's state occupancy measure and the state occupancy
-            measure for the current reward falls below this value.
-        grad_l2_eps: optimisation also terminates if the $\ell_2$ norm of the
-            MCE IRL gradient falls below this value.
-        print_interval: how often to log current loss stats (using `logging`).
-            None to disable.
-
-    Returns:
-        state occupancy measure for the final reward function. Note that `rmodel`
-        and `optimizer` will be updated in-place during optimisation.
-    """
-    # use the same device and dtype as the rmodel parameters
-    a_param = next(rmodel.parameters())
-    device = a_param.device
-    dtype = a_param.dtype
-
-    obs_mat = env.observation_matrix
-    torch_obs_mat = th.as_tensor(obs_mat, dtype=dtype, device=device)
-    # l_\infty distance between demonstrator occupancy measure (OM) and OM for
-    # soft-optimal policy w.r.t current reward (initially set to this value to
-    # prevent termination)
-    linf_delta = linf_eps + 1
-    # norm of the MCE IRL gradient (also set to this value to prevent
-    # termination)
-    grad_norm = grad_l2_eps + 1
-    # number of optimisation steps taken
-    t = 0
-    assert demo_state_om.shape == (len(obs_mat),)
-
-    while linf_delta > linf_eps and grad_norm > grad_l2_eps:
-        optimizer.zero_grad()
-
-        # get reward predicted for each state by current model, & compute
-        # expected # of times each state is visited by soft-optimal policy
-        # w.r.t that reward function
-        predicted_r = squeeze_r(rmodel(torch_obs_mat))
-        assert predicted_r.shape == (obs_mat.shape[0],)
-        predicted_r_np = predicted_r.detach().cpu().numpy()
-        _, visitations = mce_occupancy_measures(env, R=predicted_r_np)
-
-        # Forward/back/step (grads are zeroed at the top).
-        # weights_th(s) = \pi(s) - D(s)
-        weights_th = th.as_tensor(
-            visitations - demo_state_om,
-            dtype=dtype,
-            device=device,
-        )
-        # The "loss" is then:
-        #   E_\pi[r_\theta(S)] - E_D[r_\theta(S)]
-        loss = th.dot(weights_th, predicted_r)
-        # This gives the required gradient:
-        #   E_\pi[\nabla r_\theta(S)] - E_D[\nabla r_\theta(S)].
-        loss.backward()
-        optimizer.step()
-
-        # these are just for termination conditions & debug logging
-        grad_norm = tensor_iter_norm(p.grad for p in rmodel.parameters()).item()
-        linf_delta = np.max(np.abs(demo_state_om - visitations))
-        if print_interval is not None and 0 == (t % print_interval):
-            weight_norm = tensor_iter_norm(rmodel.parameters()).item()
-            logging.info(
-                "Occupancy measure error@iter % 3d: %f (||params||=%f, "
-                "||grad||=%f)"
-                % (
-                    t,
-                    linf_delta,
-                    weight_norm,
-                    grad_norm,
-                ),
-            )
-
-        t += 1
+MCEDemonstrations = Union[np.ndarray, base.AnyTransitions]
 
-    return visitations
 
+class MCEIRL(base.DemonstrationAlgorithm[types.TransitionsMinimal]):
+    """Tabular MCE IRL."""
 
-class LinearRewardModel(nn.Module):
-    """Simplest possible linear reward model."""
+    demo_state_om: Optional[np.ndarray]
 
-    def __init__(self, obs_dim: int):
-        """Construct a linear reward model, with no bias.
+    def __init__(
+        self,
+        demonstrations: Optional[MCEDemonstrations],
+        env: resettable_env.TabularModelEnv,
+        reward_net: reward_nets.RewardNet,
+        optimizer: optimizer.Optimizer,
+        discount: float = 1,
+        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,
+        *,
+        custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
+    ):
+        r"""Creates MCE IRL.
 
         Args:
-            obs_dim: size of input observations.
+            demonstrations: Demonstrations from an expert (optional). Can be a sequence
+                of trajectories, or transitions, an iterable over mappings that
+                represent a batch of transitions, or a state occupancy measure.
+                The demonstrations must have observations one-hot coded unless
+                demonstrations is a state-occupancy measure.
+            env: a tabular MDP.
+            reward_net: a neural network that computes rewards for the supplied
+                observations.
+            optimizer: an optimizer for `reward_net`.
+            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
+                are *not allowed* as we cannot discount them appropriately without
+                knowing the timestep they were drawn from.
+            linf_eps: optimisation terminates if the $l_{\infty}$ distance between
+                the demonstrator's state occupancy measure and the state occupancy
+                measure for the current reward falls below this value.
+            grad_l2_eps: optimisation also terminates if the $\ell_2$ norm of the
+                MCE IRL gradient falls below this value.
+            log_interval: how often to log current loss stats (using `logging`).
+                None to disable.
+            custom_logger: Where to log to; if None (default), creates a new logger.
         """
-        super().__init__()
-        # no bias in this model for the sake of simplicity
-        self.lin = nn.Linear(obs_dim, 1, bias=False)
-
-    def forward(self, obs: th.Tensor) -> th.Tensor:
-        return th.squeeze(self.lin(obs), 1)
-
+        self.discount = discount
+        self.env = env
+        self.demo_state_om = None
+        super().__init__(
+            demonstrations=demonstrations,
+            custom_logger=custom_logger,
+        )
 
-class MLPRewardModel(nn.Module):
-    """MLP-based reward model."""
+        self.optimizer = optimizer
+        self.reward_net = reward_net
+        self.linf_eps = linf_eps
+        self.grad_l2_eps = grad_l2_eps
+        self.log_interval = log_interval
 
-    def __init__(
-        self,
-        obs_dim: int,
-        hiddens: Iterable[int],
-        activation: Callable[[], nn.Module] = nn.Tanh,
-    ):
-        """Construct an MLP-based reward model.
+    def _set_demo_from_obs(self, obses: np.ndarray) -> None:
+        if self.discount != 1.0:
+            raise ValueError(
+                "Cannot compute discounted OM from timeless Transitions.",
+            )
+        for obs in obses:
+            if isinstance(obs, th.Tensor):
+                obs = obs.numpy()
+            self.demo_state_om[obs.astype(bool)] += 1.0
+
+    def set_demonstrations(self, demonstrations: MCEDemonstrations) -> None:
+        if isinstance(demonstrations, np.ndarray):
+            # Demonstrations are an occupancy measure
+            assert demonstrations.ndim == 1
+            self.demo_state_om = demonstrations
+        else:
+            # Demonstrations are either trajectories or transitions;
+            # we must compute occupancy measure from this.
+            self.demo_state_om = np.zeros((self.env.n_states,))
+
+            if isinstance(demonstrations, Iterable):
+                first_item = next(iter(demonstrations))
+                if isinstance(first_item, types.Trajectory):
+                    # Demonstrations are trajectories.
+                    for traj in demonstrations:
+                        # TODO(adam): vectorize?
+                        cum_discount = 1
+                        for obs in traj.obs:
+                            self.demo_state_om[obs.astype(bool)] += cum_discount
+                            cum_discount *= self.discount
+                else:
+                    # Demonstrations are a Torch DataLoader or other Mapping iterable
+                    for batch in demonstrations:
+                        self._set_demo_from_obs(batch["obs"])
+
+            elif isinstance(demonstrations, types.TransitionsMinimal):
+                self._set_demo_from_obs(demonstrations.obs)
+            else:
+                raise TypeError(
+                    f"Unsupported demonstration type {type(demonstrations)}",
+                )
+
+            self.demo_state_om /= self.demo_state_om.sum()  # normalize
+
+    def train(self, max_iter: int = 1000) -> np.ndarray:
+        """Runs MCE IRL.
 
         Args:
-            obs_dim: size of input observations.
-            hiddens: list of zero or more hidden layer sizes.
-            activation: constructs an activation to insert between hidden
-                layers (e.g. you could supply `nn.ReLU`, `nn.Tanh`, etc.`).
+            max_iter: The maximum number of iterations to train for. May terminate
+                earlier if `self.linf_eps` or `self.grad_l2_eps` thresholds are reached.
+
+        Returns:
+            State occupancy measure for the final reward function. `self.reward_net`
+            and `self.optimizer` will be updated in-place during optimisation.
         """
-        super().__init__()
-        layer_list = []
-        prev_h_size = obs_dim
-        for h_size in hiddens:
-            layer_list.append(nn.Linear(prev_h_size, h_size))
-            layer_list.append(activation())
-            prev_h_size = h_size
-        layer_list.append(nn.Linear(prev_h_size, 1))
-        self.layers = nn.Sequential(*layer_list)
-
-    def forward(self, obs: th.Tensor) -> th.Tensor:
-        return th.squeeze(self.layers(obs), 1)
+        # use the same device and dtype as the rmodel parameters
+        obs_mat = self.env.observation_matrix
+        dtype = self.reward_net.dtype
+        device = self.reward_net.device
+        torch_obs_mat = th.as_tensor(
+            obs_mat,
+            dtype=dtype,
+            device=device,
+        )
+        assert self.demo_state_om.shape == (len(obs_mat),)
+
+        for t in range(max_iter):
+            self.optimizer.zero_grad()
+
+            # get reward predicted for each state by current model, & compute
+            # expected # of times each state is visited by soft-optimal policy
+            # w.r.t that reward function
+            # TODO(adam): support not just state-only reward?
+            predicted_r = squeeze_r(self.reward_net(torch_obs_mat, None, None, None))
+            assert predicted_r.shape == (obs_mat.shape[0],)
+            predicted_r_np = predicted_r.detach().cpu().numpy()
+            _, visitations = mce_occupancy_measures(
+                self.env,
+                reward=predicted_r_np,
+                discount=self.discount,
+            )
+
+            # Forward/back/step (grads are zeroed at the top).
+            # weights_th(s) = \pi(s) - D(s)
+            weights_th = th.as_tensor(
+                visitations - self.demo_state_om,
+                dtype=dtype,
+                device=device,
+            )
+            # The "loss" is then:
+            #   E_\pi[r_\theta(S)] - E_D[r_\theta(S)]
+            loss = th.dot(weights_th, predicted_r)
+            # This gives the required gradient:
+            #   E_\pi[\nabla r_\theta(S)] - E_D[\nabla r_\theta(S)].
+            loss.backward()
+            self.optimizer.step()
+
+            # these are just for termination conditions & debug logging
+            grad_norm = util.tensor_iter_norm(
+                p.grad for p in self.reward_net.parameters()
+            ).item()
+            linf_delta = np.max(np.abs(self.demo_state_om - visitations))
+
+            if self.log_interval is not None and 0 == (t % self.log_interval):
+                weight_norm = util.tensor_iter_norm(self.reward_net.parameters()).item()
+                self.logger.record("iteration", t)
+                self.logger.record("linf_delta", linf_delta)
+                self.logger.record("weight_norm", weight_norm)
+                self.logger.record("grad_norm", grad_norm)
+                self.logger.dump(t)
+
+            if linf_delta <= self.linf_eps or grad_norm <= self.grad_l2_eps:
+                break
+
+        return visitations
diff --git a/src/imitation/data/rollout.py b/src/imitation/data/rollout.py
index d52e6c716..5fb82aadf 100644
--- a/src/imitation/data/rollout.py
+++ b/src/imitation/data/rollout.py
@@ -585,23 +585,33 @@ def rollout_and_save(
     types.save(path, trajs)
 
 
-def compute_returns(rewards: np.ndarray, gamma: float) -> float:
-    """Calculate the discounted returns from an array of undiscounted rewards.
+def discounted_sum(arr: np.ndarray, gamma: float) -> Union[np.ndarray, float]:
+    """Calculate the discounted sum of `arr`.
+
+    If `arr` is an array of rewards, then this computes the return;
+    however, it can also be used to e.g. compute discounted state
+    occupancy measures.
 
     Args:
-        rewards: array of rewards, from the current time step (first)
-            to the last timestep (last).
-        gamma: the discount factor used for calculating returns.
+        arr: 1 or 2-dimensional array to compute discounted sum over.
+            Last axis is timestep, from current time step (first) to
+            last timestep (last). First axis (if present) is batch
+            dimension.
+        gamma: the discount factor used.
 
     Returns:
-        The discounted returns (the first reward is undiscounted,
-        i.e. we start at gamma^0)
+        The discounted sum over the timestep axis. The first timestep is undiscounted,
+        i.e. we start at gamma^0.
     """
     # We want to calculate sum_{t = 0}^T gamma^t r_t, which can be
     # interpreted as the polynomial sum_{t = 0}^T r_t x^t
     # evaluated at x=gamma.
     # Compared to first computing all the powers of gamma, then
-    # multiplying with the rewards and then summing, this method
+    # multiplying with the `arr` values and then summing, this method
     # should require fewer computations and potentially be more
     # numerically stable.
-    return np.polynomial.polynomial.polyval(gamma, rewards)
+    assert arr.ndim in (1, 2)
+    if gamma == 1.0:
+        return arr.sum(axis=0)
+    else:
+        return np.polynomial.polynomial.polyval(gamma, arr)
diff --git a/src/imitation/data/types.py b/src/imitation/data/types.py
index 3a0084f5e..8fc0f98c9 100644
--- a/src/imitation/data/types.py
+++ b/src/imitation/data/types.py
@@ -116,7 +116,7 @@ def __post_init__(self):
 
 def transitions_collate_fn(
     batch: Sequence[Mapping[str, np.ndarray]],
-) -> Dict[str, Union[np.ndarray, th.Tensor]]:
+) -> Mapping[str, Union[np.ndarray, th.Tensor]]:
     """Custom `torch.utils.data.DataLoader` collate_fn for `TransitionsMinimal`.
 
     Use this as the `collate_fn` argument to `DataLoader` if using an instance of
diff --git a/src/imitation/rewards/reward_nets.py b/src/imitation/rewards/reward_nets.py
index b88318011..3080a6d99 100644
--- a/src/imitation/rewards/reward_nets.py
+++ b/src/imitation/rewards/reward_nets.py
@@ -113,6 +113,16 @@ def device(self) -> th.device:
             # if the model has no parameters, we use the CPU
             return th.device("cpu")
 
+    @property
+    def dtype(self) -> th.dtype:
+        """Heuristic to determine dtype of module."""
+        try:
+            first_param = next(self.parameters())
+            return first_param.dtype
+        except StopIteration:
+            # if the model has no parameters, default to float32
+            return th.float32
+
 
 class ShapedRewardNet(RewardNet):
     """A RewardNet consisting of a base net and a potential shaping."""
diff --git a/src/imitation/scripts/config/train_adversarial.py b/src/imitation/scripts/config/train_adversarial.py
index f88fc2cd5..a56e3d95e 100644
--- a/src/imitation/scripts/config/train_adversarial.py
+++ b/src/imitation/scripts/config/train_adversarial.py
@@ -32,7 +32,7 @@ def train_defaults():
     # Kwargs for initializing GAIL and AIRL
     algorithm_kwargs = dict(
         shared=dict(
-            expert_batch_size=1024,  # Number of expert samples per discriminator update
+            demo_batch_size=1024,  # Number of expert samples per discriminator update
             # Number of discriminator updates after each round of generator updates
             n_disc_updates_per_round=4,
         ),
@@ -115,7 +115,7 @@ def airl():
 ANT_SHARED_LOCALS = dict(
     total_timesteps=3e7,
     max_episode_steps=500,  # To match `inverse_rl` settings.
-    algorithm_kwargs=dict(shared=dict(expert_batch_size=8192)),
+    algorithm_kwargs=dict(shared=dict(demo_batch_size=8192)),
     gen_batch_size=16384,
 )
 
@@ -190,7 +190,7 @@ def seals_ant():
             n_disc_updates_per_round=16,
             # Equivalent to no replay buffer if batch size is the same
             gen_replay_buffer_capacity=16384,
-            expert_batch_size=8192,
+            demo_batch_size=8192,
         ),
         airl=dict(
             reward_net_kwargs=dict(
@@ -302,7 +302,7 @@ def fast():
     n_episodes_eval = 1
     algorithm_kwargs = dict(
         shared=dict(
-            expert_batch_size=1,
+            demo_batch_size=1,
             n_disc_updates_per_round=4,
         ),
     )
diff --git a/src/imitation/scripts/train_adversarial.py b/src/imitation/scripts/train_adversarial.py
index 03c33aa5a..d0f3c284e 100644
--- a/src/imitation/scripts/train_adversarial.py
+++ b/src/imitation/scripts/train_adversarial.py
@@ -192,7 +192,7 @@ def train_adversarial(
 
     trainer = algo_cls(
         venv=venv,
-        expert_data=expert_transitions,
+        demonstrations=expert_transitions,
         gen_algo=gen_algo,
         log_dir=log_dir,
         discrim_kwargs=final_discrim_kwargs,
diff --git a/src/imitation/scripts/train_bc.py b/src/imitation/scripts/train_bc.py
index dc21f260e..54fdf8611 100644
--- a/src/imitation/scripts/train_bc.py
+++ b/src/imitation/scripts/train_bc.py
@@ -7,7 +7,6 @@
 import torch as th
 from sacred.observers import FileStorageObserver
 from stable_baselines3.common import vec_env
-from torch.utils import data as th_data
 
 from imitation.algorithms import bc
 from imitation.data import rollout, types
@@ -105,18 +104,11 @@ def train_bc(
             )
         expert_trajs = expert_trajs[:n_expert_demos]
 
-    expert_data_trans = rollout.flatten_trajectories(expert_trajs)
-    expert_data = th_data.DataLoader(
-        expert_data_trans,
-        batch_size=batch_size,
-        shuffle=True,
-        collate_fn=types.transitions_collate_fn,
-    )
-
     model = bc.BC(
-        observation_space,
-        action_space,
-        expert_data=expert_data,
+        observation_space=observation_space,
+        action_space=action_space,
+        demonstrations=expert_trajs,
+        demo_batch_size=batch_size,
         l2_weight=l2_weight,
         optimizer_cls=optimizer_cls,
         optimizer_kwargs=optimizer_kwargs,


Test patch:
diff --git a/tests/algorithms/test_adversarial.py b/tests/algorithms/test_adversarial.py
index c9d6f6fa2..11c39f7f4 100644
--- a/tests/algorithms/test_adversarial.py
+++ b/tests/algorithms/test_adversarial.py
@@ -41,42 +41,6 @@ def _algorithm_kwargs(request):
     return dict(request.param)
 
 
-def test_train_disc_small_expert_data_warning(
-    tmpdir,
-    custom_logger,
-    _algorithm_kwargs,
-):
-    venv = util.make_vec_env(
-        "seals/CartPole-v0",
-        n_envs=1,
-        parallel=_parallel,
-    )
-
-    _algorithm_cls = _algorithm_kwargs.pop("algorithm_cls")
-    gen_algo = util.init_rl(venv, verbose=1, **_algorithm_kwargs)
-    small_data = rollout.generate_transitions(gen_algo, venv, n_timesteps=20)
-
-    with pytest.raises(ValueError, match="Transitions.*expert_batch_size"):
-        _algorithm_cls(
-            venv=venv,
-            expert_data=small_data,
-            expert_batch_size=21,
-            gen_algo=gen_algo,
-            log_dir=tmpdir,
-            custom_logger=custom_logger,
-        )
-
-    with pytest.raises(ValueError, match="expert_batch_size.*positive"):
-        _algorithm_cls(
-            venv=venv,
-            expert_data=small_data,
-            expert_batch_size=-1,
-            gen_algo=gen_algo,
-            log_dir=tmpdir,
-            custom_logger=custom_logger,
-        )
-
-
 def test_airl_fail_fast(custom_logger, tmpdir):
     venv = util.make_vec_env(
         "seals/CartPole-v0",
@@ -94,8 +58,8 @@ def test_airl_fail_fast(custom_logger, tmpdir):
     with pytest.raises(TypeError, match="AIRL needs a stochastic policy.*"):
         adversarial.AIRL(
             venv=venv,
-            expert_data=small_data,
-            expert_batch_size=20,
+            demonstrations=small_data,
+            demo_batch_size=20,
             gen_algo=gen_algo,
             log_dir=tmpdir,
             custom_logger=custom_logger,
@@ -162,8 +126,8 @@ def trainer(
     custom_logger = logger.configure(tmpdir, ["tensorboard", "stdout"])
     trainer = _algorithm_cls(
         venv=venv,
-        expert_data=expert_data,
-        expert_batch_size=expert_batch_size,
+        demonstrations=expert_data,
+        demo_batch_size=expert_batch_size,
         gen_algo=gen_algo,
         log_dir=tmpdir,
         custom_logger=custom_logger,
diff --git a/tests/algorithms/test_base.py b/tests/algorithms/test_base.py
index b357ca6fd..c19d3235a 100644
--- a/tests/algorithms/test_base.py
+++ b/tests/algorithms/test_base.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 import pytest
+import torch as th
 
 from imitation.algorithms import base
 from imitation.data import types
@@ -60,3 +61,113 @@ def test_check_fixed_horizon_flag(custom_logger):
     algo._check_fixed_horizon(trajs=gen_trajectories([42], [True]))
     algo._check_fixed_horizon(trajs=gen_trajectories([5, 42], [True, True]))
     assert algo._horizon is None
+
+
+def _make_and_iterate_loader(*args, **kwargs):
+    loader = base.make_data_loader(*args, **kwargs)
+    for batch in loader:
+        pass
+
+
+def test_make_data_loader_batch_size():
+    """Tests data loader performs batch size validation."""
+    for batch_size in [0, -1, -42]:
+        with pytest.raises(ValueError, match=".*must be positive"):
+            base.make_data_loader([], batch_size=batch_size)
+
+    # batch size = 5
+    batch_iterable = [{"obs": np.zeros((5, 2)), "acts": np.zeros((5, 1))}]
+    for wrong_batch_size in [4, 6, 42]:
+        with pytest.raises(ValueError, match="Expected batch size.*"):
+            _make_and_iterate_loader(batch_iterable, batch_size=wrong_batch_size)
+    _make_and_iterate_loader(batch_iterable, batch_size=5)
+
+    batch_iterable2 = [{"obs": np.zeros((5, 2)), "acts": np.zeros((4, 1))}]
+    with pytest.raises(ValueError, match="Expected batch size.*"):
+        _make_and_iterate_loader(batch_iterable2, batch_size=5)
+
+    batch_iterable3 = [
+        {"obs": np.zeros((5, 2)), "acts": np.zeros((5, 1))},
+        {"obs": np.zeros((6, 2)), "acts": np.zeros((5, 1))},
+    ]
+    with pytest.raises(ValueError, match="Expected batch size.*"):
+        _make_and_iterate_loader(batch_iterable3, batch_size=5)
+
+    trans = types.TransitionsMinimal(
+        obs=np.zeros((5, 2)),
+        acts=np.zeros((5, 1)),
+        infos=np.array([{}] * 5),
+    )
+    for smaller_bs in range(1, 6):
+        base.make_data_loader(trans, batch_size=smaller_bs)
+    for larger_bs in [6, 7, 42]:
+        with pytest.raises(ValueError, match=".* smaller than batch size.*"):
+            base.make_data_loader(trans, batch_size=larger_bs)
+
+
+def test_make_data_loader():
+    """Tests data loader produces same results for same input in different formats."""
+    trajs = [
+        types.Trajectory(
+            obs=np.array([0, 1]),
+            acts=np.array([100]),
+            infos=None,
+            terminal=True,
+        ),
+        types.Trajectory(
+            obs=np.array([4, 5, 6]),
+            acts=np.array([102, 103]),
+            infos=None,
+            terminal=True,
+        ),
+        types.Trajectory(
+            obs=np.array([10, 11, 12, 13]),
+            acts=np.array([104, 105, 106]),
+            infos=None,
+            terminal=False,
+        ),
+    ]
+    trans = types.Transitions(
+        obs=np.array([0, 4, 5, 10, 11, 12]),
+        acts=np.array([100, 102, 103, 104, 105, 106]),
+        next_obs=np.array([1, 5, 6, 11, 12, 13]),
+        dones=np.array([True, False, True, False, False, False]),
+        infos=np.array([{}] * 6),
+    )
+    trans_mapping = [
+        {
+            "obs": np.array([0, 4]),
+            "acts": np.array([100, 102]),
+            "next_obs": np.array([1, 5]),
+            "dones": np.array([True, False]),
+            "infos": np.array([{}, {}]),
+        },
+        {
+            "obs": np.array([5, 10]),
+            "acts": np.array([103, 104]),
+            "next_obs": np.array([6, 11]),
+            "dones": np.array([True, False]),
+            "infos": np.array([{}, {}]),
+        },
+        {
+            "obs": np.array([11, 12]),
+            "acts": np.array([105, 106]),
+            "next_obs": np.array([12, 13]),
+            "dones": np.array([False]),
+            "infos": np.array([{}, {}]),
+        },
+    ]
+
+    for data in [trajs, trans, trans_mapping]:
+        data_loader = base.make_data_loader(
+            data,
+            batch_size=2,
+            data_loader_kwargs=dict(shuffle=False, drop_last=False),
+        )
+        for batch, expected_batch in zip(data_loader, trans_mapping):
+            assert batch.keys() == expected_batch.keys()
+            for k in batch.keys():
+                v = batch[k]
+                if isinstance(v, th.Tensor):
+                    v = v.numpy()
+                assert np.all(v == expected_batch[k])
diff --git a/tests/algorithms/test_bc.py b/tests/algorithms/test_bc.py
index df69cffba..d6dfb9918 100644
--- a/tests/algorithms/test_bc.py
+++ b/tests/algorithms/test_bc.py
@@ -39,8 +39,9 @@ def __init__(self, transitions, batch_size):
         self.batch_size = batch_size
 
     def __iter__(self):
-        for start in range(0, len(self.trans), self.batch_size):
-            d = dict(obs=self.trans.obs, acts=self.trans.acts)
+        for start in range(0, len(self.trans) - self.batch_size, self.batch_size):
+            end = start + self.batch_size
+            d = dict(obs=self.trans.obs[start:end], acts=self.trans.acts[start:end])
             d = {k: th.from_numpy(v) for k, v in d.items()}
             yield d
 
@@ -54,6 +55,7 @@ def trainer(batch_size, venv, expert_data_type, custom_logger):
             trans,
             batch_size=batch_size,
             shuffle=True,
+            drop_last=True,
             collate_fn=types.transitions_collate_fn,
         )
     elif expert_data_type == "ducktyped_data_loader":
@@ -64,9 +66,10 @@ def trainer(batch_size, venv, expert_data_type, custom_logger):
         raise ValueError(expert_data_type)
 
     return bc.BC(
-        venv.observation_space,
-        venv.action_space,
-        expert_data=expert_data,
+        observation_space=venv.observation_space,
+        action_space=venv.action_space,
+        demo_batch_size=batch_size,
+        demonstrations=expert_data,
         custom_logger=custom_logger,
     )
 
@@ -74,9 +77,9 @@ def trainer(batch_size, venv, expert_data_type, custom_logger):
 def test_weight_decay_init_error(venv, custom_logger):
     with pytest.raises(ValueError, match=".*weight_decay.*"):
         bc.BC(
-            venv.observation_space,
-            venv.action_space,
-            expert_data=None,
+            observation_space=venv.observation_space,
+            action_space=venv.action_space,
+            demonstrations=None,
             optimizer_kwargs=dict(weight_decay=1e-4),
             custom_logger=custom_logger,
         )
@@ -138,20 +141,22 @@ def test_bc_data_loader_empty_iter_error(venv, no_yield_after_iter, custom_logge
 
     At one point, we entered an updateless infinite loop in this edge case.
     """
+    batch_size = 32
     rollouts = types.load(ROLLOUT_PATH)
     trans = rollout.flatten_trajectories(rollouts)
-    dummy_yield_value = dataclasses.asdict(trans[:3])
+    dummy_yield_value = dataclasses.asdict(trans[:batch_size])
 
     bad_data_loader = _DataLoaderFailsOnSecondIter(
         dummy_yield_value=dummy_yield_value,
         no_yield_after_iter=no_yield_after_iter,
     )
     trainer = bc.BC(
-        venv.observation_space,
-        venv.action_space,
+        observation_space=venv.observation_space,
+        action_space=venv.action_space,
+        demo_batch_size=batch_size,
         custom_logger=custom_logger,
     )
-    trainer.set_expert_data_loader(bad_data_loader)
+    trainer.set_demonstrations(bad_data_loader)
     with pytest.raises(AssertionError, match=".*no data.*"):
         trainer.train(n_batches=20)
 
diff --git a/tests/algorithms/test_dagger.py b/tests/algorithms/test_dagger.py
index 35080eb49..82f62b763 100644
--- a/tests/algorithms/test_dagger.py
+++ b/tests/algorithms/test_dagger.py
@@ -345,8 +345,8 @@ def test_dagger_not_enough_transitions_error(tmpdir, custom_logger):
     venv = util.make_vec_env("CartPole-v0")
     # Initialize with large batch size to ensure error down the line.
     trainer = dagger.DAggerTrainer(
-        venv,
-        tmpdir,
+        venv=venv,
+        scratch_dir=tmpdir,
         batch_size=100_000,
         custom_logger=custom_logger,
     )
diff --git a/tests/algorithms/test_density_baselines.py b/tests/algorithms/test_density_baselines.py
index 508e41742..944bb0a3c 100644
--- a/tests/algorithms/test_density_baselines.py
+++ b/tests/algorithms/test_density_baselines.py
@@ -5,38 +5,29 @@
 import numpy as np
 import pytest
 
-from imitation.algorithms.density_baselines import (
-    STATE_ACTION_DENSITY,
-    STATE_DENSITY,
-    STATE_STATE_DENSITY,
-    DensityReward,
-    DensityTrainer,
-)
+from imitation.algorithms.density import DensityAlgorithm, DensityType
 from imitation.data import rollout, types
 from imitation.policies.base import RandomPolicy
-from imitation.rewards import common
 from imitation.util import util
 
 parametrize_density_stationary = pytest.mark.parametrize(
     "density_type,is_stationary",
-    [
-        (STATE_DENSITY, True),
-        (STATE_DENSITY, False),
-        (STATE_ACTION_DENSITY, True),
-        (STATE_STATE_DENSITY, True),
-    ],
+    [(density_type, True) for density_type in DensityType]
+    + [(DensityType.STATE_DENSITY, False)],
 )
 
 
 def score_trajectories(
     trajectories: Sequence[types.Trajectory],
-    reward_fn: common.RewardFn,
+    density_reward: DensityAlgorithm,
 ):
     # score trajectories under given reward function w/o discount
     returns = []
     for traj in trajectories:
+        dones = np.zeros(len(traj), dtype=np.bool)
+        dones[-1] = True
         steps = np.arange(0, len(traj.acts))
-        rewards = reward_fn(traj.obs[:-1], traj.acts, traj.obs[1:], steps)
+        rewards = density_reward(traj.obs[:-1], traj.acts, traj.obs[1:], dones, steps)
         ret = np.sum(rewards)
         returns.append(ret)
     return np.mean(returns)
@@ -47,7 +38,7 @@ def test_density_reward(density_type, is_stationary):
     # test on Pendulum rather than Cartpole because I don't handle episodes that
     # terminate early yet (see issue #40)
     env_name = "Pendulum-v0"
-    env = util.make_vec_env(env_name, 2)
+    venv = util.make_vec_env(env_name, 2)
 
     # construct density-based reward from expert rollouts
     rollout_path = "tests/testdata/expert_models/pendulum_0/rollouts/final.pkl"
@@ -55,24 +46,24 @@ def test_density_reward(density_type, is_stationary):
     expert_trajectories_all = types.load(rollout_path)[:8]
     n_experts = len(expert_trajectories_all)
     expert_trajectories_train = expert_trajectories_all[: n_experts // 2]
-    reward_fn = DensityReward(
-        trajectories=expert_trajectories_train,
+    reward_fn = DensityAlgorithm(
+        demonstrations=expert_trajectories_train,
         density_type=density_type,
         kernel="gaussian",
-        obs_space=env.observation_space,
-        act_space=env.action_space,
+        venv=venv,
         is_stationary=is_stationary,
         kernel_bandwidth=0.2,
         standardise_inputs=True,
     )
+    reward_fn.train()
 
     # check that expert policy does better than a random policy under our reward
     # function
-    random_policy = RandomPolicy(env.observation_space, env.action_space)
+    random_policy = RandomPolicy(venv.observation_space, venv.action_space)
     sample_until = rollout.make_min_episodes(15)
     random_trajectories = rollout.generate_trajectories(
         random_policy,
-        env,
+        venv,
         sample_until=sample_until,
     )
     expert_trajectories_test = expert_trajectories_all[n_experts // 2 :]
@@ -88,15 +79,13 @@ def test_density_trainer_smoke():
     env_name = "Pendulum-v0"
     rollout_path = "tests/testdata/expert_models/pendulum_0/rollouts/final.pkl"
     rollouts = types.load(rollout_path)[:2]
-    env = util.make_vec_env(env_name, 2)
-    imitation_trainer = util.init_rl(env)
-    density_trainer = DensityTrainer(
-        env,
-        rollouts=rollouts,
-        imitation_trainer=imitation_trainer,
-        density_type=STATE_ACTION_DENSITY,
-        is_stationary=False,
-        kernel="gaussian",
+    venv = util.make_vec_env(env_name, 2)
+    rl_algo = util.init_rl(venv)
+    density_trainer = DensityAlgorithm(
+        demonstrations=rollouts,
+        venv=venv,
+        rl_algo=rl_algo,
     )
+    density_trainer.train()
     density_trainer.train_policy(n_timesteps=2)
     density_trainer.test_policy(n_trajectories=2)
diff --git a/tests/algorithms/test_tabular.py b/tests/algorithms/test_tabular.py
index c41814cfa..e56da85bc 100644
--- a/tests/algorithms/test_tabular.py
+++ b/tests/algorithms/test_tabular.py
@@ -1,22 +1,24 @@
 """Test tabular environments and tabular MCE IRL."""
 
+from typing import Any, Mapping, Type
+
 import gym
 import numpy as np
 import pytest
 import torch as th
+from stable_baselines3.common import vec_env
+from torch.optim import adam
 
-# need direct import since torch.optim.Adam references cause pytype to fail
-from torch.optim.adam import Adam
-
+from imitation.algorithms import base
 from imitation.algorithms.tabular_irl import (
-    LinearRewardModel,
-    MLPRewardModel,
-    mce_irl,
+    MCEIRL,
     mce_occupancy_measures,
     mce_partition_fh,
 )
-from imitation.envs.examples.model_envs import RandomMDP
-from imitation.envs.resettable_env import TabularModelEnv
+from imitation.data import rollout
+from imitation.envs import resettable_env
+from imitation.envs.examples import model_envs
+from imitation.rewards import reward_nets
 from imitation.util.util import tensor_iter_norm
 
 
@@ -50,7 +52,7 @@ def test_random_mdp():
         horizon = 5 * (i + 1)
         random_obs = (i % 2) == 0
         obs_dim = (i * 3 + 4) ** 2 + i
-        mdp = RandomMDP(
+        mdp = model_envs.RandomMDP(
             n_states=n_states,
             n_actions=n_actions,
             branch_factor=branch_factor,
@@ -85,10 +87,15 @@ def test_random_mdp():
         assert len(set(map(str, trajectories))) == 1
 
 
-def test_policy_om_random_mdp():
+FEW_DISCOUNT_RATES = [0.0, 0.99, 1.0]
+DISCOUNT_RATES = FEW_DISCOUNT_RATES + [0.5, 0.9]
+
+
+@pytest.mark.parametrize("discount", DISCOUNT_RATES)
+def test_policy_om_random_mdp(discount: float):
     """Test that optimal policy occupancy measure ("om") for a random MDP is sane."""
     mdp = gym.make("imitation/Random-v0")
-    V, Q, pi = mce_partition_fh(mdp)
+    V, Q, pi = mce_partition_fh(mdp, discount=discount)
     assert np.all(np.isfinite(V))
     assert np.all(np.isfinite(Q))
     assert np.all(np.isfinite(pi))
@@ -96,15 +103,19 @@ def test_policy_om_random_mdp():
     assert np.all(pi >= 0)
     assert np.allclose(np.sum(pi, axis=-1), 1)
 
-    Dt, D = mce_occupancy_measures(mdp, pi=pi)
+    Dt, D = mce_occupancy_measures(mdp, pi=pi, discount=discount)
     assert np.all(np.isfinite(D))
     assert np.any(D > 0)
     # expected number of state visits (over all states) should be equal to the
     # horizon
-    assert np.allclose(np.sum(D), mdp.horizon)
+    if discount == 1.0:
+        expected_sum = mdp.horizon
+    else:
+        expected_sum = (1 - discount ** mdp.horizon) / (1 - discount)
+    assert np.allclose(np.sum(D), expected_sum)
 
 
-class ReasonableMDP(TabularModelEnv):
+class ReasonableMDP(resettable_env.TabularModelEnv):
     observation_matrix = np.array(
         [
             [3, -5, -1, -1, -4, 5, 3, 0],
@@ -186,12 +197,13 @@ class ReasonableMDP(TabularModelEnv):
     horizon = 20
 
 
-def test_policy_om_reasonable_mdp():
+@pytest.mark.parametrize("discount", DISCOUNT_RATES)
+def test_policy_om_reasonable_mdp(discount: float):
     # MDP described above
     mdp = ReasonableMDP()
     # get policy etc. for our MDP
-    V, Q, pi = mce_partition_fh(mdp)
-    Dt, D = mce_occupancy_measures(mdp, pi=pi)
+    V, Q, pi = mce_partition_fh(mdp, discount=discount)
+    Dt, D = mce_occupancy_measures(mdp, pi=pi, discount=discount)
     assert np.all(np.isfinite(V))
     assert np.all(np.isfinite(Q))
     assert np.all(np.isfinite(pi))
@@ -201,7 +213,8 @@ def test_policy_om_reasonable_mdp():
     assert np.allclose(pi[:19, 0, 0], pi[:19, 0, 2])
     # also check that they're by far preferred to action 1 (that goes to state
     # 3, which has poor reward)
-    assert np.all(pi[:19, 0, 0] > 2 * pi[:19, 0, 1])
+    if discount > 0:
+        assert np.all(pi[:19, 0, 0] > 2 * pi[:19, 0, 1])
     # make sure that states 3 & 4 have roughly uniform policies
     pi_34 = pi[:5, 3:5]
     assert np.allclose(pi_34, np.ones_like(pi_34) / 3.0)
@@ -210,18 +223,75 @@ def test_policy_om_reasonable_mdp():
     # check that in state 1, action 2 (which goes to state 4 with certainty) is
     # better than action 0 (which only gets there with some probability), and
     # that both are better than action 1 (which always goes to the bad state).
-    assert np.all(pi[:19, 1, 2] > pi[:19, 1, 0])
-    assert np.all(pi[:19, 1, 0] > pi[:19, 1, 1])
+    if discount > 0:
+        assert np.all(pi[:19, 1, 2] > pi[:19, 1, 0])
+        assert np.all(pi[:19, 1, 0] > pi[:19, 1, 1])
     # check that Dt[0] matches our initial state dist
     assert np.allclose(Dt[0], mdp.initial_state_dist)
 
 
+def test_mce_irl_demo_formats():
+    mdp = model_envs.RandomMDP(
+        n_states=5,
+        n_actions=3,
+        branch_factor=2,
+        horizon=10,
+        random_obs=False,
+        obs_dim=None,
+        generator_seed=42,
+    )
+    venv = vec_env.DummyVecEnv([lambda: mdp])
+    trajs = rollout.generate_trajectories(
+        policy=None,
+        venv=venv,
+        sample_until=rollout.make_min_timesteps(100),
+    )
+
+    demonstrations = {
+        "trajs": trajs,
+        "trans": rollout.flatten_trajectories(trajs),
+        "data_loader": base.make_data_loader(trajs, batch_size=32),
+    }
+
+    final_counts = {}
+    for kind, demo in demonstrations.items():
+        with th.random.fork_rng():
+            th.random.manual_seed(715298)
+
+            reward_net = reward_nets.BasicRewardNet(
+                mdp.observation_space,
+                mdp.action_space,
+                use_action=False,
+                use_next_state=False,
+                use_done=False,
+                hid_sizes=[],
+            )
+            opt = adam.Adam(reward_net.parameters(), lr=1e-2)
+            mce_irl = MCEIRL(demo, mdp, reward_net, opt, linf_eps=1e-3)
+            assert np.allclose(mce_irl.demo_state_om.sum(), 1.0)
+            final_counts[kind] = mce_irl.train(max_iter=5)
+
+        # make sure weights have non-insane norm
+        assert tensor_iter_norm(reward_net.parameters()) < 1000
+
+    for cts in final_counts.values():
+        assert np.allclose(cts, final_counts["trajs"], atol=1e-3, rtol=1e-3)
+
+
 @pytest.mark.expensive
 @pytest.mark.parametrize(
     "model_class,model_kwargs",
-    [(LinearRewardModel, dict()), (MLPRewardModel, dict(hiddens=[32, 32]))],
+    [
+        (reward_nets.BasicRewardNet, dict(hid_sizes=[])),
+        (reward_nets.BasicRewardNet, dict(hid_sizes=[32, 32])),
+    ],
 )
-def test_mce_irl_reasonable_mdp(model_class, model_kwargs):
+@pytest.mark.parametrize("discount", FEW_DISCOUNT_RATES)
+def test_mce_irl_reasonable_mdp(
+    model_class: Type[reward_nets.RewardNet],
+    model_kwargs: Mapping[str, Any],
+    discount: float,
+):
     with th.random.fork_rng():
         th.random.manual_seed(715298)
 
@@ -229,13 +299,21 @@ def test_mce_irl_reasonable_mdp(model_class, model_kwargs):
         mdp = ReasonableMDP()
 
         # demo occupancy measure
-        V, Q, pi = mce_partition_fh(mdp)
-        Dt, D = mce_occupancy_measures(mdp, pi=pi)
+        V, Q, pi = mce_partition_fh(mdp, discount=discount)
+        Dt, D = mce_occupancy_measures(mdp, pi=pi, discount=discount)
 
-        rmodel = model_class(mdp.obs_dim, **model_kwargs)
-        opt = Adam(rmodel.parameters(), lr=1e-2)
-        final_counts = mce_irl(mdp, opt, rmodel, D, linf_eps=1e-3)
+        reward_net = model_class(
+            mdp.observation_space,
+            mdp.action_space,
+            use_action=False,
+            use_next_state=False,
+            use_done=False,
+            **model_kwargs,
+        )
+        opt = adam.Adam(reward_net.parameters(), lr=1e-2)
+        mce_irl = MCEIRL(D, mdp, reward_net, opt, linf_eps=1e-3, discount=discount)
+        final_counts = mce_irl.train()
 
         assert np.allclose(final_counts, D, atol=1e-3, rtol=1e-3)
         # make sure weights have non-insane norm
-        assert tensor_iter_norm(rmodel.parameters()) < 1000
+        assert tensor_iter_norm(reward_net.parameters()) < 1000
diff --git a/tests/data/test_rollout.py b/tests/data/test_rollout.py
index 122144bcc..824eb04b1 100644
--- a/tests/data/test_rollout.py
+++ b/tests/data/test_rollout.py
@@ -236,7 +236,7 @@ def test_compute_returns(gamma):
     # small numerical errors will occur because compute_returns
     # uses a somewhat different method based on evaluating
     # polynomials
-    assert abs(rollout.compute_returns(rewards, gamma) - returns) < 1e-8
+    assert abs(rollout.discounted_sum(rewards, gamma) - returns) < 1e-8
 
 
 def test_generate_trajectories_type_error():