day4/rl/policy_gradient_example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "policy_gradient_example.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/lisaong/diec/blob/master/day4/rl/policy_gradient_example.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jqQIepT-gLFZ",
        "colab_type": "text"
      },
      "source": [
        "Training an Agent using Policy Gradient on Open AI Gym.\n",
        "\n",
        "https://gist.github.com/karpathy/a4166c7fe253700972fcbc77e4ea32c5"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RLd3ritKfjmw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\" Trains an agent with (stochastic) Policy Gradients on Pong. Uses OpenAI Gym. \"\"\"\n",
        "\n",
        "import numpy as np\n",
        "import pickle\n",
        "import gym"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "szlkVIDzfdSY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# hyperparameters\n",
        "H = 200 # number of hidden layer neurons\n",
        "batch_size = 10 # every how many episodes to do a param update?\n",
        "learning_rate = 1e-4\n",
        "gamma = 0.99 # discount factor for reward\n",
        "decay_rate = 0.99 # decay factor for RMSProp leaky sum of grad^2\n",
        "resume = False # resume from previous checkpoint?\n",
        "render = False"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B1yuAnOwgZo4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# model initialization\n",
        "D = 80 * 80 # input dimensionality: 80x80 grid\n",
        "if resume:\n",
        "  model = pickle.load(open('save.p', 'rb'))\n",
        "else:\n",
        "  model = {}\n",
        "  model['W1'] = np.random.randn(H,D) / np.sqrt(D) # \"Xavier\" initialization\n",
        "  model['W2'] = np.random.randn(H) / np.sqrt(H)\n",
        "  \n",
        "grad_buffer = { k : np.zeros_like(v) for k,v in model.items() } # update buffers that add up gradients over a batch\n",
        "rmsprop_cache = { k : np.zeros_like(v) for k,v in model.items() } # rmsprop memory"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_9EwbqpNgdBp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Policy Gradient Implementation\n",
        "\n",
        "def sigmoid(x): \n",
        "  return 1.0 / (1.0 + np.exp(-x)) # sigmoid \"squashing\" function to interval [0,1]\n",
        "\n",
        "def prepro(I):\n",
        "  \"\"\" prepro 210x160x3 uint8 frame into 6400 (80x80) 1D float vector \"\"\"\n",
        "  I = I[35:195] # crop\n",
        "  I = I[::2,::2,0] # downsample by factor of 2\n",
        "  I[I == 144] = 0 # erase background (background type 1)\n",
        "  I[I == 109] = 0 # erase background (background type 2)\n",
        "  I[I != 0] = 1 # everything else (paddles, ball) just set to 1\n",
        "  return I.astype(np.float).ravel()\n",
        "\n",
        "def discount_rewards(r):\n",
        "  \"\"\" take 1D float array of rewards and compute discounted reward \n",
        "      (rewards for steps far into the future are deemed less important)\n",
        "  \"\"\"\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 # reset the sum, since this was a game boundary (pong specific!)\n",
        "    running_add = running_add * gamma + r[t] # gamma: discount factor\n",
        "                                             # higher degree of discount further back in time\n",
        "    discounted_r[t] = running_add\n",
        "  return discounted_r\n",
        "\n",
        "def policy_forward(x):\n",
        "  h = np.dot(model['W1'], x)\n",
        "  h[h<0] = 0 # ReLU nonlinearity\n",
        "  logp = np.dot(model['W2'], h)\n",
        "  p = sigmoid(logp)\n",
        "  return p, h # return probability of taking action 2, and hidden state\n",
        "\n",
        "def policy_backward(eph, epdlogp):\n",
        "  \"\"\" backward pass. (eph is array of intermediate hidden states) \n",
        "      https://medium.com/@jonathan_hui/rl-policy-gradients-explained-9b13b688b146  \n",
        "  \"\"\"\n",
        "  # h = relu(W1 * x)\n",
        "  # logp = W2 * h = W2 * relu(W1 * x)\n",
        "  dW2 = np.dot(eph.T, epdlogp).ravel()\n",
        "  dh = np.outer(epdlogp, model['W2'])\n",
        "  dh[eph <= 0] = 0 # backprop relu\n",
        "  dW1 = np.dot(dh.T, epx)\n",
        "  return {'W1':dW1, 'W2':dW2}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zg1TG5RVfpNG",
        "colab_type": "code",
        "outputId": "233c238f-110f-4d46-ae29-e57350363c30",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 662
        }
      },
      "source": [
        "# Training using OpenAI Gym's Pong preset enviroment\n",
        "\n",
        "render = False\n",
        "env = gym.make(\"Pong-v0\")\n",
        "observation = env.reset()\n",
        "prev_x = None # used in computing the difference frame\n",
        "xs,hs,dlogps,drs = [],[],[],[]\n",
        "running_reward = None\n",
        "reward_sum = 0\n",
        "episode_number = 0\n",
        "while True:\n",
        "  if render: env.render()\n",
        "\n",
        "  # preprocess the observation, set input to network to be difference image\n",
        "  cur_x = prepro(observation)\n",
        "  x = cur_x - prev_x if prev_x is not None else np.zeros(D)\n",
        "  prev_x = cur_x\n",
        "\n",
        "  # forward the policy network and sample an action from the returned probability\n",
        "  aprob, h = policy_forward(x)\n",
        "  action = 2 if np.random.uniform() < aprob else 3 # roll the dice!\n",
        "\n",
        "  # record various intermediates (needed later for backprop)\n",
        "  xs.append(x) # observation\n",
        "  hs.append(h) # hidden state\n",
        "  y = 1 if action == 2 else 0 # a \"fake label\"\n",
        "  dlogps.append(y - aprob) # grad that encourages the action that was taken to be taken (see http://cs231n.github.io/neural-networks-2/#losses if confused)\n",
        "\n",
        "  # step the environment and get new measurements\n",
        "  observation, reward, done, info = env.step(action)\n",
        "\n",
        "  reward_sum += reward\n",
        "\n",
        "  drs.append(reward) # record reward (has to be done after we call step() to get reward for previous action)\n",
        "\n",
        "  if done: # an episode finished\n",
        "    print(f'Episode {episode_number}: input: {x}, reward: {reward}')\n",
        "\n",
        "    episode_number += 1\n",
        "\n",
        "    # stack together all inputs, hidden states, action gradients, and rewards for this episode\n",
        "    epx = np.vstack(xs) # inputs\n",
        "    eph = np.vstack(hs) # hidden states\n",
        "    epdlogp = np.vstack(dlogps) # action gradients\n",
        "    epr = np.vstack(drs) # rewards\n",
        "    xs,hs,dlogps,drs = [],[],[],[] # reset array memory\n",
        "\n",
        "    # compute the discounted reward backwards through time\n",
        "    discounted_epr = discount_rewards(epr)\n",
        "\n",
        "    # standardize the rewards to be unit normal (helps control the gradient estimator variance)\n",
        "    discounted_epr -= np.mean(discounted_epr)\n",
        "    discounted_epr /= np.std(discounted_epr)\n",
        "\n",
        "    epdlogp *= discounted_epr # modulate the gradient with advantage (PG magic happens right here.)\n",
        "    grad = policy_backward(eph, epdlogp)\n",
        "    for k in model: grad_buffer[k] += grad[k] # accumulate grad over batch\n",
        "\n",
        "    # perform rmsprop parameter update every batch_size episodes\n",
        "    if episode_number % batch_size == 0:\n",
        "      for k,v in model.items():\n",
        "        g = grad_buffer[k] # gradient\n",
        "        rmsprop_cache[k] = decay_rate * rmsprop_cache[k] + (1 - decay_rate) * g**2\n",
        "        model[k] += learning_rate * g / (np.sqrt(rmsprop_cache[k]) + 1e-5)\n",
        "        grad_buffer[k] = np.zeros_like(v) # reset batch gradient buffer\n",
        "\n",
        "    # boring book-keeping\n",
        "    running_reward = reward_sum if running_reward is None else running_reward * 0.99 + reward_sum * 0.01\n",
        "    print(f'resetting env. episode reward total was {reward_sum}. running mean: {running_reward}')\n",
        "    if episode_number % 100 == 0: pickle.dump(model, open('save.p', 'wb'))\n",
        "    reward_sum = 0\n",
        "    observation = env.reset() # reset env\n",
        "    prev_x = None\n",
        "\n",
        "  #if reward != 0: # Pong has either +1 or -1 reward exactly when game ends.\n",
        "  #  print (f'ep {episode_number}: game finished, reward: {reward}')\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Episode 0: input: [0. 0. 0. ... 0. 1. 0.], reward: -1.0\n",
            "resetting env. episode reward total was -21.0. running mean: -21.0\n",
            "Episode 1: input: [0. 0. 0. ... 0. 1. 0.], reward: -1.0\n",
            "resetting env. episode reward total was -21.0. running mean: -21.0\n",
            "Episode 2: input: [0. 0. 0. ... 0. 1. 0.], reward: -1.0\n",
            "resetting env. episode reward total was -21.0. running mean: -21.0\n",
            "Episode 3: input: [0. 0. 0. ... 0. 0. 0.], reward: -1.0\n",
            "resetting env. episode reward total was -21.0. running mean: -21.0\n",
            "Episode 4: input: [0. 0. 0. ... 0. 0. 0.], reward: -1.0\n",
            "resetting env. episode reward total was -20.0. running mean: -20.99\n",
            "Episode 5: input: [0. 0. 0. ... 0. 0. 0.], reward: -1.0\n",
            "resetting env. episode reward total was -21.0. running mean: -20.990099999999998\n",
            "Episode 6: input: [0. 0. 0. ... 0. 0. 0.], reward: -1.0\n",
            "resetting env. episode reward total was -21.0. running mean: -20.990199\n",
            "Episode 7: input: [0. 0. 0. ... 0. 0. 0.], reward: -1.0\n",
            "resetting env. episode reward total was -21.0. running mean: -20.990297010000003\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-28-e1b3842c8306>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m   \u001b[0;31m# forward the policy network and sample an action from the returned probability\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 19\u001b[0;31m   \u001b[0maprob\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mh\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpolicy_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     20\u001b[0m   \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m2\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muniform\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0maprob\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;36m3\u001b[0m \u001b[0;31m# roll the dice!\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-18-aff9163f3ed2>\u001b[0m in \u001b[0;36mpolicy_forward\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mpolicy_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m   \u001b[0mh\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'W1'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m   \u001b[0mh\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mh\u001b[0m\u001b[0;34m<\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;31m# ReLU nonlinearity\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m   \u001b[0mlogp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'W2'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mh\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<__array_function__ internals>\u001b[0m in \u001b[0;36mdot\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fWkZ-UIvhoCe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}