Reduce training time and improve expert loading code in the tutorials (#810)

* Use normal env name when loading policies in the tutorial notebooks.

* Use load_policy from the utils to load a policy in the density tutorial notebook.

* Reduce the number of train steps in the FAST mode of the AIRL tutorial notebook.

Author: ernestum

docs/tutorials/1_train_bc.ipynb

@@ -47,7 +47,7 @@
     "expert = load_policy(\n",
     "    \"ppo-huggingface\",\n",
     "    organization=\"HumanCompatibleAI\",\n",
-    "    env_name=\"seals-CartPole-v0\",\n",
+    "    env_name=\"seals/CartPole-v0\",\n",
     "    venv=env,\n",
     ")"
    ]

docs/tutorials/2_train_dagger.ipynb

@@ -39,7 +39,7 @@
     "expert = load_policy(\n",
     "    \"ppo-huggingface\",\n",
     "    organization=\"HumanCompatibleAI\",\n",
-    "    env_name=\"seals-CartPole-v0\",\n",
+    "    env_name=\"seals/CartPole-v0\",\n",
     "    venv=env,\n",
     ")"
    ]

docs/tutorials/3_train_gail.ipynb

@@ -44,7 +44,7 @@
     "expert = load_policy(\n",
     "    \"ppo-huggingface\",\n",
     "    organization=\"HumanCompatibleAI\",\n",
-    "    env_name=\"seals:seals/CartPole-v0\",\n",
+    "    env_name=\"seals/CartPole-v0\",\n",
     "    venv=env,\n",
     ")"
    ]

docs/tutorials/4_train_airl.ipynb

@@ -23,7 +23,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import seals  # noqa: F401  # needed to load \"seals/\" environments\n",
     "import numpy as np\n",
     "from imitation.policies.serialize import load_policy\n",
     "from imitation.util.util import make_vec_env\n",
@@ -34,12 +33,12 @@
     "FAST = True\n",
     "\n",
     "if FAST:\n",
-    "    N_RL_TRAIN_STEPS = 300_000\n",
+    "    N_RL_TRAIN_STEPS = 100_000\n",
     "else:\n",
     "    N_RL_TRAIN_STEPS = 2_000_000\n",
     "\n",
     "venv = make_vec_env(\n",
-    "    \"seals/CartPole-v0\",\n",
+    "    \"seals:seals/CartPole-v0\",\n",
     "    rng=np.random.default_rng(SEED),\n",
     "    n_envs=8,\n",
     "    post_wrappers=[\n",
@@ -49,7 +48,7 @@
     "expert = load_policy(\n",
     "    \"ppo-huggingface\",\n",
     "    organization=\"HumanCompatibleAI\",\n",
-    "    env_name=\"seals-CartPole-v0\",\n",
+    "    env_name=\"seals/CartPole-v0\",\n",
     "    venv=venv,\n",
     ")"
    ]

docs/tutorials/7_train_density.ipynb

@@ -56,9 +56,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from imitation.policies.serialize import load_policy\n",
     "from stable_baselines3.common.policies import ActorCriticPolicy\n",
     "from stable_baselines3 import PPO\n",
-    "from huggingface_sb3 import load_from_hub\n",
     "from imitation.data import rollout\n",
     "from stable_baselines3.common.vec_env import DummyVecEnv\n",
     "from stable_baselines3.common.evaluation import evaluate_policy\n",
@@ -70,12 +70,15 @@
     "\n",
     "rng = np.random.default_rng(seed=SEED)\n",
     "env_name = \"Pendulum-v1\"\n",
-    "expert = PPO.load(\n",
-    "    load_from_hub(\"HumanCompatibleAI/ppo-Pendulum-v1\", \"ppo-Pendulum-v1.zip\")\n",
-    ").policy\n",
     "rollout_env = DummyVecEnv(\n",
     "    [lambda: RolloutInfoWrapper(gym.make(env_name)) for _ in range(N_VEC)]\n",
     ")\n",
+    "expert = load_policy(\n",
+    "    \"ppo-huggingface\",\n",
+    "    organization=\"HumanCompatibleAI\",\n",
+    "    env_name=env_name,\n",
+    "    venv=rollout_env,\n",
+    ")\n",
     "rollouts = rollout.rollout(\n",
     "    expert,\n",
     "    rollout_env,\n",