notebooks/density_baseline_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo of simple density-based imitation learning baselines\n",
    "\n",
    "This demo shows how to train a `Pendulum` agent (exciting!) with our simple density-based imitation learning baselines. `DensityTrainer` has a few interesting parameters, but the key ones are:\n",
    "\n",
    "1. `density_type`: this governs whether density is measured on $(s,s')$ pairs (`db.STATE_STATE_DENSITY`), $(s,a)$ pairs (`db.STATE_ACTION_DENSITY`), or single states (`db.STATE_DENSITY`).\n",
    "2. `is_stationary`: determines whether a separate density model is used for each time step $t$ (`False`), or the same model is used for transitions at all times (`True`).\n",
    "3. `standardise_inputs`: if `True`, each dimension of the agent state vectors will be normalised to have zero mean and unit variance over the training dataset. This can be useful when not all elements of the demonstration vector are on the same scale, or when some elements have too wide a variation to be captured by the fixed kernel width (1 for Gaussian kernel).\n",
    "4. `kernel`: changes the kernel used for non-parametric density estimation. `gaussian` and `exponential` are the best bets; see the [sklearn docs](https://scikit-learn.org/stable/modules/density.html#kernel-density) for the rest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "#%load_ext autoreload\n",
    "#%autoreload 2\n",
    "\n",
    "import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from imitation import density_baselines as db\n",
    "from imitation import util"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_id = 'Pendulum-v0'\n",
    "env = util.make_vec_env(env_id, 8)\n",
    "rollouts = util.rollout.load_trajectories(\"../tests/data/expert_models/pendulum_0/rollouts/final.pkl\")\n",
    "imitation_trainer = util.init_rl(env, learning_rate=3e-4, nminibatches=32, noptepochs=10, n_steps=2048)\n",
    "density_trainer = db.DensityTrainer(env,\n",
    "                                    rollouts=rollouts,\n",
    "                                    imitation_trainer=imitation_trainer,\n",
    "                                    density_type=db.STATE_ACTION_DENSITY,\n",
    "                                    is_stationary=False,\n",
    "                                    kernel='gaussian',\n",
    "                                    kernel_bandwidth=0.2,  # found using divination & some palm reading\n",
    "                                    standardise_inputs=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "novice_stats = density_trainer.test_policy()\n",
    "print('Novice stats (true reward function):')\n",
    "pprint.pprint(novice_stats)\n",
    "novice_stats_im = density_trainer.test_policy(true_reward=False)\n",
    "print('Novice stats (imitation reward function):')\n",
    "pprint.pprint(novice_stats_im)\n",
    "\n",
    "for i in range(100):\n",
    "    density_trainer.train_policy(100000)\n",
    "\n",
    "    good_stats = density_trainer.test_policy()\n",
    "    print(f'Trained stats (epoch {i}):')\n",
    "    pprint.pprint(good_stats)\n",
    "    novice_stats_im = density_trainer.test_policy(true_reward=False)\n",
    "    print(f'Trained stats (imitation reward function, epoch {i}):')\n",
    "    pprint.pprint(novice_stats_im)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}