7_3_TRPO_TF.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "eUwrWBxIEGT_",
    "outputId": "b9787a8f-9fe6-4553-d426-4027814f489b"
   },
   "outputs": [],
   "source": [
    "# # in google colab uncomment this\n",
    "\n",
    "import os\n",
    "\n",
    "# os.system('apt-get install -y xvfb')\n",
    "# os.system('wget https://raw.githubusercontent.com/yandexdataschool/Practical_DL/fall18/xvfb -O ../xvfb')\n",
    "# os.system('apt-get install -y python-opengl ffmpeg')\n",
    "# os.system('pip install pyglet==1.2.4')\n",
    "\n",
    "# # launch XVFB if you run on a server\n",
    "# import os\n",
    "# if type(os.environ.get(\"DISPLAY\")) is not str or len(os.environ.get(\"DISPLAY\")) == 0:\n",
    "#     !bash ../xvfb start\n",
    "#     %env DISPLAY = : 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9XgoNxq0EGUR"
   },
   "source": [
    "### Let's make a TRPO!\n",
    "\n",
    "In this notebook we will write the code of the one Trust Region Policy Optimization.\n",
    "As usually, it contains a few different parts which we are going to reproduce.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bBNgf5s8EGUU"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "collapsed": true,
    "id": "lsVSlNjYEGUe",
    "outputId": "23d41679-c270-4b16-9a30-e529923aa5ae"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Observation Space Box(6,)\n",
      "Action Space Discrete(3)\n"
     ]
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "env = gym.make(\"Acrobot-v1\")\n",
    "env.reset()\n",
    "observation_shape = env.observation_space.shape\n",
    "n_actions = env.action_space.n\n",
    "print(\"Observation Space\", env.observation_space)\n",
    "print(\"Action Space\", env.action_space)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 287
    },
    "colab_type": "code",
    "collapsed": true,
    "id": "7VUzWFQHEGUz",
    "outputId": "3998209c-054a-4364-be1f-1c7d62fceb4e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x253a0b50e10>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQsAAAD8CAYAAABgtYFHAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAADlFJREFUeJzt3X+o3Xd9x/Hna+kP3RRj22sJSSQVw9A/ttpcaqRjuFZH7cT0jxYqMoMEApsDxYFLN9gQ9ofuD1uEoYZVFofadv6goXRzJW0Z+8PaG/vD1qz2Kp29pJhIf7ghulXf++N8rh6T29xPcu/5lTwfcDif7+f7Oee8Tzl59fP9fj/3nFQVkrSa35h0AZJmg2EhqYthIamLYSGpi2EhqYthIanLSMIiybVJnkyymGTfKF5D0nhlvddZJNkAfBd4J7AEPAS8t6q+s64vJGmsRjGzuBJYrKrvV9X/ArcDu0bwOpLG6LwRPOdm4Jmh7SXgrad6wCWXXFLbtm0bQSmSlh0+fPhHVTV3po8fRVhkhb6TjnWS7AX2Arz+9a9nYWFhBKVIWpbkv9by+FEchiwBW4e2twBHTxxUVfurar6q5ufmzjjsJI3JKMLiIWB7ksuSXADcBBwcwetIGqN1PwypqpeS/BnwdWAD8LmqemK9X0fSeI3inAVVdQ9wzyieW9JkuIJTUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlKXVcMiyeeSHEvy+FDfRUnuTfJUu39t60+STyVZTPJYkitGWbyk8emZWfwjcO0JffuAQ1W1HTjUtgHeBWxvt73Ap9enTEmTtmpYVNW/A8+d0L0LONDaB4Drh/o/XwPfADYm2bRexUqanDM9Z3FpVT0L0O5f1/o3A88MjVtqfSdJsjfJQpKF48ePn2EZksZlvU9wZoW+WmlgVe2vqvmqmp+bm1vnMiSttzMNix8uH160+2OtfwnYOjRuC3D0zMuTNC3ONCwOArtbezdw11D/+9tVkZ3Ai8uHK5Jm23mrDUjyJeDtwCVJloC/AT4O3JlkD/AD4MY2/B7gOmAR+AnwgRHULGkCVg2Lqnrvy+y6ZoWxBXxwrUVJmj6u4JTUxbCQ1MWwkNTFsJDUxbCQ1MWwkNTFsJDUxbCQ1MWwkNQlg0WXEy4imXwR0tnvcFXNn+mDV13uPQ47duxgYWFh0mVIZ7VkpW+Q6OdhiKQuhoWkLoaFpC6GhaQuhoWkLoaFpC6GhaQuhoWkLoaFpC6GhaQuhoWkLoaFpC6GhaQuhoWkLoaFpC6GhaQuhoWkLoaFpC6GhaQuhoWkLoaFpC6GhaQuhoWkLquGRZKtSe5PciTJE0k+1PovSnJvkqfa/Wtbf5J8KslikseSXDHqNyFp9HpmFi8Bf15VbwJ2Ah9M8mZgH3CoqrYDh9o2wLuA7e22F/j0ulctaexWDYuqeraqvtXa/w0cATYDu4ADbdgB4PrW3gV8vga+AWxMsmndK5c0Vqd1ziLJNuAtwIPApVX1LAwCBXhdG7YZeGboYUutT9IM6w6LJK8CvgJ8uKp+fKqhK/Sd9MPHSfYmWUiycPz48d4yJE1IV1gkOZ9BUHyhqr7aun+4fHjR7o+1/iVg69DDtwBHT3zOqtpfVfNVNT83N3em9Usak56rIQFuA45U1SeHdh0Edrf2buCuof73t6siO4EXlw9XJM2u8zrGXAX8MfDtJI+0vr8EPg7cmWQP8APgxrbvHuA6YBH4CfCBda1Y0kSsGhZV9R+sfB4C4JoVxhfwwTXWJWnKuIJTUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlIXw0JSF8NCUhfDQlKXnj8kk37p8OFf/zOhHTtO+qoSnaWcWajbiUHxcn06OxkW6nKqUDAwzg2GhVbVEwYGxtnPsJDUxbCQ1MWwkNTFsNCq5llYlzGabYaFupwqDAyKc4NhoW4rhYJBce5wBadOi+Fw7nJmIamLYSGpi2EhqYthIamLYSGpi2EhqYthIamLYSGpi2EhqYthIamLYSGpi2EhqYthIanLqmGR5BVJvpnk0SRPJPlY678syYNJnkpyR5ILWv+FbXux7d822rcgaRx6ZhY/A66uqt8FLgeuTbIT+ARwS1VtB54H9rTxe4Dnq+qNwC1tnKQZt2pY1MD/tM3z262Aq4Evt/4DwPWtvatt0/Zfk8TviZdmXNc5iyQbkjwCHAPuBb4HvFBVL7UhS8Dm1t4MPAPQ9r8IXLzCc+5NspBk4fjx42t7F5JGrissqurnVXU5sAW4EnjTSsPa/UqziJN+ELOq9lfVfFXNz83N9dYraUJO62pIVb0APADsBDYmWf5avi3A0dZeArYCtP2vAZ5bj2IlTU7P1ZC5JBtb+5XAO4AjwP3ADW3YbuCu1j7Ytmn776sqf2pbmnE9X9i7CTiQZAODcLmzqu5O8h3g9iR/CzwM3NbG3wb8U5JFBjOKm0ZQt6QxWzUsquox4C0r9H+fwfmLE/t/Cty4LtVJmhqu4JTUxbCQ1MWwkNTFsJDUxbCQ1MWwkNTFsJDUxbCQ1MWwkNTFsJDUxbCQ1MWwkNTFsJDUxbCQ1MWwkNTFsJDUxbDQaVtgngXmJ12Gxqzna/UkgJMCYnl7noVJlKMxc2ahLqeaSTjLODcYFpK6GBZaVc/MwdnF2c+wkNTFsJDUxbDQqnqudnhF5OxnWEjqYlioy6lmDs4qzg0uylK3U4VC7dgxxko0Cc4sdEo5fHjSJWhKGBaSuhgWkroYFpK6GBaSuhgWkroYFpK6dIdFkg1JHk5yd9u+LMmDSZ5KckeSC1r/hW17se3fNprSJY3T6cwsPgQcGdr+BHBLVW0Hngf2tP49wPNV9UbgljZO0ozrCoskW4A/Av6hbQe4GvhyG3IAuL61d7Vt2v5r2nhJM6x3ZnEr8FHgF237YuCFqnqpbS8Bm1t7M/AMQNv/YhuvGdO7etOl3ueGVcMiybuBY1U1/MlZaaZQHfuGn3dvkoUkC8ePH+8qVtLk9MwsrgLek+Rp4HYGhx+3AhuTLP8h2hbgaGsvAVsB2v7XAM+d+KRVtb+q5qtqfm5ubk1vQtLorRoWVXVzVW2pqm3ATcB9VfU+4H7ghjZsN3BXax9s27T991XVSTMLSbNlLess/gL4SJJFBuckbmv9twEXt/6PAPvWVqKkaXBa32dRVQ8AD7T294ErVxjzU+DGdahN0hRxBaekLoaFpC6GhaQuhoWkLoaFpC6GhaQuhoWkLoaFpC6GhaQuhoWkLoaF1sTvsjh3GBZakT9bqBMZFpK6GBaSuhgWkroYFpK6GBaSuhgWkroYFpK6GBaSuhgWkroYFjqJP1uolRgWkroYFpK6GBaSuhgWkroYFpK6GBaSupzWDyPr3OAlUa3EmYWkLoaFpC6GhaQuhoWkLoaFpC6GhaQuXWGR5Okk307ySJKF1ndRknuTPNXuX9v6k+RTSRaTPJbkilG+AUnjcToziz+oqsurar5t7wMOVdV24FDbBngXsL3d9gKfXq9iJU3OWg5DdgEHWvsAcP1Q/+dr4BvAxiSb1vA6kqZA7wrOAv4tSQGfrar9wKVV9SxAVT2b5HVt7GbgmaHHLrW+Z4efMMleBjMPgJ8lefwM38MkXAL8aNJFdJqlWmG26p2lWgF+ey0P7g2Lq6rqaAuEe5P85ynGZoW+OqljEDj7AZIsDB3eTL1ZqneWaoXZqneWaoVBvWt5fNdhSFUdbffHgK8BVwI/XD68aPfH2vAlYOvQw7cAR9dSpKTJWzUskvxWklcvt4E/BB4HDgK727DdwF2tfRB4f7sqshN4cflwRdLs6jkMuRT4WpLl8V+sqn9N8hBwZ5I9wA+AG9v4e4DrgEXgJ8AHOl5j/+kWPmGzVO8s1QqzVe8s1QprrDdVJ51OkKSTuIJTUpeJh0WSa5M82VZ87lv9ESOv53NJjg1fyp3m1apJtia5P8mRJE8k+dC01pzkFUm+meTRVuvHWv9lSR5std6R5ILWf2HbXmz7t42r1qGaNyR5OMndM1DraFdaV9XEbsAG4HvAG4ALgEeBN0+4pt8HrgAeH+r7O2Bfa+8DPtHa1wH/wuBy8U7gwQnUuwm4orVfDXwXePM01txe81WtfT7wYKvhTuCm1v8Z4E9a+0+Bz7T2TcAdE/jv+xHgi8DdbXuaa30auOSEvnX7HIz1zazw5t4GfH1o+2bg5knW1OrYdkJYPAlsau1NwJOt/VngvSuNm2DtdwHvnPaagd8EvgW8lcHCpvNO/EwAXwfe1trntXEZY41bGPwpw9XA3e0f1lTW2l53pbBYt8/BpA9DXm6157T5tdWqwGqrVSeiTX3fwuD/2FNZc5vWP8JgXc69DGaWL1TVSyvU88ta2/4XgYvHVStwK/BR4Bdt+2Kmt1b41Urrw22FNKzj52DSX9jbtdpzik1N/UleBXwF+HBV/bhd6l5x6Ap9Y6u5qn4OXJ5kI4MFfm86RT0TqzXJu4FjVXU4yds76pmGz8K6r7QeNumZxays9pzq1apJzmcQFF+oqq+27qmuuapeAB5gcLy8Mcny/7iG6/llrW3/a4DnxlTiVcB7kjwN3M7gUOTWKa0VGP1K60mHxUPA9naG+QIGJ4YOTrimlUztatUMphC3AUeq6pNDu6au5iRzbUZBklcC7wCOAPcDN7xMrcvv4QbgvmoH2KNWVTdX1Zaq2sbgc3lfVb1vGmuFMa20HucJmJc5KXMdgzP43wP+agrq+RKDv5D9Pwbpu4fBsech4Kl2f1EbG+DvW+3fBuYnUO/vMZg+PgY80m7XTWPNwO8AD7daHwf+uvW/Afgmg1W//wxc2Ppf0bYX2/43TOgz8XZ+dTVkKmttdT3abk8s/1taz8+BKzgldZn0YYikGWFYSOpiWEjqYlhI6mJYSOpiWEjqYlhI6mJYSOry/2SX+95XYuZuAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "plt.imshow(env.render('rgb_array'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uyCQlVo2EGU9"
   },
   "source": [
    "### Step 1: Defining a network\n",
    "\n",
    "With all it's complexity, at it's core TRPO is yet another policy gradient method. \n",
    "\n",
    "This essentially means we're actually training a stochastic policy $ \\pi_\\theta(a|s) $. \n",
    "\n",
    "And yes, it's gonna be a neural network. So let's start by defining one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "id": "9pPGepy3EGVA"
   },
   "outputs": [],
   "source": [
    "# input tensors\n",
    "observations_ph = tf.placeholder(\n",
    "    shape=(None, observation_shape[0]), dtype=tf.float32)\n",
    "# Actions that we made\n",
    "actions_ph = tf.placeholder(shape=(None,), dtype=tf.int32)\n",
    "# \"G = r + gamma*r' + gamma^2*r'' + ...\"\n",
    "cummulative_returns_ph = tf.placeholder(shape=(None,), dtype=tf.float32)\n",
    "# Action probabilities from previous iteration\n",
    "old_probs_ph = tf.placeholder(shape=(None, n_actions), dtype=tf.float32)\n",
    "\n",
    "all_inputs = [observations_ph, actions_ph,\n",
    "              cummulative_returns_ph, old_probs_ph]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 92
    },
    "colab_type": "code",
    "id": "ciXnYXFuEGVJ",
    "outputId": "131f34f9-c448-4b68-f196-8f8303d79617"
   },
   "outputs": [],
   "source": [
    "def denselayer(name, x, out_dim, nonlinearity=None):\n",
    "    with tf.variable_scope(name):\n",
    "        if nonlinearity is None:\n",
    "            nonlinearity = tf.identity\n",
    "\n",
    "        x_shape = x.get_shape().as_list()\n",
    "\n",
    "        w = tf.get_variable('w', shape=[x_shape[1], out_dim])\n",
    "        b = tf.get_variable(\n",
    "            'b', shape=[out_dim], initializer=tf.constant_initializer(0))\n",
    "        o = nonlinearity(tf.matmul(x, w) + b)\n",
    "\n",
    "        return o\n",
    "\n",
    "\n",
    "sess = tf.InteractiveSession()\n",
    "\n",
    "nn = observations_ph\n",
    "\n",
    "lin1 = denselayer(\"dense1\", nn, 32, nonlinearity=tf.nn.relu);\n",
    "lin2 = denselayer(\"dense2\", lin1, n_actions, nonlinearity=tf.nn.log_softmax)\n",
    "\n",
    "policy_out = lin2;\n",
    "\n",
    "probs_out = tf.exp(policy_out)\n",
    "\n",
    "weights = tf.trainable_variables()\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QJlCVVV0EGVT"
   },
   "source": [
    "### Step 2: Actions and rollouts\n",
    "\n",
    "In this section, we'll define functions that take actions $ a \\sim \\pi_\\theta(a|s) $ and rollouts $ <s_0,a_0,s_1,a_1,s_2,a_2,...s_n,a_n> $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vmTba4zTEGVX"
   },
   "outputs": [],
   "source": [
    "# compile function\n",
    "\n",
    "\n",
    "def act(obs, sample=True):\n",
    "    \"\"\"\n",
    "    Samples action from policy distribution (sample = True) or takes most likely action (sample = False)\n",
    "    :param: obs - single observation vector\n",
    "    :param sample: if True, samples from \\pi, otherwise takes most likely action\n",
    "    :returns: action (single integer) and probabilities for all actions\n",
    "    \"\"\"\n",
    "\n",
    "    probs = sess.run(probs_out, feed_dict={\n",
    "                     observations_ph: obs.reshape((1, -1))})[0]\n",
    "\n",
    "    if sample:\n",
    "        action = int(np.random.choice(n_actions, p=probs))\n",
    "    else:\n",
    "        action = int(np.argmax(probs))\n",
    "\n",
    "    return action, probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 73
    },
    "colab_type": "code",
    "id": "EzCSm0cfEGVk",
    "outputId": "95a0e5ac-3dc3-47dc-d2ee-c5b3402fe6a2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sampled: [(0, array([0.44120154, 0.28181925, 0.27697915], dtype=float32)), (1, array([0.44336593, 0.27932698, 0.27730703], dtype=float32)), (0, array([0.44128194, 0.28797308, 0.27074504], dtype=float32)), (0, array([0.4459713 , 0.26784635, 0.28618237], dtype=float32)), (0, array([0.44373965, 0.2739944 , 0.282266  ], dtype=float32))]\n",
      "greedy: [(0, array([0.44161856, 0.28442848, 0.273953  ], dtype=float32)), (0, array([0.44412243, 0.28011602, 0.2757615 ], dtype=float32)), (0, array([0.4407704 , 0.2819702 , 0.27725932], dtype=float32)), (0, array([0.4409779 , 0.27940676, 0.27961534], dtype=float32)), (0, array([0.44314262, 0.27710864, 0.27974877], dtype=float32))]\n"
     ]
    }
   ],
   "source": [
    "# demo\n",
    "print(\"sampled:\", [act(env.reset()) for _ in range(5)])\n",
    "print(\"greedy:\", [act(env.reset(), sample=False) for _ in range(5)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bxOp7_57EGVz"
   },
   "source": [
    "Compute cummulative reward just like you did in vanilla REINFORCE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kRTjBJzxEGV1"
   },
   "outputs": [],
   "source": [
    "import scipy.signal\n",
    "\n",
    "\n",
    "def get_cummulative_returns(r, gamma=1):\n",
    "    \"\"\"\n",
    "    Computes cummulative discounted rewards given immediate rewards\n",
    "    G_i = r_i + gamma*r_{i+1} + gamma^2*r_{i+2} + ...\n",
    "    Also known as R(s,a).\n",
    "    \"\"\"\n",
    "    r = np.array(r)\n",
    "    assert r.ndim >= 1\n",
    "    return scipy.signal.lfilter([1], [1, -gamma], r[::-1], axis=0)[::-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "nMVh-LevEGV7",
    "outputId": "40cccf8c-4969-4cb5-ab66-26f9956c94ba"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.40049, 1.5561 , 1.729  , 0.81   , 0.9    , 1.     ])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# simple demo on rewards [0,0,1,0,0,1]\n",
    "get_cummulative_returns([0, 0, 1, 0, 0, 1], gamma=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-_ULjgZrEGWB"
   },
   "source": [
    "**Rollout**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QEecckt8EGWD"
   },
   "outputs": [],
   "source": [
    "def rollout(env, act, max_pathlength=2500, n_timesteps=50000):\n",
    "    \"\"\"\n",
    "    Generate rollouts for training.\n",
    "    :param: env - environment in which we will make actions to generate rollouts.\n",
    "    :param: act - the function that can return policy and action given observation.\n",
    "    :param: max_pathlength - maximum size of one path that we generate.\n",
    "    :param: n_timesteps - total sum of sizes of all pathes we generate.\n",
    "    \"\"\"\n",
    "    paths = []\n",
    "\n",
    "    total_timesteps = 0\n",
    "    while total_timesteps < n_timesteps:\n",
    "        obervations, actions, rewards, action_probs = [], [], [], []\n",
    "        obervation = env.reset()\n",
    "        for _ in range(max_pathlength):\n",
    "            action, policy = act(obervation)\n",
    "            obervations.append(obervation)\n",
    "            actions.append(action)\n",
    "            action_probs.append(policy)\n",
    "            obervation, reward, done, _ = env.step(action)\n",
    "            rewards.append(reward)\n",
    "            total_timesteps += 1\n",
    "            if done or total_timesteps == n_timesteps:\n",
    "                path = {\"observations\": np.array(obervations),\n",
    "                        \"policy\": np.array(action_probs),\n",
    "                        \"actions\": np.array(actions),\n",
    "                        \"rewards\": np.array(rewards),\n",
    "                        \"cumulative_returns\": get_cummulative_returns(rewards),\n",
    "                        }\n",
    "                paths.append(path)\n",
    "                break\n",
    "    return paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 310
    },
    "colab_type": "code",
    "id": "AoCMmIXBEGWN",
    "outputId": "42586574-b843-40cd-d1b7-c6fe1bce5344"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'observations': array([[ 0.99760514,  0.06916638,  0.99907676, -0.04296071,  0.0221272 ,\n",
      "        -0.04023622],\n",
      "       [ 0.99793818,  0.0641824 ,  0.99920724, -0.03981058, -0.07085715,\n",
      "         0.07027222],\n",
      "       [ 0.99958117,  0.0289394 ,  0.99984549,  0.01757803, -0.27276454,\n",
      "         0.48788278],\n",
      "       [ 0.9999199 , -0.01265711,  0.99707116,  0.07647943, -0.13326009,\n",
      "         0.08585853],\n",
      "       [ 0.99976016, -0.02190014,  0.99880106,  0.04895341,  0.04162871,\n",
      "        -0.35791806]]), 'policy': array([[0.44099003, 0.283529  , 0.27548093],\n",
      "       [0.44043428, 0.27865574, 0.28091   ],\n",
      "       [0.43189543, 0.29182473, 0.27627984],\n",
      "       [0.44264403, 0.2731225 , 0.28423348],\n",
      "       [0.46287042, 0.2599591 , 0.27717042]], dtype=float32), 'actions': array([1, 2, 0, 0, 0]), 'rewards': array([-1., -1., -1., -1., -1.]), 'cumulative_returns': array([-5., -4., -3., -2., -1.])}\n",
      "It's ok\n"
     ]
    }
   ],
   "source": [
    "paths = rollout(env, act, max_pathlength=5, n_timesteps=100)\n",
    "print(paths[-1])\n",
    "assert (paths[0]['policy'].shape == (5, n_actions))\n",
    "assert (paths[0]['cumulative_returns'].shape == (5,))\n",
    "assert (paths[0]['rewards'].shape == (5,))\n",
    "assert (paths[0]['observations'].shape == (5,)+observation_shape)\n",
    "assert (paths[0]['actions'].shape == (5,))\n",
    "print('It\\'s ok')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "b7Ht_oU5EGWS"
   },
   "source": [
    "### Step 3: loss functions\n",
    "\n",
    "Now let's define the loss functions and constraints for actual TRPO training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EIOE4E0lEGWT"
   },
   "source": [
    "The surrogate reward should be\n",
    "$$J_{surr}= {1 \\over N} \\sum\\limits_{i=0}^N \\frac{\\pi_{\\theta}(s_i, a_i)}{\\pi_{\\theta_{old}}(s_i, a_i)}A_{\\theta_{old}(s_i, a_i)}$$\n",
    "\n",
    "For simplicity, let's use cummulative returns instead of advantage for now:\n",
    "$$J'_{surr}= {1 \\over N} \\sum\\limits_{i=0}^N \\frac{\\pi_{\\theta}(s_i, a_i)}{\\pi_{\\theta_{old}}(s_i, a_i)}G_{\\theta_{old}(s_i, a_i)}$$\n",
    "\n",
    "Or alternatively, minimize the surrogate loss:\n",
    "$$ L_{surr} = - J'_{surr} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "P7OI3o0bEGWZ"
   },
   "outputs": [],
   "source": [
    "# select probabilities of chosen actions\n",
    "batch_size = tf.shape(observations_ph)[0]\n",
    "probs_all = tf.reshape(probs_out, [-1])\n",
    "probs_for_actions = tf.gather(probs_all, tf.range(\n",
    "    0, batch_size) * n_actions + actions_ph)\n",
    "old_probs_all = tf.reshape(old_probs_ph, [-1])\n",
    "old_probs_for_actions = tf.gather(\n",
    "    old_probs_all, tf.range(0, batch_size) * n_actions + actions_ph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZhlT8u_9EGWf"
   },
   "outputs": [],
   "source": [
    "# Compute surrogate loss: negative importance-sampled policy gradient\n",
    "\n",
    "L_surr = - tf.reduce_mean(cummulative_returns_ph * (probs_for_actions/old_probs_for_actions))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 128
    },
    "colab_type": "code",
    "id": "Lu3LTHgCEGWt",
    "outputId": "d72f6db5-d06f-46f8-c435-ee93b211f062"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Tools\\Python\\Anaconda\\lib\\site-packages\\tensorflow\\python\\ops\\gradients_impl.py:100: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
      "  \"Converting sparse IndexedSlices to a dense Tensor of unknown shape. \"\n"
     ]
    }
   ],
   "source": [
    "# compute and return surrogate policy gradient\n",
    "def var_shape(x):\n",
    "    return [k.value for k in x.get_shape()]\n",
    "\n",
    "\n",
    "def numel(x):\n",
    "    return np.prod(var_shape(x))\n",
    "\n",
    "\n",
    "def flatgrad(loss, var_list):\n",
    "    grads = tf.gradients(loss, var_list)\n",
    "    return tf.concat([tf.reshape(grad, [numel(v)])\n",
    "                      for (v, grad) in zip(var_list, grads)], 0)\n",
    "\n",
    "\n",
    "flat_grad_surr = flatgrad(L_surr, weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xwNzWUtAEGW-"
   },
   "source": [
    "We can ascend these gradients as long as our $pi_\\theta(a|s)$ satisfies the constraint\n",
    "$$E_{s,\\pi_{\\Theta_{t}}}\\Big[KL(\\pi(\\Theta_{t}, s) \\:||\\:\\pi(\\Theta_{t+1}, s))\\Big]< \\alpha$$\n",
    "\n",
    "\n",
    "where\n",
    "\n",
    "$$KL(p||q) = E _p log({p \\over q})$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 90
    },
    "colab_type": "code",
    "id": "12KtgaCBEGXA",
    "outputId": "15947ab8-2f5e-4d45-96c2-1828c7816f3f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kl Tensor(\"truediv_1:0\", shape=(), dtype=float32)\n",
      "probs_all Tensor(\"Exp:0\", shape=(?, 3), dtype=float32)\n",
      "log_probs_all Tensor(\"dense2/LogSoftmax:0\", shape=(?, 3), dtype=float32)\n",
      "entropy Tensor(\"truediv_2:0\", shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Compute Kullback-Leibler divergence (see formula above)\n",
    "# Note: you need to sum KL and entropy over all actions, not just the ones agent took\n",
    "old_log_probs = tf.log(old_probs_ph+1e-10)\n",
    "log_probs_all = policy_out\n",
    "\n",
    "kl = tf.reduce_sum(old_probs_ph * (old_log_probs - log_probs_all))/tf.cast(batch_size, tf.float32)\n",
    "print(\"kl\", kl)\n",
    "# Compute policy entropy\n",
    "print(\"probs_all\", probs_out)\n",
    "print(\"log_probs_all\", log_probs_all)\n",
    "\n",
    "\n",
    "entropy = tf.reduce_sum(-probs_out * log_probs_all) / tf.cast(batch_size, tf.float32)\n",
    "print(\"entropy\",entropy)\n",
    "losses = [L_surr, kl, entropy]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xtq2xb62EGXb"
   },
   "source": [
    "**Linear search**\n",
    "\n",
    "TRPO in its core involves ascending surrogate policy gradient constrained by KL divergence. \n",
    "\n",
    "In order to enforce this constraint, we're gonna use linesearch. You can find out more about it [here](https://en.wikipedia.org/wiki/Linear_search)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8Px_ldYOEGXd"
   },
   "outputs": [],
   "source": [
    "def linesearch(f, x, fullstep, max_kl):\n",
    "    \"\"\"\n",
    "    Linesearch finds the best parameters of neural networks in the direction of fullstep contrainted by KL divergence.\n",
    "    :param: f - function that returns loss, kl and arbitrary third component.\n",
    "    :param: x - old parameters of neural network.\n",
    "    :param: fullstep - direction in which we make search.\n",
    "    :param: max_kl - constraint of KL divergence.\n",
    "    :returns:\n",
    "    \"\"\"\n",
    "    max_backtracks = 10\n",
    "    loss, _, _ = f(x)\n",
    "    for stepfrac in .5**np.arange(max_backtracks):\n",
    "        xnew = x + stepfrac * fullstep\n",
    "        new_loss, kl, _ = f(xnew)\n",
    "        actual_improve = new_loss - loss\n",
    "        if kl <= max_kl and actual_improve < 0:\n",
    "            x = xnew\n",
    "            loss = new_loss\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "50nwdhFnEGXw"
   },
   "source": [
    "### Step 4: training\n",
    "In this section we construct rest parts of our computational graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EksHGVN6EGX2"
   },
   "outputs": [],
   "source": [
    "def slice_vector(vector, shapes):\n",
    "    \"\"\"\n",
    "    Slices symbolic vector into several symbolic tensors of given shapes.\n",
    "    Auxilary function used to un-flatten gradients, tangents etc.\n",
    "    :param vector: 1-dimensional symbolic vector\n",
    "    :param shapes: list or tuple of shapes (list, tuple or symbolic)\n",
    "    :returns: list of symbolic tensors of given shapes\n",
    "    \"\"\"\n",
    "    assert len(vector.get_shape()) == 1, \"vector must be 1-dimensional\"\n",
    "    start = 0\n",
    "    tensors = []\n",
    "    for shape in shapes:\n",
    "        size = np.prod(shape)\n",
    "        tensor = tf.reshape(vector[start:(start + size)], shape)\n",
    "        tensors.append(tensor)\n",
    "        start += size\n",
    "    return tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jsWdcVC3EGYC"
   },
   "outputs": [],
   "source": [
    "# intermediate grad in conjugate_gradient\n",
    "conjugate_grad_intermediate_vector = tf.placeholder(\n",
    "    dtype=tf.float32, shape=(None,))\n",
    "\n",
    "# slice flat_tangent into chunks for each weight\n",
    "weight_shapes = [sess.run(var).shape for var in weights]\n",
    "tangents = slice_vector(conjugate_grad_intermediate_vector, weight_shapes)\n",
    "\n",
    "# KL divergence where first arg is fixed\n",
    "kl_firstfixed = tf.reduce_sum((tf.stop_gradient(probs_out) * (tf.stop_gradient(\n",
    "    tf.log(probs_out)) - tf.log(probs_out)))) / tf.cast(batch_size, tf.float32)\n",
    "\n",
    "# compute fisher information matrix (used for conjugate gradients and to estimate KL)\n",
    "gradients = tf.gradients(kl_firstfixed, weights)\n",
    "gradient_vector_product = [tf.reduce_sum(\n",
    "    g[0] * t) for (g, t) in zip(gradients, tangents)]\n",
    "\n",
    "fisher_vec_prod = flatgrad(gradient_vector_product, weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4e4HzTBFEGYM"
   },
   "source": [
    "### TRPO helpers\n",
    "\n",
    "Here we define a few helper functions used in the main TRPO loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6SR6q7CHEGYO"
   },
   "source": [
    "**Conjugate gradients**\n",
    "\n",
    "Since TRPO includes contrainted optimization, we will need to solve Ax=b using conjugate gradients.\n",
    "\n",
    "In general, CG is an algorithm that solves Ax=b where A is positive-defined. A is Hessian matrix so A is positive-defined. You can find out more about them [here](https://en.wikipedia.org/wiki/Conjugate_gradient_method)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eR96VRIuEGYR"
   },
   "outputs": [],
   "source": [
    "from numpy.linalg import inv\n",
    "\n",
    "\n",
    "def conjugate_gradient(f_Ax, b, cg_iters=10, residual_tol=1e-10):\n",
    "    \"\"\"\n",
    "    This method solves system of equation Ax=b using iterative method called conjugate gradients\n",
    "    :f_Ax: function that returns Ax\n",
    "    :b: targets for Ax\n",
    "    :cg_iters: how many iterations this method should do\n",
    "    :residual_tol: epsilon for stability\n",
    "    \"\"\"\n",
    "    p = b.copy()\n",
    "    r = b.copy()\n",
    "    x = np.zeros_like(b)\n",
    "    rdotr = r.dot(r)\n",
    "    for i in range(cg_iters):\n",
    "        z = f_Ax(p)\n",
    "        v = rdotr / (p.dot(z) + 1e-8)\n",
    "        x += v * p\n",
    "        r -= v * z\n",
    "        newrdotr = r.dot(r)\n",
    "        mu = newrdotr / (rdotr + 1e-8)\n",
    "        p = r + mu * p\n",
    "        rdotr = newrdotr\n",
    "        if rdotr < residual_tol:\n",
    "            break\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 90
    },
    "colab_type": "code",
    "id": "WVIvJenjEGYg",
    "outputId": "a26cac10-84a0-49b5-d731-dc09ef6196fa"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-27.65448251  19.73062791 -10.46556453 -13.3967524   -7.91207936\n",
      "  49.35834653  44.35182242 -57.93542044]\n",
      "[-27.6542676   19.73053067 -10.46553218 -13.39670521  -7.91201668\n",
      "  49.35804765  44.35140544 -57.9349304 ]\n"
     ]
    }
   ],
   "source": [
    "# This code validates conjugate gradients\n",
    "A = np.random.rand(8, 8)\n",
    "A = np.matmul(np.transpose(A), A)\n",
    "\n",
    "\n",
    "def f_Ax(x):\n",
    "    return np.matmul(A, x.reshape(-1, 1)).reshape(-1)\n",
    "\n",
    "\n",
    "b = np.random.rand(8)\n",
    "\n",
    "w = np.matmul(np.matmul(inv(np.matmul(np.transpose(A), A)),\n",
    "                        np.transpose(A)), b.reshape((-1, 1))).reshape(-1)\n",
    "print(w)\n",
    "print(conjugate_gradient(f_Ax, b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bli-ZTcwEGYo"
   },
   "outputs": [],
   "source": [
    "# Compile a function that exports network weights as a vector\n",
    "flat_weights = tf.concat([tf.reshape(var, [-1]) for var in weights], axis=0)\n",
    "\n",
    "# ... and another function that imports vector back into network weights\n",
    "flat_weights_placeholder = tf.placeholder(tf.float32, shape=(None,))\n",
    "assigns = slice_vector(flat_weights_placeholder, weight_shapes)\n",
    "\n",
    "load_flat_weights = [w.assign(ph) for w, ph in zip(weights, assigns)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-bAmtKreEGY0"
   },
   "source": [
    "##### Step 5: Main TRPO loop\n",
    "\n",
    "Here we will train our network!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1868
    },
    "colab_type": "code",
    "id": "19DFjbMzEGY5",
    "outputId": "820c4b78-8d1e-4a12-9fb3-1b117de3f210"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "********** Iteration 1 ************\n",
      "Rollout\n",
      "Made rollout\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Tools\\Python\\Anaconda\\lib\\site-packages\\ipykernel_launcher.py:43: RuntimeWarning: invalid value encountered in sqrt\n",
      "D:\\Tools\\Python\\Anaconda\\lib\\site-packages\\ipykernel_launcher.py:16: RuntimeWarning: invalid value encountered in less_equal\n",
      "  app.launch_new_instance()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of episodes:                 156\n",
      "Average sum of rewards per episode:       -319.5897435897436\n",
      "Std of rewards per episode:               86.8916855804465\n",
      "Entropy:                                  0.93045706\n",
      "Time elapsed:                             0.57 mins\n",
      "KL between old and new distribution:      -4.794534e-10\n",
      "Surrogate loss:                           171.61148\n",
      "\n",
      "********** Iteration 2 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 311\n",
      "Average sum of rewards per episode:       -321.63870967741934\n",
      "Std of rewards per episode:               89.27781098245487\n",
      "Entropy:                                  0.87479377\n",
      "Time elapsed:                             1.06 mins\n",
      "KL between old and new distribution:      0.009837995\n",
      "Surrogate loss:                           172.69958\n",
      "\n",
      "********** Iteration 3 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 522\n",
      "Average sum of rewards per episode:       -235.97630331753555\n",
      "Std of rewards per episode:               58.06422663542982\n",
      "Entropy:                                  0.810788\n",
      "Time elapsed:                             1.68 mins\n",
      "KL between old and new distribution:      0.009954151\n",
      "Surrogate loss:                           124.79254\n",
      "\n",
      "********** Iteration 4 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 783\n",
      "Average sum of rewards per episode:       -190.57471264367817\n",
      "Std of rewards per episode:               40.652899172436925\n",
      "Entropy:                                  0.7775275\n",
      "Time elapsed:                             2.32 mins\n",
      "KL between old and new distribution:      0.009986863\n",
      "Surrogate loss:                           99.37344\n",
      "\n",
      "********** Iteration 5 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 1093\n",
      "Average sum of rewards per episode:       -160.29354838709676\n",
      "Std of rewards per episode:               34.369636219497025\n",
      "Entropy:                                  0.7524064\n",
      "Time elapsed:                             2.95 mins\n",
      "KL between old and new distribution:      0.009995325\n",
      "Surrogate loss:                           83.66674\n",
      "\n",
      "********** Iteration 6 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 1447\n",
      "Average sum of rewards per episode:       -140.24576271186442\n",
      "Std of rewards per episode:               29.802249310518377\n",
      "Entropy:                                  0.7007453\n",
      "Time elapsed:                             3.52 mins\n",
      "KL between old and new distribution:      0.009985239\n",
      "Surrogate loss:                           73.07046\n",
      "\n",
      "********** Iteration 7 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 1821\n",
      "Average sum of rewards per episode:       -132.69251336898395\n",
      "Std of rewards per episode:               27.163381642187446\n",
      "Entropy:                                  0.6348548\n",
      "Time elapsed:                             4.09 mins\n",
      "KL between old and new distribution:      0.00994857\n",
      "Surrogate loss:                           68.97993\n",
      "\n",
      "********** Iteration 8 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 2230\n",
      "Average sum of rewards per episode:       -121.2518337408313\n",
      "Std of rewards per episode:               29.272393309820554\n",
      "Entropy:                                  0.5886749\n",
      "Time elapsed:                             4.64 mins\n",
      "KL between old and new distribution:      0.0099472245\n",
      "Surrogate loss:                           64.039635\n",
      "\n",
      "********** Iteration 9 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 2635\n",
      "Average sum of rewards per episode:       -122.45925925925926\n",
      "Std of rewards per episode:               32.12508645433062\n",
      "Entropy:                                  0.5170329\n",
      "Time elapsed:                             5.21 mins\n",
      "KL between old and new distribution:      0.009742734\n",
      "Surrogate loss:                           65.27621\n",
      "\n",
      "********** Iteration 10 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 3079\n",
      "Average sum of rewards per episode:       -111.61711711711712\n",
      "Std of rewards per episode:               31.530726823708466\n",
      "Entropy:                                  0.5009117\n",
      "Time elapsed:                             5.79 mins\n",
      "KL between old and new distribution:      0.009914713\n",
      "Surrogate loss:                           60.086815\n",
      "\n",
      "********** Iteration 11 ************\n",
      "Rollout\n",
      "Made rollout\n",
      "Total number of episodes:                 3535\n",
      "Average sum of rewards per episode:       -108.65350877192982\n",
      "Std of rewards per episode:               31.534485939650157\n",
      "Entropy:                                  0.465178\n",
      "Time elapsed:                             6.34 mins\n",
      "KL between old and new distribution:      0.009925838\n",
      "Surrogate loss:                           58.7691\n",
      "Well done!\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "from itertools import count\n",
    "from collections import OrderedDict\n",
    "\n",
    "# this is hyperparameter of TRPO. It controls how big KL divergence may be between old and new policy every step.\n",
    "max_kl = 0.01\n",
    "cg_damping = 0.1  # This parameters regularize addition to\n",
    "numeptotal = 0  # this is number of episodes that we played.\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "for i in count(1):\n",
    "\n",
    "    print(\"\\n********** Iteration %i ************\" % i)\n",
    "\n",
    "    # Generating paths.\n",
    "    print(\"Rollout\")\n",
    "    paths = rollout(env, act)\n",
    "    print(\"Made rollout\")\n",
    "\n",
    "    # Updating policy.\n",
    "    observations = np.concatenate([path[\"observations\"] for path in paths])\n",
    "    actions = np.concatenate([path[\"actions\"] for path in paths])\n",
    "    returns = np.concatenate([path[\"cumulative_returns\"] for path in paths])\n",
    "    old_probs = np.concatenate([path[\"policy\"] for path in paths])\n",
    "    inputs_batch = [observations, actions, returns, old_probs]\n",
    "    feed_dict = {observations_ph: observations,\n",
    "                 actions_ph: actions,\n",
    "                 old_probs_ph: old_probs,\n",
    "                 cummulative_returns_ph: returns,\n",
    "                 }\n",
    "    old_weights = sess.run(flat_weights)\n",
    "\n",
    "    def fisher_vector_product(p):\n",
    "        \"\"\"gets intermediate grads (p) and computes fisher*vector \"\"\"\n",
    "        feed_dict[conjugate_grad_intermediate_vector] = p\n",
    "        return sess.run(fisher_vec_prod, feed_dict) + cg_damping * p\n",
    "\n",
    "    flat_grad = sess.run(flat_grad_surr, feed_dict)\n",
    "\n",
    "    stepdir = conjugate_gradient(fisher_vector_product, -flat_grad)\n",
    "    shs = .5 * stepdir.dot(fisher_vector_product(stepdir))\n",
    "    lm = np.sqrt(shs / max_kl)\n",
    "    fullstep = stepdir / lm\n",
    "\n",
    "    # Compute new weights with linesearch in the direction we found with CG\n",
    "\n",
    "    def losses_f(flat_weights):\n",
    "        feed_dict[flat_weights_placeholder] = flat_weights\n",
    "        sess.run(load_flat_weights, feed_dict)\n",
    "        return sess.run(losses, feed_dict)\n",
    "\n",
    "    new_weights = linesearch(losses_f, old_weights, fullstep, max_kl)\n",
    "    feed_dict[flat_weights_placeholder] = new_weights\n",
    "    sess.run(load_flat_weights, feed_dict)\n",
    "\n",
    "    # Report current progress\n",
    "    L_surr, kl, entropy = sess.run(losses, feed_dict)\n",
    "    episode_rewards = np.array([path[\"rewards\"].sum() for path in paths])\n",
    "\n",
    "    stats = OrderedDict()\n",
    "    numeptotal += len(episode_rewards)\n",
    "    stats[\"Total number of episodes\"] = numeptotal\n",
    "    stats[\"Average sum of rewards per episode\"] = episode_rewards.mean()\n",
    "    stats[\"Std of rewards per episode\"] = episode_rewards.std()\n",
    "    stats[\"Entropy\"] = entropy\n",
    "    stats[\"Time elapsed\"] = \"%.2f mins\" % ((time.time() - start_time)/60.)\n",
    "    stats[\"KL between old and new distribution\"] = kl\n",
    "    stats[\"Surrogate loss\"] = L_surr\n",
    "    for k, v in stats.items():\n",
    "        print(k + \": \" + \" \" * (40 - len(k)) + str(v))\n",
    "    if episode_rewards.mean() > -110:\n",
    "        print(\"Well done!\")\n",
    "        break;\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ywodBkTgEGZA"
   },
   "source": [
    "# Homework option I: better sampling (10+pts)\n",
    "\n",
    "In this section, you're invited to implement a better rollout strategy called _vine_.\n",
    "\n",
    "![img](https://s17.postimg.org/i90chxgvj/vine.png)\n",
    "\n",
    "In most gym environments, you can actually backtrack by using states. You can find a wrapper that saves/loads states in [the mcts seminar](https://github.com/yandexdataschool/Practical_RL/blob/master/yet_another_week/seminar_MCTS.ipynb).\n",
    "\n",
    "You can read more about in the [TRPO article](https://arxiv.org/abs/1502.05477) in section 5.2.\n",
    "\n",
    "The goal here is to implement such rollout policy (we recommend using tree data structure like in the seminar above).\n",
    "Then you can assign cummulative rewards similar to `get_cummulative_rewards`, but for a tree.\n",
    "\n",
    "__bonus task__ - parallelize samples using multiple cores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GbhjBL9LEGZD"
   },
   "source": [
    "# Homework option II (10+pts)\n",
    "\n",
    "Let's use TRPO to train evil robots! (pick any of two)\n",
    "* [MuJoCo robots](https://gym.openai.com/envs#mujoco)\n",
    "* [Box2d robot](https://gym.openai.com/envs/BipedalWalker-v2)\n",
    "\n",
    "The catch here is that those environments have continuous action spaces. \n",
    "\n",
    "Luckily, TRPO is a policy gradient method, so it's gonna work for any parametric $\\pi_\\theta(a|s)$. We recommend starting with gaussian policy:\n",
    "\n",
    "$$\\pi_\\theta(a|s) = N(\\mu_\\theta(s),\\sigma^2_\\theta(s)) = {1 \\over \\sqrt { 2 \\pi {\\sigma^2}_\\theta(s) } } e^{ (a - \n",
    "\\mu_\\theta(s))^2 \\over 2 {\\sigma^2}_\\theta(s) } $$\n",
    "\n",
    "In the $\\sqrt { 2 \\pi {\\sigma^2}_\\theta(s) }$ clause, $\\pi$ means ~3.1415926, not agent's policy.\n",
    "\n",
    "This essentially means that you will need two output layers:\n",
    "* $\\mu_\\theta(s)$, a dense layer with linear activation\n",
    "* ${\\sigma^2}_\\theta(s)$, a dense layer with activation tf.exp (to make it positive; like rho from bandits)\n",
    "\n",
    "For multidimensional actions, you can use fully factorized gaussian (basically a vector of gaussians).\n",
    "\n",
    "__bonus task__: compare performance of continuous action space method to action space discretization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PsxMfrGOEGZF"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "7_3_TRPO_TF.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "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": 2
}