Fix test errors

src/imitation/algorithms/adversarial/common.py

@@ -101,6 +101,8 @@ def __init__(self, adversarial_trainer, *args, **kwargs):
         """Builds TrainDiscriminatorCallback.
 
         Args:
+            adversarial_trainer: The AdversarialTrainer instance in which
+                this callback will be called.
             *args: Passed through to `callbacks.BaseCallback`.
             **kwargs: Passed through to `callbacks.BaseCallback`.
         """
@@ -276,7 +278,7 @@ def __init__(
             # Would use an identity reward fn here, but RewardFns can't see rewards.
             self.venv_wrapped = self.venv_buffering
             self.gen_callback: List[callbacks.BaseCallback] = [
-                self.disc_trainer_callback
+                self.disc_trainer_callback,
             ]
         else:
             self.venv_wrapped = reward_wrapper.RewardVecEnvWrapper(
@@ -369,7 +371,7 @@ def update_rewards_of_rollouts(self) -> None:
             buffer = self.gen_algo.rollout_buffer
             assert buffer is not None
             reward_fn_inputs = replay_buffer_wrapper._rollout_buffer_to_reward_fn_input(
-                self.gen_algo.rollout_buffer
+                self.gen_algo.rollout_buffer,
             )
             rewards = self._reward_net.predict(**reward_fn_inputs)
             rewards = rewards.reshape(buffer.rewards.shape)
@@ -380,13 +382,14 @@ def update_rewards_of_rollouts(self) -> None:
             last_dones = last_values == 0.0
             self.gen_algo.rollout_buffer.rewards[:] = rewards
             self.gen_algo.rollout_buffer.compute_returns_and_advantage(
-                th.tensor(last_values), last_dones
+                th.tensor(last_values),
+                last_dones,
             )
         elif isinstance(self.gen_algo, off_policy_algorithm.OffPolicyAlgorithm):
             buffer = self.gen_algo.replay_buffer
             assert buffer is not None
             reward_fn_inputs = replay_buffer_wrapper._replay_buffer_to_reward_fn_input(
-                buffer
+                buffer,
             )
             rewards = self._reward_net.predict(**reward_fn_inputs)
             buffer.rewards[:] = rewards.reshape(buffer.rewards.shape)
@@ -465,13 +468,15 @@ def train_disc(
 
         return train_stats
 
-    def train_gen(
+    def train_gen_with_disc(
         self,
         total_timesteps: Optional[int] = None,
         learn_kwargs: Optional[Mapping] = None,
     ) -> None:
         """Trains the generator to maximize the discriminator loss.
 
+        The discriminator is also trained after the rollouts are collected and before
+        the generator is trained.
         After the end of training populates the generator replay buffer (used in
         discriminator training) with `self.disc_batch_size` transitions.
 
@@ -502,7 +507,7 @@ def train(
     ) -> None:
         """Alternates between training the generator and discriminator.
 
-        Every "round" consists of a call to `train_gen(self.gen_train_timesteps)`,
+        Every "round" consists of a call to `train_gen_with_disc(self.gen_train_timesteps)`,
         a call to `train_disc`, and finally a call to `callback(round)`.
 
         Training ends once an additional "round" would cause the number of transitions
@@ -522,7 +527,7 @@ def train(
             f"total_timesteps={total_timesteps})!"
         )
         for r in tqdm.tqdm(range(0, n_rounds), desc="round"):
-            self.train_gen(self.gen_train_timesteps)
+            self.train_gen_with_disc(self.gen_train_timesteps)
             if callback:
                 callback(r)
             self.logger.dump(self._global_step)
@@ -610,7 +615,8 @@ def _make_disc_train_batches(
         if gen_samples is None:
             if self._gen_replay_buffer.size() == 0:
                 raise RuntimeError(
-                    "No generator samples for training. " "Call `train_gen()` first.",
+                    "No generator samples for training. "
+                    "Call `train_gen_with_disc()` first.",
                 )
             gen_samples_dataclass = self._gen_replay_buffer.sample(batch_size)
             gen_samples = types.dataclass_quick_asdict(gen_samples_dataclass)

tests/algorithms/test_adversarial.py

@@ -231,8 +231,9 @@ def test_train_gen_train_disc_no_crash(
     trainer_parametrized: common.AdversarialTrainer,
     n_updates: int = 2,
 ) -> None:
-    trainer_parametrized.train_gen(n_updates * trainer_parametrized.gen_train_timesteps)
-    trainer_parametrized.train_disc()
+    trainer_parametrized.train_gen_with_disc(
+        n_updates * trainer_parametrized.gen_train_timesteps
+    )
 
 
 @pytest.fixture