#625 specialized PebbleAgentTrainer to distinguish from old preference comparison trainer

Author: Jan Michelfeit

src/imitation/algorithms/preference_comparisons.py

@@ -33,6 +33,7 @@
 from tqdm.auto import tqdm
 
 from imitation.algorithms import base
+from imitation.algorithms.pebble.entropy_reward import PebbleStateEntropyReward
 from imitation.data import rollout, types, wrappers
 from imitation.data.types import (
     AnyPath,
@@ -329,6 +330,27 @@ def logger(self, value: imit_logger.HierarchicalLogger) -> None:
         self.algorithm.set_logger(self.logger)
 
 
+class PebbleAgentTrainer(AgentTrainer):
+    """
+    Specialization of AgentTrainer for PEBBLE training.
+    Includes unsupervised pretraining with an entropy based reward function.
+    """
+
+    reward_fn: PebbleStateEntropyReward
+
+    def __init__(
+        self,
+        *,
+        reward_fn: PebbleStateEntropyReward,
+        **kwargs,
+    ) -> None:
+        super().__init__(reward_fn=reward_fn, **kwargs)
+
+    def unsupervised_pretrain(self, steps: int, **kwargs: Any) -> None:
+        self.train(steps, **kwargs)
+        self.reward_fn.unsupervised_exploration_finish()
+
+
 def _get_trajectories(
     trajectories: Sequence[TrajectoryWithRew],
     steps: int,
@@ -1705,7 +1727,9 @@ def train(
             self.logger.log(
                 f"Pre-training agent for {unsupervised_pretrain_timesteps} timesteps"
             )
-            self.trajectory_generator.unsupervised_pretrain(unsupervised_pretrain_timesteps)
+            self.trajectory_generator.unsupervised_pretrain(
+                unsupervised_pretrain_timesteps
+            )
 
         for i, num_pairs in enumerate(preference_query_schedule):
             ##########################