Q-Learning+ASP-in-GridWorld.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random\n",
    "import PW_explorer\n",
    "from PW_explorer.run_clingo import run_clingo\n",
    "from PW_explorer.load_worlds import load_worlds\n",
    "from PW_explorer.pwe_nb_helper import ASPRules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "IPython.CodeCell.options_default.highlight_modes['prolog'] = {'reg':[/^%%(clingo|dlv)/]};"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%load_ext PWE_NB_Extension"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We show how we can use a combination of Answer Set Programming (ASP) and Python for Reinforcement Learning, specifically the Q-Learning alogorithm. We leverage the modelling capabilities of ASP to define, model and drive the environment and its state, and use Python to drive the \"learning\". This allows us to achieve the best of both worlds and create a truly modular and efficient learning procedure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can encode a small GridsWorld instance in ASP as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%clingo --donot-run --load_combined_input_to grid_worlds_env_rules --donot-display_input\n",
    "\n",
    "% GRIDS World\n",
    "\n",
    "#const min_x=0.\n",
    "#const max_x=5.\n",
    "#const min_y=0.\n",
    "#const max_y=5.\n",
    "#const min_timestep=0.\n",
    "#const max_timestep=1.\n",
    "\n",
    "time(min_timestep..max_timestep).\n",
    "valid_x(min_x..max_x).\n",
    "valid_y(min_y..max_y).\n",
    "\n",
    "% schema state(X, Y, TIME)\n",
    "% temporal state(_,_,T)\n",
    "\n",
    "%% state(0,0,0).  %% COMES FROM PYTHON LATER\n",
    ":- state(X1,Y1,T), state(X2,Y2,T), valid_x(X1), valid_x(X2), valid_y(Y1), valid_y(Y2), time(T), X1!=X2.\n",
    ":- state(X1,Y1,T), state(X2,Y2,T), valid_x(X1), valid_x(X2), valid_y(Y1), valid_y(Y2), time(T), Y1!=Y2.\n",
    "\n",
    "% schema reward(REWARD)\n",
    "reward(100) :- state(X,Y,_), reward_state(X,Y).\n",
    "\n",
    "% schema barrier(X,Y)\n",
    "reward_state(4, 4).\n",
    "barrier(3,0).\n",
    "barrier(1,4).\n",
    "barrier(4,5).\n",
    "\n",
    "%% :- not reward(100).\n",
    "\n",
    "\n",
    "% schema action(ACTION, TIME)\n",
    "% temporal action(_,T)\n",
    "%% ACTION COMES FROM PYTHON LATER\n",
    "%action(up, 0).\n",
    "%% 1 {action(up, T) ; action(down, T); action(left, T) ; action(right, T)} 1 :- time(T).\n",
    "\n",
    "state(X, Y, T+1) :- state(X, Y2, T), Y2=Y-1, action(up, T), valid_x(X), valid_y(Y), valid_x(X), valid_y(Y2), time(T), time(T+1), not barrier(X,Y).\n",
    "state(X, Y, T+1) :- state(X, Y2, T), Y2=Y+1, action(down, T), valid_x(X), valid_y(Y), valid_x(X), valid_y(Y2), time(T), time(T+1), not barrier(X,Y).\n",
    "state(X, Y, T+1) :- state(X2, Y, T), X2=X+1, action(left, T), valid_x(X), valid_y(Y), valid_x(X2), valid_y(Y), time(T), time(T+1), not barrier(X,Y).\n",
    "state(X, Y, T+1) :- state(X2, Y, T), X2=X-1, action(right, T), valid_x(X), valid_y(Y), valid_x(X2), valid_y(Y), time(T), time(T+1), not barrier(X,Y).\n",
    "\n",
    "%:- not reward(100).\n",
    "\n",
    "#show state/3.\n",
    "#show action/2.\n",
    "#show reward/1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now define the learning procedure in Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def egreedy_policy(q_values, state, epsilon=0.1):\n",
    "    if np.random.rand() < epsilon:\n",
    "        return random.choice(ACTIONS)\n",
    "    else:\n",
    "        return ACTIONS[np.argmax(q_values[state[0]][state[1]])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "ACTIONS = ['left', 'right', 'up', 'down']\n",
    "class Grids_World_Env:\n",
    "    \n",
    "    asp_file = 'grids_world.lp4'\n",
    "    env_rules = [grid_worlds_env_rules] #ASPRules(asp_file).splitlines()\n",
    "    \n",
    "    @staticmethod\n",
    "    def _get_next_state(state, dfs):\n",
    "        next_state = state\n",
    "        state_df = dfs['state_3']\n",
    "        state_df = state_df[state_df['pw']==1]\n",
    "        state_df = state_df[state_df['TIME']==1]\n",
    "        if len(state_df) > 0:\n",
    "            next_state = (int(state_df.iloc[0]['X']), int(state_df.iloc[0]['Y']))\n",
    "        return next_state\n",
    "    \n",
    "    @staticmethod\n",
    "    def _get_reward(dfs):\n",
    "        reward = 0\n",
    "        if 'reward_1' in dfs:\n",
    "            reward_df = dfs['reward_1']\n",
    "            reward_df = reward_df[reward_df['pw'] == 1]\n",
    "            if len(reward_df) > 0:\n",
    "                reward = int(reward_df.iloc[0]['REWARD'])\n",
    "        return reward\n",
    "    \n",
    "    @staticmethod\n",
    "    def get_next_state_and_reward(state, action):\n",
    "        curr_rules = Grids_World_Env.env_rules + ['action({},0).'.format(action), \n",
    "                                                  'state({},{},0).'.format(state[0], state[1])]\n",
    "        \n",
    "        asp_soln, meta_data = run_clingo(curr_rules)\n",
    "        dfs, rels, pws = load_worlds(asp_soln, meta_data, 'clingo', silent=True)\n",
    "        next_state = Grids_World_Env._get_next_state(state, dfs)\n",
    "        reward = Grids_World_Env._get_reward(dfs)\n",
    "        return next_state, reward\n",
    "    \n",
    "    @staticmethod\n",
    "    def is_game_over(state, next_state, reward):\n",
    "        return reward > 0\n",
    "    \n",
    "    @staticmethod\n",
    "    def step(state, action):\n",
    "        next_state, reward = Grids_World_Env.get_next_state_and_reward(state, action)\n",
    "        game_over = Grids_World_Env.is_game_over(state, next_state, reward)\n",
    "        return next_state, reward, game_over"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can initilize the experiment as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "num_states = 6\n",
    "num_actions = 4\n",
    "learning_rate = 0.1\n",
    "num_episodes_for_exploration = 20\n",
    "exploration_epsilon = 0.9\n",
    "exploitation_epsilon = 0.2\n",
    "gamma = 0.9\n",
    "num_episodes_per_iter = 10\n",
    "\n",
    "env = Grids_World_Env()\n",
    "q_values = np.zeros((num_states, num_states, num_actions))\n",
    "total_episodes = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we have the Q-Learning procedure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 1 Done\n",
      "reward: 100\n",
      "num_steps: 42\n",
      "Episode 2 Done\n",
      "reward: 100\n",
      "num_steps: 362\n",
      "Episode 3 Done\n",
      "reward: 100\n",
      "num_steps: 41\n",
      "Episode 4 Done\n",
      "reward: 100\n",
      "num_steps: 92\n",
      "Episode 5 Done\n",
      "reward: 100\n",
      "num_steps: 81\n",
      "Episode 6 Done\n",
      "reward: 100\n",
      "num_steps: 118\n",
      "Episode 7 Done\n",
      "reward: 100\n",
      "num_steps: 40\n",
      "Episode 8 Done\n",
      "reward: 100\n",
      "num_steps: 10\n",
      "Episode 9 Done\n",
      "reward: 100\n",
      "num_steps: 101\n",
      "Episode 10 Done\n",
      "reward: 100\n",
      "num_steps: 39\n"
     ]
    }
   ],
   "source": [
    "for ep in range(num_episodes_per_iter):\n",
    "    total_episodes += 1\n",
    "    state = (0,0)\n",
    "    num_steps = 0\n",
    "    done = False\n",
    "    # While episode is not over\n",
    "    while not done:\n",
    "        num_steps += 1\n",
    "        # print(\"state:\", state)\n",
    "        # Choose action        \n",
    "        action = egreedy_policy(q_values, state, \n",
    "                                epsilon=exploitation_epsilon \n",
    "                                if total_episodes > num_episodes_for_exploration \n",
    "                                else exploration_epsilon)\n",
    "        # print(\"action:\", action)\n",
    "        # Do the action\n",
    "        next_state, reward, done = Grids_World_Env.step(state, action)\n",
    "        # Update q_values        \n",
    "        td_target = reward + gamma * np.max(q_values[next_state[0]][next_state[1]])\n",
    "        td_error = td_target - q_values[state[0]][state[1]][ACTIONS.index(action)]\n",
    "        q_values[state[0]][state[1]][ACTIONS.index(action)] += learning_rate * td_error\n",
    "        # Update state\n",
    "        state = next_state\n",
    "        if done:\n",
    "            print(\"Episode {} Done\".format(total_episodes))\n",
    "            print(\"reward:\", reward)\n",
    "            print(\"num_steps:\", num_steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Testing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have learned a Q-value table, we can freeze it, and then try it out as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "State: (0, 0)\n",
      "Action: right\n",
      "Next State: (1, 0)\n",
      "Reward: 0\n",
      "\n",
      "State: (1, 0)\n",
      "Action: right\n",
      "Next State: (2, 0)\n",
      "Reward: 0\n",
      "\n",
      "State: (2, 0)\n",
      "Action: up\n",
      "Next State: (2, 1)\n",
      "Reward: 0\n",
      "\n",
      "State: (2, 1)\n",
      "Action: up\n",
      "Next State: (2, 2)\n",
      "Reward: 0\n",
      "\n",
      "State: (2, 2)\n",
      "Action: up\n",
      "Next State: (2, 3)\n",
      "Reward: 0\n",
      "\n",
      "State: (2, 3)\n",
      "Action: up\n",
      "Next State: (2, 4)\n",
      "Reward: 0\n",
      "\n",
      "State: (2, 4)\n",
      "Action: right\n",
      "Next State: (3, 4)\n",
      "Reward: 0\n",
      "\n",
      "State: (3, 4)\n",
      "Action: right\n",
      "Next State: (4, 4)\n",
      "Reward: 100\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Simulating based on Q-Value table\n",
    "state = (0,0)\n",
    "is_game_over = False\n",
    "while not is_game_over:\n",
    "    print(\"State:\", state)\n",
    "    action = ACTIONS[np.argmax(q_values[state[0]][state[1]])]\n",
    "    print(\"Action:\", action)\n",
    "    next_state, reward = Grids_World_Env.get_next_state_and_reward(state, action)\n",
    "    print(\"Next State:\", next_state)\n",
    "    print(\"Reward:\", reward)\n",
    "    is_game_over = Grids_World_Env.is_game_over(state, next_state, reward)\n",
    "    state = next_state\n",
    "    print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can easily display the \"optimal\" learned policy as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>right</td>\n",
       "      <td>right</td>\n",
       "      <td>up</td>\n",
       "      <td>left</td>\n",
       "      <td>up</td>\n",
       "      <td>up</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>right</td>\n",
       "      <td>right</td>\n",
       "      <td>up</td>\n",
       "      <td>up</td>\n",
       "      <td>left</td>\n",
       "      <td>up</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>right</td>\n",
       "      <td>right</td>\n",
       "      <td>up</td>\n",
       "      <td>up</td>\n",
       "      <td>up</td>\n",
       "      <td>left</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>right</td>\n",
       "      <td>right</td>\n",
       "      <td>up</td>\n",
       "      <td>up</td>\n",
       "      <td>left</td>\n",
       "      <td>left</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>up</td>\n",
       "      <td>left</td>\n",
       "      <td>right</td>\n",
       "      <td>right</td>\n",
       "      <td>left</td>\n",
       "      <td>left</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>right</td>\n",
       "      <td>right</td>\n",
       "      <td>down</td>\n",
       "      <td>down</td>\n",
       "      <td>left</td>\n",
       "      <td>left</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       0      1      2      3     4     5\n",
       "0  right  right     up   left    up    up\n",
       "1  right  right     up     up  left    up\n",
       "2  right  right     up     up    up  left\n",
       "3  right  right     up     up  left  left\n",
       "4     up   left  right  right  left  left\n",
       "5  right  right   down   down  left  left"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "policy = [[ACTIONS[np.argmax(q_values[i][j])] for i in range(num_states)] for j in range(num_states)]\n",
    "pd.DataFrame(policy)"
   ]
  },
  {
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}