HW1_Blank.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "HW1 -- Blank.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/naji-s/DeepRL/blob/master/HW1_Blank.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O3lRSKm9LZFQ",
        "colab_type": "text"
      },
      "source": [
        "# 10-703: Homework 1 - Behavior Cloning, DAGGER, CMA-ES, and GAIL\n",
        "\n",
        "You will implement this assignment right here in this Colab notebook. Colab is a Jupyter notebook that runs in the cloud. If you haven't used Colab before, we recommend checking out the following tutorial:\n",
        "https://colab.sandbox.google.com/notebooks/welcome.ipynb.\n",
        "Note that all cells modify the same global state, so imported packages as well as functions and variables declared in one cell will be accessible in other cells.\n",
        "\n",
        "\n",
        "To get started, click the ``Open in Playground`` button in the upper right. Then click the ``Copy to Drive`` button in the upper center to save a copy in your Google drive. In the future, you will be able to find the notebook by looking in your Google drive folder.\n",
        "\n",
        "Now, you're ready start coding. You will want to run each cell in this notebook by clicking the \"play\" button to the left of the cell (or using [ctrl -> enter]. Look for ``WRITE CODE HERE'' to identify places where you need to write some code. Each section involves writing 3 - 10 lines of code. \n",
        "\n",
        "When you're done, copy plots genetated by your code into the solution boxes in the submission LaTeX file released with this assignment. In addition to uploading your PDF submission to GradeScope (one per group), you should also upload your code. To do this, explort the code from this notebook (File -> Export.py), and then upload the .py file to GradeScope.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HDrq4jSWWYx1",
        "colab_type": "text"
      },
      "source": [
        "# Preliminaries\n",
        "In these first few cells, you will implement some compoments that will be used for all problems."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WIXzVYlFLL-m",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "### Setup: Import Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MgHGKD-_iB3r",
        "colab_type": "code",
        "outputId": "5ad0f247-936f-45f0-f0a1-b88cdcaa788c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "from collections import OrderedDict \n",
        "import gym\n",
        "import keras\n",
        "from keras.models import Sequential\n",
        "from keras.layers import Dense\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import subprocess\n",
        "import random\n",
        "import tensorflow as tf"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Using TensorFlow backend.\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5YiASXZnSH7C",
        "colab_type": "text"
      },
      "source": [
        "### Make the TF Model\n",
        "We'll use the same architecture for each of the problems. By implementing a function that creates the model here, you won't need to implement it again for each problem."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Cr7A-CqASErb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_model():\n",
        "  model = Sequential()      \n",
        "  # WRITE CODE HERE\n",
        "  # Add layers to the model:\n",
        "  # a fully connected layer with 10 units\n",
        "  # a tanh activation\n",
        "  # another fully connected layer with 2 units (the number of actions)\n",
        "  # a softmax activation (so the output is a proper distribution)\n",
        "  model.compile(loss='categorical_crossentropy',\n",
        "                     optimizer=tf.train.AdamOptimizer(),\n",
        "                     metrics=['accuracy'])\n",
        "  \n",
        "  # We expect the model to have four weight variables (a kernel and bias for\n",
        "  # both layers)\n",
        "  assert len(model.weights) == 4, 'Model should have 4 weights.'\n",
        "  return model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c0z4oMxZSgq6",
        "colab_type": "text"
      },
      "source": [
        "### Test the model\n",
        "To confirm that the model is correct, we'll use it to solve a binary classification problem. The target function $f: \\mathbb{R}^4 \\rightarrow {0, 1}$ indicates whether the sum of the vector coordinates is positive:\n",
        "$$f(x) = \\delta \\left(\\sum_{i=1}^4 x_i > 0 \\right)$$\n",
        "\n",
        "You should achieve an accuracy of at least 98%."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y1iJByvzSp7_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = make_model()\n",
        "for t in range(20):\n",
        "  X = np.random.normal(size=(1000, 4))  # some random data\n",
        "  is_positive = np.sum(X, axis=1) > 0  # A simple binary function\n",
        "  Y = np.zeros((1000, 2))\n",
        "  Y[np.arange(1000), is_positive.astype(int)] = 1  # one-hot labels\n",
        "  history = model.fit(X, Y, epochs=10, batch_size=256, verbose=0)\n",
        "  loss = history.history['loss'][-1]\n",
        "  acc = history.history['acc'][-1]\n",
        "  print('(%d) loss= %.3f; accuracy = %.1f%%' % (t, loss, 100 * acc))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sw_HvyFnWXa5",
        "colab_type": "text"
      },
      "source": [
        "### Interacting with the Gym\n",
        "Implement the function below for gathering an episode (a \"rollout\"). The environment we will use will implement the OpenAI Gym interface. For documentation, please see the link below:\n",
        "http://gym.openai.com/docs/#environments"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dmRoHiliWdJf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def action_to_one_hot(env, action):\n",
        "    action_vec = np.zeros(env.action_space.n)\n",
        "    action_vec[action] = 1\n",
        "    return action_vec    \n",
        "      \n",
        "      \n",
        "def generate_episode(env, policy):\n",
        "  \"\"\"Collects one rollout from the policy in an environment. The environment\n",
        "  should implement the OpenAI Gym interface. A rollout ends when done=True. The\n",
        "  number of states and actions should be the same, so you should not include\n",
        "  the final state when done=True.\n",
        "\n",
        "  Args:\n",
        "    env: an OpenAI Gym environment.\n",
        "    policy: a keras model\n",
        "  Returns:\n",
        "    states: a list of states visited by the agent.\n",
        "    actions: a list of actions taken by the agent. While the original actions\n",
        "      are discrete, it will be helpful to use a one-hot encoding. The actions\n",
        "      that you return should be one-hot vectors (use action_to_one_hot())\n",
        "    rewards: the reward received by the agent at each step.\n",
        "  \"\"\"\n",
        "  done = False\n",
        "  state = env.reset()\n",
        "\n",
        "  states = []\n",
        "  actions = []\n",
        "  rewards = []\n",
        "  while not done:\n",
        "      # WRITE CODE HERE\n",
        "  return np.array(states), np.array(actions), np.array(rewards)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s9mCrbZDXVvI",
        "colab_type": "text"
      },
      "source": [
        "### Test the data collection\n",
        "Run the following cell and make sure you see \"Test passed!\""
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WCo0B_aDXZfX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create the environment.\n",
        "env = gym.make('CartPole-v0')\n",
        "policy = make_model()\n",
        "states, actions, rewards = generate_episode(env, policy)\n",
        "assert len(states) == len(actions), 'Number of states and actions should be equal.'\n",
        "assert len(actions) == len(rewards), 'Number of actions and rewards should be equal.'\n",
        "assert len(actions[0]) == 2, 'Actions should use one-hot encoding.'\n",
        "print('Test passed!')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k5GAtN56n9hk",
        "colab_type": "text"
      },
      "source": [
        "### Download the expert policy\n",
        "Run the cell below to upload the expert policy (`expert.h5`) to the Colab runtime."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "91G9CWHJtVdR",
        "colab_type": "code",
        "outputId": "c179bb70-79be-4565-ea44-7eb70fcc41d5",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "!wget https://raw.githubusercontent.com/cmudeeprl/703website/master/assets/homework/hw1/expert.h5"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Status: OK\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w8mAIl5xLc6e",
        "colab_type": "text"
      },
      "source": [
        "## Problem 1: Behavior Cloning and DAGGER (50 pt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oYIziDr-VUG7",
        "colab_type": "text"
      },
      "source": [
        "### Implementing Behavior Cloning and DAGGER\n",
        "To implement behavior cloning and DAGGER, fill in the missing blocks of code below. The provided code loads an expert model upon creation of the `Imitation` class. The function `generate_behavior_cloning_data()` fills in `self._train_states` and `self._train_actions` with states and actions from a single episode. Later, when implementing DAGGER, you will finish implementing `generate_dagger_data()`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qn2jQXBNh2WQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Imitation():\n",
        "\n",
        "    def __init__(self, env, num_episodes):\n",
        "        self.env = env\n",
        "        self.expert = tf.keras.models.load_model('expert.h5')\n",
        "        self.num_episodes = num_episodes\n",
        "        \n",
        "        \n",
        "        self.model = make_model()\n",
        "        \n",
        "    def generate_behavior_cloning_data(self):\n",
        "        self._train_states = []\n",
        "        self._train_actions = []\n",
        "        for _ in range(self.num_episodes):\n",
        "            states, actions, rewards = generate_episode(self.env, self.expert)\n",
        "            self._train_states.extend(states)\n",
        "            self._train_actions.extend(actions)\n",
        "        self._train_states = np.array(self._train_states)\n",
        "        self._train_actions = np.array(self._train_actions)\n",
        "        \n",
        "    def generate_dagger_data(self):\n",
        "        # WRITE CODE HERE\n",
        "        # You should collect states and actions from the student policy\n",
        "        # (self.model), and then relabel the actions using the expert policy.\n",
        "        # This method does not return anything.\n",
        "        \n",
        "    def train(self, num_epochs=200):\n",
        "        \"\"\"Trains the model on training data generated by the expert policy.\n",
        "        Args:\n",
        "          env: The environment to run the expert policy on.\n",
        "          num_epochs: number of epochs to train on the data generated by the expert.\n",
        "        Return:\n",
        "          loss: (float) final loss of the trained policy.\n",
        "          acc: (float) final accuracy of the trained policy\n",
        "        \"\"\"\n",
        "        # WRITE CODE HERE\n",
        "        return loss, acc\n",
        "\n",
        "\n",
        "    def evaluate(self, policy, n_episodes=50):\n",
        "        rewards = []\n",
        "        for i in range(n_episodes):\n",
        "            _, _, r = generate_episode(self.env, policy)\n",
        "            rewards.append(sum(r))\n",
        "        r_mean = np.mean(rewards)\n",
        "        return r_mean"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vu7FcPOkAz-c",
        "colab_type": "text"
      },
      "source": [
        "### Experiment: Student vs Expert\n",
        "In the next two cells, you will compare the performance of the expert policy\n",
        "to the imitation policies obtained via behavior cloning and DAGGER."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sRMPf6r2itw3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Uncomment one of the two lines below to select whether to run behavior\n",
        "# cloning or dagger\n",
        "mode = 'behavior cloning'\n",
        "# mode = 'dagger'\n",
        "\n",
        "num_episodes = 100  # Leave this fixed for now. You will experiment with\n",
        "                    # changing it later.\n",
        "num_iterations = 10  # Number of training iterations. Use a small number\n",
        "                     # (e.g., 10) for debugging, and then try a larger number\n",
        "                     # (e.g., 100).\n",
        "\n",
        "# Create the environment.\n",
        "env = gym.make('CartPole-v0')\n",
        "im = Imitation(env, num_episodes)\n",
        "expert_reward = im.evaluate(im.expert)\n",
        "print('Expert reward: %.2f' % expert_reward)\n",
        "\n",
        "loss_vec = []\n",
        "acc_vec = []\n",
        "imitation_reward_vec = []\n",
        "for t in range(num_iterations):\n",
        "  if mode == 'behavior cloning':\n",
        "    im.generate_behavior_cloning_data()\n",
        "  elif mode == 'dagger':\n",
        "    im.generate_dagger_data()\n",
        "  else:\n",
        "    raise ValueError('Unknown mode: %s' % mode)\n",
        "  loss, acc = im.train(num_epochs=1)\n",
        "  imitation_reward = im.evaluate(im.model)\n",
        "  loss_vec.append(loss)\n",
        "  acc_vec.append(acc)\n",
        "  imitation_reward_vec.append(imitation_reward)\n",
        "  print('(%d) loss = %.3f; accuracy = %.2f; reward = %.1f' % (t, loss, acc, imitation_reward))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wacKZfLU1oAC",
        "colab_type": "text"
      },
      "source": [
        "### Plot the results\n",
        "After saving your plots by running `plt.savefig(FILENAME)`, you can download them by navigating to the `Files` tab on the left, and then right-clicking on each filename and selecting `Download`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fcytZUZYmrzc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "### Plot the results\n",
        "plt.figure(figsize=(12, 3))\n",
        "plt.subplot(131)\n",
        "plt.title('Reward')\n",
        "plt.plot(imitation_reward_vec, label='imitation')\n",
        "plt.hlines(expert_reward, 0, len(imitation_reward_vec), label='expert')\n",
        "plt.xlabel('iterations')\n",
        "plt.ylabel('return')\n",
        "plt.legend()\n",
        "plt.ylim([0, None])\n",
        "\n",
        "plt.subplot(132)\n",
        "plt.title('Loss')\n",
        "plt.plot(loss_vec)\n",
        "plt.xlabel('iterations')\n",
        "plt.ylabel('loss')\n",
        "\n",
        "plt.subplot(133)\n",
        "plt.title('Accuracy')\n",
        "plt.plot(acc_vec)\n",
        "plt.xlabel('iterations')\n",
        "plt.ylabel('accuracy')\n",
        "plt.tight_layout()\n",
        "plt.savefig('student_vs_expert_%s.png' % mode, dpi=300)\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d4W6BRtNBGyv",
        "colab_type": "text"
      },
      "source": [
        "### Experiment: How much expert data is needed?\n",
        "This question studies how the amount of expert data effects the performance. You will run the same experiment as above, each time varying the number of expert episodes collected at each iteration. Use values of 1, 10, 50, and 100. You can keep the number of iterations fixed at 100."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3sE8xQFW3ZbL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "random_seeds = 5\n",
        "# Dictionary mapping number of expert trajectories to a list of rewards.\n",
        "# Each is the result of running with a different random seed.\n",
        "reward_data = OrderedDict({  \n",
        "    1: [],\n",
        "    10: [],\n",
        "    50: [],\n",
        "    100: []\n",
        "})\n",
        "accuracy_data = OrderedDict({  \n",
        "    1: [],\n",
        "    10: [],\n",
        "    50: [],\n",
        "    100: []\n",
        "})\n",
        "loss_data = OrderedDict({  \n",
        "    1: [],\n",
        "    10: [],\n",
        "    50: [],\n",
        "    100: []\n",
        "})\n",
        "for num_episodes in [1, 10, 50, 100]:\n",
        "  for t in range(random_seeds):\n",
        "    print('num_episodes: %s; seed: %d' % (num_episodes, t))\n",
        "    # WRITE CODE HERE\n",
        "    # Hint: The code here should be nearly identical to code after the\n",
        "    # \"Student vs Expert\" cell. Feel free to copy and paste."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ec4HEd163il1",
        "colab_type": "text"
      },
      "source": [
        "Plot the reward, loss, and accuracy for each, remembering to label each line."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iQzT0nPc5Odm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "keys = [1, 10, 50, 100]\n",
        "plt.figure(figsize=(12, 4))\n",
        "for (index, (data, name)) in enumerate(zip([reward_data, accuracy_data, loss_data],\n",
        "                                           ['reward', 'accuracy', 'loss'])):\n",
        "  plt.subplot(1, 3, index + 1)\n",
        "  plt.plot(data.keys(), data.values(), '-o')\n",
        "  plt.xlabel('number of expert trajectories', fontsize=16)\n",
        "  plt.ylabel('name', fontsize=16)\n",
        "plt.savefig('expert_data_%s.png' % mode, dpi=300)\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l12OnOMTL214",
        "colab_type": "text"
      },
      "source": [
        "## Problem 2: CMA-ES (25 pt)\n",
        "In this section, you will implement CMA-ES, a black-box optimization algorithm. You will then use CMA-ES to solve a reinforcement learning problem. In particular, the function you will maximize is\n",
        "$$\\max_\\theta J(\\theta) = E_{\\pi_\\theta} \\left[ \\sum_t r(s_t, a_t) \\right]$$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jdGOCHD6NZsr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class CMAES:\n",
        "    def __init__(self, env, L, n, p, sigma, noise, reward_fn=None):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          env: environment with which to evaluate different weights\n",
        "          L: number of episodes used to evaluate a given set of weights\n",
        "          n: number of members (weights) in each generation\n",
        "          p: proportion of members used to update the mean and covariance\n",
        "          sigma: initial std\n",
        "          noise: additional noise to add to covariance\n",
        "          reward_fn: if specified, this reward function is used instead of the \n",
        "            default environment reward. This will be used in Problem 3, when the\n",
        "            reward function will come from the discriminator. The reward\n",
        "            function should be a function of the state and action.\n",
        "        \"\"\"\n",
        "\n",
        "        self.env = env\n",
        "        self.model = make_model()\n",
        "        # Let d be the dimension of the 1d vector of weights.\n",
        "        self.d = sum(int(np.product(w.shape)) for w in self.model.weights)\n",
        "        self.mu = np.zeros(self.d)\n",
        "        self.S = sigma**2 * np.eye(self.d)\n",
        "\n",
        "        self.L = L\n",
        "        self.n = n\n",
        "        self.p = p\n",
        "        self.noise = noise\n",
        "        self.reward_fn = reward_fn\n",
        "\n",
        "\n",
        "    def populate(self):\n",
        "        \"\"\"\n",
        "        Populate a generation using the current estimates of mu and S\n",
        "        \"\"\"\n",
        "        self.population = []\n",
        "        # WRITE CODE HERE\n",
        "\n",
        "    def set_weight(self, member):\n",
        "        ind = 0\n",
        "        weights = []\n",
        "        for w in self.model.weights:\n",
        "          if len(w.shape) > 1:\n",
        "            mat = member[ind:ind+int(w.shape[0]*w.shape[1])]\n",
        "            mat = np.reshape(mat, w.shape)\n",
        "            ind += int(w.shape[0]*w.shape[1])\n",
        "          else:\n",
        "            mat = member[ind:ind+int(w.shape[0])]\n",
        "            ind += int(w.shape[0])\n",
        "          weights.append(mat)\n",
        "\n",
        "        self.model.set_weights(weights)\n",
        "        \n",
        "\n",
        "    def evaluate(self, member, num_episodes):\n",
        "        \"\"\"\n",
        "        Evaluate a set of weights by interacting with the environment and\n",
        "        return the average total reward over num_episodes.\n",
        "        \"\"\"\n",
        "        self.set_weight(member)\n",
        "        return self.evaluate_policy(self.model, num_episodes)\n",
        "    \n",
        "    def evaluate_policy(self, policy, num_episodes):\n",
        "        episode_rewards = []\n",
        "        for episode in range(num_episodes):\n",
        "          # WRITE CODE HERE\n",
        "        return np.mean(episode_rewards)\n",
        "\n",
        "    def train(self):\n",
        "        \"\"\"\n",
        "        Perform an iteration of CMA-ES by evaluating all members of the\n",
        "        current population and then updating mu and S with the top self.p\n",
        "        proportion of members. Note that self.populate automatically deletes\n",
        "        all the old members, so you don't need to worry about deleting the\n",
        "        \"bad\" members.\n",
        "\n",
        "        \"\"\"\n",
        "        self.populate()\n",
        "        # WRITE CODE HERE\n",
        "\n",
        "\n",
        "        best_r = self.evaluate(best_member, 10)\n",
        "        mu_r = self.evaluate(self.mu, 10)\n",
        "        return mu_r, best_r\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tzbc2TqJNk-I",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "iterations = 10  # Use 10 for debugging, and then try 200 once you've got it working.\n",
        "pop_size_vec = [50]             # Start with a population size of 50. Once that\n",
        "# pop_size_vec = [20, 50, 100]  # works, try varying the population size.\n",
        "                              \n",
        "data = {pop_size: [] for pop_size in pop_size_vec}\n",
        "\n",
        "for pop_size in pop_size_vec:\n",
        "  print('Population size:', pop_size)\n",
        "  env = gym.make('CartPole-v0')\n",
        "  optimizer = CMAES(env,\n",
        "                    L=1,  # number of episodes for evaluation\n",
        "                    n=pop_size,  # population size\n",
        "                    p=0.25,  # proportion of population to keep\n",
        "                    sigma=10,  # initial std dev\n",
        "                    noise=0.25)  # noise\n",
        "\n",
        "  for t in range(iterations):\n",
        "      mu_r, best_r = optimizer.train()\n",
        "      data[pop_size].append((mu_r, best_r))\n",
        "      print('(%d) avg member rew = %.2f; best member rew = %.2f' % (t, mu_r, best_r))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kiL2DqEEjZ2i",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "plt.figure(figsize=(12, 4))\n",
        "plt.subplot(121)  # Left plot will show performance vs number of iterations\n",
        "for pop_size, values in data.items():\n",
        "  mu_r = np.array(values)[:, 1]  # Use the performance of the best point\n",
        "  x = np.arange(len(mu_r)) + 1\n",
        "  plt.plot(x, mu_r, label=str(pop_size))\n",
        "  plt.ylabel('average return', fontsize=16)\n",
        "  plt.xlabel('num. iterations', fontsize=16)\n",
        "\n",
        "plt.subplot(122)  # Right plot will show performance vs number of points evaluated\n",
        "for pop_size, values in data.items():\n",
        "  mu_r = np.array(values)[:, 1]  # Use the performance of the best point\n",
        "  x = pop_size * (np.arange(len(mu_r)) + 1)\n",
        "  plt.plot(x, mu_r, label=str(pop_size))\n",
        "  plt.ylabel('average return', fontsize=16)\n",
        "  plt.xlabel('num. points evaluated', fontsize=16)\n",
        "\n",
        "plt.legend()\n",
        "plt.tight_layout()\n",
        "plt.savefig('cmaes_pop_size.png')\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fEYzLBqXL6e-",
        "colab_type": "text"
      },
      "source": [
        "## Problem 3: GAIL (25 pt)\n",
        "For this problem, we will only condition the discriminator on the state, not the action."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XvQ0lfe4L8b2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class GAIL(object):\n",
        "  \n",
        "  def __init__(self, env):\n",
        "    self.env = env\n",
        "    self.expert = tf.keras.models.load_model('expert.h5')\n",
        "    self.model = make_model()\n",
        "    self.discriminator = make_model()\n",
        "    self.cmaes = CMAES(env,\n",
        "                  L=1,  # number of episodes for evaluation\n",
        "                  n=20,  # population size\n",
        "                  p=0.25,  # proportion of population to keep\n",
        "                  sigma=10,  # initial std dev\n",
        "                  noise=0.25,  # noise\n",
        "                  reward_fn=self._reward_fn)\n",
        "    self.expert_states = []\n",
        "    self.student_states = []\n",
        "    \n",
        "  def _reward_fn(self, s, a):\n",
        "    \"\"\"Log probability that state is from expert.\"\"\"\n",
        "    del a\n",
        "    p_expert = self.discriminator.predict(s[None])[0][0]\n",
        "    return np.log(p_expert)\n",
        "    \n",
        "  def collect_data(self, num_episodes):\n",
        "    \"\"\"Collect data from the expert and imitation policy. After the initial\n",
        "    iteration, there is no need to collect new data from the expert, as the\n",
        "    expert policy never changes.\n",
        "    \"\"\"\n",
        "    collect_expert = len(self.expert_states) == 0\n",
        "    self.student_states = []\n",
        "    for _ in range(num_episodes):\n",
        "      # WRITE CODE HERE\n",
        "      # Collect data from the expert policy\n",
        "      # Collect data from the student policy\n",
        "  \n",
        "  def train_discriminator(self):\n",
        "    # WRITE CODE HERE\n",
        "    \n",
        "    assert Y.shape[1] == 2  # Use a 1-hot encoding for the labels\n",
        "    assert np.all(np.sum(Y, axis=1) == 1)\n",
        "    history = self.discriminator.fit(X, Y, epochs=10, batch_size=256, verbose=0)\n",
        "    loss = history.history['loss'][-1]\n",
        "    acc = history.history['acc'][-1]\n",
        "    return loss, acc\n",
        "  \n",
        "  def train_policy(self):\n",
        "    mu_r, best_r = self.cmaes.train()\n",
        "    return mu_r, best_r"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xDbeWwvQI5FG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "### Implement the total variation distance to compare two policies\n",
        "def get_x_position_histogram(states):\n",
        "  x_vec = [s[0] for s in states]  # The x position is the first coordinate\n",
        "  bins = np.linspace(-2.4, 2.4, 11)  # Need 11 edges to make 10 bins\n",
        "  hist, _ = np.histogram(x_vec, bins=bins, density=True)\n",
        "  return hist\n",
        "\n",
        "def TV_distance(expert_states, student_states):\n",
        "  expert_hist = get_x_position_histogram(expert_states)\n",
        "  student_hist = get_x_position_histogram(student_states)\n",
        "  return 0.5 * np.sum(np.abs(expert_hist - student_hist))\n",
        "\n",
        "\n",
        "def evaluate(gail):\n",
        "  \"\"\"Evaluate the policy learned by GAIL according to three metrics:\n",
        "    1. Environment reward. We want this number to be large (~100)\n",
        "    2. How well it fools the discriminator. In particular, we compute the\n",
        "      discriminator's prediction that the policy is the expert. The policy\n",
        "      tries to increase this number, while the discriminator tries to decrease\n",
        "      it. We expect it to be around 40% - 60%\n",
        "    3. Total variation distance between the student and the expert, along the\n",
        "      X axis. We want this number to be small (~0).\"\"\"\n",
        "  rewards_vec = []\n",
        "  p_expert_vec = []\n",
        "  for _ in range(10):\n",
        "    states, actions, rewards = generate_episode(gail.env, gail.cmaes.model)\n",
        "    rewards_vec.append(np.sum(rewards))\n",
        "    log_p_expert = [gail._reward_fn(s, a) for (s, a) in zip(states, actions)]\n",
        "    p_expert = np.exp(log_p_expert)\n",
        "    p_expert_vec.append(np.mean(p_expert))\n",
        "  tv_dist = TV_distance(gail.expert_states, gail.student_states)\n",
        "  return np.mean(p_expert_vec), np.mean(rewards_vec), tv_dist"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J331EVBnbguh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "discriminator_update_period = 10\n",
        "\n",
        "env = gym.make('CartPole-v0')\n",
        "gail = GAIL(env)\n",
        "for t in range(100):\n",
        "  # WRITE CODE HERE\n",
        "\n",
        "  p_expert, avg_r, tv_dist = evaluate(gail)\n",
        "  print('(%d) Policy: p(expert) = %.2f%% ; reward = %.1f' % (t, 100.0 * p_expert, avg_r))\n",
        "  print('(%d) Total variation distance = %.2f' % (t, tv_dist))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jVeqRhMDgPqW",
        "colab_type": "text"
      },
      "source": [
        "# You're Done!"
      ]
    }
  ]
}