Add support to download rollouts from huggingface (#629)

* Add support to download rollouts from huggingface

* Add check for rollout_path

* Fix type error

* Add test cases

* Incorporate reviewer's suggestions

* Change pkl to npz

* Add input validation for rollout_type

* Fix error

* Update src/imitation/scripts/common/demonstrations.py

Co-authored-by: Adam Gleave <adam@gleave.me>

Co-authored-by: Adam Gleave <adam@gleave.me>

Author: taufeeque9

src/imitation/policies/serialize.py

@@ -230,3 +230,14 @@ def _on_step(self) -> bool:
         output_dir = self.policy_dir / f"{self.num_timesteps:012d}"
         save_stable_model(output_dir, self.model)
         return True
+
+
+def load_rollouts_from_huggingface(
+    algo_name: str,
+    env_name: str,
+    organization: str = "HumanCompatibleAI",
+) -> str:
+    model_name = hfsb3.ModelName(algo_name, hfsb3.EnvironmentName(env_name))
+    repo_id = hfsb3.ModelRepoId(organization, model_name)
+    filename = hfsb3.load_from_hub(repo_id, "rollouts.npz")
+    return filename

src/imitation/scripts/ingredients/demonstrations.py

@@ -1,12 +1,15 @@
 """Ingredient for scripts learning from demonstrations."""
 
 import logging
-from typing import Optional, Sequence
+import pathlib
+import warnings
+from typing import Optional, Sequence, Union
 
 import numpy as np
 import sacred
 
 from imitation.data import rollout, types
+from imitation.policies import serialize
 from imitation.scripts.ingredients import environment, expert
 from imitation.scripts.ingredients import logging as logging_ingredient
 
@@ -23,7 +26,9 @@
 
 @demonstrations_ingredient.config
 def config():
-    # path to file containing rollouts. If None, they are sampled from the expert.
+    rollout_type = "local"
+    # path to file containing rollouts. If rollout_path is None
+    # and rollout_type is local, they are sampled from the expert.
     rollout_path = None
     n_expert_demos = None  # Num demos used or sampled. None loads every demo possible.
     locals()  # quieten flake8
@@ -38,10 +43,39 @@ def fast():
 
 @demonstrations_ingredient.capture
 def get_expert_trajectories(
+    rollout_type: str,
     rollout_path: str,
 ) -> Sequence[types.Trajectory]:
+    """Loads expert demonstrations.
+
+    Args:
+        rollout_type: Can be either `local` to load rollouts from the disk or to
+            generate them locally or of the format `{algo}-huggingface` to load
+            from the huggingface hub of expert trained using `{algo}`.
+        rollout_path: A path containing a pickled sequence of `types.Trajectory`.
+
+    Returns:
+        The expert trajectories.
+
+    Raises:
+        ValueError: if `rollout_type` is not "local" or of the form {algo}-huggingface.
+    """
+    if rollout_type.endswith("-huggingface"):
+        if rollout_path is not None:
+            warnings.warn(
+                "Ignoring `rollout_path` since `rollout_type` is set to download the "
+                "rollouts from the huggingface-hub. If you want to load the rollouts "
+                'from disk, set `rollout_type`="local" and the path in `rollout_path`.',
+                RuntimeWarning,
+            )
+        rollout_path = _download_expert_rollouts(rollout_type)
+    elif rollout_type != "local":
+        raise ValueError(
+            "`rollout_type` can either be `local` or of the form `{algo}-huggingface`.",
+        )
+
     if rollout_path is not None:
-        return load_expert_trajs()
+        return load_local_expert_trajs(rollout_path)
     else:
         return generate_expert_trajs()
 
@@ -77,11 +111,11 @@ def generate_expert_trajs(
 
 
 @demonstrations_ingredient.capture
-def load_expert_trajs(
-    rollout_path: str,
+def load_local_expert_trajs(
+    rollout_path: Union[str, pathlib.Path],
     n_expert_demos: Optional[int],
 ) -> Sequence[types.Trajectory]:
-    """Loads expert demonstrations.
+    """Loads expert demonstrations from a local path.
 
     Args:
         rollout_path: A path containing a pickled sequence of `types.Trajectory`.
@@ -105,3 +139,14 @@ def load_expert_trajs(
         expert_trajs = expert_trajs[:n_expert_demos]
         logger.info(f"Truncated to {n_expert_demos} expert trajectories")
     return expert_trajs
+
+
+@demonstrations_ingredient.capture(prefix="expert")
+def _download_expert_rollouts(rollout_type, loader_kwargs):
+    assert rollout_type.endswith("-huggingface")
+    algo_name = rollout_type.split("-")[0]
+    return serialize.load_rollouts_from_huggingface(
+        algo_name,
+        env_name=loader_kwargs["env_name"],
+        organization=loader_kwargs["organization"],
+    )

tests/scripts/test_scripts.py

@@ -313,6 +313,54 @@ def test_train_bc_main_with_none_demonstrations_raises_value_error(tmpdir):
         )
 
 
+def test_train_bc_main_with_demonstrations_from_huggingface(tmpdir):
+    train_imitation.train_imitation_ex.run(
+        command_name="bc",
+        named_configs=["seals_cartpole"] + ALGO_FAST_CONFIGS["imitation"],
+        config_updates=dict(
+            logging=dict(log_root=tmpdir),
+            demonstrations=dict(rollout_type="ppo-huggingface"),
+        ),
+    )
+
+
+def test_train_bc_main_with_demonstrations_raises_error_on_wrong_huggingface_format(
+    tmpdir,
+):
+    with pytest.raises(
+        ValueError,
+        match="`rollout_type` can either be `local` or of the form .*-huggingface.S*",
+    ):
+        train_imitation.train_imitation_ex.run(
+            command_name="bc",
+            named_configs=["seals_cartpole"] + ALGO_FAST_CONFIGS["imitation"],
+            config_updates=dict(
+                logging=dict(log_root=tmpdir),
+                demonstrations=dict(rollout_type="huggingface-ppo"),
+            ),
+        )
+
+
+def test_train_bc_main_with_demonstrations_warns_setting_rollout_type(
+    tmpdir,
+):
+    with pytest.warns(
+        RuntimeWarning,
+        match="Ignoring `rollout_path` .*",
+    ):
+        train_imitation.train_imitation_ex.run(
+            command_name="bc",
+            named_configs=["seals_cartpole"] + ALGO_FAST_CONFIGS["imitation"],
+            config_updates=dict(
+                logging=dict(log_root=tmpdir),
+                demonstrations=dict(
+                    rollout_type="ppo-huggingface",
+                    rollout_path="path",
+                ),
+            ),
+        )
+
+
 @pytest.fixture(
     params=[
         "expert_from_path",