imitation_learning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import copy\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards_per_ep = pd.read_csv('data/reward_log.txt',header=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards_per_ep = rewards_per_ep[0].apply(lambda x: int(x.split()[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading expert_2022-04-18T15-08-29.txt\n",
      "reading expert_2022-04-18T15-22-57.txt\n",
      "reading expert_2022-04-18T15-26-55.txt\n",
      "reading expert_2022-04-18T15-31-12.txt\n",
      "reading expert_2022-04-18T15-34-37.txt\n",
      "reading expert_2022-04-18T15-37-46.txt\n",
      "reading expert_2022-04-18T15-40-20.txt\n",
      "reading expert_2022-04-18T15-43-02.txt\n",
      "reading expert_2022-04-18T15-46-45.txt\n",
      "reading expert_2022-04-18T15-49-47.txt\n",
      "reading expert_2022-04-18T17-45-40.txt\n",
      "reading expert_2022-04-18T17-52-24.txt\n",
      "reading expert_2022-04-18T17-54-38.txt\n",
      "reading expert_2022-04-18T17-58-00.txt\n",
      "reading expert_2022-04-18T18-00-32.txt\n",
      "reading expert_2022-04-18T18-03-22.txt\n",
      "reading expert_2022-04-18T18-05-46.txt\n",
      "reading expert_2022-04-18T18-08-26.txt\n",
      "reading expert_2022-04-18T18-11-25.txt\n",
      "reading expert_2022-04-18T18-14-07.txt\n"
     ]
    }
   ],
   "source": [
    "dfs = []\n",
    "for exp_f in os.listdir('data'):\n",
    "    if not exp_f.startswith('expert') or not exp_f.endswith('.txt'):\n",
    "        continue\n",
    "\n",
    "    print('reading',exp_f)\n",
    "    f = open(f'data/{exp_f}')\n",
    "    data = json.load(f)\n",
    "    f.close()\n",
    "\n",
    "    df = []\n",
    "    for k, v in data.items():\n",
    "        o, a = v\n",
    "        o.append(a)\n",
    "        df.append(o)\n",
    "\n",
    "    df = pd.DataFrame(df).iloc[:,:129]\n",
    "    dfs.append(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    129\n",
       "dtype: int64"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.concat(dfs).isna().sum().value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "      <th>9</th>\n",
       "      <th>...</th>\n",
       "      <th>119</th>\n",
       "      <th>120</th>\n",
       "      <th>121</th>\n",
       "      <th>122</th>\n",
       "      <th>123</th>\n",
       "      <th>124</th>\n",
       "      <th>125</th>\n",
       "      <th>126</th>\n",
       "      <th>127</th>\n",
       "      <th>128</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 129 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   0    1    2    3    4    5    6    7    8    9    ...  119  120  121  122  \\\n",
       "0    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "1    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "2    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "3    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "4    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "\n",
       "   123  124  125  126  127  128  \n",
       "0    2   66  240  146  215    3  \n",
       "1    2   66  240  146  215    3  \n",
       "2    2   66  240  146  215    3  \n",
       "3    2   66  240  146  215    3  \n",
       "4    2   66  240  146  215    3  \n",
       "\n",
       "[5 rows x 129 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train = pd.concat(dfs)\n",
    "train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "Int64Index: 23484 entries, 0 to 1177\n",
      "Columns: 129 entries, 0 to 128\n",
      "dtypes: int64(129)\n",
      "memory usage: 23.3 MB\n"
     ]
    }
   ],
   "source": [
    "train.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = train.drop(128, axis=1)\n",
    "y = train[128]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "y = pd.get_dummies(y)\n",
    "y.insert(0, 0, 0)\n",
    "y.insert(5, 5, 0)\n",
    "y.insert(6, 6, 0)\n",
    "y.insert(7, 7, 0)\n",
    "y.insert(8, 8, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "      <th>9</th>\n",
       "      <th>...</th>\n",
       "      <th>118</th>\n",
       "      <th>119</th>\n",
       "      <th>120</th>\n",
       "      <th>121</th>\n",
       "      <th>122</th>\n",
       "      <th>123</th>\n",
       "      <th>124</th>\n",
       "      <th>125</th>\n",
       "      <th>126</th>\n",
       "      <th>127</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>112</td>\n",
       "      <td>114</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>88</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>240</td>\n",
       "      <td>146</td>\n",
       "      <td>215</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 128 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   0    1    2    3    4    5    6    7    8    9    ...  118  119  120  121  \\\n",
       "0    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "1    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "2    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "3    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "4    0  112  114  115    0    3   88   88   88   88  ...    0    0    0    0   \n",
       "\n",
       "   122  123  124  125  126  127  \n",
       "0    0    2   66  240  146  215  \n",
       "1    0    2   66  240  146  215  \n",
       "2    0    2   66  240  146  215  \n",
       "3    0    2   66  240  146  215  \n",
       "4    0    2   66  240  146  215  \n",
       "\n",
       "[5 rows x 128 columns]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   0  1  2  3  4  5  6  7  8\n",
       "0  0  0  0  1  0  0  0  0  0\n",
       "1  0  0  0  1  0  0  0  0  0\n",
       "2  0  0  0  1  0  0  0  0  0\n",
       "3  0  0  0  1  0  0  0  0  0\n",
       "4  0  0  0  1  0  0  0  0  0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(23484, 128)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DDQN(nn.Module):\n",
    "    \"\"\"\n",
    "    The Double Deep Q-Network has as input a state s and\n",
    "    outputs the state-action values q(s,a_1), ..., q(s,a_n) for all n actions.\n",
    "    :param: state_dim: for input layer\n",
    "    :param: hidden_dim: for every hidden layer\n",
    "    :param: action_dim: for output layer\n",
    "    \"\"\"\n",
    "    def __init__(self, action_dim, state_dim, hidden_dim):\n",
    "        super().__init__()\n",
    "\n",
    "        self.online = nn.Sequential(\n",
    "            nn.Linear(state_dim, hidden_dim*2),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(hidden_dim*2, hidden_dim),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(hidden_dim, action_dim)\n",
    "        )\n",
    "\n",
    "        self.target = copy.deepcopy(self.online)\n",
    "\n",
    "        # Q_target parameters are frozen.\n",
    "        for p in self.target.parameters():\n",
    "            p.requires_grad = False\n",
    "\n",
    "    def forward(self, input, model):\n",
    "        \"\"\"\n",
    "        When doing update by forward, it takes:\n",
    "        :param: input: all state of each observation\n",
    "        :param: model: online or target\n",
    "        :return: Q_values of all actions given state from online/target\n",
    "        \"\"\"\n",
    "\n",
    "        if model == \"online\":\n",
    "            return self.online(input)\n",
    "        elif model == \"target\":\n",
    "            return self.target(input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, epochs, X_train, X_test, y_train, y_test, batch_size=32):\n",
    "\n",
    "    model = model\n",
    "    \n",
    "    loss_fn = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(net.parameters(), lr=0.00001)\n",
    "    \n",
    "    train_loss = [loss_fn(model(X_train, 'online'), y_train).item()]\n",
    "    print(f'epoch: {0}, train loss = {train_loss[0]: .4f}')\n",
    "    \n",
    "    running_loss = 0.0\n",
    "    for epoch in range(epochs):\n",
    "        \n",
    "        # random the order of X_train and y_train for batching\n",
    "        permutation = torch.randperm(X_train.size()[0])\n",
    "        \n",
    "        avg_train_loss = []\n",
    "        avg_test_loss = []\n",
    "        \n",
    "        for i in range(0, X_train.size()[0], batch_size):\n",
    "            \n",
    "            indices = permutation[i:i+batch_size]\n",
    "            batch_x, batch_y = X_train[indices], y_train[indices]\n",
    "            \n",
    "            # init optimizer\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # forward pass and loss\n",
    "            batch_y_hat = model(batch_x, 'online')\n",
    "            loss = loss_fn(batch_y_hat, batch_y)\n",
    "\n",
    "            # backward pass\n",
    "            loss.backward()\n",
    "\n",
    "            # update\n",
    "            optimizer.step()\n",
    "            \n",
    "            avg_train_loss.append(loss.item())\n",
    "        \n",
    "        train_loss.append(np.mean(avg_train_loss))\n",
    "        y_pred = model(X_test, 'online')\n",
    "        \n",
    "        val_loss = loss_fn(y_pred, y_test)\n",
    "        avg_test_loss.append(np.mean(val_loss.item()))\n",
    "\n",
    "        if (epoch + 1) % 10 == 0:\n",
    "            print(f'epoch: {epoch+1}, train loss = {loss.item(): .4f}, test loss = {val_loss.item(): .4f}')\n",
    "    \n",
    "    return avg_train_loss, avg_test_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, shuffle=True, stratify=y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = torch.tensor(X_train.values, dtype=torch.float32).to(device)\n",
    "y_train = torch.tensor(y_train.values, dtype=torch.float32).to(device)\n",
    "X_test = torch.tensor(X_test.values, dtype=torch.float32).to(device)\n",
    "y_test = torch.tensor(y_test.values, dtype=torch.float32).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = DDQN(9,128,128).to(device)\n",
    "\n",
    "retrain = False # True\n",
    "if retrain:\n",
    "    checkpoint = torch.load(\"data/expert_CE.chkpt\", map_location=torch.device('cpu'))\n",
    "    net.load_state_dict(checkpoint['model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, train loss =  0.8282\n",
      "epoch: 10, train loss =  0.9480, test loss =  0.9155\n",
      "epoch: 20, train loss =  0.7107, test loss =  0.8950\n",
      "epoch: 30, train loss =  0.7097, test loss =  0.8696\n",
      "epoch: 40, train loss =  0.7362, test loss =  0.8583\n",
      "epoch: 50, train loss =  0.5293, test loss =  0.8453\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "([0.6679730415344238,\n",
       "  0.718706488609314,\n",
       "  0.5056966543197632,\n",
       "  0.6972368955612183,\n",
       "  0.5345399975776672,\n",
       "  0.6854649782180786,\n",
       "  0.6328055262565613,\n",
       "  0.5927356481552124,\n",
       "  0.8564898371696472,\n",
       "  0.5288715362548828,\n",
       "  0.7811293601989746,\n",
       "  0.7239629626274109,\n",
       "  0.6966859698295593,\n",
       "  0.7225309014320374,\n",
       "  0.6797051429748535,\n",
       "  0.7152587175369263,\n",
       "  0.5407190322875977,\n",
       "  0.757960319519043,\n",
       "  0.6365983486175537,\n",
       "  0.6549027562141418,\n",
       "  0.6250765323638916,\n",
       "  0.44561904668807983,\n",
       "  0.6260679364204407,\n",
       "  0.4876243472099304,\n",
       "  0.5606369972229004,\n",
       "  0.4195913076400757,\n",
       "  0.678602933883667,\n",
       "  0.5023676156997681,\n",
       "  0.5965218544006348,\n",
       "  0.5387930274009705,\n",
       "  0.6152819991111755,\n",
       "  0.6889351606369019,\n",
       "  0.6937064528465271,\n",
       "  0.49560749530792236,\n",
       "  0.3879130780696869,\n",
       "  0.47004279494285583,\n",
       "  0.8159167766571045,\n",
       "  0.776059627532959,\n",
       "  0.5666105151176453,\n",
       "  0.4622281789779663,\n",
       "  0.5599296689033508,\n",
       "  0.8033310174942017,\n",
       "  0.7847638130187988,\n",
       "  0.5619091391563416,\n",
       "  0.7469055652618408,\n",
       "  0.743781566619873,\n",
       "  0.5651001930236816,\n",
       "  0.6953495144844055,\n",
       "  0.7926130890846252,\n",
       "  0.4584934115409851,\n",
       "  0.6348080039024353,\n",
       "  0.6852505803108215,\n",
       "  0.9928567409515381,\n",
       "  0.5479781627655029,\n",
       "  0.8088350296020508,\n",
       "  0.8065288662910461,\n",
       "  0.4993239939212799,\n",
       "  0.6829478144645691,\n",
       "  0.6249183416366577,\n",
       "  0.6641284227371216,\n",
       "  0.8058773279190063,\n",
       "  0.5399921536445618,\n",
       "  0.542636513710022,\n",
       "  0.7715076208114624,\n",
       "  0.7988998293876648,\n",
       "  0.6589305400848389,\n",
       "  0.5089260339736938,\n",
       "  0.6224508285522461,\n",
       "  0.7622687816619873,\n",
       "  0.6092082858085632,\n",
       "  0.506279468536377,\n",
       "  0.5392698049545288,\n",
       "  0.6577326059341431,\n",
       "  0.7362910509109497,\n",
       "  0.7026139497756958,\n",
       "  0.46278879046440125,\n",
       "  0.7745322585105896,\n",
       "  0.6954734325408936,\n",
       "  0.73801589012146,\n",
       "  0.7931790351867676,\n",
       "  0.5998133420944214,\n",
       "  0.7120112180709839,\n",
       "  0.7861539125442505,\n",
       "  0.7666073441505432,\n",
       "  0.43668195605278015,\n",
       "  0.4960062503814697,\n",
       "  0.7676936388015747,\n",
       "  0.5664645433425903,\n",
       "  0.4267916679382324,\n",
       "  0.6714338064193726,\n",
       "  0.8331422209739685,\n",
       "  0.6329124569892883,\n",
       "  0.5204025506973267,\n",
       "  0.7472861409187317,\n",
       "  0.7258728742599487,\n",
       "  0.4959946572780609,\n",
       "  0.7547574639320374,\n",
       "  0.7008714079856873,\n",
       "  0.438413143157959,\n",
       "  0.7750170826911926,\n",
       "  0.6109123229980469,\n",
       "  0.7633733749389648,\n",
       "  0.5613332986831665,\n",
       "  0.4743717610836029,\n",
       "  0.6955931782722473,\n",
       "  0.7158265709877014,\n",
       "  0.5909603834152222,\n",
       "  0.7989821434020996,\n",
       "  0.8927644491195679,\n",
       "  0.5238025188446045,\n",
       "  0.7609840631484985,\n",
       "  0.4926724433898926,\n",
       "  0.5154805779457092,\n",
       "  0.7610455751419067,\n",
       "  0.4321059584617615,\n",
       "  0.8246748447418213,\n",
       "  0.6277711391448975,\n",
       "  0.600191593170166,\n",
       "  0.6660836935043335,\n",
       "  0.7468252778053284,\n",
       "  0.7400929927825928,\n",
       "  0.5554651618003845,\n",
       "  0.8402838110923767,\n",
       "  0.7713428735733032,\n",
       "  0.5058744549751282,\n",
       "  0.980905294418335,\n",
       "  0.7241131067276001,\n",
       "  0.6592737436294556,\n",
       "  0.7031880617141724,\n",
       "  0.7608472108840942,\n",
       "  0.8504870533943176,\n",
       "  0.6073609590530396,\n",
       "  0.7976498603820801,\n",
       "  0.43097901344299316,\n",
       "  0.6811061501502991,\n",
       "  0.6222100257873535,\n",
       "  0.7035086750984192,\n",
       "  0.7126248478889465,\n",
       "  0.536466658115387,\n",
       "  0.7605433464050293,\n",
       "  0.6604015827178955,\n",
       "  0.931780219078064,\n",
       "  0.841742992401123,\n",
       "  0.5456141233444214,\n",
       "  0.7489231824874878,\n",
       "  0.8655914068222046,\n",
       "  0.5555460453033447,\n",
       "  0.5823910236358643,\n",
       "  0.5669419765472412,\n",
       "  0.8132622241973877,\n",
       "  0.8676180243492126,\n",
       "  0.6987923979759216,\n",
       "  0.6193059682846069,\n",
       "  0.5284513831138611,\n",
       "  0.8195589780807495,\n",
       "  0.8823904991149902,\n",
       "  0.6881762742996216,\n",
       "  0.6911250948905945,\n",
       "  0.5694552063941956,\n",
       "  0.9970006942749023,\n",
       "  0.5043207406997681,\n",
       "  0.7917052507400513,\n",
       "  0.777404248714447,\n",
       "  0.5625333786010742,\n",
       "  0.7075948715209961,\n",
       "  0.6865279674530029,\n",
       "  0.6426073312759399,\n",
       "  0.6333290934562683,\n",
       "  0.5649492144584656,\n",
       "  0.6709468960762024,\n",
       "  1.007678508758545,\n",
       "  0.4348369836807251,\n",
       "  0.4444088935852051,\n",
       "  0.6205757856369019,\n",
       "  0.4425514340400696,\n",
       "  0.48085999488830566,\n",
       "  0.7351557016372681,\n",
       "  0.4956820607185364,\n",
       "  0.7270943522453308,\n",
       "  0.715868353843689,\n",
       "  0.5590075850486755,\n",
       "  0.43325892090797424,\n",
       "  0.8282678723335266,\n",
       "  0.6965516805648804,\n",
       "  0.8546784520149231,\n",
       "  0.555433988571167,\n",
       "  0.6863144040107727,\n",
       "  0.6123737096786499,\n",
       "  0.6017913818359375,\n",
       "  0.41499269008636475,\n",
       "  0.5259023904800415,\n",
       "  0.7260187268257141,\n",
       "  0.7849626541137695,\n",
       "  0.8872387409210205,\n",
       "  0.7862933874130249,\n",
       "  0.5840222239494324,\n",
       "  0.6206648945808411,\n",
       "  0.8102579116821289,\n",
       "  0.7788963317871094,\n",
       "  0.613917887210846,\n",
       "  0.6979961395263672,\n",
       "  0.572178304195404,\n",
       "  0.8244991302490234,\n",
       "  0.6548786163330078,\n",
       "  0.5873939990997314,\n",
       "  0.7254446744918823,\n",
       "  0.9653437733650208,\n",
       "  0.759374737739563,\n",
       "  0.5171614289283752,\n",
       "  0.32272011041641235,\n",
       "  0.837933361530304,\n",
       "  0.6188645958900452,\n",
       "  0.41715407371520996,\n",
       "  0.7493668794631958,\n",
       "  0.5907366275787354,\n",
       "  0.6667819619178772,\n",
       "  0.5028605461120605,\n",
       "  0.5883191227912903,\n",
       "  0.8833215236663818,\n",
       "  0.6712597012519836,\n",
       "  0.5533852577209473,\n",
       "  0.7958792448043823,\n",
       "  0.7272927761077881,\n",
       "  0.673011064529419,\n",
       "  0.6546318531036377,\n",
       "  0.9539963006973267,\n",
       "  0.4996926188468933,\n",
       "  0.6919965147972107,\n",
       "  0.5413985252380371,\n",
       "  0.6813732981681824,\n",
       "  0.8426346778869629,\n",
       "  0.6521037817001343,\n",
       "  0.7331539392471313,\n",
       "  0.7079511284828186,\n",
       "  0.968743085861206,\n",
       "  0.5268877148628235,\n",
       "  0.47935691475868225,\n",
       "  0.9056294560432434,\n",
       "  0.8899890780448914,\n",
       "  0.8770710825920105,\n",
       "  0.4842104911804199,\n",
       "  0.6190370917320251,\n",
       "  0.519032895565033,\n",
       "  0.7137086987495422,\n",
       "  0.38603299856185913,\n",
       "  0.5106644034385681,\n",
       "  0.5938458442687988,\n",
       "  0.5802732706069946,\n",
       "  0.4535992741584778,\n",
       "  0.6691631078720093,\n",
       "  0.823567807674408,\n",
       "  0.6926376819610596,\n",
       "  0.36795952916145325,\n",
       "  0.8717048168182373,\n",
       "  0.8272625207901001,\n",
       "  0.5438469648361206,\n",
       "  0.8865727186203003,\n",
       "  0.5687763094902039,\n",
       "  0.9251924157142639,\n",
       "  0.5745810866355896,\n",
       "  0.49115127325057983,\n",
       "  0.8341981172561646,\n",
       "  0.7110648155212402,\n",
       "  0.5810550451278687,\n",
       "  0.5767058730125427,\n",
       "  0.6192591190338135,\n",
       "  0.7689361572265625,\n",
       "  0.5432673096656799,\n",
       "  0.7847689390182495,\n",
       "  0.8939043879508972,\n",
       "  0.7635911107063293,\n",
       "  0.6646199822425842,\n",
       "  0.5290762186050415,\n",
       "  0.5288593769073486,\n",
       "  0.513836145401001,\n",
       "  0.8071932196617126,\n",
       "  0.7876113653182983,\n",
       "  0.7848787307739258,\n",
       "  0.724026083946228,\n",
       "  0.781032383441925,\n",
       "  0.7963818311691284,\n",
       "  0.667943000793457,\n",
       "  0.5941826105117798,\n",
       "  0.9081788063049316,\n",
       "  0.6032267808914185,\n",
       "  0.6688550710678101,\n",
       "  0.5170092582702637,\n",
       "  0.9276499152183533,\n",
       "  0.9252760410308838,\n",
       "  0.739649772644043,\n",
       "  0.51218181848526,\n",
       "  0.7651658654212952,\n",
       "  0.7485915422439575,\n",
       "  0.695164680480957,\n",
       "  0.7796391248703003,\n",
       "  0.6081773638725281,\n",
       "  0.5321725606918335,\n",
       "  0.6123350858688354,\n",
       "  0.6964508295059204,\n",
       "  0.5332782864570618,\n",
       "  0.6077353358268738,\n",
       "  0.5873988270759583,\n",
       "  0.5251379013061523,\n",
       "  0.6568476557731628,\n",
       "  0.8283793926239014,\n",
       "  0.714030921459198,\n",
       "  0.7368108034133911,\n",
       "  0.8433695435523987,\n",
       "  0.6787470579147339,\n",
       "  0.719834566116333,\n",
       "  0.5753135085105896,\n",
       "  0.6805344820022583,\n",
       "  0.8831870555877686,\n",
       "  0.6303929090499878,\n",
       "  0.5804747343063354,\n",
       "  0.8664489388465881,\n",
       "  0.6025905013084412,\n",
       "  0.6242923140525818,\n",
       "  0.6650658249855042,\n",
       "  0.6748972535133362,\n",
       "  0.5285370945930481,\n",
       "  0.7867289185523987,\n",
       "  0.46636590361595154,\n",
       "  0.5377182364463806,\n",
       "  0.6752861142158508,\n",
       "  0.7290814518928528,\n",
       "  0.6508784294128418,\n",
       "  0.6774572730064392,\n",
       "  0.7959795594215393,\n",
       "  0.8311444520950317,\n",
       "  0.7553033828735352,\n",
       "  0.5571292638778687,\n",
       "  0.7107932567596436,\n",
       "  0.9075932502746582,\n",
       "  0.7535876035690308,\n",
       "  0.7962010502815247,\n",
       "  0.783844530582428,\n",
       "  0.5764805674552917,\n",
       "  0.6985397934913635,\n",
       "  0.6875860691070557,\n",
       "  0.5124273300170898,\n",
       "  0.6630516052246094,\n",
       "  0.5401342511177063,\n",
       "  0.42889508605003357,\n",
       "  0.7309895157814026,\n",
       "  0.6357121467590332,\n",
       "  0.5805703997612,\n",
       "  0.9231235980987549,\n",
       "  0.8413530588150024,\n",
       "  0.5351256728172302,\n",
       "  0.6324571371078491,\n",
       "  0.6698279976844788,\n",
       "  0.7320876717567444,\n",
       "  0.49708932638168335,\n",
       "  0.5380868315696716,\n",
       "  0.5098959803581238,\n",
       "  0.6163938045501709,\n",
       "  0.3855106830596924,\n",
       "  0.5157565474510193,\n",
       "  0.5767414569854736,\n",
       "  0.7115359902381897,\n",
       "  0.8271310329437256,\n",
       "  0.7212449312210083,\n",
       "  0.6104134321212769,\n",
       "  0.8000843524932861,\n",
       "  0.5707542896270752,\n",
       "  0.6922533512115479,\n",
       "  0.5606265068054199,\n",
       "  0.8958009481430054,\n",
       "  0.7167202234268188,\n",
       "  0.48367705941200256,\n",
       "  0.7175196409225464,\n",
       "  0.6355206370353699,\n",
       "  0.5925367474555969,\n",
       "  0.5206478834152222,\n",
       "  0.7424992322921753,\n",
       "  0.8181039094924927,\n",
       "  0.7441390752792358,\n",
       "  0.7403674721717834,\n",
       "  0.6144565939903259,\n",
       "  0.6445305943489075,\n",
       "  0.6715077757835388,\n",
       "  0.5270345211029053,\n",
       "  0.8546101450920105,\n",
       "  0.7077576518058777,\n",
       "  0.5304980278015137,\n",
       "  0.6431035995483398,\n",
       "  0.588070809841156,\n",
       "  0.38723331689834595,\n",
       "  0.6354002952575684,\n",
       "  0.6906026005744934,\n",
       "  0.9307663440704346,\n",
       "  0.6642323732376099,\n",
       "  0.8794179558753967,\n",
       "  0.7749402523040771,\n",
       "  0.6740888357162476,\n",
       "  0.6364889144897461,\n",
       "  0.864301860332489,\n",
       "  0.8634672164916992,\n",
       "  0.41103965044021606,\n",
       "  0.8381129503250122,\n",
       "  0.5250168442726135,\n",
       "  0.7505064010620117,\n",
       "  0.6601781845092773,\n",
       "  0.608674168586731,\n",
       "  0.7011000514030457,\n",
       "  0.5566994547843933,\n",
       "  0.5891622304916382,\n",
       "  0.4296528398990631,\n",
       "  0.5814027786254883,\n",
       "  0.6674760580062866,\n",
       "  0.7071434259414673,\n",
       "  0.7284191846847534,\n",
       "  0.5439858436584473,\n",
       "  0.6186201572418213,\n",
       "  0.8778467774391174,\n",
       "  0.7396705746650696,\n",
       "  0.912279486656189,\n",
       "  0.8385826349258423,\n",
       "  0.7512702345848083,\n",
       "  0.469036340713501,\n",
       "  0.585473895072937,\n",
       "  0.5045090913772583,\n",
       "  0.4754844605922699,\n",
       "  0.7833174467086792,\n",
       "  0.5630461573600769,\n",
       "  0.7213994860649109,\n",
       "  0.6901018023490906,\n",
       "  0.7405483722686768,\n",
       "  0.7207720279693604,\n",
       "  0.74113529920578,\n",
       "  0.722615122795105,\n",
       "  0.5842742323875427,\n",
       "  0.7892348766326904,\n",
       "  0.8726791739463806,\n",
       "  0.5485762357711792,\n",
       "  0.6854499578475952,\n",
       "  0.767641007900238,\n",
       "  0.8695201277732849,\n",
       "  0.4879511296749115,\n",
       "  0.519585907459259,\n",
       "  0.4315932095050812,\n",
       "  0.4856487214565277,\n",
       "  0.517146110534668,\n",
       "  0.6628235578536987,\n",
       "  0.7890191078186035,\n",
       "  0.6211727261543274,\n",
       "  0.9714040756225586,\n",
       "  0.5510913729667664,\n",
       "  0.7545786499977112,\n",
       "  0.6756659746170044,\n",
       "  0.6625136137008667,\n",
       "  0.7669345140457153,\n",
       "  0.6861569881439209,\n",
       "  0.817661702632904,\n",
       "  0.9072751998901367,\n",
       "  0.8731000423431396,\n",
       "  0.6833115816116333,\n",
       "  0.6262120008468628,\n",
       "  0.5638856887817383,\n",
       "  0.5421105623245239,\n",
       "  0.8364357948303223,\n",
       "  0.6137890815734863,\n",
       "  0.7927445769309998,\n",
       "  0.7850750684738159,\n",
       "  0.6002069711685181,\n",
       "  0.8689020872116089,\n",
       "  0.4909431040287018,\n",
       "  0.5483546257019043,\n",
       "  0.5482177734375,\n",
       "  0.8347494006156921,\n",
       "  0.7494511604309082,\n",
       "  0.6559104323387146,\n",
       "  0.7770793437957764,\n",
       "  0.6468007564544678,\n",
       "  0.6147613525390625,\n",
       "  0.8319859504699707,\n",
       "  0.6753239035606384,\n",
       "  0.8357244729995728,\n",
       "  0.5862772464752197,\n",
       "  0.8529380559921265,\n",
       "  0.686281681060791,\n",
       "  0.8335499167442322,\n",
       "  0.8277930021286011,\n",
       "  0.3398565351963043,\n",
       "  0.7196618914604187,\n",
       "  0.49511730670928955,\n",
       "  0.6826152205467224,\n",
       "  0.6633034944534302,\n",
       "  0.7822718620300293,\n",
       "  0.4484030604362488,\n",
       "  0.7252223491668701,\n",
       "  0.7470133900642395,\n",
       "  0.8012734651565552,\n",
       "  0.8245422840118408,\n",
       "  0.8910475373268127,\n",
       "  0.7198299169540405,\n",
       "  0.8005461692810059,\n",
       "  0.4841037094593048,\n",
       "  0.35202282667160034,\n",
       "  0.7723510265350342,\n",
       "  0.7386221885681152,\n",
       "  0.6491512656211853,\n",
       "  0.7064854502677917,\n",
       "  0.6520504355430603,\n",
       "  0.671328067779541,\n",
       "  0.6413508057594299,\n",
       "  0.6489764451980591,\n",
       "  0.5830612182617188,\n",
       "  0.6832699179649353,\n",
       "  0.7207316160202026,\n",
       "  0.5832312107086182,\n",
       "  0.863646388053894,\n",
       "  0.5293495059013367],\n",
       " [0.8452965021133423])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train(net, 50, X_train, X_test, y_train, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.tensor(X.values, dtype=torch.float32).to(device)\n",
    "y = torch.tensor(y.values, dtype=torch.float32).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7201058268547058"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_fn = nn.CrossEntropyLoss()\n",
    "y_pred = net(X, 'online')\n",
    "loss_fn(y_pred, y).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save trained weight\n",
    "\n",
    "net.target.load_state_dict(net.online.state_dict())\n",
    "torch.save(dict(model=net.state_dict(), exploration_rate=1.0), \"data/expert_CE_1.chkpt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Checking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.05048644542694092"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = DDQN(9,128,128).to(device)\n",
    "loss_fn = nn.SmoothL1Loss()\n",
    "checkpoint = torch.load(\"data/expert1.chkpt\", map_location=torch.device('cpu'))\n",
    "model.load_state_dict(checkpoint['model'])\n",
    "y_pred = model(X, 'online')\n",
    "loss_fn(y_pred, y).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.02409677766263485"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = DDQN(9,128,128).to(device)\n",
    "loss_fn = nn.SmoothL1Loss()\n",
    "checkpoint = torch.load(\"data/expert2.chkpt\", map_location=torch.device('cpu'))\n",
    "model.load_state_dict(checkpoint['model'])\n",
    "y_pred = model(X, 'online')\n",
    "loss_fn(y_pred, y).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.013653799891471863"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = DDQN(9,128,128).to(device)\n",
    "loss_fn = nn.SmoothL1Loss()\n",
    "checkpoint = torch.load(\"data/expert4.chkpt\", map_location=torch.device('cpu'))\n",
    "model.load_state_dict(checkpoint['model'])\n",
    "y_pred = model(X, 'online')\n",
    "loss_fn(y_pred, y).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.008327632211148739"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = DDQN(9,128,128).to(device)\n",
    "loss_fn = nn.SmoothL1Loss()\n",
    "checkpoint = torch.load(\"data/expert5.chkpt\", map_location=torch.device('cpu'))\n",
    "model.load_state_dict(checkpoint['model'])\n",
    "y_pred = model(X, 'online')\n",
    "loss_fn(y_pred, y).item()"
   ]
  },
  {
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}