{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hackathon 11\n",
    "\n",
    "Topics:\n",
    "- Keras on TensorFlow\n",
    "- Policy Gradients\n",
    "\n",
    "Today's hackathon will be short. I want to communicate the surprising simplicity of policy gradients.\n",
    "\n",
    "This is all setup in a IPython notebook so you can run any code you want to experiment with. Feel free to edit any cell, or add some to run your own code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Today, we'll use another deep learning framework called Keras. Keras is a python library designed to facilitate rapid prototyping and iteration of deep networks (with the drawback of some inflexibility). It can run atop TensorFlow, Theano, or most recently, CNTK. Our kernel is running it with TensorFlow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import PIL    # Python Image Library (for resizing and greyscaling images)\n",
    "import gym\n",
    "import tqdm   # make your loops show a smart progress meter\n",
    "\n",
    "import keras\n",
    "from keras import Sequential\n",
    "from keras.layers import Dense, Flatten, Reshape\n",
    "from keras.layers import Conv2D, MaxPooling2D\n",
    "from keras.optimizers import Adam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we'll set the constants we'll use later and create the Gym environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "env = gym.make('Breakout-v0').unwrapped\n",
    "\n",
    "# Fixed constants\n",
    "ACTION_SHAPE = env.action_space.n\n",
    "OBS_SHAPE = env.observation_space.shape\n",
    "\n",
    "# Training constants\n",
    "INPUT_SHAPE = (84, 84, 1)\n",
    "LEARNING_RATE = 0.001\n",
    "EPISODES = 100\n",
    "GAMMA = 0.99\n",
    "BATCH_SIZE = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we’re going to define a policy network that implements our player (or “agent”). This network will take the state of the game and decide what we should do (move LEFT or RIGHT, FIRE, or NOP). We'll use a simple neural network that takes the raw image pixels (7,056 numbers total (84\\*84\\*1)), and produces a vector indicating the probability of taking each action. Note that it is standard to use a stochastic policy. Every iteration we will sample from this distribution to get the actual move. The reason for this will become more clear once we talk about training.\n",
    "\n",
    "The core data structure of Keras is a **model**, a way to organize layers. The simplest type of model is the `Sequential` model, a linear stack of layers. It's a Python object with an `add` function to add a new layer. Finally, we declare the optimizer and call `compile` with the loss function (Keras has a number of these standard functions built in to be called by name). Then, training is as simple as calling `fit` or `train_on_batch` and passing in data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Sequential()\n",
    "model.add(Conv2D(32, (6, 6), activation=\"relu\", input_shape=INPUT_SHAPE,\n",
    "                 strides=(3, 3), padding=\"same\", kernel_initializer=\"he_uniform\"))\n",
    "model.add(Flatten())\n",
    "model.add(Dense(64, activation=\"relu\", kernel_initializer=\"he_uniform\"))\n",
    "model.add(Dense(32, activation=\"relu\", kernel_initializer=\"he_uniform\"))\n",
    "model.add(Dense(ACTION_SHAPE, activation='softmax'))\n",
    "opt = Adam(lr=LEARNING_RATE)\n",
    "model.compile(loss='categorical_crossentropy', optimizer=opt)\n",
    "model.summary()  #this gives the fancy summary below"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll preprocess frames by reducing their size to (84, 84), and converting them to grayscale (so we can use just 1 channel)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def preprocess_observation(observation):\n",
    "    img = PIL.Image.fromarray(observation)\n",
    "    img = img.resize(INPUT_SHAPE[:2]).convert('L')  # resize and convert to grayscale\n",
    "    preprocessed_observation = np.expand_dims(np.array(img), axis=2)\n",
    "    return preprocessed_observation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we continue with code, we'll take a detour to talk about the concept of policy gradients, moving from supervised learning. This is from Andrej Karpathy's blog post, [Pong to Pixels](http://karpathy.github.io/2016/05/31/rl/). See his blog for a deeper discussion of policy gradients.\n",
    "\n",
    "> **Supervised Learning.** Before we dive into the Policy Gradients solution I’d like to remind you briefly about supervised learning because, as we’ll see, RL is very similar. Refer to the diagram below. In ordinary supervised learning we would feed an image to the network and get some probabilities, e.g. for two classes UP and DOWN. I’m showing log probabilities (-1.2, -0.36) for UP and DOWN instead of the raw probabilities (30% and 70% in this case) because we always optimize the log probability of the correct label (this makes math nicer, and is equivalent to optimizing the raw probability because log is monotonic). Now, in supervised learning we would have access to a label. For example, we might be told that the correct thing to do right now is to go UP (label 0). In an implementation we would enter gradient of 1.0 on the log probability of UP and run backprop to compute the gradient vector $\\nabla_W\\log p(y=UP∣x)$. This gradient would tell us how we should change every one of our million parameters to make the network slightly more likely to predict UP. For example, one of the million parameters in the network might have a gradient of -2.1, which means that if we were to increase that parameter by a small positive amount (e.g. 0.001), the log probability of UP would decrease by 2.1 * 0.001 (decrease due to the negative sign). If we then did a parameter update then, yay, our network would now be slightly more likely to predict UP when it sees a very similar image in the future.\n",
    "\n",
    "<img src=\"http://karpathy.github.io/assets/rl/sl.png\">\n",
    "\n",
    "> **Policy Gradients**. Okay, but what do we do if we do not have the correct label in the Reinforcement Learning setting? Here is the Policy Gradients solution (again refer to diagram below). Our policy network calculated probability of going UP as 30% (logprob -1.2) and DOWN as 70% (logprob -0.36). We will now sample an action from this distribution; E.g. suppose we sample DOWN, and we will execute it in the game. At this point notice one interesting fact: We could immediately fill in a gradient of 1.0 for DOWN as we did in supervised learning, and find the gradient vector that would encourage the network to be slightly more likely to do the DOWN action in the future. So we can immediately evaluate this gradient and that’s great, but the problem is that at least for now we do not yet know if going DOWN is good. But the critical point is that that’s okay, because we can simply wait a bit and see! For example in Pong we could wait until the end of the game, then take the reward we get (either +1 if we won or -1 if we lost), and enter that scalar as the gradient for the action we have taken (DOWN in this case).\n",
    "\n",
    "<img src=\"http://karpathy.github.io/assets/rl/rl.png\">\n",
    "\n",
    "> And that’s it: we have a stochastic policy that samples actions and then actions that happen to eventually lead to good outcomes get encouraged in the future, and actions taken that lead to bad outcomes get discouraged. Also, the reward does not even need to be +1 or -1 if we win the game eventually. It can be an arbitrary measure of some kind of eventual quality. For example if things turn out really well it could be 10.0, which we would then enter as the gradient instead of -1 to start off backprop. That’s the beauty of neural nets; Using them can feel like cheating: You’re allowed to have 1 million parameters embedded in 1 teraflop of compute and you can make it do arbitrary things with SGD. It shouldn’t work, but amusingly we live in a universe where it does.\n",
    "\n",
    "<img src=\"http://karpathy.github.io/assets/rl/pg.png\">\n",
    "> A visualization of the score function gradient estimator. **Left:** A gaussian distribution and a few samples from it (blue dots). On each blue dot we also plot the gradient of the log probability with respect to the gaussian's mean parameter. The arrow indicates the direction in which the mean of the distribution should be nudged to increase the probability of that sample. **Middle:** Overlay of some score function giving -1 everywhere except +1 in some small regions (note this can be an arbitrary and not necessarily differentiable scalar-valued function). The arrows are now color coded because due to the multiplication in the update we are going to average up all the green arrows, and the negative of the red arrows. **Right:** after parameter update, the green arrows and the reversed red arrows nudge us to left and towards the bottom. Samples from this distribution will now have a higher expected score, as desired."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to collect an entire episode of rewards, we need to accumulate and discount them (as in all reinforcement learning problems) through time. We'll define a quick funtion for that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def discount_rewards(r, gamma):\n",
    "    discounted_r = np.zeros_like(r)\n",
    "    running_add = 0\n",
    "    for t in reversed(range(0, r.size)):\n",
    "        if r[t] != 0: running_add = 0\n",
    "        running_add = running_add * gamma + r[t]\n",
    "        discounted_r[t] = running_add\n",
    "    return discounted_r"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, a function which will play one episode and accumulate the observations and rewards for training. Also, notice that we're using difference frames to allow the observations to capture motion in the observations, and that the rewards are normalized to improve training performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def play_episode(max_frames=1000):\n",
    "    frames, action_probs, actions, rewards = [], [], [], []\n",
    "    prev_obs = np.zeros(INPUT_SHAPE)\n",
    "    obs = preprocess_observation(env.reset())\n",
    "    done = False\n",
    "\n",
    "    frame_count = 0\n",
    "    while not done and frame_count < max_frames:\n",
    "        frame_count += 1\n",
    "        # calculate action probs and act\n",
    "        a_probs = model.predict(np.expand_dims(obs, 0), batch_size=1).flatten()\n",
    "        action = np.random.choice(ACTION_SHAPE, 1, p=a_probs)\n",
    "        raw_obs, reward, done, _ = env.step(action)\n",
    "        \n",
    "        # process raw observation\n",
    "        obs = preprocess_observation(raw_obs) - prev_obs\n",
    "        prev_obs = obs\n",
    "        \n",
    "        # record experience\n",
    "        frames.append(obs)\n",
    "        rewards.append(reward)\n",
    "        action_probs.append(a_probs)\n",
    "        actions.append(action)\n",
    "    \n",
    "    # collect episode experience, discounting and normalizing rewards\n",
    "    ep_frames = np.vstack(frames)\n",
    "    ep_dlogp = np.squeeze(np.eye(ACTION_SHAPE)[actions]) - np.array(action_probs)\n",
    "    ep_rewards = discount_rewards(np.vstack(rewards), GAMMA)\n",
    "    ep_rewards -= np.mean(ep_rewards)\n",
    "    if np.std(ep_rewards) != 0:\n",
    "        ep_rewards /= np.std(ep_rewards)\n",
    "    ep_dlogp *= ep_rewards\n",
    "    \n",
    "    return ep_frames, ep_dlogp, np.array(action_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_x, train_y = [],[]\n",
    "# loop through frames, tracking progress with tqdm\n",
    "for ep in tqdm.tqdm(range(EPISODES)):\n",
    "    ep_obs, ep_dlogp, action_probs = play_episode()\n",
    "    train_x.append(ep_obs)\n",
    "    train_y.append(ep_dlogp)\n",
    "    if ep % BATCH_SIZE == 0:\n",
    "        input_tr_y = action_probs + LEARNING_RATE * np.squeeze(np.vstack(train_y))\n",
    "        model.train_on_batch(np.vstack(train_x).reshape((-1,)+INPUT_SHAPE), input_tr_y)\n",
    "        train_x, train_y = [],[]\n",
    "        # Checkpoint\n",
    "        model.save('./policy_gradient_model')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (keras)",
   "language": "python",
   "name": "keras"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
