Update adversarial algorithm

Author: taufeeque9

src/imitation/algorithms/adversarial/common.py

@@ -2,15 +2,17 @@
 import abc
 import dataclasses
 import logging
-from typing import Callable, Iterable, Iterator, Mapping, Optional, Type, overload
+from typing import Callable, Iterable, Iterator, List, Mapping, Optional, Type, overload
 
 import numpy as np
 import torch as th
 import torch.utils.tensorboard as thboard
 import tqdm
 from stable_baselines3.common import (
     base_class,
+    callbacks,
     distributions,
+    off_policy_algorithm,
     on_policy_algorithm,
     policies,
     vec_env,
@@ -20,6 +22,7 @@
 
 from imitation.algorithms import base
 from imitation.data import buffer, rollout, types, wrappers
+from imitation.policies import replay_buffer_wrapper
 from imitation.rewards import reward_nets, reward_wrapper
 from imitation.util import logger, networks, util
 
@@ -92,6 +95,47 @@ def compute_train_stats(
     }
 
 
+class TrainDiscriminatorCallback(callbacks.BaseCallback):
+    """Callback for training discriminator after collecting rollouts."""
+
+    def __init__(self, adversarial_trainer, *args, **kwargs):
+        """Builds TrainDiscriminatorCallback.
+
+        Args:
+            *args: Passed through to `callbacks.BaseCallback`.
+            **kwargs: Passed through to `callbacks.BaseCallback`.
+        """
+        self.adversarial_trainer = adversarial_trainer
+        self.gen_ctx_manager = None
+        super().__init__(*args, **kwargs)
+
+    def _on_step(self) -> bool:
+        return True
+
+    def _on_rollout_end(self) -> None:
+        gen_trajs, ep_lens = self.adversarial_trainer.venv_buffering.pop_trajectories()
+        self.adversarial_trainer._check_fixed_horizon(ep_lens)
+        gen_samples = rollout.flatten_trajectories_with_rew(gen_trajs)
+        self.adversarial_trainer._gen_replay_buffer.store(gen_samples)
+
+        for _ in range(self.adversarial_trainer.n_disc_updates_per_round):
+            with networks.training(self.adversarial_trainer.reward_train):
+                # switch to training mode (affects dropout, normalization)
+                self.adversarial_trainer.train_disc()
+
+        # update the rollouts with the reward of the latest discriminator
+        self.adversarial_trainer.update_rewards_of_rollouts()
+
+        # This is a hacky way to enable logger.accumulate_means for generator
+        # This is done to avoid nested loggers of discriminator and generator
+        self.gen_ctx_manager = self.adversarial_trainer.logger.accumulate_means("gen")
+        self.gen_ctx_manager.__enter__()
+
+    def _on_training_end(self) -> None:
+        assert self.gen_ctx_manager is not None
+        self.gen_ctx_manager.__exit__(None, None, None)
+
+
 class AdversarialTrainer(base.DemonstrationAlgorithm[types.Transitions]):
     """Base class for adversarial imitation learning algorithms like GAIL and AIRL."""
 
@@ -228,16 +272,22 @@ def __init__(
 
         self.venv_buffering = wrappers.BufferingWrapper(self.venv)
 
+        self.disc_trainer_callback = TrainDiscriminatorCallback(self)
         if debug_use_ground_truth:
             # Would use an identity reward fn here, but RewardFns can't see rewards.
             self.venv_wrapped = self.venv_buffering
-            self.gen_callback = None
+            self.gen_callback: List[callbacks.BaseCallback] = [
+                self.disc_trainer_callback
+            ]
         else:
             self.venv_wrapped = reward_wrapper.RewardVecEnvWrapper(
                 self.venv_buffering,
                 reward_fn=self.reward_train.predict_processed,
             )
-            self.gen_callback = self.venv_wrapped.make_log_callback()
+            self.gen_callback = [
+                self.venv_wrapped.make_log_callback(),
+                self.disc_trainer_callback,
+            ]
         self.venv_train = self.venv_wrapped
 
         self.gen_algo.set_env(self.venv_train)
@@ -314,6 +364,34 @@ def _next_expert_batch(self) -> Mapping:
         assert self._endless_expert_iterator is not None
         return next(self._endless_expert_iterator)
 
+    def update_rewards_of_rollouts(self) -> None:
+        """Updates the rewards of the rollouts using the latest discriminator."""
+        if isinstance(self.gen_algo, on_policy_algorithm.OnPolicyAlgorithm):
+            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
+            )
+            rewards = self._reward_net.predict(**reward_fn_inputs)
+            rewards = rewards.reshape(buffer.rewards.shape)
+            last_values = buffer.advantages[-1] - buffer.rewards[-1] + buffer.values[-1]
+            last_values = last_values / buffer.gamma
+            # here we assume that the actual last_values cannot exactly be 0.0 and so if
+            # last_values is 0.0 then we know that the episode terminated
+            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
+            )
+        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
+            )
+            rewards = self._reward_net.predict(**reward_fn_inputs)
+            buffer.rewards[:] = rewards.reshape(buffer.rewards.shape)
+
     def train_disc(
         self,
         *,
@@ -452,10 +530,6 @@ def train(
         )
         for r in tqdm.tqdm(range(0, n_rounds), desc="round"):
             self.train_gen(self.gen_train_timesteps)
-            for _ in range(self.n_disc_updates_per_round):
-                with networks.training(self.reward_train):
-                    # switch to training mode (affects dropout, normalization)
-                    self.train_disc()
             if callback:
                 callback(r)
             self.logger.dump(self._global_step)

src/imitation/policies/replay_buffer_wrapper.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 from gymnasium import spaces
-from stable_baselines3.common.buffers import ReplayBuffer
+from stable_baselines3.common.buffers import ReplayBuffer, RolloutBuffer
 from stable_baselines3.common.type_aliases import ReplayBufferSamples
 
 from imitation.rewards.reward_function import RewardFn
@@ -23,6 +23,48 @@ def _samples_to_reward_fn_input(
     )
 
 
+def _rollout_buffer_to_reward_fn_input(
+    buffer: RolloutBuffer,
+) -> Mapping[str, np.ndarray]:
+    """Convert a sample from a rollout buffer to a numpy array."""
+    assert buffer.observations is not None
+    assert buffer.actions is not None
+    obs = buffer.observations
+    next_obs = obs[1:]
+    next_obs = np.concatenate([next_obs, obs[-1:]], axis=0)  # last obs not available
+    actions = buffer.actions
+    dones = buffer.episode_starts
+    dones = np.roll(dones, -1, axis=0)
+    dones[-1] = np.ones_like(dones[-1])  # last dones not available
+
+    return dict(
+        state=obs.reshape(-1, *obs.shape[2:]),
+        action=actions.reshape(-1, *actions.shape[2:]),
+        next_state=next_obs.reshape(-1, *next_obs.shape[2:]),
+        done=dones.reshape(-1),
+    )
+
+
+def _replay_buffer_to_reward_fn_input(
+    buffer: ReplayBuffer,
+) -> Mapping[str, np.ndarray]:
+    """Convert a sample from a replay buffer to a numpy array."""
+    assert buffer.observations is not None
+    assert buffer.next_observations is not None
+    assert buffer.actions is not None
+    obs = buffer.observations
+    next_obs = buffer.next_observations
+    actions = buffer.actions
+    dones = buffer.dones
+
+    return dict(
+        state=obs.reshape(-1, *obs.shape[2:]),
+        action=actions.reshape(-1, *actions.shape[2:]),
+        next_state=next_obs.reshape(-1, *next_obs.shape[2:]),
+        done=dones.reshape(-1),
+    )
+
+
 class ReplayBufferRewardWrapper(ReplayBuffer):
     """Relabel the rewards in transitions sampled from a ReplayBuffer."""