one_player_env.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "from gym.spaces import Box, Discrete\n",
    "\n",
    "import numpy as np\n",
    "from collections import deque"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LightCycleEnv(gym.Env):\n",
    "    def __init__(self, grid_width=50, grid_height=50, trail_length=20) -> None:\n",
    "        \"\"\"\n",
    "        Game takes place on a 50 x 50 grid\n",
    "        An agent is a snake with a length of 20 (snake occupies 20 spaces)\n",
    "        The agent can move in one of 4 directions or not at all\n",
    "        Each agent has a constant speed and is moving in 1 direction\n",
    "\n",
    "        action_space : none, left, right, up, down\n",
    "        observation_space (i.e. state space): x position, y position, current direction (1=left, 2=right, 3=up, 4=down)\n",
    "\n",
    "        Note that (0, 0) is the top-left corner, so up is decreasing y-coordinate, while down is increasing y-coordinate\n",
    "        \"\"\"\n",
    "        self.grid_width = grid_width\n",
    "        self.grid_height = grid_height\n",
    "        self.trail_length = trail_length\n",
    "\n",
    "        self.body = deque()\n",
    "\n",
    "        for i in range(0, trail_length):\n",
    "            self.body.append([int(self.grid_width/2), i]) # (x, y) pair for each space occupied\n",
    "        \n",
    "        self.state = np.concatenate([np.array(self.body).flatten(), [2]], dtype='int32') # include direction, initial is right\n",
    "        \n",
    "        low = np.concatenate([np.zeros(shape=(self.state.shape[0]-1,)), [1]])\n",
    "        high = np.concatenate([self.grid_width*np.ones(shape=self.state.shape[0]-1,), [4]]) # non-box region not support at the moment\n",
    "        self.action_space = Discrete(5)\n",
    "        self.observation_space = Box(low=low, high=high, dtype='int32') # essentially a multi-dimensional discrete space\n",
    "    \n",
    "    def reset(self):\n",
    "        \"\"\"\n",
    "        Resets the playing board\n",
    "        Should be called before starting the game\n",
    "        \"\"\"\n",
    "        self.body = deque()\n",
    "\n",
    "        for i in range(0, self.trail_length):\n",
    "            self.body.append([int(self.grid_width/2), i]) # (x, y) pair for each space occupied\n",
    "        \n",
    "        self.state = np.concatenate([np.array(self.body).flatten(), [2]], dtype='int32') # include direction, initial is right\n",
    "        return self.state\n",
    "\n",
    "    def step(self, action):\n",
    "        \"\"\"\n",
    "        Move the agent in the direction specified by action\n",
    "        Punish the agent for hitting an invalid area (severe) or making an\n",
    "        useless action (light, i.e. trying to go left when it's already going left)\n",
    "\n",
    "        \n",
    "        \"\"\"\n",
    "        reward = 0\n",
    "        done = False\n",
    "        info = {}\n",
    "\n",
    "        dir = self.state[-1] # current direction\n",
    "\n",
    "        # moving left\n",
    "        if dir == 1:\n",
    "            # continue left\n",
    "            if action < 3: \n",
    "                # unnecessary action that doesn't do anything   \n",
    "                if action != 0:\n",
    "                    reward -= 1\n",
    "                \n",
    "                # move left\n",
    "                new_head = [self.body[-1][0], self.body[-1][1]-1]\n",
    "                if new_head[1] < 0:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into left wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [1]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "            # go up\n",
    "            if action == 3:\n",
    "                new_head = [self.body[-1][0]-1, self.body[-1][1]]\n",
    "                if new_head[0] < 0: \n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into top wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [3]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "            # go down\n",
    "            if action == 4:\n",
    "                new_head = [self.body[-1][0]+1, self.body[-1][1]]\n",
    "                if new_head[0] >= self.grid_height: \n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into bottom wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [4]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "\n",
    "        # moving right\n",
    "        if dir == 2:\n",
    "            # continue right\n",
    "            if action < 3: \n",
    "                # unnecessary action that doesn't do anything   \n",
    "                if action != 0:\n",
    "                    reward -= 1\n",
    "                \n",
    "                # move right\n",
    "                new_head = [self.body[-1][0], self.body[-1][1]+1]\n",
    "                if new_head[1] >= self.grid_width:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into right wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [2]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "            # go up\n",
    "            if action == 3:\n",
    "                new_head = [self.body[-1][0]-1, self.body[-1][1]]\n",
    "                if new_head[0] < 0: \n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into top wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [3]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "            # go down\n",
    "            if action == 4:\n",
    "                new_head = [self.body[-1][0]+1, self.body[-1][1]]\n",
    "                if new_head[0] >= self.grid_height: \n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into bottom wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [4]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "        # moving up\n",
    "        if dir == 3:\n",
    "            # continue up\n",
    "            if action > 2 or action == 0:\n",
    "                # unnecessary action\n",
    "                if action > 2:\n",
    "                    reward -= 1\n",
    "                \n",
    "                new_head = [self.body[-1][0]-1, self.body[-1][1]]\n",
    "                if new_head[0] < 0: \n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into top wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [3]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "            # go left\n",
    "            if action == 1:\n",
    "                new_head = [self.body[-1][0], self.body[-1][1]-1]\n",
    "                if new_head[1] < 0:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into left wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [1]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "            # go right\n",
    "            if action == 2:\n",
    "                new_head = [self.body[-1][0], self.body[-1][1]+1]\n",
    "                if new_head[1] >= self.grid_width:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into right wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [2]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "        # moving down\n",
    "        if dir == 4:\n",
    "            # continue down\n",
    "            if action > 2 or action == 0:\n",
    "                # unnecessary action\n",
    "                if action > 2:\n",
    "                    reward -= 1\n",
    "                \n",
    "                new_head = [self.body[-1][0]+1, self.body[-1][1]]\n",
    "                if new_head[0] >= self.grid_height: \n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into bottom wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [4]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "            # go left\n",
    "            if action == 1:\n",
    "                new_head = [self.body[-1][0], self.body[-1][1]-1]\n",
    "                if new_head[1] < 0:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into left wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [1]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "            # go right\n",
    "            if action == 2:\n",
    "                new_head = [self.body[-1][0]+1, self.body[-1][1]]\n",
    "                if new_head[1] >= self.grid_width:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into right wall'\n",
    "                elif new_head in self.body:\n",
    "                    reward -= 100\n",
    "                    done = True\n",
    "                    info['halt_state'] = 'ran into self'\n",
    "\n",
    "                self.body.popleft()\n",
    "                self.body.append(new_head)\n",
    "\n",
    "                self.state = np.concatenate([np.array(self.body).flatten(), [2]], dtype='int32')\n",
    "                return self.state, reward, done, info\n",
    "\n",
    "\n",
    "    def render(self):\n",
    "        \"\"\"\n",
    "        Print a visual representation of the board to standard output\n",
    "        \"\"\"\n",
    "        for r in range(0, self.grid_height):\n",
    "            if r > 0:\n",
    "                print()\n",
    "            for c in range(0, self.grid_width):\n",
    "                if [r, c] in self.body:\n",
    "                    print('X', end=' ')\n",
    "                else:\n",
    "                    print('O', end=' ')\n",
    "\n",
    "        print()\n",
    "\n",
    "    def create_grid(self):\n",
    "        \"\"\"\n",
    "        Return a (grid_width, grid_height) array with O's for open spaces and X's for snake spaces\n",
    "        \"\"\"\n",
    "        grid = []\n",
    "\n",
    "        for r in range(0, self.grid_width):\n",
    "            grid.append([])\n",
    "            for c in range(0, self.grid_height):\n",
    "                if [r, c] in self.body:\n",
    "                    grid[r].append('O')\n",
    "                else:\n",
    "                    grid[r].append('X')\n",
    "        \n",
    "        return grid\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "deque([[2, 3], [2, 4], [1, 4], [0, 4], [-1, 4]])\n",
      "ran into top wall\n"
     ]
    }
   ],
   "source": [
    "from IPython.display import clear_output\n",
    "import time\n",
    "\n",
    "delay = 0.33\n",
    "env = LightCycleEnv(grid_width=10, grid_height=10, trail_length=5)\n",
    "\n",
    "env.render()\n",
    "time.sleep(delay)\n",
    "\n",
    "done = False\n",
    "\n",
    "while not done:\n",
    "    clear_output(wait=True)\n",
    "    \n",
    "    action = env.action_space.sample()\n",
    "    state, action, done, info = env.step(action)\n",
    "\n",
    "    if done:\n",
    "        print(env.body)\n",
    "        print(info['halt_state'])\n",
    "        break\n",
    "\n",
    "    env.render()\n",
    "\n",
    "    time.sleep(delay)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.24900412559509277 seconds\n",
      "Completed 10000 steps\n"
     ]
    }
   ],
   "source": [
    "# test speed on a large game\n",
    "\n",
    "start = time.time()\n",
    "\n",
    "env = LightCycleEnv(grid_width=100, grid_height=100, trail_length=25)\n",
    "\n",
    "for i in range(10000):\n",
    "    action = env.action_space.sample()\n",
    "    state, action, done, info = env.step(action)\n",
    "\n",
    "    # if done:\n",
    "    #     break\n",
    "\n",
    "print(time.time() - start, 'seconds')\n",
    "print(f\"Completed {i+1} steps\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.6 ('venv': venv)",
   "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.9.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ff86e6f48323df11d88af547ef1833adb93dfb0691af991d287ce5989bfe6dd7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}