Neural Network Approximation of Policy.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import json\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"value_function.json\", 'r') as f:\n",
    "    value_function = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "S-------O    0.426296\n",
       "-S------O    0.382835\n",
       "--S-----O    0.413625\n",
       "---S----O    0.440745\n",
       "-------SO    0.471797\n",
       "----S---O    0.500000\n",
       "-----S--O    0.497570\n",
       "------S-O    0.309953\n",
       "S---O--SO    0.101673\n",
       "-S--O--SO    0.000000\n",
       "Name: value, dtype: float64"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vf = pd.Series(value_function, name='value')\n",
    "vf.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'S-------O'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vf.index[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_key(key):\n",
    "    return np.stack((\n",
    "        (np.array(tuple(key)) == 'S').reshape(3, 3), \n",
    "        (np.array(tuple(key)) == 'O').reshape(3, 3)\n",
    "    ), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 3, 3)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = convert_key('S-------O')\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ True, False, False],\n",
       "        [False, False, False],\n",
       "        [False, False, False]],\n",
       "\n",
       "       [[False, False, False],\n",
       "        [False, False, False],\n",
       "        [False, False,  True]]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ True, False, False],\n",
       "        [False, False, False],\n",
       "        [False, False, False]],\n",
       "\n",
       "       [[False, False, False],\n",
       "        [False, False, False],\n",
       "        [False, False,  True]]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vf.index.map(convert_key)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5442, 2, 3, 3)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = np.stack(vf.index.map(convert_key)).astype(np.float32)\n",
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5442,)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = vf.values.astype(np.float32)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[[1., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.]],\n",
       " \n",
       "        [[0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 1.]]], dtype=float32), 0.42629614)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X[0], y[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First try regular 'unshaped' data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5442, 18)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Xd = X.reshape((-1,18))\n",
    "Xd.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 1.],\n",
       "        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 1.],\n",
       "        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 1.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 1.]], dtype=float32),\n",
       " Index(['S-------O', '-S------O', '--S-----O', '---S----O', '-------SO'], dtype='object'))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Xd[0:5], vf.index[0:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class args:\n",
    "    \n",
    "    seed = 0\n",
    "    no_cuda = False\n",
    "    \n",
    "torch.manual_seed(args.seed)\n",
    "use_cuda = not args.no_cuda and torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Policy(nn.Module):\n",
    "    \n",
    "    def __init__(self, dim=[18, 200, 1], act_funcs=[F.relu, F.relu]):\n",
    "        \n",
    "        super(Policy, self).__init__()\n",
    "        \n",
    "        self.dim = dim\n",
    "        self.layers = []\n",
    "        self.act_funcs = act_funcs\n",
    "        for i in range(1, len(dim)):\n",
    "            self.layers.append(nn.Linear(dim[i-1], dim[i]))\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        for layer, act_func in zip(self.layers, self.act_funcs):\n",
    "            x = self.act_func(layer(x))\n",
    "        y_pred = x\n",
    "        return y_pred\n",
    "    \n",
    "    #def __repr__(self):\n",
    "    #    \n",
    "    #    return \"Policy(dim=%s)\" % (self.dim.__repr__())\n",
    "\n",
    "\n",
    "class SimplePolicy(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(SimplePolicy, self).__init__()\n",
    "        self.affine1 = nn.Linear(18, 1024)\n",
    "        self.affine2 = nn.Linear(1024, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.affine1(x))\n",
    "        y_pred = self.affine2(x)\n",
    "        return y_pred\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model = Policy()\n",
    "model = SimplePolicy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SimplePolicy(\n",
       "  (affine1): Linear(in_features=18, out_features=1024, bias=True)\n",
       "  (affine2): Linear(in_features=1024, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5442, 4353)"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = len(Xd)\n",
    "split = int(0.8*m)\n",
    "m, split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((4353, 18), (4353, 1))"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Xd_train, Xd_test = Xd[0:split], Xd[split:]\n",
    "y_train, y_test = y[0:split].reshape((-1, 1)), y[split:].reshape((-1, 1))\n",
    "\n",
    "Xd_train.shape, y_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(args, model, device, data, target, criterion, optimizer):\n",
    "    model.train()\n",
    "    data, target = data.to(device), target.to(device)\n",
    "    optimizer.zero_grad()\n",
    "    output = model(data)\n",
    "    loss = criterion(output, target)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if i % 100 == 99:\n",
    "        print('loss:', loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(args, model, device, data, target, criterion):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        output = model(data)\n",
    "        test_loss = criterion(output, target)\n",
    "    \n",
    "    return test_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SimplePolicy()\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "eps = np.finfo(np.float32).eps.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train\n",
      "loss: 0.0045851063914597034\n",
      "loss: 0.00250298622995615\n",
      "loss: 0.0018828449537977576\n",
      "loss: 0.0015443058218806982\n",
      "loss: 0.001320302370004356\n"
     ]
    }
   ],
   "source": [
    "input_data = torch.from_numpy(Xd_train)\n",
    "target = torch.from_numpy(y_train)\n",
    "print(\"Train\")\n",
    "for i in range(500):\n",
    "    train(args, model, device, input_data, target, criterion, optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test\n",
      "loss: 0.06868447363376617\n"
     ]
    }
   ],
   "source": [
    "input_data = torch.from_numpy(Xd_test)\n",
    "target = torch.from_numpy(y_test)\n",
    "test_loss = test(args, model, device, input_data, target, criterion)\n",
    "print('Test\\nloss:', float(test_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[([0.4124775230884552], [0.6355000138282776]),\n",
       " ([0.3514310121536255], [0.0]),\n",
       " ([0.400850385427475], [1.0]),\n",
       " ([0.4477371275424957], [0.0]),\n",
       " ([0.4631620943546295], [0.0]),\n",
       " ([0.4977399408817291], [1.0]),\n",
       " ([0.49888715147972107], [1.0]),\n",
       " ([0.28567180037498474], [0.6355000138282776]),\n",
       " ([0.11942158639431], [0.2922925353050232]),\n",
       " ([-0.0009653307497501373], [0.5])]"
      ]
     },
     "execution_count": 130,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_data = torch.from_numpy(Xd_train)\n",
    "list(zip(model(input_data[:10]).detach().numpy().tolist(), target[:10].detach().numpy().tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}