HW2_Policy_Graident.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Automatically reload changes to external code\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this assignment, you will solve a classic control problem - CartPole using policy gradient methods.\n",
    "\n",
    "First, you will implement the \"vanilla\" policy gradient method, i.e., a method that repeatedly computes **unbiased** estimates $\\hat{g}$ of $\\nabla_{\\theta} E[\\sum_t r_t]$ and takes gradient ascent steps $\\theta \\rightarrow \\theta + \\epsilon \\hat{g}$ so as to increase the total rewards collected in each episode. To make sure our code can solve multiple MDPs with different policy parameterizations, provided code follows an OOP manner and represents MDP and Policy as classes.\n",
    "\n",
    "The following code constructs an instance of the MDP using OpenAI gym."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import gym\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from policy_gradient import util\n",
    "from policy_gradient.policy import CategoricalPolicy\n",
    "from policy_gradient.baselines.linear_feature_baseline import LinearFeatureBaseline\n",
    "\n",
    "np.random.seed(0)\n",
    "tf.set_random_seed(0)\n",
    "\n",
    "# CartPole-v0 is a MDP with finite state and action space. \n",
    "# In this environment, A pendulum is attached by an un-actuated joint to a cart, \n",
    "# and the goal is to prevent it from falling over. You can apply a force of +1 or -1 to the cart.\n",
    "# A reward of +1 is provided for every timestep that the pendulum remains upright. \n",
    "# To visualize CartPole-v0, please see https://gym.openai.com/envs/CartPole-v0\n",
    "\n",
    "env = gym.make('CartPole-v0')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 1: construct a neural network to represent policy\n",
    "\n",
    "Make sure you know how to construct neural network using tensorflow.\n",
    "\n",
    "1. Open **homework2/policy_gradient/policy.py**.\n",
    "2. Follow the instruction of Problem 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 2: compute the surrogate loss\n",
    "\n",
    "If there are $N$ episodes in an iteration, then for $i$ th episode we define $R_t^i = \\sum_{{t^′}=t}^T \\gamma^{{t^′}-t}r(s_{t^′}, a_{t^′})$ as the accumulated discounted rewards from timestep $t$ to the end of that episode, where $\\gamma$ is the discount rate.\n",
    "\n",
    "The pseudocode for the REINFORCE algorithm is as below:\n",
    "\n",
    "1. Initialize policy $\\pi$ with parameter $\\theta_1$.\n",
    "2. For iteration $k = 1, 2, ...$:\n",
    "    * Sample N episodes $\\tau_1, \\tau_2, ..., \\tau_N$ under the current policy $\\theta_k$, where $\\tau_i =(s_i^t,a_i^t,R_i^t)_{t=0}^{T−1}$. Note that the last state is dropped since no action is taken after observing the last state.\n",
    "    * Compute the empirical policy gradient using formula: $$\\hat{g} = E_{\\pi_\\theta}[\\nabla_{\\theta} log\\pi_\\theta(a_t^i | s_t^i) R_t^i]$$\n",
    "    * Take a gradient step: $\\theta_{k+1} = \\theta_k + \\epsilon \\hat{g}$.\n",
    "    \n",
    "    \n",
    "Note that we can transform the policy gradient formula as\n",
    "\n",
    "$$\\hat{g} = \\nabla_{\\theta} \\frac{1}{(NT)}(\\sum_{i=1}^N \\sum_{t=0}^T log\\pi_\\theta(a_t^i | s_t^i) R_t^i)$$\n",
    "\n",
    "and $L(\\theta) = \\frac{1}{(NT)}(\\sum_{i=1}^N \\sum_{t=0}^T log\\pi_\\theta(a_t^i | s_t^i) R_t^i)$ is called the surrogate loss. \n",
    "\n",
    "We can first construct the computation graph for $L(\\theta)$, and then take its gradient as the empirical policy gradient.\n",
    "\n",
    "\n",
    "1. Open **homework2/policy_gradient/policy.py**.\n",
    "2. Follow the instruction of Problem 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sess = tf.Session()\n",
    "\n",
    "# Construct a neural network to represent policy which maps observed state to action. \n",
    "in_dim = util.flatten_space(env.observation_space)\n",
    "out_dim = util.flatten_space(env.action_space)\n",
    "hidden_dim = 8\n",
    "\n",
    "opt = tf.train.AdamOptimizer(learning_rate=0.01)\n",
    "policy = CategoricalPolicy(in_dim, out_dim, hidden_dim, opt, sess)\n",
    "\n",
    "sess.run(tf.initialize_all_variables())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem 3\n",
    "\n",
    "Implement a function that computes the accumulated discounted rewards of each timestep _t_ from _t_ to the end of the episode.\n",
    "\n",
    "For example:\n",
    "\n",
    "```python\n",
    "rewards = [1, 1, 1]\n",
    "discount_rate = 0.99\n",
    "util.discount_cumsum(rewards, discount_rate)\n",
    "```\n",
    "\n",
    "should return:\n",
    "\n",
    "`array([ 2.9701,  1.99  ,  1.    ])`\n",
    "\n",
    "1. Open **homework/policy_gradient/util.py**.\n",
    "2. Implement the commented function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem 4\n",
    "\n",
    "Use baseline to reduce the variance of our gradient estimate.\n",
    "\n",
    "1. Fill in the function `process_paths` of class `PolicyOptimizer` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class PolicyOptimizer(object):\n",
    "    def __init__(self, env, policy, baseline, n_iter, n_episode, path_length,\n",
    "        discount_rate=.99):\n",
    "\n",
    "        self.policy = policy\n",
    "        self.baseline = baseline\n",
    "        self.env = env\n",
    "        self.n_iter = n_iter\n",
    "        self.n_episode = n_episode\n",
    "        self.path_length = path_length\n",
    "        self.discount_rate = discount_rate\n",
    "\n",
    "    def sample_path(self):\n",
    "        obs = []\n",
    "        actions = []\n",
    "        rewards = []\n",
    "        ob = self.env.reset()\n",
    "\n",
    "        for _ in range(self.path_length):\n",
    "            a = self.policy.act(ob.reshape(1, -1))\n",
    "            next_ob, r, done, _ = self.env.step(a)\n",
    "            obs.append(ob)\n",
    "            actions.append(a)\n",
    "            rewards.append(r)\n",
    "            ob = next_ob\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "        return dict(\n",
    "            observations=np.array(obs),\n",
    "            actions=np.array(actions),\n",
    "            rewards=np.array(rewards),\n",
    "        )\n",
    "\n",
    "    def process_paths(self, paths):\n",
    "        for p in paths:\n",
    "            if self.baseline != None:\n",
    "                b = self.baseline.predict(p)\n",
    "            else:\n",
    "                b = 0\n",
    "            \n",
    "            # `p[\"rewards\"]` is a matrix contains the rewards of each timestep in a sample path\n",
    "            r = util.discount_cumsum(p[\"rewards\"], self.discount_rate)\n",
    "            \n",
    "            \"\"\"\n",
    "            Problem 4:\n",
    "\n",
    "            1. Variable `b` is the reward predicted by our baseline\n",
    "            2. Use it to reduce variance and then assign the result to the variable `a`\n",
    "\n",
    "            Sample solution should be only 1 line.\n",
    "            \"\"\"\n",
    "            # YOUR CODE HERE >>>>>>\n",
    "            # a = ???\n",
    "            # <<<<<<<<\n",
    "\n",
    "            p[\"returns\"] = r\n",
    "            p[\"baselines\"] = b\n",
    "            p[\"advantages\"] = (a - a.mean()) / (a.std() + 1e-8) # normalize\n",
    "\n",
    "        obs = np.concatenate([ p[\"observations\"] for p in paths ])\n",
    "        actions = np.concatenate([ p[\"actions\"] for p in paths ])\n",
    "        rewards = np.concatenate([ p[\"rewards\"] for p in paths ])\n",
    "        advantages = np.concatenate([ p[\"advantages\"] for p in paths ])\n",
    "\n",
    "        return dict(\n",
    "            observations=obs,\n",
    "            actions=actions,\n",
    "            rewards=rewards,\n",
    "            advantages=advantages,\n",
    "        )\n",
    "\n",
    "    def train(self):\n",
    "        for i in range(1, self.n_iter + 1):\n",
    "            paths = []\n",
    "            for _ in range(self.n_episode):\n",
    "                paths.append(self.sample_path())\n",
    "            data = self.process_paths(paths)\n",
    "            loss = self.policy.train(data[\"observations\"], data[\"actions\"], data[\"advantages\"])\n",
    "            avg_return = np.mean([sum(p[\"rewards\"]) for p in paths])\n",
    "            print(\"Iteration {}: Average Return = {}\".format(i, avg_return))\n",
    "            \n",
    "            # CartPole-v0 defines \"solving\" as getting average reward of 195.0 over 100 consecutive trials.\n",
    "            if avg_return >= 195:\n",
    "                print(\"Solve at {} iterations, which equals {} episodes.\".format(i, i*100))\n",
    "                break\n",
    "\n",
    "            if self.baseline != None:\n",
    "                self.baseline.fit(paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "n_iter = 200\n",
    "n_episode = 100\n",
    "path_length = 200\n",
    "discount_rate = 0.99\n",
    "baseline = LinearFeatureBaseline(env.spec)\n",
    "\n",
    "po = PolicyOptimizer(env, policy, baseline, n_iter, n_episode, path_length,\n",
    "                     discount_rate)\n",
    "\n",
    "# Train the policy optimizer\n",
    "po.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Verify your solutions\n",
    "\n",
    "if you solve the problems 1~4 correctly, your will solve CartPole with roughly ~ 80 iterations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem 5\n",
    "Replacing line \n",
    "\n",
    "`baseline = LinearFeatureBaseline(env.spec)` \n",
    "\n",
    "with \n",
    "\n",
    "`baseline = None`\n",
    "\n",
    "can remove the baseline.\n",
    "\n",
    "Modify the code to compare the variance and performance before and after adding baseline.\n",
    "Then, write a report about your findings. (with figures is better)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem 6\n",
    "\n",
    "In function process_paths of class `PolicyOptimizer`, why we need to normalize the advantages? i.e., what's the usage of this line:\n",
    "\n",
    "`p[\"advantages\"] = (a - a.mean()) / (a.std() + 1e-8)`\n",
    "\n",
    "Include the answer in your report."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}