BipedalWalker_COGS_workshop.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "BipedalWalker_COGS_workshop.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/turing-club/info/blob/master/BipedalWalker_COGS_workshop.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bxhu8z6rjs32",
        "colab_type": "text"
      },
      "source": [
        "# Bipedal Walker - reinforcement learning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ys3KESS-8Vlc",
        "colab_type": "text"
      },
      "source": [
        "## Introduction\n",
        "We've just shown you how reinforcement learning can be used to let a computer learn a psychology task. Now, we're going to help a robot learn how to walk! \n",
        "\n",
        "The process by which the robot will learn the task is based on the exact same principles as the process by which a computer learns to do a go/no-go task. \n",
        "\n",
        "Don't be intimidated by all of the code below; we'll take you through the most important parts!\n",
        "\n",
        "Source: https://github.com/mayurmadnani/BipedalWalker"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n-2Im14N9Hlb",
        "colab_type": "text"
      },
      "source": [
        "## Installation\n",
        "We'll be using OpenAI Gym, a framework that provides ton of games for the AIs we create to learn and play. You can see all of the available environments here: https://gym.openai.com/envs/#classic_control\n",
        "\n",
        "The environment we will use is BipedalWalker-v2. This is a simulation of a robot that will try to learn to walk on two legs and reach the other side of the screen - all by itself!\n",
        "\n",
        "Let's phrase the robot's challenge in the terms we've been using so far:\n",
        "The robot's **stimulus/environment** is a **grassy field**.\n",
        "The robot's **decision** will be how to align its legs to stand up and walk.\n",
        "Depending on the robot's decision, the **outcome** will either be a reward (a few points for walking forward) or a punishment (-100 points for falling down).\n",
        "\n",
        "The robot knows:\n",
        "- How fast it's going\n",
        "- How its feet touch the ground\n",
        "- How far away the ground is from its head (using a sensor of the kind found in self-driving cars)\n",
        "- Whether it's been rewarded or punished for its actions\n",
        "\n",
        "Let's install some software that will let us render the robot's little world into a video."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H18Dq-YjR7k7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install gym==0.12.1\n",
        "!apt-get update\n",
        "!apt-get -qq -y install xvfb freeglut3-dev ffmpeg> /dev/null\n",
        "!apt-get install xvfb\n",
        "!pip install pyvirtualdisplay\n",
        "!pip -q install pyglet\n",
        "!pip -q install pyopengl"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qhG5aATO-6cw",
        "colab_type": "text"
      },
      "source": [
        "Now, let's install the physics engine that defines the world in which our robot lives."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Eok_4dD9mFE0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!apt-get install swig x11-utils\n",
        "!pip install box2d box2d-kengz\n",
        "!pip install pybullet"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wiJyZxmB_ADe",
        "colab_type": "text"
      },
      "source": [
        "\"Hello world!\" Let's set up our robot's world."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8KjX9UtlSWKB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Start virtual display\n",
        "from pyvirtualdisplay import Display\n",
        "display = Display(visible=0, size=(1024, 768))\n",
        "display.start()\n",
        "import os\n",
        "os.environ[\"DISPLAY\"] = \":\" + str(display.display) + \".\" + str(display.screen)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q9UYPMYakrpL",
        "colab_type": "text"
      },
      "source": [
        "# Augmented Random Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kIJ1takz_JqZ",
        "colab_type": "text"
      },
      "source": [
        "**Augmented random search** is a policy that the robot will use to explore the world (and learn to walk). It will start out by trying out random movements, watch the rewards and punishments it gets for each action, and using the information it has about the world to improve its performance."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Azjqd6JOVcy7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import gym\n",
        "from gym import wrappers"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gx4x_-_Q_pCA",
        "colab_type": "text"
      },
      "source": [
        "**Hyperparameters** are variables that define how fast our robot's brain will work. These parameters require a lot of tuning in order to get a good model that will enable the robot to learn efficently."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gRtzFlKFVpAb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class HP():\n",
        "    # Hyperparameters\n",
        "    def __init__(self,\n",
        "                 nb_steps=1000,\n",
        "                 episode_length=2000,\n",
        "                 learning_rate=0.02,\n",
        "                 num_deltas=16,\n",
        "                 num_best_deltas=16,\n",
        "                 noise=0.03,\n",
        "                 seed=1,\n",
        "                 env_name='BipedalWalker-v2',\n",
        "                 record_every=50):\n",
        "\n",
        "        self.nb_steps = nb_steps\n",
        "        self.episode_length = episode_length\n",
        "        self.learning_rate = learning_rate\n",
        "        self.num_deltas = num_deltas\n",
        "        self.num_best_deltas = num_best_deltas\n",
        "        assert self.num_best_deltas <= self.num_deltas\n",
        "        self.noise = noise\n",
        "        self.seed = seed\n",
        "        self.env_name = env_name\n",
        "        self.record_every = record_every"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aESkk4BlAQwA",
        "colab_type": "text"
      },
      "source": [
        "Have you taken PSYC 218? If so, this next block of code will make a lot of sense to you. We want to normalize all of the robot's behaviours (the size of its step, speed of walking, etc.) so that we can directly compare parameters with different units. To do this, we'll convert the robot's movement parameters to z-scores!\n",
        "\n",
        "It mostly happens in the following line:\n",
        "```\n",
        "return (inputs - obs_mean) / obs_std\n",
        "```\n",
        "This is (current position - mean position) / standard deviation, which is the formula for calculating a z-score!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qcwrRlAlVpI8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Normalizer():\n",
        "    # Normalizes the inputs\n",
        "    def __init__(self, nb_inputs):\n",
        "        self.n = np.zeros(nb_inputs)\n",
        "        self.mean = np.zeros(nb_inputs)\n",
        "        self.mean_diff = np.zeros(nb_inputs)\n",
        "        self.var = np.zeros(nb_inputs)\n",
        "\n",
        "    def observe(self, x):\n",
        "        self.n += 1.0\n",
        "        last_mean = self.mean.copy()\n",
        "        self.mean += (x - self.mean) / self.n\n",
        "        self.mean_diff += (x - last_mean) * (x - self.mean)\n",
        "        self.var = (self.mean_diff / self.n).clip(min = 1e-2)\n",
        "\n",
        "    def normalize(self, inputs):\n",
        "        obs_mean = self.mean\n",
        "        obs_std = np.sqrt(self.var)\n",
        "        return (inputs - obs_mean) / obs_std"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Avy4K-WgBVxp",
        "colab_type": "text"
      },
      "source": [
        "Here, we evaluate how well our policy (the robot's program of action) is working, and update its behavioural model accordingly.\n",
        "\n",
        "Focus on this line:\n",
        "\n",
        "\n",
        "```\n",
        "self.theta += self.hp.learning_rate / (self.hp.num_best_deltas * sigma_rewards) * step\n",
        "```\n",
        "**self.theta** is the weight assigned to a particular behaviour. From a neuroscience perspective, you can think of it as a long-term potentiation or inhibition (strengthening or weakening) of a neural connection. If the behaviour is beneficial, then you want to strengthen the neurons driving that behaviour (and vice versa).\n",
        "\n",
        "**self.hp.learning_rate** is the speed at which the model learns. If the model learns too fast, it may overshoot the optimal setup for making a good decision (i.e. making too large of a step and falling over). \n",
        "\n",
        "**self.hp.num_best_deltas** represents the robot's past experiences. Here, the robot is 'thinking' about the decisions it made in the past, and the optimal deltas - the optimal shifts in strategy it has made.\n",
        "\n",
        "**sigma_rewards** is the sum of the rewards the robot has gained for making specific behaviours of that type before. If the action resulted in the robot making progress towards the goal in the past, it's assigned a positive reward. If the action resulted in the robot falling, it's assigned a punishment (negative reward).\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SAakK-txVpQR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Policy():\n",
        "    def __init__(self, input_size, output_size, hp):\n",
        "        self.theta = np.zeros((output_size, input_size))\n",
        "        self.hp = hp\n",
        "\n",
        "    def evaluate(self, input, delta = None, direction = None):\n",
        "        if direction is None:\n",
        "            return self.theta.dot(input)\n",
        "        elif direction == \"+\":\n",
        "            return (self.theta + self.hp.noise * delta).dot(input)\n",
        "        elif direction == \"-\":\n",
        "            return (self.theta - self.hp.noise * delta).dot(input)\n",
        "\n",
        "    def sample_deltas(self):\n",
        "        return [np.random.randn(*self.theta.shape) for _ in range(self.hp.num_deltas)]\n",
        "\n",
        "    def update(self, rollouts, sigma_rewards):\n",
        "        # sigma_rewards is the standard deviation of the rewards\n",
        "        step = np.zeros(self.theta.shape)\n",
        "        for r_pos, r_neg, delta in rollouts:\n",
        "            step += (r_pos - r_neg) * delta\n",
        "        self.theta += self.hp.learning_rate / (self.hp.num_best_deltas * sigma_rewards) * step"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CasihAg1DEFe",
        "colab_type": "text"
      },
      "source": [
        "Now, it's time to define how our robot is going to learn. We're putting together the world, normalizer, and policy that we defined in the previous cells. \n",
        "\n",
        "In **explore**, we're telling the robot to go forth into this brave new world, try out various strategies according to its policies, and note down the rewards and punishments it gets. \n",
        "\n",
        "In **train**, it uses the rewards and punishments that it's faced in **explore** to update its behaviour for the next strategy that it will try."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DgzIRF71VxNm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ARSTrainer():\n",
        "    def __init__(self,\n",
        "                 hp=None,\n",
        "                 input_size=None,\n",
        "                 output_size=None,\n",
        "                 normalizer=None,\n",
        "                 policy=None,\n",
        "                 monitor_dir=None):\n",
        "\n",
        "        self.hp = hp or HP()\n",
        "        np.random.seed(self.hp.seed)\n",
        "        self.env = gym.make(self.hp.env_name)\n",
        "        if monitor_dir is not None:\n",
        "            should_record = lambda i: self.record_video\n",
        "            self.env = wrappers.Monitor(self.env, monitor_dir, video_callable=should_record, force=True)\n",
        "        self.hp.episode_length = self.env.spec.timestep_limit or self.hp.episode_length\n",
        "        self.input_size = input_size or self.env.observation_space.shape[0]\n",
        "        self.output_size = output_size or self.env.action_space.shape[0]\n",
        "        self.normalizer = normalizer or Normalizer(self.input_size)\n",
        "        self.policy = policy or Policy(self.input_size, self.output_size, self.hp)\n",
        "        self.record_video = False\n",
        "\n",
        "    # Explore the policy on one specific direction and over one episode\n",
        "    def explore(self, direction=None, delta=None):\n",
        "        state = self.env.reset()\n",
        "        done = False\n",
        "        num_plays = 0.0\n",
        "        sum_rewards = 0.0\n",
        "        while not done and num_plays < self.hp.episode_length:\n",
        "            self.normalizer.observe(state)\n",
        "            state = self.normalizer.normalize(state)\n",
        "            action = self.policy.evaluate(state, delta, direction)\n",
        "            state, reward, done, _ = self.env.step(action)\n",
        "            reward = max(min(reward, 1), -1)\n",
        "            sum_rewards += reward\n",
        "            num_plays += 1\n",
        "        return sum_rewards\n",
        "\n",
        "    def train(self):\n",
        "        for step in range(self.hp.nb_steps):\n",
        "            # initialize the random noise deltas and the positive/negative rewards\n",
        "            deltas = self.policy.sample_deltas()\n",
        "            positive_rewards = [0] * self.hp.num_deltas\n",
        "            negative_rewards = [0] * self.hp.num_deltas\n",
        "\n",
        "            # play an episode each with positive deltas and negative deltas, collect rewards\n",
        "            for k in range(self.hp.num_deltas):\n",
        "                positive_rewards[k] = self.explore(direction=\"+\", delta=deltas[k])\n",
        "                negative_rewards[k] = self.explore(direction=\"-\", delta=deltas[k])\n",
        "                \n",
        "            # Compute the standard deviation of all rewards\n",
        "            sigma_rewards = np.array(positive_rewards + negative_rewards).std()\n",
        "\n",
        "            # Sort the rollouts by the max(r_pos, r_neg) and select the deltas with best rewards\n",
        "            scores = {k:max(r_pos, r_neg) for k,(r_pos,r_neg) in enumerate(zip(positive_rewards, negative_rewards))}\n",
        "            order = sorted(scores.keys(), key = lambda x:scores[x], reverse = True)[:self.hp.num_best_deltas]\n",
        "            rollouts = [(positive_rewards[k], negative_rewards[k], deltas[k]) for k in order]\n",
        "\n",
        "            # Update the policy\n",
        "            self.policy.update(rollouts, sigma_rewards)\n",
        "\n",
        "            # Only record video during evaluation, every n steps\n",
        "            if step % self.hp.record_every == 0:\n",
        "                self.record_video = True\n",
        "            # Play an episode with the new weights and print the score\n",
        "            reward_evaluation = self.explore()\n",
        "            print('Step: ', step, 'Reward: ', reward_evaluation)\n",
        "            self.record_video = False"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "09Sl4COGKCGS",
        "colab_type": "text"
      },
      "source": [
        "This block of code just sets up a directory to which we'll save the videos of the robot's behaviour."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p2aIv0lMV3_F",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def mkdir(base, name):\n",
        "    path = os.path.join(base, name)\n",
        "    if not os.path.exists(path):\n",
        "        os.makedirs(path)\n",
        "    return path"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bfO9KDnYKLYb",
        "colab_type": "text"
      },
      "source": [
        "We've defined all of the parameters of the model - now, it's time to unleash the robot upon the little world we've created for it!\n",
        "\n",
        "This step will begin the robot's training. Please be patient - it takes a while for the robot to be able to walk. In the meantime, you can watch the robot's rewards and punishments. When you feel that you've trained the robot enough, click the play button again to stop the training."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IVuDOFFVV9xZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "ENV_NAME = 'BipedalWalker-v2'\n",
        "\n",
        "videos_dir = mkdir('.', 'videos')\n",
        "monitor_dir = mkdir(videos_dir, ENV_NAME)\n",
        "\n",
        "hp = HP(env_name=ENV_NAME)\n",
        "trainer = ARSTrainer(hp=hp, monitor_dir=monitor_dir)\n",
        "trainer.train()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4C_waL3_lKXu",
        "colab_type": "text"
      },
      "source": [
        "# Download the episodes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LbtWobZFKxpo",
        "colab_type": "text"
      },
      "source": [
        "Congratulations, the robot has (hopefully) learned to walk a bit! To see how it performed, run the following code blocks to download a video of the robot's behaviour. That's it!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5xvR-vXbBLB1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!ls videos/{ENV_NAME}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MYitauj1SePX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import files\n",
        "import glob\n",
        "\n",
        "for file in glob.glob(\"videos/{}/openaigym.video.*.mp4\".format(ENV_NAME)):\n",
        "  files.download(file)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}