path: root/revised_snake_q_table.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "85da5df2-c926-417c-bd7b-d214ad31ebe1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pygame 2.5.1 (SDL 2.28.2, Python 3.11.5)\n",
      "Hello from the pygame community. https://www.pygame.org/contribute.html\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from collections import namedtuple\n",
    "from IPython.core.debugger import Pdb\n",
    "\n",
    "from GameEngine import multiplayer\n",
    "Point = namedtuple('Point', 'x, y')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d264ae4a-380c-47b6-9b29-c6fe16c6399c",
   "metadata": {},
   "source": [
    "### New Game Implementation\n",
    "\n",
    "I have an improved game implementation which allows for multiplayer snake games, as well as simplified training. This notebook will go over both of these, including an implementation of q-table learning, as well as a match between a manually filled out q-table and a learned one.\n",
    "\n",
    "Let's start by initializing the engine object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2a382240-906d-474f-94c0-9af1a5de97ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# defines game window size and block size, in pixels\n",
    "WINDOW_WIDTH = 640\n",
    "WINDOW_HEIGHT = 480\n",
    "GAME_UNITS = 80"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "976eff80-c50a-492e-a49b-975d9905e274",
   "metadata": {},
   "outputs": [],
   "source": [
    "game_engine = multiplayer.Playfield(window_width=WINDOW_WIDTH,\n",
    "                                    window_height=WINDOW_HEIGHT,\n",
    "                                    units=GAME_UNITS,\n",
    "                                    g_speed=35,\n",
    "                                    s_size=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23de2fca-d31a-497f-9a0c-845c568e5df7",
   "metadata": {},
   "source": [
    "Here is a run-down of the current functions available to programs utilizing the game engine:\n",
    "\n",
    "**add_player**: Returns the player's number to the callee (between 0-3, for a total of 4 players). This number can be used with other functions to index that player's state.\n",
    "\n",
    "**get_heads_tails_and_goal**: Returns an array of player heads (in order of player number) the locations of all snake tails, as well as the goal location. Each is stored in an array of named tuples.\n",
    "\n",
    "**get_viable_actions**: Given a player's id, returns a list of integers corresponding to actions which will not immediately result in the snake's death.\n",
    "0 = UP\n",
    "1 = RIGHT\n",
    "2 = DOWN\n",
    "3 = LEFT\n",
    "\n",
    "**start_game**: Initializes goal, player, score, and playfield objects. Disables the ability to add new players. Enables use of player_advance function.\n",
    "\n",
    "**stop_game**: Sets the game_state to false, allowing new players to be added.\n",
    "\n",
    "**cleanup**: Quits the pygame window.\n",
    "\n",
    "**player_advance**: Given an array corresponding to each player's action (integers), returns a list of collision results, and updates the internal game state.\n",
    "\n",
    "**toggle_draw**: Turns off/on the game UI for faster training.\n",
    "\n",
    "Let's do a test of these functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "64470c1a-ddc6-4ce1-b6d4-f1e17e3d0176",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Game starting with 1 players.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p1 = game_engine.add_player()\n",
    "game_engine.start_game()\n",
    "p1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "28d7fa43-4419-4940-9d4f-d6bd5ccc11ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "viable_actions = game_engine.get_viable_actions(p1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ba45b03a-9c42-48e9-a259-5d01e667157d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<CollisionType.NONE: 2>]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "game_engine.player_advance([np.random.choice(viable_actions)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df7170ae-e95d-404b-b632-e8e9de46d69e",
   "metadata": {},
   "source": [
    "If you looked at the UI for this last statement, you should have seen that the game moved the snake (yellow) in a random direction away from immediate death."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "144bc5ff-756c-4e07-a855-4020a4474d52",
   "metadata": {},
   "source": [
    "### State-sensing methods, creating and reading a q-table\n",
    "Now, we can start redesigning some functions used to allow the snake to play intelligently. We'll use a multi-dimensional numpy array to store the rewards corresponding to each state and action. This is called a q-function, or a q-table in this case.\n",
    "\n",
    "How many states do I need? Seeing how the new **get_viable_actions** method already prevents the snake from choosing life-ending moves, the snake is no longer tasked with learning or memorizing it.\n",
    "\n",
    "The snake doesneed to be able to interpret progress towards the goal, so I will reinclude one state-sensing to sense the goal direction. I need only 8 states (with entries for each four actions) to represent our game now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "88a69876-fbc7-4dc1-a471-d8449fada4e4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0.]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "goal_relations = 8\n",
    "actions = 4\n",
    "q = np.zeros((goal_relations,\n",
    "              actions))\n",
    "q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3e15744e-1251-4315-8a4c-ed5788d04478",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sense_goal(head, goal):\n",
    "    '''\n",
    "    maps head and goal location onto an\n",
    "    integer corresponding to approx location\n",
    "    '''\n",
    "    diffs = Point(goal.x - head.x, goal.y - head.y)\n",
    "\n",
    "    if diffs.x == 0 and diffs.y <  0:\n",
    "        return 0\n",
    "    if diffs.x >  0 and diffs.y <  0:\n",
    "        return 1\n",
    "    if diffs.x >  0 and diffs.y == 0:\n",
    "        return 2\n",
    "    if diffs.x >  0 and diffs.y >  0:\n",
    "        return 3\n",
    "    if diffs.x == 0 and diffs.y >  0:\n",
    "        return 4\n",
    "    if diffs.x <  0 and diffs.y >  0:\n",
    "        return 5\n",
    "    if diffs.x <  0 and diffs.y == 0:\n",
    "        return 6\n",
    "    return 7"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "addf716b-892c-4f7f-b71f-c6af8779dff7",
   "metadata": {},
   "source": [
    "I will use the getter provided by my engine, which queries various statistics about all agents in the game:\n",
    "1. An array of head positions\n",
    "2. An array of all tail locations\n",
    "3. The goal location"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7b5f0d57-8e26-4f95-b7ea-a6810936ad5d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([Point(x=0, y=160)], [0, Point(x=0, y=160)], Point(x=0, y=320))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "game_engine.get_heads_tails_and_goal()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3fd47ce-55fe-4d2f-9147-8848193f7ca1",
   "metadata": {},
   "source": [
    "Now to use sense_goal as an index into our q table:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "85e2bab1-c98b-400e-be41-a47ccd4bc163",
   "metadata": {},
   "outputs": [],
   "source": [
    "def index_actions(q, id):\n",
    "    '''\n",
    "    given q, player_id, an array of heads,\n",
    "    and the goal position,\n",
    "    indexes into the corresponding expected\n",
    "    reward of each action\n",
    "    '''\n",
    "    heads, tails, goal = game_engine.get_heads_tails_and_goal()\n",
    "    state = sense_goal(heads[id], goal)\n",
    "    return state, q[state, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33ae53a8-989a-410c-a04f-2ac76561ce21",
   "metadata": {},
   "source": [
    "Returning state here simplifies some logic later when I train the agent. It will be passed along to my next function, but it can be ignored for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3808c200-2f67-43c4-a80f-c4c17dcfeacf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0., 0., 0., 0.])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_, rewards = index_actions(q, p1)\n",
    "rewards"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a2ef7f7-f6f7-4610-8e98-1d389327f3e8",
   "metadata": {},
   "source": [
    "In our learning agent, these actions will obviously be associated with different expected rewards. But it is not enough to take the best reward, because the positions of the hazards have not been accounted for. I chose to implement a replacement argmin/max function to select actions from this table, which generates new actions in order from highest expected reward to lowest expected reward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a172e347-75b7-4b0a-8dcc-07a6ba04f77e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def argmin_gen(rewards):\n",
    "    rewards = rewards.copy()\n",
    "    for i in range(rewards.size):\n",
    "        best_action = np.argmin(rewards)\n",
    "        rewards[best_action] = float(\"inf\")\n",
    "        yield best_action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9674f225-e7df-4551-baa8-29eb23fcc1d5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n",
      "2\n",
      "3\n"
     ]
    }
   ],
   "source": [
    "for action in argmin_gen(rewards):\n",
    "    print(action)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9bed101-661c-4e61-b8ad-94bbc1900e03",
   "metadata": {},
   "source": [
    "How will we use this? If the action generated is not a viable action, we will take the next best action.\n",
    "\n",
    "What if no actions are viable? Then the agent has boxed itself in, and it doesn't matter what action we choose.\n",
    "\n",
    "Previously, I reset the snake if it got stuck in a learning-loop. I will instead use epsilon, as it gives me a bit more control over training. Here is my greedy-action selector function, combining the work of all the previous code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d3d65397-62a9-47b6-9c84-808282656f80",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pick_greedy_action(q, id, epsilon):\n",
    "    viable_actions = game_engine.get_viable_actions(id)\n",
    "    state, rewards = index_actions(q, id)\n",
    "\n",
    "    if np.random.uniform() < epsilon:\n",
    "        return (state, np.random.choice(viable_actions)) if viable_actions.size > 0 else (state, 0)\n",
    "    for action in argmin_gen(rewards):\n",
    "        if action in viable_actions:\n",
    "            return (state, action)\n",
    "    return (state, 0) # death"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2169ac2b-df83-4f53-8a83-b35a2d1a521a",
   "metadata": {},
   "source": [
    "We'll set up epsilon to decay over our 500-step test..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a5932d47-adbe-46de-b91e-582e40faf369",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_steps = 200\n",
    "epsilon = 1\n",
    "final_epsilon = 0.001\n",
    "epsilon_decay =  np.exp(np.log(final_epsilon) / (n_steps))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee7b066a-330d-4cdd-bbb2-9d2a7ad2ceb4",
   "metadata": {},
   "source": [
    "And watch the snake explore:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6beff583-e32a-4c15-8fcb-f5b5d45ad548",
   "metadata": {},
   "outputs": [],
   "source": [
    "for step in range(n_steps):\n",
    "    _, p1_action = pick_greedy_action(q, p1, epsilon)\n",
    "    game_engine.player_advance([p1_action])\n",
    "    epsilon *= epsilon_decay"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ee2c2e6-933d-460e-b2fc-f5bf1f22e381",
   "metadata": {},
   "source": [
    "This snake obviously has no prior knowledge of how to earn the most reward, but it still does remarkably well because it is not allowed to die. It behaves as expected, favoring up and right when it is not forced to choose a random action.\n",
    "\n",
    "Our q_table only has 32 values as a result of removing the 16 danger states... It would be incredibly easy to manually select reward values to fill our q_table with..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7af51359-872c-4d1b-b178-e27cf86eb3cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_q = np.array([[-10., -2., 0., -2.],\n",
    "                 [-5., -5., 0., 0.],\n",
    "                 [-2., -10., 2., 0.],\n",
    "                 [0., -5., -5., 0.],\n",
    "                 [0., -2., -10., -2.],\n",
    "                 [0., 0., -5., -5.],\n",
    "                 [-2., 0., -2., -10.],\n",
    "                 [-5., 0., 0., -5.]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "4da4d318-f7e0-412b-8545-8fd346e167b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "epsilon = 0\n",
    "for step in range(n_steps):\n",
    "    _, p1_action = pick_greedy_action(set_q, p1, epsilon)\n",
    "    game_engine.player_advance([p1_action])\n",
    "    epsilon *= epsilon_decay"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c36ab97-2ca0-4468-8d4c-ebd1e4deec23",
   "metadata": {},
   "source": [
    "And the snake already plays optimally, no learning required.\n",
    "\n",
    "Now that we have these methods, I will create functions to allow the snake to learn by its own, and then pair it off against the q-table I just built."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b9968af-0ec2-4b92-a19d-50912703dd4a",
   "metadata": {},
   "source": [
    "### Learning and Temporal Difference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce537e44-ac8c-4f09-b89d-a330f13277da",
   "metadata": {},
   "source": [
    "I will be using the temporal difference equation as the key learning element of my reinforcement function. In theory, it allows me adjust the expected reward of a state to agree with its observed successor. In practice, it will allow out agent to take actions that previously led it closer to the goal.\n",
    "\n",
    "In order to use this equation, I simply need to create a function that takes the current state/action/outcome, and the previous state/action, as this will be updated in cases where the agent did not reach the goal.\n",
    "\n",
    "When the agent does reach the goal, I will manually set that state and action to the best reward, 0. Remember that the q-table is initialized with zeros, meaning untravelled actions are pre-assigned good rewards. Both this and epsilon will encourage exploration.\n",
    "\n",
    "Here is the complete function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a0d3942a-af6a-41f3-be74-167e3abaae0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_q(q, old_state_action, new_state_action, outcome, lr=0.05):\n",
    "    if outcome == multiplayer.CollisionType.GOAL:\n",
    "        q[new_state_action[0], new_state_action[1]] = 0\n",
    "    else:\n",
    "        td_error = -1 + q[new_state_action[0], new_state_action[1]] - q[old_state_action[0], old_state_action[1]]\n",
    "        q[old_state_action[0], old_state_action[1]] += lr * td_error"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01b21e01-174e-4fdd-ad70-dcc1e6483fb2",
   "metadata": {},
   "source": [
    "Now all that is needed is the training loop. I have high expectations for this agent, so I will only allow it 1500 moves to train itself! Here is where the outputs of pick_greedy_action come in handy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "95a2ec72-e30c-4730-a876-21f054d3727f",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_steps = 1500\n",
    "epsilon = 1\n",
    "final_epsilon = 0.001\n",
    "epsilon_decay =  np.exp(np.log(final_epsilon) / (n_steps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a71dc022-5e51-46f8-bfed-37b95243fc5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "p1_old_s_a = pick_greedy_action(q, p1, epsilon) # state, action\n",
    "game_engine.player_advance([p1_old_s_a[1]])\n",
    "\n",
    "for step in range(n_steps):\n",
    "    p1_new_s_a = pick_greedy_action(q, p1, epsilon) # state, action\n",
    "    outcome = game_engine.player_advance([p1_new_s_a[1]])\n",
    "\n",
    "    update_q(q, p1_old_s_a, p1_new_s_a, outcome)\n",
    "\n",
    "    epsilon *= epsilon_decay\n",
    "    p1_old_s_a = p1_new_s_a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6cbf429-c790-4bb9-8775-ed067844ab4e",
   "metadata": {},
   "source": [
    "The results look promising. Here is everything it learned:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ad12d31a-a1ec-45af-89b0-f66ca375b524",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-5.61286955, -0.92562893, -0.14868602, -0.77936816],\n",
       "       [-4.87612819, -3.84302796, -0.9194685 , -0.46585524],\n",
       "       [-0.65546246, -4.77095018, -0.49342682,  0.        ],\n",
       "       [-1.58528747, -6.27145881, -2.59719179, -0.51455945],\n",
       "       [-0.48715079, -1.43421385, -6.12413612, -0.7260854 ],\n",
       "       [-0.78394669, -1.70799532, -3.10876847, -6.45896201],\n",
       "       [-0.798533  , -0.41896323, -1.30809359, -3.63055269],\n",
       "       [-2.96333643, -0.98546563, -2.05542179, -6.07365687]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2414728-c36c-45d6-8f2d-18e78e482054",
   "metadata": {},
   "source": [
    "### Multiplayer Demonstration, Saving Tables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24c73f20-2853-4ba7-a24e-22c5a4b8da5e",
   "metadata": {},
   "source": [
    "The most entertaining way to test the success of my implementation is pair the agents q and set_q against each other. I will first stop and set up a new game:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a93f0fe5-2bc4-45d7-ba31-2306efaa9806",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Game over!\n",
      "Game starting with 2 players.\n"
     ]
    }
   ],
   "source": [
    "game_engine.stop_game()\n",
    "p2 = game_engine.add_player()\n",
    "game_engine.start_game()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47a1adde-d95d-444e-9688-1290764f8cfa",
   "metadata": {},
   "source": [
    "Now, I can simply call player advance with both player's actions in order, and the engine will handle the rest. I will define a new game loop, similar to the previous one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d5b5d089-cb66-47de-b352-36c13fd7dead",
   "metadata": {},
   "outputs": [],
   "source": [
    "epsilon = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "4de326e6-82dc-48df-8920-d28d0154fbd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for step in range(n_steps):\n",
    "    # p1\n",
    "    _, p1_action = pick_greedy_action(set_q, p1, epsilon)\n",
    "\n",
    "    # p2\n",
    "    p2_new_s_a = pick_greedy_action(q, p2, epsilon) # state, action\n",
    "    \n",
    "    game_engine.player_advance([p1_action, p2_new_s_a[1]])\n",
    "\n",
    "    epsilon *= epsilon_decay\n",
    "    p2_old_s_a = p2_new_s_a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "820d7e5f-c3e9-4dae-82ce-3c9188d8a8d8",
   "metadata": {},
   "source": [
    "The learned agent normally plays significantly worse than the artificially-learned one, which is okay given I hardly spent time optimizing the number of steps and learning rate. I plan to compare both of the agents again my neural-network approach, so the last this I will do is save the q_tables to a file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "95227c89-e7db-4923-9160-dacfa1cf4af8",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('superior_qt.npy', set_q)\n",
    "np.save('inferior_qt.npy', q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9eec33d8-9a65-426a-8ad5-8ffb2cbe2541",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}