Dino_Classes.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numpy.random as rnd\n",
    "import random\n",
    "import tensorflow as tf\n",
    "\n",
    "from functools import reduce\n",
    "\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "import base64, json, re, time, threading\n",
    "import multiprocessing\n",
    "\n",
    "from scipy.misc import imresize\n",
    "%run Dino_Server.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Memory:\n",
    "\n",
    "    def __init__(self, size):\n",
    "        self.size = size\n",
    "        self.mem = np.ndarray((size,5), dtype=object)\n",
    "        self.iter = 0\n",
    "        self.current_size = 0\n",
    "\n",
    "    def remember(self, state1, action, reward, state2, crashed):\n",
    "        self.mem[self.iter,:] = state1, action, reward, state2, crashed\n",
    "        self.iter = (self.iter + 1) % self.size\n",
    "        self.current_size = min(self.current_size + 1, self.size)\n",
    "\n",
    "    def sample(self, n):\n",
    "        n = min(self.current_size, n)\n",
    "        random_idx = random.sample(list(range(self.current_size)), n)\n",
    "        sample = self.mem[random_idx]\n",
    "        return (np.stack(sample[:,i], axis=0) for i in range(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DDQNAgent:\n",
    "\n",
    "    def __init__(self, session, num_actions, width, height, path, writer=None):\n",
    "        self.path_checkpoints = path\n",
    "        self.session = session\n",
    "        self.num_actions = num_actions\n",
    "        self.memory_size = 10000\n",
    "        self.explore_prob = 1.\n",
    "        self.explore_min = 0.01\n",
    "        self.explore_decay = 0.997\n",
    "        self.batch_size = 32\n",
    "        self.discount = .95\n",
    "        self.memory = Memory(self.memory_size)\n",
    "        self.main_dqn = DQN(session, height, width, num_actions, \"main\", writer)\n",
    "        self.target_dqn = DQN(session, height, width, num_actions, \"target\", None)\n",
    "        self.session.run(tf.global_variables_initializer())\n",
    "\n",
    "        self.update_target_network()\n",
    "        self.saver = tf.train.Saver()\n",
    "\n",
    "    def act(self, state):\n",
    "        \"\"\"\n",
    "        return an action and a boolean.\n",
    "        The returned boolean: - False: action generated by the DQN\n",
    "                              - True: random action (exploration)\n",
    "        \"\"\"\n",
    "        if self.explore_prob > 0 and rnd.rand() <= self.explore_prob:\n",
    "            # explore\n",
    "            return rnd.randint(self.num_actions), True\n",
    "\n",
    "        return self.main_dqn.get_action(state), False\n",
    "\n",
    "    def remember(self, state, action, reward, state_next, crashed):\n",
    "        self.memory.remember(state, action, reward, state_next, crashed)\n",
    "\n",
    "    def replay(self, cnt):\n",
    "        if self.memory.current_size < self.batch_size:\n",
    "            return\n",
    "\n",
    "        print(\"...Training...\")\n",
    "        states, actions, rewards, states_next, crashes = self.memory.sample(self.batch_size)\n",
    "        target = rewards\n",
    "        # add Q value of next state to not terminal states (not crashed)\n",
    "        target[~crashes] += self.discount * self.target_dqn.get_action_and_q(states_next[~crashes])[1]\n",
    "        self.main_dqn.train(states, actions, target, cnt)\n",
    "\n",
    "    def explore_less(self):\n",
    "        self.explore_prob = max(self.explore_min, self.explore_prob * self.explore_decay)\n",
    "\n",
    "    def update_target_network(self):\n",
    "        self.target_dqn.tranfer_variables_from(self.main_dqn)\n",
    "\n",
    "    def save(self, cnt):\n",
    "        save_path = self.saver.save(self.session, self.path_checkpoints + \"rex.ckpt\", global_step=cnt)\n",
    "        print(\"Model saved in file: %s\" % save_path)\n",
    "\n",
    "    def load(self, checkpoint_name):\n",
    "        self.saver.restore(self.session, checkpoint_name)\n",
    "        print(\"Model restored:\", checkpoint_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DQN:\n",
    "\n",
    "    def __init__(self, session, height, width, num_actions, name, writer=None):\n",
    "        self.num_actions = num_actions\n",
    "        self.height = height\n",
    "        self.width = width\n",
    "        self.name = name\n",
    "        self.vars = []\n",
    "        self.session = session\n",
    "\n",
    "        self.summary_ops = []\n",
    "        self._create_network()\n",
    "        self.writer = writer\n",
    "\n",
    "    def linear(self, x, output_size, name, activation_fn=tf.nn.relu):\n",
    "        shape = x.get_shape().as_list()\n",
    "\n",
    "        with tf.variable_scope(name):\n",
    "            w = tf.Variable(tf.random_normal([shape[1], output_size], stddev=.02), dtype=tf.float32, name='w')\n",
    "            b = tf.Variable(tf.zeros([output_size]), name='b')\n",
    "            out = tf.nn.bias_add(tf.matmul(x, w), b)\n",
    "\n",
    "            if activation_fn != None:\n",
    "                out =  activation_fn(out)\n",
    "\n",
    "        return out, w, b\n",
    "    \n",
    "    def conv2d(self, x, output_dim, kernel_shape, stride, name):\n",
    "        stride = [1, stride[0], stride[1], 1]\n",
    "\n",
    "        with tf.variable_scope(name):\n",
    "            w = tf.Variable(tf.truncated_normal(kernel_shape, mean=0, stddev=.1), dtype=tf.float32, name=\"w\")\n",
    "            conv = tf.nn.conv2d(x, w, stride, \"VALID\")\n",
    "            b = tf.Variable(tf.constant(0.1, shape=[output_dim]), name=\"b\")\n",
    "            out = tf.nn.bias_add(conv, b)\n",
    "            out = tf.nn.relu(out)\n",
    "\n",
    "        return out, w, b\n",
    "    \n",
    "    def max_pool_2x2(self, x, kernel_shape, name):\n",
    "        ksize = [1, *kernel_shape, 1]\n",
    "        strides = [1, *kernel_shape, 1]\n",
    "        return tf.nn.max_pool(x, ksize, strides, padding='SAME', name=name)\n",
    "\n",
    "    def get_action_and_q(self, states):\n",
    "        \"\"\"\n",
    "        returns array:\n",
    "            array[0]: actions: is a array of length len(state) with the action with the highest score\n",
    "            array[1]: q value: is a array of length len(state) with the Q-value belonging to the action\n",
    "        \"\"\"\n",
    "        states = states.reshape(-1, 4, self.height, self.width)\n",
    "        return self.session.run([self.a, self.Q], {self.state: states})\n",
    "\n",
    "    def get_action(self, states):\n",
    "        \"\"\"\n",
    "        returns action(s),\n",
    "            - if states contains only a single state then we return the optimal action as an integer,\n",
    "            - if states contains an array of states then we return the optimal action for each state of the array\n",
    "        \"\"\"\n",
    "        states = states.reshape(-1, 4, self.height, self.width)\n",
    "        num_states = states.shape[0]\n",
    "        actions = self.session.run(self.a, {self.state: states})\n",
    "        return actions[0] if num_states == 1 else actions\n",
    "\n",
    "    def train(self, states, actions, targets, cnt):\n",
    "        states = states.reshape(-1, 4, self.height, self.width)\n",
    "        feed_dict = {self.state: states, self.actions: actions, self.Q_target: targets}\n",
    "        summary,_ = self.session.run([tf.summary.merge(self.summary_ops), self.minimize], feed_dict)\n",
    "        if self.writer: self.writer.add_summary(summary, global_step=cnt)\n",
    "\n",
    "    def tranfer_variables_from(self, other):\n",
    "        \"\"\"\n",
    "            Builds the operations required to transfer the values of the variables\n",
    "            from other to self\n",
    "        \"\"\"\n",
    "        ops = []\n",
    "        for var_self, var_other in zip(self.vars, other.vars):\n",
    "            ops.append(var_self.assign(var_other.value()))\n",
    "\n",
    "        self.session.run(ops)\n",
    "\n",
    "\n",
    "    def _create_network(self):\n",
    "\n",
    "        with tf.variable_scope(self.name):\n",
    "            self.state =  tf.placeholder(shape=[None, 4, self.height, self.width],dtype=tf.float32)\n",
    "            self.state_perm = tf.transpose(self.state, perm=[0, 2, 3, 1])\n",
    "            self.summary_ops.append(tf.summary.image(\"states\", self.state[:, 0, :, :][..., tf.newaxis], max_outputs=10))\n",
    "\n",
    "            conv1, w1, b1 = self.conv2d(self.state_perm, 32, [8, 8, 4, 32], [4, 4], \"conv1\")\n",
    "            max_pool = self.max_pool_2x2(conv1, [2, 2], \"maxpool\")\n",
    "            conv2, w2, b2 = self.conv2d(max_pool, 64, [4, 4, 32, 64], [2, 2], \"conv2\")\n",
    "            conv3, w3, b3 = self.conv2d(conv2, 64, [3, 3, 64, 64], [1, 1], \"conv3\")\n",
    "            self.vars += [w1, b1, w2, b2, w3, b3]\n",
    "\n",
    "            shape = conv3.get_shape().as_list()\n",
    "            conv3_flat = tf.reshape(conv3, [-1, reduce(lambda x, y: x * y, shape[1:])])\n",
    "\n",
    "            value_hid, w4, b4 = self.linear(conv3_flat, 512, \"value_hid\")\n",
    "            adv_hid, w5, b5 = self.linear(conv3_flat, 512, \"adv_hid\")\n",
    "\n",
    "            value, w6, b6 = self.linear(value_hid, 1, \"value\", activation_fn=None)\n",
    "            advantage, w7, b7 = self.linear(adv_hid, self.num_actions, \"advantage\", activation_fn=None)\n",
    "            self.vars += [w4, b4, w5, b5, w6, b6, w7, b7]\n",
    "\n",
    "            self.Qs = value + (advantage - tf.reduce_mean(advantage, axis=1, keepdims=True))\n",
    "\n",
    "            # action with highest Q values\n",
    "            self.a = tf.argmax(self.Qs, 1)\n",
    "            self.Q = tf.reduce_max(self.Qs, 1)\n",
    "            tf.summary.scalar(\"Q\", self.Q)\n",
    "\n",
    "            self.Q_target = tf.placeholder(shape=[None], dtype=tf.float32)\n",
    "            self.actions = tf.placeholder(shape=[None], dtype=tf.int32)\n",
    "            actions_onehot = tf.one_hot(self.actions, self.num_actions, on_value=1., off_value=0., axis=1, dtype=tf.float32)\n",
    "\n",
    "            Q_tmp = tf.reduce_sum(tf.multiply(self.Qs, actions_onehot), axis=1)\n",
    "            loss = tf.reduce_mean(tf.square(self.Q_target - Q_tmp))\n",
    "            self.summary_ops.append(tf.summary.scalar(\"loss\", loss))\n",
    "            optimizer = tf.train.AdamOptimizer()\n",
    "            self.minimize = optimizer.minimize(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Action:\n",
    "    UP = 0\n",
    "    DOWN = 1\n",
    "    FORWARD = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Environment:\n",
    "    \n",
    "    actions = {Action.UP:'UP', Action.FORWARD:'FORTH', Action.DOWN:'DOWN'}\n",
    "\n",
    "    def __init__(self, host, port, debug=False):\n",
    "        self.debug = debug\n",
    "        self.queue = multiprocessing.Queue()\n",
    "        self.game_client = None\n",
    "        self.server = WebsocketServer(port, host=host)\n",
    "        self.server.set_fn_new_client(self.new_client)\n",
    "        self.server.set_fn_message_received(self.new_message)\n",
    "        thread = threading.Thread(target = self.server.run_forever)\n",
    "        thread.daemon = True\n",
    "        thread.start()\n",
    "\n",
    "    def new_client(self, client, server):\n",
    "        if self.debug: print(\"GameAgent: Game just connected\")\n",
    "        self.game_client = client\n",
    "        self.server.send_message(self.game_client, \"Connection to Game Agent Established\");\n",
    "\n",
    "    def new_message(self, client, server, message):\n",
    "        if self.debug: print(\"GameAgent: Incoming data from game\")\n",
    "        data = json.loads(message)\n",
    "        image, crashed = data['world'], data['crashed']\n",
    "\n",
    "        image = re.sub('data:image/png;base64,', '',image)\n",
    "        # convert image from base64 decoding to np array\n",
    "        image = np.array(Image.open(BytesIO(base64.b64decode(image))))\n",
    "\n",
    "        crashed = True if crashed in ['True', 'true'] else False\n",
    "\n",
    "        self.queue.put((image, crashed))\n",
    "\n",
    "    def start_game(self):\n",
    "        \"\"\"\n",
    "        Starts the game and lets the TRex run for half a second and then returns the initial state.\n",
    "\n",
    "        return the initial state of the game (np.array, reward, crashed).\n",
    "        \"\"\"\n",
    "        # game not start when browser is not ready\n",
    "        while self.game_client is None:\n",
    "            time.sleep(1)\n",
    "\n",
    "        self.server.send_message(self.game_client, \"START\");\n",
    "        time.sleep(4)\n",
    "        return self.get_state(Action.FORWARD)\n",
    "\n",
    "    def refresh_game(self):\n",
    "        time.sleep(0.5)\n",
    "        print(\"...refreshing game...\")\n",
    "        self.server.send_message(self.game_client, \"REFRESH\");\n",
    "        time.sleep(1)\n",
    "\n",
    "    def do_action(self, action):\n",
    "        \"\"\"\n",
    "        Performs action and returns the updated status\n",
    "\n",
    "        param action:  Must come from the class Action.\n",
    "                        The only allowed actions are Action.UP, Action.Down and Action.FORWARD.\n",
    "        return the image of the game after performing the action, the reward (after the action) and\n",
    "                        whether the TRex crashed or not.\n",
    "        \"\"\"\n",
    "        if action != Action.FORWARD:\n",
    "            # noting needs to be send when the action is going forward\n",
    "            self.server.send_message(self.game_client, self.actions[action]);\n",
    "\n",
    "        time.sleep(.05)\n",
    "        return self.get_state(action)\n",
    "\n",
    "    def get_state(self, action):\n",
    "        self.server.send_message(self.game_client, \"STATE\");\n",
    "\n",
    "        image, crashed = self.queue.get()\n",
    "\n",
    "        if crashed:\n",
    "            reward = -100.\n",
    "        else:\n",
    "            if action == Action.UP:\n",
    "                reward = -5.\n",
    "            elif action == Action.DOWN:\n",
    "                reward = -3.\n",
    "            else:\n",
    "                reward = 1.\n",
    "\n",
    "        return image, reward, crashed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Preprocessor:\n",
    "\n",
    "    def __init__(self, width, height):\n",
    "        self.width = width\n",
    "        self.height = height\n",
    "\n",
    "    def process(self, frame):\n",
    "        roi_height, roi_width = frame.shape[0], int(frame.shape[1] * .68)\n",
    "        processed = np.zeros((roi_height, roi_width))\n",
    "\n",
    "        roi = frame[:, :roi_width, 0]\n",
    "        all_obstacles_idx = roi > 50\n",
    "        processed[all_obstacles_idx] = 1\n",
    "        unharmful_obstacles_idx = roi > 200\n",
    "        processed[unharmful_obstacles_idx] = 0\n",
    "\n",
    "        processed = imresize(processed, (self.height, self.width, 1))\n",
    "        processed = processed / 255.0\n",
    "        return processed\n",
    "\n",
    "    def get_initial_state(self, first_frame):\n",
    "        self.state = np.array([first_frame, first_frame, first_frame, first_frame])\n",
    "        return self.state\n",
    "\n",
    "    def get_updated_state(self, next_frame):\n",
    "        self.state =  np.array([*self.state[-3:], next_frame])\n",
    "        return self.state"
   ]
  },
  {
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}