#625 merge pebble to train_preference_comparisons.py and configure only through sacred

Author: Jan Michelfeit

src/imitation/algorithms/preference_comparisons.py

@@ -344,6 +344,10 @@ def __init__(
         reward_fn: PebbleStateEntropyReward,
         **kwargs,
     ) -> None:
+        if not isinstance(reward_fn, PebbleStateEntropyReward):
+            raise ValueError(
+                f"{self.__class__.__name__} expects {PebbleStateEntropyReward.__name__} reward function"
+            )
         super().__init__(reward_fn=reward_fn, **kwargs)
 
     def unsupervised_pretrain(self, steps: int, **kwargs: Any) -> None:

src/imitation/scripts/config/train_preference_comparisons.py

@@ -1,8 +1,10 @@
 """Configuration for imitation.scripts.train_preference_comparisons."""
 
 import sacred
+import stable_baselines3 as sb3
 
 from imitation.algorithms import preference_comparisons
+from imitation.policies import base
 from imitation.scripts.common import common, reward, rl, train
 
 train_preference_comparisons_ex = sacred.Experiment(
@@ -15,7 +17,6 @@
     ],
 )
 
-
 MUJOCO_SHARED_LOCALS = dict(rl=dict(rl_kwargs=dict(ent_coef=0.1)))
 ANT_SHARED_LOCALS = dict(
     total_timesteps=int(3e7),
@@ -61,6 +62,26 @@ def train_defaults():
     query_schedule = "hyperbolic"
 
 
+@train_preference_comparisons_ex.named_config
+def pebble():
+    # fraction of total_timesteps for training before preference gathering
+    unsupervised_agent_pretrain_frac = 0.05
+    pebble_nearest_neighbor_k = 5
+
+    rl = {
+        "rl_cls": sb3.SAC,
+        "batch_size": 256,  # batch size for RL algorithm
+        "rl_kwargs": {"batch_size": None},  # make sure to set batch size to None
+    }
+    train = {
+        "policy_cls": base.SAC1024Policy,  # noqa: F841
+    }
+    common = {"env_name": "MountainCarContinuous-v0"}
+    allow_variable_horizon = True
+
+    locals()  # quieten flake8
+
+
 @train_preference_comparisons_ex.named_config
 def cartpole():
     common = dict(env_name="CartPole-v1")
@@ -121,6 +142,7 @@ def fast():
     total_timesteps = 50
     total_comparisons = 5
     initial_comparison_frac = 0.2
+    unsupervised_agent_pretrain_frac = 0.2
     num_iterations = 1
     fragment_length = 2
     reward_trainer_kwargs = {

src/imitation/scripts/train_preference_comparisons.py

@@ -82,6 +82,8 @@ def train_preference_comparisons(
     allow_variable_horizon: bool,
     checkpoint_interval: int,
     query_schedule: Union[str, type_aliases.Schedule],
+    unsupervised_agent_pretrain_frac: Optional[float],
+    pebble_nearest_neighbor_k: Optional[int],
 ) -> Mapping[str, Any]:
     """Train a reward model using preference comparisons.
 
@@ -141,6 +143,11 @@ def train_preference_comparisons(
             be allocated to each iteration. "hyperbolic" and "inverse_quadratic"
             apportion fewer queries to later iterations when the policy is assumed
             to be better and more stable.
+        unsupervised_agent_pretrain_frac: fraction of total_timesteps for which the
+                agent will be trained without preference gathering (and reward model
+                training)
+        pebble_nearest_neighbor_k: Parameter for state entropy computation (for PEBBLE
+            training only)
 
     Returns:
         Rollout statistics from trained policy.
@@ -244,6 +251,7 @@ def train_preference_comparisons(
             custom_logger=custom_logger,
             allow_variable_horizon=allow_variable_horizon,
             query_schedule=query_schedule,
+            unsupervised_agent_pretrain_frac=unsupervised_agent_pretrain_frac,
         )
 
         def save_callback(iteration_num):