mario_dqn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import datetime\n",
    "import math\n",
    "import pickle\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "import gym_super_mario_bros\n",
    "\n",
    "from nes_py.wrappers import JoypadSpace\n",
    "from gym_super_mario_bros import actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_shape = (84, 84)\n",
    "learning_rate = 0.00025\n",
    "discount_factor = 0.99\n",
    "update_target_network_interval = 20\n",
    "action_set = gym_super_mario_bros.actions.COMPLEX_MOVEMENT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model():\n",
    "    model = tf.keras.Sequential()\n",
    "    # Conv2D with data_format='channels_first' doesn't have CPU kernel without MKL.\n",
    "    # That's why I convert NCHW to NHWC here.  It probably has performance cost.\n",
    "    model.add(tf.keras.layers.Permute((2, 3, 1), input_shape=(4,) + state_shape))\n",
    "    model.add(tf.keras.layers.Conv2D(32, (8, 8), strides=(4, 4), activation='relu', \n",
    "                               kernel_initializer='he_normal',\n",
    "                               bias_initializer='ones'))\n",
    "    model.add(tf.keras.layers.Conv2D(64, (4, 4), strides=(2, 2), activation='relu', \n",
    "                               kernel_initializer='he_normal',\n",
    "                               bias_initializer='ones'))\n",
    "    model.add(tf.keras.layers.Conv2D(64, (3, 3), strides=(1, 1), activation='relu', \n",
    "                               kernel_initializer='he_normal',\n",
    "                               bias_initializer='ones'))\n",
    "    model.add(tf.keras.layers.Flatten())\n",
    "    model.add(tf.keras.layers.Dense(512, activation='relu', \n",
    "                              kernel_initializer='he_normal',\n",
    "                              bias_initializer='ones'))\n",
    "    model.add(tf.keras.layers.Dense(len(action_set), kernel_initializer='he_normal',\n",
    "                              bias_initializer='ones'))\n",
    "    model.compile(optimizer = tf.keras.optimizers.Adam(lr = learning_rate),\n",
    "                  loss = 'mse', metrics=['mse'])\n",
    "    return model\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "permute (Permute)            (None, 84, 84, 4)         0         \n",
      "_________________________________________________________________\n",
      "conv2d (Conv2D)              (None, 20, 20, 32)        8224      \n",
      "_________________________________________________________________\n",
      "conv2d_1 (Conv2D)            (None, 9, 9, 64)          32832     \n",
      "_________________________________________________________________\n",
      "conv2d_2 (Conv2D)            (None, 7, 7, 64)          36928     \n",
      "_________________________________________________________________\n",
      "flatten (Flatten)            (None, 3136)              0         \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (None, 512)               1606144   \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 12)                6156      \n",
      "=================================================================\n",
      "Total params: 1,690,284\n",
      "Trainable params: 1,690,284\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model = create_model()\n",
    "model.summary()\n",
    "target_model = create_model()\n",
    "target_model.set_weights(model.get_weights())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "logdir = \"logs/scalars/\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_epoch = 0\n",
    "starting_episode = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "continuing @ episode 137 epoch 385000\n"
     ]
    }
   ],
   "source": [
    "with open(\"models/state.txt\", 'rb') as f:\n",
    "    starting_episode, current_epoch, logdir = pickle.load(f)\n",
    "    model = tf.keras.models.load_model('models/mario0')\n",
    "    target_model.set_weights(model.get_weights())\n",
    "    print ('continuing @ episode {} epoch {}'.format(starting_episode, current_epoch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir, update_freq=5991)\n",
    "file_writer = tf.summary.create_file_writer(logdir + \"/metrics\")\n",
    "file_writer.set_as_default()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_epoch_end(epoch, logs=None):\n",
    "    if tensorboard_callback.histogram_freq and epoch % tensorboard_callback.histogram_freq == 0:\n",
    "      tensorboard_callback._log_weights(epoch)\n",
    "tensorboard_callback.on_epoch_end = on_epoch_end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_target_network(\n",
    "    episode, update_target_network_interval, main_network, target_network):\n",
    "    if ((episode+1) % update_target_network_interval) == 0:\n",
    "        target_network.set_weights(main_network.get_weights())\n",
    "    return target_network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym_super_mario_bros.make('SuperMarioBros-v0')\n",
    "env = JoypadSpace(env, action_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_action(epsilon, state):\n",
    "    if (np.random.random() <= epsilon):\n",
    "        return np.random.choice(len(action_set))\n",
    "    else:\n",
    "        state = tf.reshape(tf.concat(state, axis=0), (1,4,) + state_shape)\n",
    "        return np.argmax(model.predict(state))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def greyscale(state):\n",
    "    return tf.image.rgb_to_grayscale([state])[0]\n",
    "\n",
    "def resize(state):\n",
    "    return tf.compat.v1.image.resize_images([state], (state_shape[0], state_shape[1]))[0]\n",
    "\n",
    "def normalize_image(state):\n",
    "    return tf.image.per_image_standardization(state)\n",
    "\n",
    "def downsample(state):\n",
    "    state = resize(state)\n",
    "    state = greyscale(state)\n",
    "    state = normalize_image(state)\n",
    "    return tf.cast(tf.reshape(state, (1,) + state_shape), tf.dtypes.float32)\n",
    "\n",
    "def normalize_reward(reward):\n",
    "    return reward / 5.0 # rewards are in the [-15,15] range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_bellman_target(discount_factor, reward, model, state_next, done):\n",
    "    if done:\n",
    "        return reward\n",
    "    return (reward + discount_factor * np.max(model.predict(state_next)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_start = 0.4\n",
    "beta_frames = 1000 \n",
    "beta_by_frame = lambda frame_idx: min(1.0, beta_start + frame_idx * (1.0 - beta_start) / beta_frames)\n",
    "#TODO: rename frame_idx to episode_id\n",
    "alpha_start = 0.6\n",
    "alpha_frames = 1000 \n",
    "alpha_by_frame = lambda frame_idx: min(1.0, alpha_start + frame_idx * (1.0 - alpha_start) / alpha_frames)\n",
    "\n",
    "epsilon_start = 0.99\n",
    "epsilon_final = 0.01\n",
    "epsilon_decay = 500\n",
    "epsilon_by_frame = lambda frame_idx: epsilon_final + (epsilon_start - epsilon_final) * math.exp(-1. * frame_idx / epsilon_decay)\n",
    "\n",
    "alpha = alpha_by_frame(starting_episode)\n",
    "beta = beta_by_frame(starting_episode)\n",
    "epsilon = epsilon_by_frame(starting_episode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PrioritizedBuffer(object):\n",
    "    def __init__(self, capacity):\n",
    "        self.capacity = capacity\n",
    "        self.position = 0\n",
    "        self.buffer = []\n",
    "        self.priorities = np.zeros((capacity,), dtype=np.float32)\n",
    "    \n",
    "    def append(self, frame):#state, action, reward, state_next, done, info):\n",
    "        max_priority = self.priorities.max() if self.buffer else 1.0\n",
    "        \n",
    "        if len(self.buffer) < self.capacity:\n",
    "            self.buffer.append(frame)\n",
    "        else:\n",
    "            self.buffer[self.position] = frame\n",
    "        \n",
    "        self.priorities[self.position] = max_priority\n",
    "        self.position = (self.position + 1) % self.capacity\n",
    "    \n",
    "    def sample(self, batch_size, alpha=0.6, beta=0.4):\n",
    "        if len(self.buffer) == self.capacity:\n",
    "            priorities = self.priorities\n",
    "        else:\n",
    "            priorities = self.priorities[:self.position]\n",
    "        \n",
    "        probabilities = priorities ** alpha\n",
    "        probabilities /= probabilities.sum()\n",
    "        \n",
    "        indices = np.random.choice(len(self.buffer), batch_size, p=probabilities)\n",
    "        samples = [self.buffer[idx] for idx in indices]\n",
    "        \n",
    "        total = len(self.buffer)\n",
    "        weights = (total * probabilities[indices]) ** (-beta)\n",
    "        weights /= weights.max()\n",
    "        weights = np.array(weights, dtype=np.float32)\n",
    "        \n",
    "        states, actions, rewards, states_next, dones, infos = [], [], [], [], [], []\n",
    "        for frame in samples:            \n",
    "            batch = list(zip(*frame))\n",
    "            states.append(tf.reshape(tf.concat([s for s in batch[0]], axis=0), (1,4,) + state_shape))\n",
    "            actions.append(batch[1][-1])\n",
    "            rewards.append(batch[2][-1])\n",
    "            states_next.append(tf.reshape(tf.concat([s for s in batch[3]], axis=0), (1,4,) + state_shape))\n",
    "            dones.append(batch[4][-1])\n",
    "            infos.append(batch[5][-1])\n",
    "        \n",
    "        return states, actions, rewards, states_next, dones, infos, indices, weights\n",
    "    \n",
    "    def update_priorities(self, batch_indices, batch_priorities):\n",
    "        for idx, priority in zip(batch_indices, batch_priorities):\n",
    "            self.priorities[idx] = priority\n",
    "            \n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.buffer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_from_replay_buffer_and_train_model(replay_buffer, batch_size, \n",
    "                                              model, target_model, discount_factor):\n",
    "    global current_epoch\n",
    "    if(len(replay_buffer) >= batch_size):\n",
    "        (states, actions, rewards, states_next, dones, infos, \n",
    "         indices, weights) = replay_buffer.sample(batch_size, alpha, beta)\n",
    "        \n",
    "        stacked_states = np.empty(shape=(0,4,) + state_shape)\n",
    "        for state in states:\n",
    "            stacked_states = tf.concat((stacked_states, state), axis=0)\n",
    "\n",
    "        target_q_values = model.predict(stacked_states)\n",
    "        td_errors = np.zeros((len(states),), dtype=np.float32)\n",
    "        for i in range(len(states)):\n",
    "            updated_target = compute_bellman_target(\n",
    "                discount_factor, rewards[i], target_model, states_next[i], dones[i])\n",
    "            td_errors[i] = updated_target - target_q_values[i, actions[i]]\n",
    "            target_q_values[i, actions[i]] = updated_target\n",
    "\n",
    "        callbacks = []\n",
    "        if current_epoch % 1000 == 0:\n",
    "            def summarize_q_values(epoch, logs): \n",
    "                # TODO: move this function out of the way.\n",
    "                tf.summary.scalar('target_q_values', data=tf.reduce_mean(target_q_values), step=epoch)\n",
    "                tf.summary.scalar('epsilon', data=epsilon, step=epoch)\n",
    "                x, y, t, c = zip(*[(i['x_pos'],i['y_pos'],i['time'],i['cumulative_reward']) for i in infos])\n",
    "                tf.summary.scalar('x_pos', data=tf.reduce_mean(x), step=epoch)\n",
    "                tf.summary.scalar('y_pos', data=tf.reduce_mean(y), step=epoch)\n",
    "                tf.summary.scalar('cumulative_reward', data=tf.reduce_mean(c), step=epoch)\n",
    "                tf.summary.scalar('time', data=tf.reduce_mean(t), step=epoch)\n",
    "                tf.summary.scalar('td_errors', data=tf.reduce_mean(td_errors), step=epoch)\n",
    "            summarize = tf.keras.callbacks.LambdaCallback(on_epoch_begin=summarize_q_values)\n",
    "            callbacks=[tensorboard_callback, summarize]\n",
    "\n",
    "        model.fit(stacked_states, target_q_values, sample_weight=weights,\n",
    "                  epochs=current_epoch + 1, initial_epoch=current_epoch,\n",
    "                  verbose=False, callbacks=callbacks)\n",
    "        td_errors = (np.power(td_errors, 2) * weights) + 1e-5\n",
    "        replay_buffer.update_priorities(indices, td_errors)\n",
    "    current_epoch += 1\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "replay_buffer_size = 1000000\n",
    "replay_buffer = PrioritizedBuffer(replay_buffer_size)\n",
    "overlapping_buffer = collections.deque(maxlen=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/image_ops_impl.py:1511: div (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Deprecated in favor of operator or tf.math.divide.\n",
      "WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (0.267066). Check your callbacks.\n"
     ]
    }
   ],
   "source": [
    "for episode in range(starting_episode, 33333):\n",
    "    state = env.reset()\n",
    "    for _ in range(np.random.randint(0, 133)):\n",
    "        state, _, _, _ = env.step(0)\n",
    "    state = downsample(state)\n",
    "    frame_states = [state, state, state, state]\n",
    "    \n",
    "    done = False\n",
    "    cumulative_reward = 0\n",
    "    for __ in range(3333):\n",
    "        if done:\n",
    "            break\n",
    "        action = select_action(epsilon, frame_states)\n",
    "        \n",
    "        total_reward = 0\n",
    "        for _ in range(4):\n",
    "            state_next, reward, done, info = env.step(action)\n",
    "            total_reward += reward\n",
    "            if done: break\n",
    "            \n",
    "        state_next, reward = downsample(state_next), normalize_reward(total_reward)\n",
    "        cumulative_reward += total_reward\n",
    "        info['cumulative_reward'] = cumulative_reward\n",
    "        overlapping_buffer.append((state, action, reward, state_next, done, info))\n",
    "        \n",
    "        if len(overlapping_buffer) == overlapping_buffer.maxlen:\n",
    "            frame = [s for s in overlapping_buffer]\n",
    "            frame_states =  [s[0] for s in overlapping_buffer]\n",
    "            replay_buffer.append(frame)\n",
    "        \n",
    "        model = sample_from_replay_buffer_and_train_model(\n",
    "            replay_buffer, batch_size, model, target_model, discount_factor)\n",
    "        state = state_next\n",
    "        \n",
    "        if current_epoch % 30000 == 0:\n",
    "            model.save('models/mario0')\n",
    "            model = tf.keras.models.load_model('models/mario0')\n",
    "            with open(\"models/state.txt\", 'wb') as f:\n",
    "                pickle.dump((episode, current_epoch, logdir), f)\n",
    "        \n",
    "    target_network = update_target_network(\n",
    "            episode, update_target_network_interval, model, target_model)\n",
    "    \n",
    "    alpha = alpha_by_frame(episode)\n",
    "    beta = beta_by_frame(episode)\n",
    "    epsilon = epsilon_by_frame(episode)"
   ]
  },
  {
   "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}