6-Vanilla-Policy-Gradient.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise 6: Vanilla Policy Gradient (VPG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import tensorflow as tf\n",
    "import tensorflow.contrib.slim as slim\n",
    "from tensorflow.python.framework import ops\n",
    "\n",
    "from unityenv import UnityEnvironment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "y = 0.99 # Discount rate.\n",
    "total_episodes = 20000 #Set total number of episodes to train agent on.\n",
    "update_frequency = 10 # How many episodes before updating model.\n",
    "learning_rate = 1e-3 # Agent learning rate.\n",
    "hidden_units = 128 # Number of units in hidden layer.\n",
    "summary_path = './summaries/vpg' # Where to save summary statistics.\n",
    "model_path = './models/vpg' # Where to save model checkpoints.\n",
    "load_model = False # Whether to load a saved model.\n",
    "train_model = True # Whether to train the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the Unity environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "env_config = {}\n",
    "env = UnityEnvironment(file_name=\"2DBall\", train_model=train_model, worker_num=5, config=env_config)\n",
    "print(str(env))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examine state space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "_, state = env.reset()\n",
    "print(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "State (s) is a vector whose values corresponds to:\n",
    "* Platform rotation\n",
    "* Ball X position\n",
    "* Ball Y position\n",
    "* Ball X velocity\n",
    "* Ball Y velocity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Actor Critic Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def discount_rewards(r, gamma):\n",
    "    \"\"\" \n",
    "    function from karpathy.github.io/2016/05/31/rl/\n",
    "    take 1D float array of rewards and compute discounted reward \n",
    "    \"\"\"\n",
    "    discounted_r = np.zeros_like(r)\n",
    "    running_add = 0\n",
    "    for t in reversed(range(0, r.size)):\n",
    "        running_add = running_add * gamma + r[t]\n",
    "        discounted_r[t] = running_add\n",
    "    return discounted_r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class VPGAgent():\n",
    "    def __init__(self, lr, s_size,a_size,h_size):\n",
    "        #These lines established the feed-forward part of the network. The agent takes a state and produces an action.\n",
    "        self.state_in= tf.placeholder(shape=[None,s_size],dtype=tf.float32)\n",
    "        self.batch_size = tf.placeholder(shape=None, dtype=tf.int32)\n",
    "        hidden = slim.fully_connected(self.state_in,h_size,biases_initializer=None,activation_fn=tf.nn.elu)\n",
    "        self.out = slim.fully_connected(hidden,a_size,activation_fn=tf.nn.softmax,biases_initializer=None)\n",
    "        self.output = self.out * (0.9) + 0.1/a_size\n",
    "        self.chosen_action = tf.argmax(self.output,1)\n",
    "\n",
    "        #The next six lines establish the training proceedure. We feed the reward and chosen action into the network\n",
    "        #to compute the loss, and use it to update the network.\n",
    "        self.reward_holder = tf.placeholder(shape=[None],dtype=tf.float32)\n",
    "        self.action_holder = tf.placeholder(shape=[None],dtype=tf.int32)\n",
    "        self.actions = slim.one_hot_encoding(self.action_holder, a_size)\n",
    "        \n",
    "        self.responsible_outputs = tf.reduce_sum(self.output * self.actions, axis=1)\n",
    "        self.advantage = self.reward_holder\n",
    "        \n",
    "        self.loss = -tf.reduce_mean(tf.log(self.responsible_outputs)*self.advantage)\n",
    "        \n",
    "        tvars = tf.trainable_variables()\n",
    "        self.gradient_holders = []\n",
    "        for idx,var in enumerate(tvars):\n",
    "            placeholder = tf.placeholder(tf.float32,name=str(idx)+'_holder')\n",
    "            self.gradient_holders.append(placeholder)\n",
    "        \n",
    "        self.gradients = tf.gradients(self.loss,tvars)\n",
    "        \n",
    "        optimizer = tf.train.AdamOptimizer(learning_rate=lr)\n",
    "        self.update_batch = optimizer.apply_gradients(zip(self.gradient_holders,tvars))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Clear the Tensorflow graph.\n",
    "tf.reset_default_graph() \n",
    "\n",
    "#Load the agent.\n",
    "myAgent = VPGAgent(lr=learning_rate, s_size=env.state_space_size, \n",
    "                          a_size=env.action_space_size, h_size=hidden_units) \n",
    "\n",
    "if not os.path.exists(summary_path):\n",
    "    os.makedirs(summary_path)\n",
    "\n",
    "if not os.path.exists(model_path):\n",
    "    os.makedirs(model_path)\n",
    "\n",
    "init = tf.global_variables_initializer()\n",
    "saver = tf.train.Saver()\n",
    "\n",
    "# Launch the tensorflow graph\n",
    "with tf.Session() as sess:\n",
    "    sess.run(init)\n",
    "    if load_model == True:\n",
    "        print('Loading Model...')\n",
    "        ckpt = tf.train.get_checkpoint_state(model_path)\n",
    "        saver.restore(sess,ckpt.model_checkpoint_path)\n",
    "    summary_writer = tf.summary.FileWriter(summary_path)\n",
    "    total_reward = []\n",
    "    total_length = []\n",
    "        \n",
    "    gradBuffer = sess.run(tf.trainable_variables())\n",
    "    for ix,grad in enumerate(gradBuffer):\n",
    "        gradBuffer[ix] = grad * 0\n",
    "    \n",
    "    i = 0\n",
    "    while i < total_episodes:\n",
    "        _, state = env.reset()\n",
    "        running_reward = 0\n",
    "        j = 0\n",
    "        done = False\n",
    "        ep_history = []\n",
    "        while not done:\n",
    "            j += 1\n",
    "            # Probabilistically pick an action given our network outputs.\n",
    "            a_dist = sess.run(myAgent.output, \n",
    "                              feed_dict={myAgent.state_in:[state], myAgent.batch_size: 1})\n",
    "            a = np.random.choice(a_dist[0],p=a_dist[0])\n",
    "            action = np.argmax(a_dist == a)\n",
    "            _, state_1, reward, done = env.step(action, 0) \n",
    "            ep_history.append([state, action, reward, state_1])\n",
    "            state = state_1\n",
    "            running_reward += reward\n",
    "            if done and train_model:\n",
    "                # Update the network.\n",
    "                ep_history = np.array(ep_history)\n",
    "                ep_history[:,2] = discount_rewards(ep_history[:,2], y)\n",
    "                feed_dict={myAgent.reward_holder:ep_history[:,2],\n",
    "                        myAgent.action_holder:ep_history[:,1],\n",
    "                        myAgent.state_in:np.vstack(ep_history[:,0]),\n",
    "                        myAgent.batch_size: len(ep_history)}\n",
    "                p_loss, grads = sess.run([myAgent.loss, myAgent.gradients], feed_dict=feed_dict)\n",
    "                for idx,grad in enumerate(grads):\n",
    "                    gradBuffer[idx] += grad\n",
    "\n",
    "                if i % update_frequency == 0 and i != 0:\n",
    "                    feed_dict = dictionary = dict(zip(myAgent.gradient_holders, gradBuffer))\n",
    "                    _ = sess.run(myAgent.update_batch, feed_dict=feed_dict)\n",
    "                    for ix,grad in enumerate(gradBuffer):\n",
    "                        gradBuffer[ix] = grad * 0\n",
    "                \n",
    "            total_reward.append(running_reward)\n",
    "            total_length.append(j)\n",
    "\n",
    "        # Write summary statistics to tensorboard.\n",
    "        if i % 50 == 0 and i != 0:\n",
    "            summary = tf.Summary()\n",
    "            summary.value.add(tag='Info/Reward', simple_value=float(np.mean(total_reward[-50:])))\n",
    "            summary.value.add(tag='Info/Policy Loss', simple_value=float(p_loss))\n",
    "            summary_writer.add_summary(summary, i)\n",
    "            summary_writer.flush()\n",
    "            print (\"Mean Reward: {}\".format(np.mean(total_reward[-50:])))\n",
    "        # Save agent's model\n",
    "        if i % 1000 == 0 and i != 0:\n",
    "            saver.save(sess, model_path+'/model-'+str(i)+'.cptk')\n",
    "            print(\"Saved Model\")\n",
    "        i += 1\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}