Support asynchronous human preference gathering in RLHP implementation

ci/clean_notebooks.py

@@ -18,6 +18,7 @@ class UncleanNotebookError(Exception):
     "metadata": {"do": "constant", "value": dict()},
     "source": {"do": "keep"},
     "id": {"do": "keep"},
+    "attachments": {"do": "constant", "value": {}},
 }
 
 code_structure: Dict[str, Dict[str, Any]] = {
@@ -76,7 +77,8 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
             if key not in structure[cell["cell_type"]]:
                 if check_only:
                     raise UncleanNotebookError(
-                        f"Notebook {file} has unknown cell key {key}",
+                        f"Notebook {file} has unknown cell key {key} for cell type "
+                        + f"{cell['cell_type']}",
                     )
                 del cell[key]
                 was_dirty = True
@@ -108,7 +110,12 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
 
 
 def parse_args():
-    """Parse command-line arguments."""
+    """Parse command-line arguments.
+
+    Returns:
+        parser: The parser object.
+        args: The parsed arguments.
+    """
     # if the argument --check has been passed, check if the notebooks are clean
     # otherwise, clean them in-place
     parser = argparse.ArgumentParser()
@@ -125,7 +132,14 @@ def parse_args():
 
 
 def get_files(input_paths: List):
-    """Build list of files to scan from list of paths and files."""
+    """Build list of files to scan from list of paths and files.
+
+    Args:
+        input_paths: List of paths and files to scan.
+
+    Returns:
+        files: List of files to scan.
+    """
     files = []
     for file in input_paths:
         if file.is_dir():

docs/index.rst

@@ -86,6 +86,7 @@ If you use ``imitation`` in your research project, please cite our paper to help
    tutorials/4_train_airl
    tutorials/5_train_preference_comparisons
    tutorials/5a_train_preference_comparisons_with_cnn
+   tutorials/5b_train_preference_comparisons_with_synchronous_human_feedback
    tutorials/6_train_mce
    tutorials/7_train_density
    tutorials/8_train_sqil

docs/tutorials/5b_train_preference_comparisons_with_synchronous_human_feedback.ipynb

@@ -0,0 +1,251 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[download this notebook here](https://github.com/HumanCompatibleAI/imitation/blob/master/docs/tutorials/5_train_preference_comparisons.ipynb)\n",
+    "# Learning a Reward Function using Preference Comparisons with Synchronous Human Feedback\n",
+    "\n",
+    "You can request human feedback via synchronous CLI or Notebook interactions as well. The setup is only slightly different than it would be with a synthetic preference gatherer."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Here's the starting setup. The major differences from the synthetic setup are indicated with comments"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pathlib\n",
+    "import random\n",
+    "import tempfile\n",
+    "from imitation.algorithms import preference_comparisons\n",
+    "from imitation.rewards.reward_nets import BasicRewardNet\n",
+    "from imitation.util import video_wrapper\n",
+    "from imitation.util.networks import RunningNorm\n",
+    "from imitation.util.util import make_vec_env\n",
+    "from imitation.policies.base import FeedForward32Policy, NormalizeFeaturesExtractor\n",
+    "import gym\n",
+    "from stable_baselines3 import PPO\n",
+    "import numpy as np\n",
+    "\n",
+    "# Add a temporary directory for video recordings of trajectories. Unfortunately Jupyter\n",
+    "# won't play videos outside the current directory, so we have to put them here. We'll\n",
+    "# delete them at the end of the script.\n",
+    "video_dir = tempfile.mkdtemp(dir=\".\", prefix=\"videos_\")\n",
+    "\n",
+    "rng = np.random.default_rng(0)\n",
+    "\n",
+    "# Add a video wrapper to the environment. This will record videos of the agent's\n",
+    "# trajectories so we can review them later.\n",
+    "venv = make_vec_env(\n",
+    "    \"Pendulum-v1\",\n",
+    "    rng=rng,\n",
+    "    post_wrappers={\n",
+    "        \"VideoWrapper\": video_wrapper.video_wrapper_factory(\n",
+    "            pathlib.Path(video_dir), single_video=False\n",
+    "        )\n",
+    "    },\n",
+    ")\n",
+    "\n",
+    "reward_net = BasicRewardNet(\n",
+    "    venv.observation_space, venv.action_space, normalize_input_layer=RunningNorm\n",
+    ")\n",
+    "\n",
+    "fragmenter = preference_comparisons.RandomFragmenter(\n",
+    "    warning_threshold=0,\n",
+    "    rng=rng,\n",
+    ")\n",
+    "\n",
+    "querent = preference_comparisons.PreferenceQuerent()\n",
+    "\n",
+    "# This gatherer will show the user (you!) pairs of trajectories and ask it to choose\n",
+    "# which one is better. It will then use the user's feedback to train the reward network.\n",
+    "gatherer = preference_comparisons.SynchronousHumanGatherer(video_dir=video_dir)\n",
+    "\n",
+    "preference_model = preference_comparisons.PreferenceModel(reward_net)\n",
+    "reward_trainer = preference_comparisons.BasicRewardTrainer(\n",
+    "    preference_model=preference_model,\n",
+    "    loss=preference_comparisons.CrossEntropyRewardLoss(),\n",
+    "    epochs=3,\n",
+    "    rng=rng,\n",
+    ")\n",
+    "\n",
+    "agent = PPO(\n",
+    "    policy=FeedForward32Policy,\n",
+    "    policy_kwargs=dict(\n",
+    "        features_extractor_class=NormalizeFeaturesExtractor,\n",
+    "        features_extractor_kwargs=dict(normalize_class=RunningNorm),\n",
+    "    ),\n",
+    "    env=venv,\n",
+    "    seed=0,\n",
+    "    n_steps=2048 // venv.num_envs,\n",
+    "    batch_size=64,\n",
+    "    ent_coef=0.0,\n",
+    "    learning_rate=0.0003,\n",
+    "    n_epochs=10,\n",
+    ")\n",
+    "\n",
+    "trajectory_generator = preference_comparisons.AgentTrainer(\n",
+    "    algorithm=agent,\n",
+    "    reward_fn=reward_net,\n",
+    "    venv=venv,\n",
+    "    exploration_frac=0.0,\n",
+    "    rng=rng,\n",
+    ")\n",
+    "\n",
+    "pref_comparisons = preference_comparisons.PreferenceComparisons(\n",
+    "    trajectory_generator,\n",
+    "    reward_net,\n",
+    "    num_iterations=5,\n",
+    "    fragmenter=fragmenter,\n",
+    "    preference_querent=querent,\n",
+    "    preference_gatherer=gatherer,\n",
+    "    reward_trainer=reward_trainer,\n",
+    "    fragment_length=100,\n",
+    "    transition_oversampling=1,\n",
+    "    initial_comparison_frac=0.1,\n",
+    "    allow_variable_horizon=False,\n",
+    "    initial_epoch_multiplier=1,\n",
+    ")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We're going to train with only 20 comparisons to make it faster for you to evaluate. The videos will appear in-line in this notebook for you to watch, and a text input will appear for you to choose one."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pref_comparisons.train(\n",
+    "    total_timesteps=5_000,  # For good performance this should be 1_000_000\n",
+    "    total_comparisons=20,  # For good performance this should be 5_000\n",
+    ")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "From this point onward, this notebook is the same as [the synthetic gatherer notebook](5_train_preference_comparisons.ipynb).\n",
+    "\n",
+    "After we trained the reward network using the preference comparisons algorithm, we can wrap our environment with that learned reward."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from imitation.rewards.reward_wrapper import RewardVecEnvWrapper\n",
+    "\n",
+    "\n",
+    "learned_reward_venv = RewardVecEnvWrapper(venv, reward_net.predict)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can train an agent, that only sees those learned reward."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from stable_baselines3 import PPO\n",
+    "from stable_baselines3.ppo import MlpPolicy\n",
+    "\n",
+    "learner = PPO(\n",
+    "    policy=MlpPolicy,\n",
+    "    env=learned_reward_venv,\n",
+    "    seed=0,\n",
+    "    batch_size=64,\n",
+    "    ent_coef=0.0,\n",
+    "    learning_rate=0.0003,\n",
+    "    n_epochs=10,\n",
+    "    n_steps=64,\n",
+    ")\n",
+    "learner.learn(1000)  # Note: set to 100000 to train a proficient expert"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Then we can evaluate it using the original reward."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from stable_baselines3.common.evaluation import evaluate_policy\n",
+    "\n",
+    "reward, _ = evaluate_policy(learner.policy, venv, 10)\n",
+    "print(reward)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# clean up the videos we made\n",
+    "import shutil\n",
+    "\n",
+    "shutil.rmtree(video_dir)"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "439158cd89905785fcc749928062ade7bfccc3f087fab145e5671f895c635937"
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

setup.py

@@ -55,6 +55,8 @@
         "wandb==0.12.21",
         "setuptools_scm~=7.0.5",
         "pre-commit>=2.20.0",
+        "types-requests~=2.31.0.1",
+        "requests-mock~=1.11.0",
     ]
     + PARALLEL_REQUIRE
     + ATARI_REQUIRE
@@ -209,6 +211,7 @@ def get_local_version(version: "ScmVersion", time_format="%Y%m%d") -> str:
         "huggingface_sb3~=3.0",
         "optuna>=3.0.1",
         "datasets>=2.8.0",
+        "opencv-python",  # TODO: specify version
     ],
     tests_require=TESTS_REQUIRE,
     extras_require={