turtlebot_circle_dqn_clean.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from collections import deque\n",
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "import time\n",
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"   # comment this to use GPUs (if available)\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"          # if having multiple GPUs, explicitly define which GPU to use\n",
    "                                                 # tf has a nasty habit of taking up ALL available memory\n",
    "\n",
    "# gym stuff\n",
    "import gym\n",
    "import gym_gazebo\n",
    "\n",
    "# keras stuff\n",
    "from keras.models import Sequential\n",
    "from keras.layers import Dense\n",
    "from keras.optimizers import Adam\n",
    "\n",
    "import action_heatmap\n",
    "\n",
    "model_save_path = './models/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make environment and set parameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "SEED = 0\n",
    "MAX_PATH_LENGTH = 5000\n",
    "NUM_EPISODES = 1000\n",
    "\n",
    "# Make the turtlebot-in-empty-world environment\n",
    "env = gym.make('GazeboEmptyTurtlebotLidarNn-v0')\n",
    "env.reset()\n",
    "env._max_episode_steps = MAX_PATH_LENGTH\n",
    "\n",
    "np.random.seed(SEED)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an experience replay buffer\n",
    "Has the ability to store state transitions and return randomly sampled transitions. Vital for breaking correlations between data for training a neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from collections import namedtuple, deque\n",
    "\n",
    "class ReplayMemory:\n",
    "    def __init__(self, capacity = 10000):\n",
    "        self.memory = deque(maxlen=capacity) # max size for replay buffer\n",
    "        self.Transition = namedtuple('Transition',\n",
    "                        ('state', 'action', 'reward', 'next_state', 'done'))    \n",
    "    \n",
    "    def push(self, *args):\n",
    "        \"\"\"Saves a transition.\"\"\"\n",
    "        self.memory.append(self.Transition(*args))\n",
    "        \n",
    "    def sample(self, batch_size):\n",
    "        if batch_size <= self.__len__():\n",
    "            return random.sample(self.memory, batch_size)\n",
    "        else:\n",
    "            print('Tried to sample more samples than are in buffer')\n",
    "            return -1\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define a DQN class\n",
    "\n",
    "Has methods for creating the main and target networks, sampling actions, and syncing weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class QNetwork:\n",
    "    def __init__(self, state_size, action_size, activation_type, learning_rate, \n",
    "                 alpha_decay, gamma, batch_size):\n",
    "        self.state_size = state_size\n",
    "        self.action_size = action_size\n",
    "        self.learning_rate = learning_rate\n",
    "        self.alpha_decay = alpha_decay\n",
    "        self.activation_type = activation_type\n",
    "        self.gamma = gamma\n",
    "        self.replay_buffer = ReplayMemory()\n",
    "        self.batch_size = batch_size\n",
    "        \n",
    "    def create_model(self, hidden_layers):\n",
    "        model = Sequential()\n",
    "        if len(hidden_layers) == 0:\n",
    "            model.add(Dense(self.action_size, input_shape=self.state_size, activation=\"linear\"))\n",
    "        else:\n",
    "            for i, hid in enumerate(hidden_layers):\n",
    "                if i == 0:\n",
    "                    model.add(Dense(hid, input_dim=state_size, activation=self.activation_type))\n",
    "                else:\n",
    "                    model.add(Dense(hid, activation=self.activation_type))\n",
    "            model.add(Dense(action_size, activation='linear'))\n",
    "        optimizer = Adam(lr=self.learning_rate, decay=self.alpha_decay)\n",
    "        model.compile(loss='mse', optimizer=optimizer)\n",
    "        return model\n",
    "        \n",
    "    def init_models(self, hidden_layers):\n",
    "        self.mainDQN = self.create_model(hidden_layers)\n",
    "        self.targetDQN = self.create_model(hidden_layers)\n",
    "        self.sync_weights()\n",
    "    \n",
    "    def model_descriptions(self):\n",
    "        print('Main DQN model architecture --->')\n",
    "        self.mainDQN.summary()\n",
    "        \n",
    "        print('Target DQN model architecture --->')\n",
    "        self.targetDQN.summary()\n",
    "        \n",
    "    def sync_weights(self):\n",
    "        self.targetDQN.set_weights(self.mainDQN.get_weights())\n",
    "        \n",
    "    def save_weights(self):\n",
    "        self.mainDQN.save('./models/turtlebot_dqn_2.h5')\n",
    "        \n",
    "    def action_sampler(self, state, exploration_rate):\n",
    "        if epsilon > np.random.rand():\n",
    "            \n",
    "            # Make a random action\n",
    "            action = np.random.randint(0, self.action_size)\n",
    "        else:\n",
    "            \n",
    "            # Get action from Q-network\n",
    "            Qs = self.mainDQN.predict(state)[0]\n",
    "            action = np.argmax(Qs)\n",
    "        return action\n",
    "    \n",
    "    def train(self):\n",
    "        transitions = self.replay_buffer.sample(self.batch_size)\n",
    "        batch = self.replay_buffer.Transition(*zip(*transitions))\n",
    "        x_batch = np.squeeze(np.asarray(batch.state))\n",
    "        y_batch = self.mainDQN.predict(x_batch)\n",
    "        for i, sample in enumerate(transitions):\n",
    "            y_batch[i, sample.action] = sample.reward + self.gamma * \\\n",
    "            np.amax(self.targetDQN.predict(sample.next_state)[0]) * (1 - sample.done)\n",
    "   \n",
    "        self.sync_weights()\n",
    "        \n",
    "        self.mainDQN.fit(x_batch, y_batch, epochs=1, verbose=0)\n",
    "        \n",
    "        self.save_weights()\n",
    "        \n",
    "               "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set hyperparameters and initialize DQN\n",
    "\n",
    "Print out network architecture for both models to make sure everything is per plan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "learning_rate = 0.01\n",
    "state_size = 2\n",
    "action_size = 21\n",
    "activation_type = 'relu'\n",
    "hidden_size = [100, 100]\n",
    "alpha_decay = 0.01\n",
    "gamma = 0.99\n",
    "batch_size = 64\n",
    "learn_start = 64\n",
    "\n",
    "DQN = QNetwork(state_size, action_size, activation_type, learning_rate, alpha_decay, gamma, batch_size)\n",
    "DQN.init_models(hidden_size)\n",
    "DQN.model_descriptions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start training!\n",
    "\n",
    "Epsilon is set to decay at an exponential rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Exploration parameters\n",
    "min_epsilon = 0.01           # minimum exploration probability\n",
    "decay_rate = 20/NUM_EPISODES\n",
    "\n",
    "returns = np.zeros(NUM_EPISODES)\n",
    "\n",
    "for ep in range(1, NUM_EPISODES):\n",
    "    # Start new episode\n",
    "    state = env.reset()\n",
    "    state = np.reshape(state, [1, state.size])\n",
    "\n",
    "    total_reward = 0\n",
    "    episode_step = 0\n",
    "    done = False\n",
    "\n",
    "    epsilon = min_epsilon + (1.0 - min_epsilon)*np.exp(-decay_rate*ep)\n",
    "\n",
    "    while not done:\n",
    "        episode_step += 1\n",
    "        \n",
    "        action = DQN.action_sampler(state, epsilon)\n",
    "        \n",
    "        # Take action, get new state and reward\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "        next_state = np.reshape(next_state, [1, state.size])\n",
    "\n",
    "        total_reward += reward\n",
    "        \n",
    "        # add to replay buffer\n",
    "        DQN.replay_buffer.push(state, action, reward, next_state, done)\n",
    "        state = next_state\n",
    "    \n",
    "    # Progress updates\n",
    "    returns[ep]=total_reward\n",
    "\n",
    "    # Plot policy learning process\n",
    "    if ep % 10 == 0:\n",
    "        ax = action_heatmap.main(env, DQN.mainDQN)\n",
    "        fig = ax.get_figure()\n",
    "        fig.savefig('./heatmaps/action_heatmap_ep_' + str(ep) +'.png')\n",
    "    \n",
    "    print('Return recorded for episode ' + str(ep) + ' : ', total_reward)\n",
    "    if ep % 10 == 0 or ep == NUM_EPISODES-1: #print out gradual improvement\n",
    "        print('Episode {}'.format(ep-10),'..{}'.format(ep),'  Avg reward: {}'.format(sum(returns[ep-10:ep])/10),'  Epsilon: {:.4f}'.format(epsilon))\n",
    "    \n",
    "    # Train network!\n",
    "    if len(DQN.replay_buffer) >= learn_start:\n",
    "        DQN.train()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# plot average returns\n",
    "returns_over_each_episode = []\n",
    "x = []\n",
    "for i in range(ep):\n",
    "    returns_over_each_episode.append(sum(returns[1*i:1*(i+1)-1])/1)\n",
    "    x.append((i+1)*1)\n",
    "plt.plot(returns[:ep], '.-r')\n",
    "plt.ylabel('Average Returns per Episode')\n",
    "plt.xlabel('Episodes')\n",
    "plt.show()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ignore the rest of this file\n",
    "\n",
    "Called only for testing helper scripts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import importlib\n",
    "importlib.reload(action_heatmap)\n",
    "print(env)\n",
    "ax = action_heatmap.main(env, DQN.mainDQN)\n",
    "fig = ax.get_figure()\n",
    "fig.savefig('./heatmaps/action_heatmap_ep_' + str(ep) +'.png')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    " %matplotlib qt\n",
    "import reward_surface\n",
    "import importlib\n",
    "importlib.reload(reward_surface)\n",
    "ax = reward_surface.main(env)\n",
    "fig = ax.get_figure()"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}