Sketching_with_RL_2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Sketching with RL 2.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/astrivedi/CSCI2270/blob/master/Sketching_with_RL_2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eg_dCVlzEsVq",
        "outputId": "aed7a188-8f63-4b67-9f94-2406fd258d37"
      },
      "source": [
        "!pip install torch tqdm"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (1.9.0+cu102)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (4.41.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch) (3.7.4.3)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QWOl6SM9Ed4r"
      },
      "source": [
        "# Importing\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import random\n",
        "from tqdm import tqdm\n",
        "\n",
        "# Simple feedforward network with tanh activations\n",
        "class FFN(nn.Module):\n",
        "    def __init__(self, in_dim, out_dim, hidden_dim=32, hidden_layers=2):\n",
        "        super().__init__()\n",
        "        \n",
        "        seq = []\n",
        "        # Input layer\n",
        "        seq.append(nn.Linear(in_dim, hidden_dim))\n",
        "        seq.append(nn.Tanh())\n",
        "        # Hidden layers\n",
        "        for _ in range(hidden_layers-1):\n",
        "            seq.append(nn.Linear(hidden_dim, hidden_dim))\n",
        "            seq.append(nn.Tanh())\n",
        "        # Output layer\n",
        "        seq.append(nn.Linear(hidden_dim, out_dim))\n",
        "        \n",
        "        self.model = nn.Sequential(*seq)\n",
        "    \n",
        "    def forward(self, x):\n",
        "        return self.model(x)\n",
        "\n",
        "# Define Filler class which will fill holes\n",
        "class Filler():\n",
        "    def __init__(self):\n",
        "        # Evaluation/training mode\n",
        "        self.evaluation_mode = False\n",
        "        # Define torch device\n",
        "        self.device = 'cpu'\n",
        "        # List of networks for each hole\n",
        "        self.holes = []\n",
        "        # prob_action history\n",
        "        self.history = []\n",
        "        # Accumulated reward per program run\n",
        "        self.reward = 0\n",
        "    \n",
        "    def evaluate(self):\n",
        "        self.evaluation_mode = True\n",
        "    \n",
        "    def add_hole(self, size):\n",
        "        self.holes.append(FFN(size, 2).to(self.device))\n",
        "        # Return hole number\n",
        "        return len(self.holes)-1\n",
        "    \n",
        "    def less_than_equal_assert(self, a, b):\n",
        "        if a <= b:\n",
        "            self.reward += 1\n",
        "        else:\n",
        "            self.reward += -1\n",
        "    \n",
        "    def except_assert(self):\n",
        "        self.reward -= 10\n",
        "    \n",
        "    def reset(self):\n",
        "        self.history = []\n",
        "        self.reward = 0\n",
        "    \n",
        "    def sample(self, hole_number, *inputs):\n",
        "        state = torch.tensor(inputs, device=self.device).float().view(1, -1)\n",
        "        net = self.holes[hole_number]\n",
        "        \n",
        "        if self.evaluation_mode:\n",
        "            _, ind = torch.max(net(state[0]), dim=0) # Take most likely action (no sampling)\n",
        "            return bool(ind)\n",
        "        else :\n",
        "            p = torch.softmax(net(state), dim=1)\n",
        "            dist = torch.distributions.Categorical(p)\n",
        "            action = dist.sample()\n",
        "            self.history.append(p[0,action])\n",
        "            return bool(action)\n",
        "    \n",
        "    def train(self, function, \\\n",
        "              lr=1e-3, baseline_lr=0.1, num_train_steps=500):\n",
        "        parameters = []\n",
        "        for net in self.holes:\n",
        "            parameters += list(net.parameters())\n",
        "        optim = torch.optim.Adam(parameters, lr=lr)\n",
        "\n",
        "        baseline = torch.zeros(1).to(self.device)\n",
        "\n",
        "        for train_step in tqdm(range(num_train_steps)):\n",
        "            # Reset history and reward\n",
        "            self.reset()\n",
        "            \n",
        "            # Call function (random inputs generated internally to function)\n",
        "            try:\n",
        "                function()\n",
        "            except:\n",
        "                self.except_assert()\n",
        "            \n",
        "            if len(self.history) != 0:\n",
        "                # Update baseline\n",
        "                baseline = (1-baseline_lr)*baseline + baseline_lr*self.reward\n",
        "                # Update holes\n",
        "                loss = 0\n",
        "                for p_action in self.history:\n",
        "                    loss += -torch.log(p_action)*(self.reward-baseline)\n",
        "\n",
        "                optim.zero_grad()\n",
        "                loss.backward()\n",
        "                optim.step()    "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "De2zGc0BEi29",
        "outputId": "658c2474-68e5-4ba3-8830-9b23e7a3c91b"
      },
      "source": [
        "filler = Filler()\n",
        "hole = filler.add_hole(1)\n",
        "\n",
        "# The program to fill\n",
        "def absolute_value(n):\n",
        "    if filler.sample(hole, n):\n",
        "        return n\n",
        "    else:\n",
        "        return -n\n",
        "    \n",
        "def test_abs():\n",
        "    n = random.randint(-10, 10)\n",
        "    abs_n = absolute_value(n)\n",
        "    filler.less_than_equal_assert(0, abs_n)\n",
        "\n",
        "filler.train(test_abs)\n",
        "\n",
        "# Put in eval mode\n",
        "filler.evaluate()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 500/500 [00:00<00:00, 551.62it/s]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MlLxikNSEmQR",
        "outputId": "0788dedd-abf9-4ad5-b9e6-c6b7183dea47"
      },
      "source": [
        "for x in range(-10,11):\n",
        "    print('Absolute value of ' + str(x) + ' is ' + str(absolute_value(x)))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Absolute value of -10 is 10\n",
            "Absolute value of -9 is 9\n",
            "Absolute value of -8 is 8\n",
            "Absolute value of -7 is 7\n",
            "Absolute value of -6 is 6\n",
            "Absolute value of -5 is 5\n",
            "Absolute value of -4 is 4\n",
            "Absolute value of -3 is 3\n",
            "Absolute value of -2 is 2\n",
            "Absolute value of -1 is 1\n",
            "Absolute value of 0 is 0\n",
            "Absolute value of 1 is 1\n",
            "Absolute value of 2 is 2\n",
            "Absolute value of 3 is 3\n",
            "Absolute value of 4 is 4\n",
            "Absolute value of 5 is 5\n",
            "Absolute value of 6 is 6\n",
            "Absolute value of 7 is 7\n",
            "Absolute value of 8 is 8\n",
            "Absolute value of 9 is 9\n",
            "Absolute value of 10 is 10\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WggRLqX6EotH",
        "outputId": "25ad2286-1762-4f52-edc9-9d1015730496"
      },
      "source": [
        "filler = Filler()\n",
        "hole0 = filler.add_hole(2)\n",
        "hole1 = filler.add_hole(2)\n",
        "\n",
        "# Program with holes\n",
        "def partition(l, low, high):\n",
        "    pivot = l[high]\n",
        "    swap_index = low\n",
        "    \n",
        "    for t in range(low, high+1):\n",
        "        if filler.sample(hole0, l[t], pivot): # l[t] < pivot:\n",
        "            l[swap_index], l[t] = l[t], l[swap_index]\n",
        "            swap_index += 1\n",
        "    l[swap_index], l[high] = l[high], l[swap_index]\n",
        "    \n",
        "    return swap_index\n",
        "\n",
        "def quick_sort(l, low, high, depth=0):\n",
        "    # Bound recursion depth\n",
        "    if depth > len(l):\n",
        "        return\n",
        "    \n",
        "    if filler.sample(hole1, low, high): # low < high\n",
        "        p = partition(l, low, high)\n",
        "        quick_sort(l, low, p - 1, depth=depth+1)\n",
        "        quick_sort(l, p + 1, high, depth=depth+1)   \n",
        "\n",
        "# Test on random lists\n",
        "def test_quick_sort(max_length=10, val_range=(-10,10)):\n",
        "    length = random.randint(0, max_length)\n",
        "    l = [random.randint(val_range[0], val_range[1]) for _ in range(length)]\n",
        "\n",
        "    quick_sort(l, 0, len(l)-1)\n",
        "    \n",
        "    # Assert that the list is sorted\n",
        "    if len(l) > 1:\n",
        "        for t in range(0, len(l)-1):\n",
        "            filler.less_than_equal_assert(l[t], l[t+1]) #assert(l[t] <= l[t+1])\n",
        "\n",
        "filler.train(test_quick_sort, num_train_steps=4000)\n",
        "\n",
        "filler.evaluate()            "
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 4000/4000 [00:28<00:00, 142.25it/s]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5YVur-kIEq3g",
        "outputId": "53aeca66-fd9a-4fdc-c337-0456f2fcd51c"
      },
      "source": [
        "for _ in range(10):\n",
        "    length = random.randint(0, 10)\n",
        "    l = [random.randint(-10, 10) for _ in range(length)]\n",
        "    print(l, end='')\n",
        "    quick_sort(l, 0, len(l)-1)\n",
        "    print(' ==> ', end='')\n",
        "    print(l)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[0, 1, -8, -10, 3, 5, -6] ==> [-10, -8, -6, 0, 1, 3, 5]\n",
            "[-6, 3, -2, -8, 1, 1] ==> [-8, -6, -2, 1, 1, 3]\n",
            "[] ==> []\n",
            "[-9] ==> [-9]\n",
            "[1, -3, -3] ==> [-3, -3, 1]\n",
            "[] ==> []\n",
            "[-7, 4] ==> [-7, 4]\n",
            "[2, -8] ==> [-8, 2]\n",
            "[-4, -1, -9, 6, 8, -10] ==> [-10, -9, -4, -1, 6, 8]\n",
            "[-1, -1, 6, -2] ==> [-2, -1, -1, 6]\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}