TicTacToe-RL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tic-tac-toe solver to illustrate Q Learning\n",
    "\n",
    "Q (and Deep-Q) Learning primer at https://www.nervanasys.com/demystifying-deep-reinforcement-learning/\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pylab inline\n",
    "\n",
    "import matplotlib.cm as cm\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "from collections import defaultdict\n",
    "from six.moves import zip_longest\n",
    "import numpy as np\n",
    "import random, sys, pickle, os, time\n",
    "\n",
    "from IPython.core.debugger import set_trace\n",
    "#set_trace()\n",
    "\n",
    "# For some reason Notebook doesn't like this...\n",
    "#from builtins import input\n",
    "# Hacky py3 backwards compatibility\n",
    "try:\n",
    "    input = raw_input\n",
    "except NameError:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0,
     27,
     30,
     33,
     38,
     42,
     48,
     52,
     65,
     69
    ],
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class BoardState(object):\n",
    "    # 0 1 2\n",
    "    # 3 4 5\n",
    "    # 6 7 8\n",
    "    winning_spots = np.array([\n",
    "        [0, 1, 2], [3, 4, 5], [6, 7, 8], # Horizontal\n",
    "        [0, 3, 6], [1, 4, 7], [2, 5, 8], # Vertical\n",
    "        [0, 4, 8], [2, 4, 6]             # Diagonal\n",
    "        ])\n",
    "    \n",
    "    board_format = '\\n'.join([\n",
    "        ' {} | {} | {} ',\n",
    "        '---+---+---',\n",
    "        ' {} | {} | {} ',\n",
    "        '---+---+---',\n",
    "        ' {} | {} | {} ',\n",
    "        ])\n",
    "\n",
    "    def __init__(self, prev=None, action=None):\n",
    "        if prev is not None:\n",
    "            self.marks = prev.marks.copy()\n",
    "            self.marks[action] = prev.active_player\n",
    "            self.active_player = 'X' if prev.active_player == 'O' else 'O'\n",
    "        else:\n",
    "            self.active_player = 'X'\n",
    "            self.marks = np.array(['_']*9)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return ''.join(self.marks) + ',' + self.active_player\n",
    "\n",
    "    def __str__(self):\n",
    "        return BoardState.board_format.format(*self.marks)\n",
    "\n",
    "    def __eq__(self, other):\n",
    "        return isinstance(other, self.__class__) \\\n",
    "            and np.array_equal(self.marks, other.marks) \\\n",
    "            and self.active_player == other.active_player\n",
    "\n",
    "    def __hash__(self):\n",
    "        return hash(repr(self))\n",
    "\n",
    "    @staticmethod\n",
    "    def from_repr(s):\n",
    "        out = BoardState()\n",
    "        out.active_player = s[-1]\n",
    "        out.marks = np.array(list(s[:-2]))\n",
    "        return out\n",
    "\n",
    "    def render(self):\n",
    "        print(self.__str__())\n",
    "\n",
    "    # (action:int) -> (next_state:BoardState, reward:float, done:bool)\n",
    "    def step(self, action:int):\n",
    "        # Construct next_state by applying action to the current board\n",
    "        # (placing 0 or X on square# `action` depending on whose turn it is).        \n",
    "        next_state = BoardState(self, action)\n",
    "\n",
    "        # Score the resulting board by performing a static evaluation:\n",
    "        #    -1  if the action is an illegal move (attempting to mark a nonempty cell), else\n",
    "        #    +1  if the action wins the game, else\n",
    "        #     0  if it completes the board without winning (tie)\n",
    "        reward = -1.0 if self.marks[action] != '_' \\\n",
    "            else +1.0 if next_state.check_win(self.active_player) \\\n",
    "            else  0.0\n",
    "\n",
    "        done = next_state.is_full()  or  reward != 0.0\n",
    "        \n",
    "        return (next_state, reward, done)\n",
    "\n",
    "    def check_win(self, player):\n",
    "        slices = self.marks[BoardState.winning_spots]\n",
    "        return (slices == player).all(axis=1).any()\n",
    "\n",
    "    def is_full(self):\n",
    "        return (self.marks != '_').all()\n",
    "    \n",
    "    def draw(self, agent):\n",
    "        fig = figure(figsize=[3,3])\n",
    "        ax = fig.add_subplot(111)\n",
    "\n",
    "        def draw_cell(pos, mark, val):\n",
    "            y, x = divmod(pos, 3)\n",
    "            \n",
    "            # If the game has been won, give cells in winning line(s) a red background\n",
    "            slices = self.marks[self.winning_spots]\n",
    "            O = (slices == 'O').all(axis=1)\n",
    "            X = (slices == 'X').all(axis=1)\n",
    "            winningSliceIndices = np.append( O.nonzero(), X.nonzero() )\n",
    "            winningSquares = np.unique( self.winning_spots[ winningSliceIndices ] )           \n",
    "            \n",
    "            if pos in winningSquares:\n",
    "                ax.add_patch(patches.Rectangle((x,y), 1, 1, ec='none', fc='red'))\n",
    "                            \n",
    "            if mark == 'X':\n",
    "                ax.plot([x+.2, x+.8], [y+.8, y+.2], 'k', lw=2.0)\n",
    "                ax.plot([x+.2, x+.8], [y+.2, y+.8], 'k', lw=2.0)\n",
    "            elif mark == 'O':\n",
    "                ax.add_patch(patches.Circle((x+.5,y+.5), .35, ec='k', fc='none', lw=2.0))\n",
    "            else:\n",
    "                # Colour empty squares according to predicted value\n",
    "                color = cm.viridis((val+1)/2.)\n",
    "                ax.add_patch(patches.Rectangle((x,y), 1, 1, ec='none', fc=color))\n",
    "                ax.text(x+.5 , y+.5, '%.2f'%val    , ha='center', va='center') \n",
    "                ax.text(x+.08, y+.12,  '%d'%(pos+1), ha='center', va='center') \n",
    "\n",
    "        for i in range(9):\n",
    "            draw_cell(i, self.marks[i], agent.Q_read(i,self))\n",
    "\n",
    "        ax.set_position([0,0,1,1])\n",
    "        ax.set_axis_off()\n",
    "\n",
    "        ax.set_xlim(0,3)\n",
    "        ax.set_ylim(3,0)\n",
    "\n",
    "        for x in range(1,3):\n",
    "            ax.plot([x, x], [0,3], 'k', lw=2.0)\n",
    "            ax.plot([0,3], [x, x], 'k', lw=2.0)\n",
    "        show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0,
     37,
     49,
     64,
     70
    ],
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# The agent is a computer player.\n",
    "# It builds a Q_table, e.g.\n",
    "#\n",
    "#       Q_tables[4][('XO_X___OX','O')] = +1.0\n",
    "#\n",
    "# This says that placing a O at location 4 (i.e. center) on the board:\n",
    "#       X O -\n",
    "#       X - -\n",
    "#       - O X\n",
    "# ... will evaluate a score (or Quality Q) of +1.0 (i.e. a win)\n",
    "#\n",
    "# Use Q_read & Q_writeOrUpdate to access.\n",
    "\n",
    "\n",
    "class TabularAgent(object):\n",
    "    def __init__(self, num_actions, alpha=0.75, gamma=1.0, epsilon=1.0, default_Q=0):\n",
    "        self.alpha = alpha\n",
    "        self.gamma = gamma\n",
    "        self.epsilon = epsilon\n",
    "        self.default_Q = default_Q\n",
    "        self.num_actions = num_actions  # should be 9 (9 squares so 9 possible actions)\n",
    "        self.Q_tables = [{} for _ in range(self.num_actions)]\n",
    "        \n",
    "    # For a given board, return the action# that predicts the highest Q\n",
    "    def max_action(self, state):\n",
    "        # type: (BoardState) -> int\n",
    "        predictions = [self.Q_read(ndx, state) for ndx in range(self.num_actions)]\n",
    "        return np.argmax(predictions)\n",
    "\n",
    "    # Choose a random action (0-8) with probability epsilon, \n",
    "    # or the optimal action with probability 1-epsilon\n",
    "    def choose_action(self, state):\n",
    "        # type: (BoardState) -> int\n",
    "        if random.random() > self.epsilon:\n",
    "            return random.choice(range(self.num_actions))\n",
    "        return self.max_action(state)\n",
    "\n",
    "    # Get Q-value for a particular action on a given board(-state)\n",
    "    def Q_read(self, nAction, state):\n",
    "        # type: (int, BoardState) -> float\n",
    "        return self.Q_tables[nAction].get(state, self.default_Q)\n",
    "\n",
    "    # Update Q-value for a particular state+action pair\n",
    "    # (creating a new entry if necessary)\n",
    "    def Q_writeOrUpdate(self, nAction, state, new_Q):\n",
    "        # type: (int, BoardState, float) -> None\n",
    "        buf = self.Q_tables[nAction]\n",
    "        if state in buf:\n",
    "            buf[state] = (1-self.alpha)*buf[state] + self.alpha*new_Q\n",
    "        else:\n",
    "            buf[state] = new_Q\n",
    "\n",
    "    def train(self, history):\n",
    "        raise NotImplementedError()\n",
    "\n",
    "# - - - - - - - - - - - - - - - - - - - - - \n",
    "# train: (history:[(action:int, reward:float, state:BoardState)]) -> (none)\n",
    "#     Use history (from the perspective of one player over the duration of one game) to update Q-table.\n",
    "#     For each (state, action) pair, replace the entry (reward) in the Q-table with a one-step lookahead prediction.\n",
    "#     i.e. \n",
    "#         - Apply the action to the state resulting in a new state s'\n",
    "#         - from s' determine Q of the best action\n",
    "#         - Q(state, action) <- reward(state, action) + gamma * Q(s', best action)\n",
    "\n",
    "class MonteCarloAgent(TabularAgent):\n",
    "    def train(self, history):\n",
    "        if len(history) == 0:\n",
    "            return\n",
    "        prev_action, return_, _ = history[-1]\n",
    "        for (action, reward, state) in reversed(history[:-1]):\n",
    "            self.Q_writeOrUpdate(prev_action, state, return_)\n",
    "\n",
    "            prev_action = action\n",
    "            if reward is not None:\n",
    "                return_ += reward\n",
    "                \n",
    "class TemporalDifferenceAgent(TabularAgent):\n",
    "    def new_val(self, history, ndx):\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def train(self, history):\n",
    "        if len(history) == 1:\n",
    "            return\n",
    "        for i in range(len(history)-2):\n",
    "            (_, _, state), (action, _, _) = history[i:i+2]\n",
    "            self.Q_writeOrUpdate(action, state, self.new_val(history, i))\n",
    "            \n",
    "        (_, _, state), (action, reward, _) = history[-2:]\n",
    "        \n",
    "        self.Q_writeOrUpdate(action, state, reward)\n",
    "\n",
    "# - - - - - - - - - - - - - - - - - - - - - \n",
    "# One-step Q-table lookahead reward prediction\n",
    "\n",
    "class QLearningAgent(TemporalDifferenceAgent):\n",
    "    def new_val(self, history, ndx):\n",
    "        (_, _, state), (action, reward, next_state)                      = history[ndx:ndx+2]\n",
    "        next_action = self.max_action(next_state)\n",
    "        return reward + self.gamma * self.Q_read(next_action, next_state)\n",
    "\n",
    "class SarsaAgent(TemporalDifferenceAgent):\n",
    "    def new_val(self, history, ndx):\n",
    "        (_, _, state), (action, reward, next_state), (next_action, _, _) = history[ndx:ndx+3]\n",
    "        return reward + self.gamma * self.Q_read(next_action, next_state)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TRAIN (Computer vs. Computer)\n",
    "episodes = 10000\n",
    "fname = 'tictac.txt'\n",
    "\n",
    "def progressbar(callback, iters, refresh_rate=2.0):\n",
    "    prev_clock = time.time()\n",
    "    start_clock = prev_clock\n",
    "\n",
    "    for i in range(iters):\n",
    "        callback(i)\n",
    "        curr_clock = time.time()\n",
    "        if (curr_clock-prev_clock)*refresh_rate >= 1:\n",
    "            sys.stdout.write('\\r[ %s / %s ]' % (i, iters))\n",
    "            sys.stdout.flush()\n",
    "            prev_clock = curr_clock\n",
    "\n",
    "    clearstr = ' '*len('[ %s / %s ]' % (iters, iters))\n",
    "    sys.stdout.write('\\r%s\\r' % clearstr)\n",
    "    sys.stdout.flush()\n",
    "\n",
    "    return time.time() - start_clock\n",
    "\n",
    "# Assumes zero-sum, two-player, sequential-turn game\n",
    "def train_episode(agent, state=None):\n",
    "    if state is None:\n",
    "        # Start at a random previously encountered state\n",
    "        keys = list( agent.Q_tables[random.randint(0,8)] )  # list of keys for Q-table dict\n",
    "        if len(keys) > 0:\n",
    "            state = random.choice(keys)\n",
    "        else:\n",
    "            state = BoardState()\n",
    "\n",
    "    # Play out a game, recording each (action, reward, state) tuple.\n",
    "    first_player = state.active_player\n",
    "    history = [(None, None, state)]\n",
    "\n",
    "    while True:\n",
    "        action = agent.choose_action(state)\n",
    "        state, reward, done = state.step(action)\n",
    "        history.append((action, reward, state))\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "    # Split history into a separate history for each player.\n",
    "    #    history stores things like [(None, None, s1), (p1a1, p1r1, s2), (p2a1, p2r1, s3), (p1a2, p1r2, s4), ...]\n",
    "    #    player_history transforms that to [(None, None, s1), (p1a1, p1r1-p2r1, s3), (p1a2, p1r2-p2r2, s5), ...]\n",
    "    #    You subtract the reward given to the other player because of the assumption of it being a zero-sum game.\n",
    "    #    (Think: relative reward)\n",
    "    def player_history(history):\n",
    "        # e.g.  grouped('ABCDEFG', 3, 'x') --> 'ABC' 'DEF' 'Gxx'\n",
    "        def grouped(iterable, n, fillvalue=None):\n",
    "            \"Collect data into fixed-length chunks or blocks\"\n",
    "            # https://docs.python.org/2/library/itertools.html#recipes\n",
    "            args = [iter(iterable)] * n\n",
    "            return zip_longest(fillvalue=fillvalue, *args)\n",
    "\n",
    "        out = [(None, None, history[0][2])]\n",
    "        for (action, reward, state), (_, other_reward, other_state) \\\n",
    "                                in grouped(history[1:], 2, (None,)*3):\n",
    "            if other_reward is None:\n",
    "                out.append((action, reward, state))\n",
    "            else:\n",
    "                out.append((action, reward-other_reward, other_state))\n",
    "        return out\n",
    "    \n",
    "    first_history = player_history(history)\n",
    "    second_history = player_history(history[1:])\n",
    "    \n",
    "    # Update Q-tables\n",
    "    agent.train(first_history)\n",
    "    agent.train(second_history)\n",
    "\n",
    "    \n",
    "if os.path.isfile(fname):\n",
    "    print('Loading agent from %s...' % fname)\n",
    "    agent = pickle.load(open(fname, 'rb'))\n",
    "else:\n",
    "    agent = QLearningAgent(num_actions=9, epsilon=0.8, default_Q=2)\n",
    "\n",
    "#init_state = BoardState() # Always start from actual inital state\n",
    "init_state = None # Random restarts\n",
    "\n",
    "print('Training for %d episodes...' % episodes)\n",
    "progressbar(lambda x: train_episode(agent, init_state), episodes)\n",
    "\n",
    "print('Saving agent to %s...' % fname)\n",
    "pickle.dump(agent, open(fname, 'wb'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def play_vs_human(agent, state=None):\n",
    "    if state is None:\n",
    "        state = BoardState()    \n",
    "    state.draw(agent)\n",
    "    \n",
    "    # Flip a coin for who goes first\n",
    "    compToMove = random.random() > 0.5\n",
    "    \n",
    "    while True:\n",
    "        if compToMove:\n",
    "            state, reward, done = state.step(agent.choose_action(state))\n",
    "        else:\n",
    "            move = int(input('Choose your move [1-9]: ')) - 1\n",
    "            state, reward, done = state.step(move)\n",
    "            \n",
    "        state.draw(agent)\n",
    "        \n",
    "        if done:\n",
    "            if compToMove:\n",
    "                s = 'Tie!' if reward == 0  else   'I win!' if reward > 0  else   'I lose!'\n",
    "            else:\n",
    "                s = 'Tie!' if reward == 0  else 'You win!' if reward > 0  else 'You lose!'\n",
    "                \n",
    "            input(s + \"\\nPress Enter to play again...\")\n",
    "            return\n",
    "            \n",
    "        compToMove = not compToMove"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PLAY\n",
    "fname = 'tictac.txt'\n",
    "\n",
    "agent = pickle.load(open(fname, 'rb'))\n",
    "agent.epsilon = 0.99\n",
    "try:\n",
    "    while True:\n",
    "        play_vs_human(agent)\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}