#625 rename unsupervised_agent_pretrain_frac parameter

Author: Jan Michelfeit

src/imitation/algorithms/preference_comparisons.py

@@ -1506,7 +1506,7 @@ def __init__(
         transition_oversampling: float = 1,
         initial_comparison_frac: float = 0.1,
         initial_epoch_multiplier: float = 200.0,
-        initial_agent_pretrain_frac: float = 0.05,
+        unsupervised_agent_pretrain_frac: float = 0.05,
         custom_logger: Optional[imit_logger.HierarchicalLogger] = None,
         allow_variable_horizon: bool = False,
         rng: Optional[np.random.Generator] = None,
@@ -1556,7 +1556,7 @@ def __init__(
             initial_epoch_multiplier: before agent training begins, train the reward
                 model for this many more epochs than usual (on fragments sampled from a
                 random agent).
-            initial_agent_pretrain_frac: fraction of total_timesteps for which the
+            unsupervised_agent_pretrain_frac: fraction of total_timesteps for which the
                 agent will be trained without preference gathering (and reward model
                 training)
             custom_logger: Where to log to; if None (default), creates a new logger.
@@ -1657,7 +1657,7 @@ def __init__(
         self.fragment_length = fragment_length
         self.initial_comparison_frac = initial_comparison_frac
         self.initial_epoch_multiplier = initial_epoch_multiplier
-        self.initial_agent_pretrain_frac = initial_agent_pretrain_frac
+        self.unsupervised_agent_pretrain_frac = unsupervised_agent_pretrain_frac
         self.num_iterations = num_iterations
         self.transition_oversampling = transition_oversampling
         if callable(query_schedule):
@@ -1691,7 +1691,7 @@ def train(
         print(f"Query schedule: {preference_query_schedule}")
 
         (
-            agent_pretrain_timesteps,
+            unsupervised_pretrain_timesteps,
             timesteps_per_iteration,
             extra_timesteps,
         ) = self._compute_timesteps(total_timesteps)
@@ -1703,9 +1703,9 @@ def train(
         ###################################################
         with self.logger.accumulate_means("agent"):
             self.logger.log(
-                f"Pre-training agent for {agent_pretrain_timesteps} timesteps"
+                f"Pre-training agent for {unsupervised_pretrain_timesteps} timesteps"
             )
-            self.trajectory_generator.unsupervised_pretrain(agent_pretrain_timesteps)
+            self.trajectory_generator.unsupervised_pretrain(unsupervised_pretrain_timesteps)
 
         for i, num_pairs in enumerate(preference_query_schedule):
             ##########################
@@ -1782,11 +1782,11 @@ def _preference_gather_schedule(self, total_comparisons):
         return schedule
 
     def _compute_timesteps(self, total_timesteps: int) -> Tuple[int, int, int]:
-        agent_pretrain_timesteps = int(
-            total_timesteps * self.initial_agent_pretrain_frac
+        unsupervised_pretrain_timesteps = int(
+            total_timesteps * self.unsupervised_agent_pretrain_frac
         )
         timesteps_per_iteration, extra_timesteps = divmod(
-            total_timesteps - agent_pretrain_timesteps,
+            total_timesteps - unsupervised_pretrain_timesteps,
             self.num_iterations,
         )
-        return agent_pretrain_timesteps, timesteps_per_iteration, extra_timesteps
+        return unsupervised_pretrain_timesteps, timesteps_per_iteration, extra_timesteps