6_1_reinforce_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# REINFORCE in pytorch\n",
    "\n",
    "Just like we did before for q-learning, this time we'll design a pytorch network to learn `CartPole-v0` via policy gradient (REINFORCE).\n",
    "\n",
    "Most of the code in this notebook is taken from approximate qlearning, so you'll find it more or less familiar and even simpler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x1af53c55fd0>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXgAAAD8CAYAAAB9y7/cAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAAEolJREFUeJzt3X+s3fV93/Hnq5hAlmQ1hAty/WMmjbeGTouhd8QR00QhbYH9MJWaCTY1KEK6mUSkRI22QietiTSkVlrDFq1DcQuNM2UhjCTDQ6wpc4iq/BGInTiOjUO5SZz41h42C5Bk0dhM3vvjfG44Ncf3Ht8fvr6fPB/S0fl+P+fz/Z73Bw6v+72f+/1wUlVIkvrzMytdgCRpeRjwktQpA16SOmXAS1KnDHhJ6pQBL0mdWraAT3JDkqeTTCe5c7neR5I0WpbjPvgk5wF/AfwKMAN8Gbi1qp5a8jeTJI20XFfwVwPTVfWtqvq/wAPA9mV6L0nSCGuW6bzrgSND+zPA207X+ZJLLqnNmzcvUymStPocPnyY5557Los5x3IF/Kii/spcUJIpYApg06ZN7NmzZ5lKkaTVZ3JyctHnWK4pmhlg49D+BuDocIeq2lFVk1U1OTExsUxlSNJPr+UK+C8DW5JcnuQ1wC3ArmV6L0nSCMsyRVNVJ5O8F/gccB5wf1UdXI73kiSNtlxz8FTVo8Cjy3V+SdLcXMkqSZ0y4CWpUwa8JHXKgJekThnwktQpA16SOmXAS1KnDHhJ6pQBL0mdMuAlqVMGvCR1yoCXpE4Z8JLUKQNekjplwEtSpwx4SeqUAS9JnTLgJalTi/rKviSHgR8ALwMnq2oyycXAp4DNwGHgn1TV84srU5J0ppbiCv6Xq2prVU22/TuB3VW1Bdjd9iVJZ9lyTNFsB3a27Z3AzcvwHpKkeSw24Av4syR7k0y1tsuq6hhAe750ke8hSVqARc3BA9dU1dEklwKPJfnGuAe2HwhTAJs2bVpkGZKkUy3qCr6qjrbn48BngauBZ5OsA2jPx09z7I6qmqyqyYmJicWUIUkaYcEBn+R1Sd4wuw38KnAA2AXc1rrdBjy82CIlSWduMVM0lwGfTTJ7nv9cVX+a5MvAg0luB74LvHPxZUqSztSCA76qvgW8dUT7/wKuX0xRkqTFcyWrJHXKgJekThnwktQpA16SOmXAS1KnDHhJ6pQBL0mdMuAlqVMGvCR1yoCXpE4Z8JLUKQNekjplwEtSpwx4SeqUAS9JnTLgJalTBrwkdcqAl6ROGfCS1Kl5Az7J/UmOJzkw1HZxkseSPNOeL2rtSfKRJNNJ9ie5ajmLlySd3jhX8B8Dbjil7U5gd1VtAXa3fYAbgS3tMQXcuzRlSpLO1LwBX1V/DnzvlObtwM62vRO4eaj94zXwJWBtknVLVawkaXwLnYO/rKqOAbTnS1v7euDIUL+Z1vYqSaaS7Emy58SJEwssQ5J0Okv9R9aMaKtRHatqR1VNVtXkxMTEEpchSVpowD87O/XSno+39hlg41C/DcDRhZcnSVqohQb8LuC2tn0b8PBQ+7va3TTbgBdnp3IkSWfXmvk6JPkkcC1wSZIZ4HeB3wMeTHI78F3gna37o8BNwDTwI+Ddy1CzJGkM8wZ8Vd16mpeuH9G3gDsWW5QkafFcySpJnTLgJalTBrwkdcqAl6ROGfCS1CkDXpI6ZcBLUqcMeEnqlAEvSZ0y4CWpUwa8JHXKgJekThnwktQpA16SOmXAS1KnDHhJ6pQBL0mdMuAlqVPzBnyS+5McT3JgqO2DSf4yyb72uGnotbuSTCd5OsmvLVfhkqS5jXMF/zHghhHt91TV1vZ4FCDJFcAtwC+2Y/5jkvOWqlhJ0vjmDfiq+nPge2OebzvwQFW9VFXfBqaBqxdRnyRpgRYzB//eJPvbFM5FrW09cGSoz0xre5UkU0n2JNlz4sSJRZQhSRploQF/L/DzwFbgGPAHrT0j+taoE1TVjqqarKrJiYmJBZYhSTqdBQV8VT1bVS9X1Y+BP+KVaZgZYONQ1w3A0cWVKElaiAUFfJJ1Q7u/DszeYbMLuCXJBUkuB7YATy6uREnSQqyZr0OSTwLXApckmQF+F7g2yVYG0y+HgfcAVNXBJA8CTwEngTuq6uXlKV2SNJd5A76qbh3RfN8c/e8G7l5MUZKkxXMlqyR1yoCXpE4Z8JLUKQNekjplwEtSpwx4SerUvLdJSj8N9u54z8j2X5r66FmuRFo6XsFLUqcMeEnqlAEvSZ0y4CWpUwa8JHXKgJekThnwktQpA16SOmXAS1KnDHhJ6pQBL0mdmjfgk2xM8niSQ0kOJnlfa784yWNJnmnPF7X2JPlIkukk+5NctdyDkCS92jhX8CeBD1TVW4BtwB1JrgDuBHZX1RZgd9sHuBHY0h5TwL1LXrUkaV7zBnxVHauqr7TtHwCHgPXAdmBn67YTuLltbwc+XgNfAtYmWbfklUuS5nRGc/BJNgNXAk8Al1XVMRj8EAAubd3WA0eGDptpbaeeayrJniR7Tpw4ceaVS5LmNHbAJ3k98Gng/VX1/bm6jmirVzVU7aiqyaqanJiYGLcMSdKYxgr4JOczCPdPVNVnWvOzs1Mv7fl4a58BNg4dvgE4ujTlSpLGNc5dNAHuAw5V1YeHXtoF3Na2bwMeHmp/V7ubZhvw4uxUjiTp7BnnK/uuAX4T+HqSfa3td4DfAx5McjvwXeCd7bVHgZuAaeBHwLuXtGJJ0ljmDfiq+iKj59UBrh/Rv4A7FlmXJGmRXMkqSZ0y4CWpUwa8JHXKgJekThnwktQpA16SOmXAS1KnDHhJ6pQBL0mdMuAlqVMGvCR1yoCXpE4Z8NJp/NLUR1e6BGlRDHhJ6pQBL0mdMuAlqVMGvCR1yoCXpE6N86XbG5M8nuRQkoNJ3tfaP5jkL5Psa4+bho65K8l0kqeT/NpyDkCSNNo4X7p9EvhAVX0lyRuAvUkea6/dU1X/drhzkiuAW4BfBH4O+B9J/mZVvbyUhUuS5jbvFXxVHauqr7TtHwCHgPVzHLIdeKCqXqqqbwPTwNVLUawkaXxnNAefZDNwJfBEa3pvkv1J7k9yUWtbDxwZOmyGuX8gSJKWwdgBn+T1wKeB91fV94F7gZ8HtgLHgD+Y7Tri8Bpxvqkke5LsOXHixBkXLkma21gBn+R8BuH+iar6DEBVPVtVL1fVj4E/4pVpmBlg49DhG4Cjp56zqnZU1WRVTU5MTCxmDJKkEca5iybAfcChqvrwUPu6oW6/Dhxo27uAW5JckORyYAvw5NKVLEkaxzh30VwD/Cbw9ST7WtvvALcm2cpg+uUw8B6AqjqY5EHgKQZ34NzhHTSSdPbNG/BV9UVGz6s/OscxdwN3L6IuSdIiuZJVkjplwEtSpwx4SeqUAS9JnTLgJalTBrwkdcqAl6ROGfCS1CkDXpI6ZcBLUqcMeEnqlAEvSZ0y4CWpUwa8upVk7MdyHC+tNANekjo1zhd+SD8V/tvRqZ9s/6Of27GClUhLwyt4ib8a7qP2pdXIgJekTo3zpdsXJnkyydeSHEzyodZ+eZInkjyT5FNJXtPaL2j70+31zcs7BEnSKONcwb8EXFdVbwW2Ajck2Qb8PnBPVW0Bngdub/1vB56vqjcD97R+0jnt1Dl35+DVg3G+dLuAH7bd89ujgOuAf9radwIfBO4FtrdtgIeA/5Ak7TzSOWnyPTuAV0L9QytXirRkxrqLJsl5wF7gzcAfAt8EXqiqk63LDLC+ba8HjgBU1ckkLwJvBJ473fn37t3rvcRa1fz86lw0VsBX1cvA1iRrgc8CbxnVrT2P+qS/6uo9yRQwBbBp0ya+853vjFWwNK6zGbr+gqqlNjk5uehznNFdNFX1AvAFYBuwNsnsD4gNwNG2PQNsBGiv/yzwvRHn2lFVk1U1OTExsbDqJUmnNc5dNBPtyp0krwXeARwCHgd+o3W7DXi4be9q+7TXP+/8uySdfeNM0awDdrZ5+J8BHqyqR5I8BTyQ5N8AXwXua/3vA/5TkmkGV+63LEPdkqR5jHMXzX7gyhHt3wKuHtH+f4B3Lkl1kqQFcyWrJHXKgJekThnwktQp/3fB6pY3b+mnnVfwktQpA16SOmXAS1KnDHhJ6pQBL0mdMuAlqVMGvCR1yoCXpE4Z8JLUKQNekjplwEtSpwx4SeqUAS9JnTLgJalT43zp9oVJnkzytSQHk3yotX8sybeT7GuPra09ST6SZDrJ/iRXLfcgJEmvNs7/D/4l4Lqq+mGS84EvJvnv7bV/UVUPndL/RmBLe7wNuLc9S5LOonmv4Gvgh233/PaY65sUtgMfb8d9CVibZN3iS5UknYmx5uCTnJdkH3AceKyqnmgv3d2mYe5JckFrWw8cGTp8prVJks6isQK+ql6uqq3ABuDqJH8buAv4BeDvAhcDv926Z9QpTm1IMpVkT5I9J06cWFDxkqTTO6O7aKrqBeALwA1VdaxNw7wE/Alwdes2A2wcOmwDcHTEuXZU1WRVTU5MTCyoeEnS6Y1zF81EkrVt+7XAO4BvzM6rJwlwM3CgHbILeFe7m2Yb8GJVHVuW6iVJpzXOXTTrgJ1JzmPwA+HBqnokyeeTTDCYktkH/PPW/1HgJmAa+BHw7qUvW5I0n3kDvqr2A1eOaL/uNP0LuGPxpUmSFsOVrJLUKQNekjplwEtSpwx4SeqUAS9JnTLgJalTBrwkdcqAl6ROGfCS1CkDXpI6ZcBLUqcMeEnqlAEvSZ0y4CWpUwa8JHXKgJekThnwktQpA16SOmXAS1Knxg74JOcl+WqSR9r+5UmeSPJMkk8leU1rv6DtT7fXNy9P6ZKkuZzJFfz7gEND+78P3FNVW4Dngdtb++3A81X1ZuCe1k+SdJaNFfBJNgD/APjjth/gOuCh1mUncHPb3t72aa9f3/pLks6iNWP2+3fAvwTe0PbfCLxQVSfb/gywvm2vB44AVNXJJC+2/s8NnzDJFDDVdl9KcmBBIzj3XcIpY+9Er+OCfsfmuFaXv5Fkqqp2LPQE8wZ8kn8IHK+qvUmunW0e0bXGeO2VhkHRO9p77KmqybEqXmV6HVuv44J+x+a4Vp8ke2g5uRDjXMFfA/zjJDcBFwJ/ncEV/doka9pV/AbgaOs/A2wEZpKsAX4W+N5CC5QkLcy8c/BVdVdVbaiqzcAtwOer6p8BjwO/0brdBjzctne1fdrrn6+qV13BS5KW12Lug/9t4LeSTDOYY7+vtd8HvLG1/xZw5xjnWvCvIKtAr2PrdVzQ79gc1+qzqLHFi2tJ6pMrWSWpUyse8EluSPJ0W/k6znTOOSXJ/UmOD9/mmeTiJI+1Vb6PJbmotSfJR9pY9ye5auUqn1uSjUkeT3IoycEk72vtq3psSS5M8mSSr7Vxfai1d7Eyu9cV50kOJ/l6kn3tzpJV/1kESLI2yUNJvtH+W3v7Uo5rRQM+yXnAHwI3AlcAtya5YiVrWoCPATec0nYnsLut8t3NK3+HuBHY0h5TwL1nqcaFOAl8oKreAmwD7mj/blb72F4CrquqtwJbgRuSbKOfldk9rzj/5araOnRL5Gr/LAL8e+BPq+oXgLcy+He3dOOqqhV7AG8HPje0fxdw10rWtMBxbAYODO0/Daxr2+uAp9v2R4FbR/U71x8M7pL6lZ7GBvw14CvA2xgslFnT2n/yuQQ+B7y9ba9p/bLStZ9mPBtaIFwHPMJgTcqqH1er8TBwySltq/qzyOCW82+f+s99Kce10lM0P1n12gyviF3NLquqYwDt+dLWvirH2359vxJ4gg7G1qYx9gHHgceAbzLmymxgdmX2uWh2xfmP2/7YK845t8cFg8WSf5Zkb1sFD6v/s/gm4ATwJ21a7Y+TvI4lHNdKB/xYq147surGm+T1wKeB91fV9+fqOqLtnBxbVb1cVVsZXPFeDbxlVLf2vCrGlaEV58PNI7quqnENuaaqrmIwTXFHkr8/R9/VMrY1wFXAvVV1JfC/mfu28jMe10oH/Oyq11nDK2JXs2eTrANoz8db+6oab5LzGYT7J6rqM625i7EBVNULwBcY/I1hbVt5DaNXZnOOr8yeXXF+GHiAwTTNT1actz6rcVwAVNXR9nwc+CyDH8yr/bM4A8xU1RNt/yEGgb9k41rpgP8ysKX9pf81DFbK7lrhmpbC8GreU1f5vqv9NXwb8OLsr2LnmiRhsGjtUFV9eOilVT22JBNJ1rbt1wLvYPCHrVW9Mrs6XnGe5HVJ3jC7DfwqcIBV/lmsqv8JHEnyt1rT9cBTLOW4zoE/NNwE/AWDedB/tdL1LKD+TwLHgP/H4Cfs7QzmMncDz7Tni1vfMLhr6JvA14HJla5/jnH9PQa//u0H9rXHTat9bMDfAb7axnUA+Net/U3Ak8A08F+AC1r7hW1/ur3+ppUewxhjvBZ4pJdxtTF8rT0OzubEav8stlq3Anva5/G/Ahct5bhcySpJnVrpKRpJ0jIx4CWpUwa8JHXKgJekThnwktQpA16SOmXAS1KnDHhJ6tT/Bxqnf/9gmrVaAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "env = gym.make(\"CartPole-v0\").env\n",
    "env.reset()\n",
    "\n",
    "plt.imshow(env.render(\"rgb_array\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the network for REINFORCE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For REINFORCE algorithm, we'll need a model that predicts action probabilities given states. Let's define such a model below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_actions = env.action_space.n\n",
    "state_dim = env.observation_space.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build a simple neural network that predicts policy logits. \n",
    "# Keep it simple: CartPole isn't worth deep architectures.\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(state_dim[0], 100),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(100, n_actions)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.5562, 0.4438]], grad_fn=<SoftmaxBackward>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s = env.reset()\n",
    "t = torch.tensor([s], dtype=torch.float32)\n",
    "t.shape\n",
    "F.softmax(model.forward(t), dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_probs(states):\n",
    "    \"\"\" \n",
    "    Predict action probabilities given states.\n",
    "    :param states: numpy array of shape [batch, state_shape]\n",
    "    :returns: numpy array of shape [batch, n_actions]\n",
    "    \"\"\"\n",
    "    # convert states, compute logits, use softmax to get probability\n",
    "    torch_states = torch.as_tensor(states, dtype=torch.float32)\n",
    "    return F.softmax(model(torch_states), dim=1).detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_states = np.array([env.reset() for _ in range(5)])\n",
    "test_probas = predict_probs(test_states)\n",
    "assert isinstance(\n",
    "    test_probas, np.ndarray), \"you must return np array and not %s\" % type(test_probas)\n",
    "assert tuple(test_probas.shape) == (\n",
    "    test_states.shape[0], env.action_space.n), \"wrong output shape: %s\" % np.shape(test_probas)\n",
    "assert np.allclose(np.sum(test_probas, axis=1),\n",
    "                   1), \"probabilities do not sum to 1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Play the game\n",
    "\n",
    "We can now use our newly built agent to play the game."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_session(t_max=1000):\n",
    "    \"\"\" \n",
    "    play a full session with REINFORCE agent and train at the session end.\n",
    "    returns sequences of states, actions andrewards\n",
    "    \"\"\"\n",
    "    # arrays to record session\n",
    "    states, actions, rewards = [], [], []\n",
    "    s = env.reset()\n",
    "\n",
    "    for t in range(t_max):\n",
    "        # action probabilities array aka pi(a|s)\n",
    "        action_probs = predict_probs(np.array([s]))[0]\n",
    "\n",
    "        # Sample action with given probabilities.\n",
    "        a = np.random.choice(n_actions, p=action_probs)\n",
    "        new_s, r, done, info = env.step(a)\n",
    "\n",
    "        # record session history to train later\n",
    "        states.append(s)\n",
    "        actions.append(a)\n",
    "        rewards.append(r)\n",
    "\n",
    "        s = new_s\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "    return states, actions, rewards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test it\n",
    "states, actions, rewards = generate_session()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewards"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing cumulative rewards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_cumulative_rewards(rewards,  # rewards at each step\n",
    "                           gamma=0.99  # discount for reward\n",
    "                           ):\n",
    "    \"\"\"\n",
    "    take a list of immediate rewards r(s,a) for the whole session \n",
    "    compute cumulative returns (a.k.a. G(s,a) in Sutton '16)\n",
    "    G_t = r_t + gamma*r_{t+1} + gamma^2*r_{t+2} + ...\n",
    "\n",
    "    The simple way to compute cumulative rewards is to iterate from last to first time tick\n",
    "    and compute G_t = r_t + gamma*G_{t+1} recurrently\n",
    "\n",
    "    You must return an array/list of cumulative rewards with as many elements as in the initial rewards.\n",
    "    \"\"\"\n",
    "    G = [rewards[-1]]\n",
    "    \n",
    "    for r in rewards[-2::-1]:\n",
    "        G.append(r + gamma * G[-1])\n",
    "    return G[::-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "looks good!\n"
     ]
    }
   ],
   "source": [
    "get_cumulative_rewards(rewards)\n",
    "assert len(get_cumulative_rewards(list(range(100)))) == 100\n",
    "assert np.allclose(get_cumulative_rewards([0, 0, 1, 0, 0, 1, 0], gamma=0.9), [\n",
    "                   1.40049, 1.5561, 1.729, 0.81, 0.9, 1.0, 0.0])\n",
    "assert np.allclose(get_cumulative_rewards(\n",
    "    [0, 0, 1, -2, 3, -4, 0], gamma=0.5), [0.0625, 0.125, 0.25, -1.5, 1.0, -4.0, 0.0])\n",
    "assert np.allclose(get_cumulative_rewards(\n",
    "    [0, 0, 1, 2, 3, 4, 0], gamma=0), [0, 0, 1, 2, 3, 4, 0])\n",
    "print(\"looks good!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Loss function and updates\n",
    "\n",
    "We now need to define objective and update over policy gradient.\n",
    "\n",
    "Our objective function is\n",
    "\n",
    "$$ J \\approx  { 1 \\over N } \\sum  _{s_i,a_i} \\pi_\\theta (a_i | s_i) \\cdot G(s_i,a_i) $$\n",
    "\n",
    "\n",
    "Following the REINFORCE algorithm, we can define our objective as follows: \n",
    "\n",
    "$$ \\hat J \\approx { 1 \\over N } \\sum  _{s_i,a_i} log \\pi_\\theta (a_i | s_i) \\cdot G(s_i,a_i) $$\n",
    "\n",
    "When you compute gradient of that function over network weights $ \\theta $, it will become exactly the policy gradient.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_one_hot(y_tensor, ndims):\n",
    "    \"\"\" helper: take an integer vector and convert it to 1-hot matrix. \"\"\"\n",
    "    y_tensor = y_tensor.type(torch.LongTensor).view(-1, 1)\n",
    "    y_one_hot = torch.zeros(\n",
    "        y_tensor.size()[0], ndims).scatter_(1, y_tensor, 1)\n",
    "    return y_one_hot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your code: define optimizers\n",
    "optimizer = torch.optim.Adam(model.parameters(), 1e-3)\n",
    "\n",
    "\n",
    "def train_on_session(states, actions, rewards, gamma=0.99, entropy_coef=1e-2):\n",
    "    \"\"\"\n",
    "    Takes a sequence of states, actions and rewards produced by generate_session.\n",
    "    Updates agent's weights by following the policy gradient above.\n",
    "    Please use Adam optimizer with default parameters.\n",
    "    \"\"\"\n",
    "\n",
    "    # cast everything into torch tensors\n",
    "    states = torch.tensor(states, dtype=torch.float32)\n",
    "    actions = torch.tensor(actions, dtype=torch.int32)\n",
    "    cumulative_returns = np.array(get_cumulative_rewards(rewards, gamma))\n",
    "    cumulative_returns = torch.tensor(cumulative_returns, dtype=torch.float32)\n",
    "\n",
    "    # predict logits, probas and log-probas using an agent.\n",
    "    logits = model(states)\n",
    "    probs = F.softmax(logits, -1)\n",
    "    log_probs = F.log_softmax(logits, -1)\n",
    "\n",
    "    assert all(isinstance(v, torch.Tensor) for v in [logits, probs, log_probs]), \\\n",
    "        \"please use compute using torch tensors and don't use predict_probs function\"\n",
    "\n",
    "    # select log-probabilities for chosen actions, log pi(a_i|s_i)\n",
    "    log_probs_for_actions = torch.sum(\n",
    "        log_probs * to_one_hot(actions, env.action_space.n), dim=1)\n",
    "    J_hat = torch.mean(log_probs_for_actions * cumulative_returns)\n",
    "    \n",
    "    # Compute loss here. Don't forgen entropy regularization with `entropy_coef` \n",
    "    entropy = -(probs * log_probs).sum(-1).mean()\n",
    "    loss = - J_hat - entropy_coef * entropy\n",
    "\n",
    "    # Gradient descent step\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # technical: return session rewards to print them later\n",
    "    return np.sum(rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The actual training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:24.110\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:25.350\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:47.590\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:55.780\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:72.520\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:156.350\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:132.210\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:277.970\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:162.590\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:359.230\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:256.970\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:232.730\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:272.090\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:260.710\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:298.500\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:269.110\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:167.300\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:161.170\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean reward:528.940\nYou Win!\n"
     ]
    }
   ],
   "source": [
    "for i in range(100):\n",
    "    rewards = [train_on_session(*generate_session())\n",
    "               for _ in range(100)]  # generate new sessions\n",
    "    print(\"mean reward:%.3f\" % (np.mean(rewards)))\n",
    "    if np.mean(rewards) > 500:\n",
    "        print(\"You Win!\")  # but you can train even further\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "# record sessions\n",
    "import gym.wrappers\n",
    "env = gym.wrappers.Monitor(gym.make(\"CartPole-v0\"),\n",
    "                           directory=\"videos\", force=True)\n",
    "sessions = [generate_session() for _ in range(100)]\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<video width=\"640\" height=\"480\" controls>\n",
       "  <source src=\"./videos/openaigym.video.0.5816.video000064.mp4\" type=\"video/mp4\">\n",
       "</video>\n"
      ],
      "text/plain": [
       "\n",
       "<video width=\"640\" height=\"480\" controls>\n",
       "  <source src=\"./videos/openaigym.video.0.5816.video000064.mp4\" type=\"video/mp4\">\n",
       "</video>\n"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# show video\n",
    "from IPython.display import HTML\n",
    "import os\n",
    "\n",
    "video_names = list(\n",
    "    filter(lambda s: s.endswith(\".mp4\"), os.listdir(\"./videos/\")))\n",
    "\n",
    "HTML(\"\"\"\n",
    "<video width=\"640\" height=\"480\" controls>\n",
    "  <source src=\"{}\" type=\"video/mp4\">\n",
    "</video>\n",
    "\"\"\".format(\"./videos/\"+video_names[-1]))  # this may or may not be the _last_ video. Try other indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}