2.0-quadrotor-train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "import pickle\n",
    "import random\n",
    "import os\n",
    "import math\n",
    "import types\n",
    "import uuid\n",
    "import time\n",
    "from copy import copy\n",
    "from collections import defaultdict, Counter\n",
    "\n",
    "import numpy as np\n",
    "from gym.envs.classic_control import rendering\n",
    "from pyglet.window import key as pygkey\n",
    "import gym\n",
    "from gym import spaces, wrappers\n",
    "\n",
    "import dill\n",
    "import tempfile\n",
    "import tensorflow as tf\n",
    "from tensorflow.contrib import rnn\n",
    "import zipfile\n",
    "\n",
    "import baselines.common.tf_util as U\n",
    "\n",
    "from baselines import logger\n",
    "from baselines.common.schedules import LinearSchedule\n",
    "from baselines.deepq.models import mlp as deepq_mlp\n",
    "from baselines.deepq import learn as deepq_learn\n",
    "from baselines.deepq.replay_buffer import ReplayBuffer, PrioritizedReplayBuffer\n",
    "from baselines.deepq.simple import ActWrapper\n",
    "\n",
    "from scipy.special import logsumexp\n",
    "\n",
    "from pyquaternion import Quaternion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas\n",
    "from matplotlib.figure import Figure\n",
    "import matplotlib as mpl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data_dir = os.path.join('data', 'quadrotor')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "create idealized ARDrone environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "latest_vicon_state = np.zeros(7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def init_latest_vicon_state():\n",
    "  global latest_vicon_state\n",
    "  latest_vicon_state = np.zeros(7)\n",
    "  latest_vicon_state[-2] = 1\n",
    "  latest_vicon_state[-1] = np.random.random() * 2 * math.pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "init_latest_vicon_state()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_act_dim = 9\n",
    "n_obs_dim = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ARDrone(gym.Env):\n",
    "  metadata = {'render.modes': ['human']}\n",
    "  \n",
    "  def __init__(self):      \n",
    "    self.action_space = spaces.Discrete(n_act_dim)\n",
    "    obs_low = -np.array([np.inf] * n_obs_dim)\n",
    "    obs_high = -obs_low\n",
    "    self.observation_space = spaces.Box(obs_low, obs_high)\n",
    "    \n",
    "    self.img_height = 368\n",
    "    self.img_width = 640\n",
    "    self.ground_contact_z_thresh = 0.5\n",
    "    self.max_pos = np.array([2, 2, 2.5])\n",
    "    self.max_ep_len = 300\n",
    "    self.start_z_thresh = 2\n",
    "    \n",
    "    self.goal_coords = self._init_goal_coords()\n",
    "    self.goal_twist = self._init_goal_twist()\n",
    "    self.terrain = self._init_terrain()\n",
    "    \n",
    "    self.goal_dist_thresh = 0.3\n",
    "    self.goal_twist_thresh = 0.1\n",
    "    \n",
    "    self.translation = None\n",
    "    self.quaternion = None\n",
    "    self.prev_pos = None\n",
    "    self.prev_obs = None\n",
    "    self.prev_time = None\n",
    "    self.prev_shaping = None\n",
    "    self.curr_step = None\n",
    "    self.latest_zs = np.zeros(10)\n",
    "    self.rollout = None\n",
    "    self.ep_start_time = None\n",
    "    \n",
    "    self.action_duration = 1 / 15\n",
    "    \n",
    "    self.viewer = None\n",
    "        \n",
    "    self.set_transform()\n",
    "    \n",
    "  def _init_terrain(self):\n",
    "    return np.ones((10, 10))\n",
    "    \n",
    "  def _init_goal_coords(self):\n",
    "    return np.random.random(2)*2-1\n",
    "  \n",
    "  def _init_goal_twist(self):\n",
    "    return -math.pi + np.random.random()*2*math.pi\n",
    "    \n",
    "  def _get_vicon_pos(self):\n",
    "    return copy(latest_vicon_state)\n",
    "    \n",
    "  def set_transform(self, pos=None):\n",
    "    if pos is None:\n",
    "      pos = self._get_vicon_pos()\n",
    "    self.translation = -pos[:3]\n",
    "    self.quaternion = Quaternion(pos[3:])\n",
    "    \n",
    "  def _transform_pos(self, pos):\n",
    "    return pos\n",
    "    \n",
    "  def _obs(self):\n",
    "    pos = self._transform_pos(self._get_vicon_pos())\n",
    "    pos = np.concatenate((pos[:3], pos[6:7]))\n",
    "    curr_time = time.time()\n",
    "    if self.prev_pos is not None:\n",
    "      vel = (pos - self.prev_pos) / (curr_time - self.prev_time)\n",
    "    else:\n",
    "      vel = np.zeros(pos.shape)\n",
    "    self.prev_pos = pos\n",
    "    self.prev_time = curr_time\n",
    "        \n",
    "    rot_ang = -pos[3]\n",
    "    rot = np.array([[np.cos(rot_ang), -np.sin(rot_ang)], [np.sin(rot_ang), np.cos(rot_ang)]])\n",
    "    delta_xy_to_goal = rot.dot(self.goal_coords - pos[:2])\n",
    "        \n",
    "    return np.concatenate((pos, vel, delta_xy_to_goal))\n",
    "  \n",
    "  def _at_site(self):\n",
    "    pos = self._transform_pos(self._get_vicon_pos())\n",
    "    at_loc = np.linalg.norm(pos[:2] - self.goal_coords) <= self.goal_dist_thresh\n",
    "    at_rot = abs(pos[6] - self.goal_twist) <= self.goal_twist_thresh\n",
    "    return at_loc, at_rot\n",
    "  \n",
    "  def _exec_action(self, action, vel_unit=0.75):\n",
    "    vel = {\n",
    "      'linear_x': 0,\n",
    "      'linear_y': 0,\n",
    "      'linear_z': 0,\n",
    "      'angular_x': 0,\n",
    "      'angular_y': 0,\n",
    "      'angular_z': 0\n",
    "    }\n",
    "\n",
    "    if action == 0:\n",
    "      vel['linear_x'] = vel_unit\n",
    "    elif action == 1:\n",
    "      vel['linear_x'] = -vel_unit\n",
    "    elif action == 2:\n",
    "      vel['linear_y'] = vel_unit\n",
    "    elif action == 3:\n",
    "      vel['linear_y'] = -vel_unit\n",
    "    elif action == 4:\n",
    "      vel['linear_z'] = vel_unit\n",
    "    elif action == 5:\n",
    "      vel['linear_z'] = -vel_unit\n",
    "    elif action == 6:\n",
    "      vel['angular_z'] = vel_unit\n",
    "    elif action == 7:\n",
    "      vel['angular_z'] = -vel_unit\n",
    "    elif action == NOOP:\n",
    "      pass\n",
    "    else:\n",
    "      raise ValueError\n",
    "      \n",
    "    global latest_vicon_state\n",
    "    \n",
    "    pos = self._transform_pos(self._get_vicon_pos())\n",
    "    rot_ang = pos[6]\n",
    "    rot = np.array([[np.cos(rot_ang), -np.sin(rot_ang)], [np.sin(rot_ang), np.cos(rot_ang)]])\n",
    "    vx, vy = rot.dot(np.array([vel['linear_x'], vel['linear_y']]))\n",
    "    latest_vicon_state[0] += vx * self.action_duration\n",
    "    latest_vicon_state[1] += vy * self.action_duration\n",
    "    \n",
    "    latest_vicon_state[2] += vel['linear_z'] * self.action_duration\n",
    "    latest_vicon_state[6] += vel['angular_z'] * self.action_duration\n",
    "    if latest_vicon_state[6] < 0:\n",
    "      latest_vicon_state[6] += 2 * math.pi\n",
    "    latest_vicon_state[6] = latest_vicon_state[6] % (2 * math.pi)\n",
    "    \n",
    "  def _out_of_bounds(self):\n",
    "    pos = self._transform_pos(self._get_vicon_pos())\n",
    "    return (np.abs(pos[:3]) >= self.max_pos).any()\n",
    "  \n",
    "  def _step(self, action):\n",
    "    self._exec_action(action)\n",
    "    \n",
    "    obs = self._obs()\n",
    "        \n",
    "    dist_to_goal = np.linalg.norm(obs[-2:])\n",
    "    if dist_to_goal < 0.5*self.goal_dist_thresh:\n",
    "      shaping = -100*obs[2]\n",
    "    else:\n",
    "      shaping = -100*dist_to_goal\n",
    "\n",
    "    r = shaping\n",
    "    if self.prev_shaping is not None:\n",
    "      r -= self.prev_shaping\n",
    "    self.prev_shaping = shaping\n",
    "    \n",
    "    oob = self._out_of_bounds()\n",
    "    timeout = self.curr_step >= self.max_ep_len\n",
    "    on_ground = obs[2] <= self.ground_contact_z_thresh\n",
    "    done = oob or timeout or on_ground\n",
    "    \n",
    "    self.curr_step += 1\n",
    "    self.latest_zs[:-1] = self.latest_zs[1:]\n",
    "    self.latest_zs[-1] = obs[2]\n",
    "    \n",
    "    at_loc, at_rot = self._at_site()\n",
    "        \n",
    "    info = {}\n",
    "    if done:\n",
    "      self._land()\n",
    "      info['duration'] = time.time() - self.ep_start_time\n",
    "      if oob or (on_ground and not at_loc):\n",
    "        r = -100\n",
    "      elif at_loc and on_ground:\n",
    "        r = 100\n",
    "    \n",
    "    if self.prev_obs is not None:\n",
    "      self.rollout.append((self.prev_obs, action, r, obs, done))\n",
    "    self.prev_obs = obs\n",
    "    \n",
    "    return obs, r, done, info\n",
    "    \n",
    "  def _takeoff(self):\n",
    "    global latest_vicon_state\n",
    "    latest_vicon_state[2] = self.start_z_thresh\n",
    "    \n",
    "  def _land(self):\n",
    "    init_latest_vicon_state()\n",
    "    \n",
    "  def _reset(self):\n",
    "    self._land()\n",
    "    \n",
    "    self.prev_pos = None\n",
    "    self.prev_time = None\n",
    "    self.prev_shaping = None\n",
    "    self.prev_obs = None\n",
    "    self.curr_step = 0\n",
    "    self.rollout = []\n",
    "    \n",
    "    self.goal_coords = self._init_goal_coords()\n",
    "    self.goal_twist = self._init_goal_twist()\n",
    "    self.terrain = self._init_terrain()\n",
    "    \n",
    "    self._takeoff()\n",
    "        \n",
    "    self.ep_start_time = time.time()\n",
    "    return self._step(NOOP)[0]\n",
    "  \n",
    "  def _render(self, mode='human', close=False):\n",
    "    if close:\n",
    "      if self.viewer is not None:\n",
    "        self.viewer.close()\n",
    "        self.viewer = None\n",
    "      return\n",
    "    \n",
    "    if self.viewer is None:\n",
    "      self.viewer = rendering.SimpleImageViewer()\n",
    "    \n",
    "    fig = plt.figure(figsize=(10, 10))\n",
    "    canvas = FigureCanvas(fig)\n",
    "    ax = plt.axes(xlim=(-0.5, self.terrain.shape[0] - 0.5), ylim=(-0.5, self.terrain.shape[1] - 0.5))\n",
    "    pursuer_pos = self._transform_pos(self._get_vicon_pos())\n",
    "    evader_loc = (self.goal_coords + self.max_pos[:2]) / (2 * self.max_pos[:2])\n",
    "    evader_loc[0] = evader_loc[0] * self.terrain.shape[0]-0.5\n",
    "    evader_loc[1] = evader_loc[1] * self.terrain.shape[1]-0.5\n",
    "    ax.scatter([evader_loc[0]], [evader_loc[1]], s=500, color='red', linewidth=0, alpha=0.75)\n",
    "    pursuer_size = 500 + 10000 * pursuer_pos[2] / self.max_pos[2]\n",
    "    pursuer_loc = (pursuer_pos[:2] + self.max_pos[:2]) / (2 * self.max_pos[:2])\n",
    "    pursuer_loc[0] = pursuer_loc[0] * self.terrain.shape[0]-0.5\n",
    "    pursuer_loc[1] = pursuer_loc[1] * self.terrain.shape[1]-0.5\n",
    "    ax.scatter([pursuer_loc[0]], [pursuer_loc[1]], s=pursuer_size, c='blue', linewidth=0, alpha=0.75)\n",
    "        \n",
    "    ax.set_title('%d seconds remaining' % (int(20 * (1 - ((self.curr_step if self.curr_step is not None else 0) / self.max_ep_len)))),\n",
    "                fontsize=30)\n",
    "    ax.set_xlabel('z = %0.2f' % (pursuer_pos[2] / self.max_pos[2]),\n",
    "                 fontsize=30)\n",
    "    \n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "    agg = canvas.switch_backends(FigureCanvas)\n",
    "    agg.draw()\n",
    "    width, height = fig.get_size_inches() * fig.get_dpi()\n",
    "    self.viewer.imshow(np.fromstring(agg.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3))\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "env = ARDrone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "max_ep_len = env.max_ep_len\n",
    "n_training_episodes = 300"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "max_timesteps = max_ep_len * (1 if load_pretrained_full_pilot else n_training_episodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def run_ep(policy, env, max_ep_len=max_ep_len, render=False, pilot_is_human=False):\n",
    "    if pilot_is_human:\n",
    "      global human_agent_action\n",
    "      human_agent_action = init_human_action()\n",
    "    obs = env.reset()\n",
    "    done = False\n",
    "    totalr = 0.\n",
    "    trajectory = [obs]\n",
    "    actions = []\n",
    "    for step_idx in range(max_ep_len+1):\n",
    "        if done:\n",
    "            break\n",
    "        action = policy(obs[None, :])\n",
    "        obs, r, done, info = env.step(action)\n",
    "        actions.append(action)\n",
    "        trajectory.append(obs)\n",
    "        if render:\n",
    "          env.render()\n",
    "        totalr += r\n",
    "    outcome = r if r % 100 == 0 else 0\n",
    "    return totalr, outcome, trajectory, actions, info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def noop_pilot_policy(obs):\n",
    "  return NOOP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "init_human_action = lambda: NOOP\n",
    "human_agent_action = init_human_action()\n",
    "\n",
    "action_of_key = {\n",
    "  pygkey.RIGHT: 0,\n",
    "  pygkey.LEFT: 1,\n",
    "  pygkey.UP: 2,\n",
    "  pygkey.DOWN: 3,\n",
    "  pygkey.B: 4,\n",
    "  pygkey.SPACE: 5,\n",
    "  pygkey.A: 6,\n",
    "  pygkey.D: 7\n",
    "}\n",
    "\n",
    "non_orient_keys = [pygkey.UP, pygkey.DOWN, pygkey.LEFT, pygkey.RIGHT, pygkey.SPACE, pygkey.B]\n",
    "\n",
    "def key_press(key, mod):\n",
    "  k = int(key)\n",
    "  if k in action_of_key:\n",
    "    global human_agent_action\n",
    "    human_agent_action = action_of_key[k]\n",
    "\n",
    "def key_release(key, mod):\n",
    "  k = int(key)\n",
    "  if k in action_of_key:\n",
    "    global human_agent_action\n",
    "    human_agent_action = 8\n",
    "      \n",
    "def human_pilot_policy(obs):\n",
    "  return human_agent_action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "env.render()\n",
    "env.unwrapped.viewer.window.on_key_press = key_press\n",
    "env.unwrapped.viewer.window.on_key_release = key_release"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "run_ep(twist_pilot_policy, env, render=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_eval_eps = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "full_pilot_eval = list(zip(*[run_ep(full_pilot_policy, env, render=False) for _ in range(n_eval_eps)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "twist_pilot_eval = list(zip(*[run_ep(twist_pilot_policy, env, render=False) for _ in range(n_eval_eps)]))\n",
    "oracle_pilot_eval = list(zip(*[run_ep(oracle_pilot_policy, env, render=False) for _ in range(n_eval_eps)]))\n",
    "noop_pilot_eval = list(zip(*[run_ep(noop_pilot_policy, env, render=False) for _ in range(n_eval_eps)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(full_pilot_eval[0]), Counter(full_pilot_eval[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(twist_pilot_eval[0]), Counter(twist_pilot_eval[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(oracle_pilot_eval[0]), Counter(oracle_pilot_eval[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(noop_pilot_eval[0]), Counter(noop_pilot_eval[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def save_tf_vars(scope, path):\n",
    "  sess = U.get_session()\n",
    "  saver = tf.train.Saver([v for v in tf.global_variables() if v.name.startswith(scope + '/')])\n",
    "  saver.save(sess, save_path=path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def load_tf_vars(scope, path):\n",
    "  sess = U.get_session()\n",
    "  saver = tf.train.Saver([v for v in tf.global_variables() if v.name.startswith(scope + '/')])\n",
    "  saver.restore(sess, path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "train assistive copilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_training_episodes = 500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "make_q_func = lambda: deepq_mlp([64, 64])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "copilot_dqn_learn_kwargs = {\n",
    "  'lr': 1e-4,\n",
    "  'exploration_fraction': 0.1,\n",
    "  'exploration_final_eps': 0.02,\n",
    "  'target_network_update_freq': 3000,\n",
    "  'print_freq': 100,\n",
    "  'num_cpu': 5,\n",
    "  'gamma': 0.99\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_co_env(pilot_policy, **extras):\n",
    "  env = ARDrone()\n",
    "  env.unwrapped.pilot_policy = pilot_policy\n",
    "  return env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def co_build_act(make_obs_ph, q_func, num_actions, scope=\"deepq\", reuse=None, using_control_sharing=True):\n",
    "  with tf.variable_scope(scope, reuse=reuse):\n",
    "    observations_ph = U.ensure_tf_input(make_obs_ph(\"observation\"))\n",
    "    if using_control_sharing:\n",
    "      pilot_action_ph = tf.placeholder(tf.int32, (), name='pilot_action')\n",
    "      pilot_tol_ph = tf.placeholder(tf.float32, (), name='pilot_tol')\n",
    "    else:\n",
    "      eps = tf.get_variable(\"eps\", (), initializer=tf.constant_initializer(0))\n",
    "      stochastic_ph = tf.placeholder(tf.bool, (), name=\"stochastic\")\n",
    "      update_eps_ph = tf.placeholder(tf.float32, (), name=\"update_eps\")\n",
    "\n",
    "    q_values = q_func(observations_ph.get(), num_actions, scope=\"q_func\")\n",
    "\n",
    "    batch_size = tf.shape(q_values)[0]\n",
    "\n",
    "    if using_control_sharing:\n",
    "      q_values -= tf.reduce_min(q_values, axis=1)\n",
    "      opt_actions = tf.argmax(q_values, axis=1, output_type=tf.int32)\n",
    "      opt_q_values = tf.reduce_max(q_values, axis=1)\n",
    "\n",
    "      batch_idxes = tf.reshape(tf.range(0, batch_size, 1), [batch_size, 1])\n",
    "      reshaped_batch_size = tf.reshape(batch_size, [1])\n",
    "\n",
    "      pi_actions = tf.tile(tf.reshape(pilot_action_ph, [1]), reshaped_batch_size)\n",
    "      pi_act_idxes = tf.concat([batch_idxes, tf.reshape(pi_actions, [batch_size, 1])], axis=1)\n",
    "      pi_act_q_values = tf.gather_nd(q_values, pi_act_idxes)\n",
    "      \n",
    "      actions = tf.where(pi_act_q_values >= (1 - pilot_tol_ph) * opt_q_values, pi_actions, opt_actions)\n",
    "      \n",
    "      act = U.function(inputs=[\n",
    "        observations_ph, pilot_action_ph, pilot_tol_ph\n",
    "      ],\n",
    "                       outputs=[actions])\n",
    "    else:\n",
    "      deterministic_actions = tf.argmax(q_values, axis=1)\n",
    "\n",
    "      random_actions = tf.random_uniform(tf.stack([batch_size]), minval=0, maxval=num_actions, dtype=tf.int64)\n",
    "      chose_random = tf.random_uniform(tf.stack([batch_size]), minval=0, maxval=1, dtype=tf.float32) < eps\n",
    "      stochastic_actions = tf.where(chose_random, random_actions, deterministic_actions)\n",
    "\n",
    "      output_actions = tf.cond(stochastic_ph, lambda: stochastic_actions, lambda: deterministic_actions)\n",
    "      update_eps_expr = eps.assign(tf.cond(update_eps_ph >= 0, lambda: update_eps_ph, lambda: eps))\n",
    "      act = U.function(inputs=[observations_ph, stochastic_ph, update_eps_ph],\n",
    "                       outputs=[output_actions],\n",
    "                       givens={update_eps_ph: -1.0, stochastic_ph: True},\n",
    "                       updates=[update_eps_expr])\n",
    "    return act"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def co_build_train(make_obs_ph, q_func, num_actions, optimizer, grad_norm_clipping=None, gamma=1.0,\n",
    "    double_q=True, scope=\"deepq\", reuse=None, using_control_sharing=True):\n",
    "    act_f = co_build_act(make_obs_ph, q_func, num_actions, scope=scope, reuse=reuse, using_control_sharing=using_control_sharing)\n",
    "\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "        # set up placeholders\n",
    "        obs_t_input = U.ensure_tf_input(make_obs_ph(\"obs_t\"))\n",
    "        act_t_ph = tf.placeholder(tf.int32, [None], name=\"action\")\n",
    "        rew_t_ph = tf.placeholder(tf.float32, [None], name=\"reward\")\n",
    "        obs_tp1_input = U.ensure_tf_input(make_obs_ph(\"obs_tp1\"))\n",
    "        done_mask_ph = tf.placeholder(tf.float32, [None], name=\"done\")\n",
    "        importance_weights_ph = tf.placeholder(tf.float32, [None], name=\"weight\")\n",
    "\n",
    "        obs_t_input_get = obs_t_input.get()\n",
    "        obs_tp1_input_get = obs_tp1_input.get()\n",
    "\n",
    "        # q network evaluation\n",
    "        q_t = q_func(obs_t_input_get, num_actions, scope='q_func', reuse=True)  # reuse parameters from act\n",
    "        q_func_vars = U.scope_vars(U.absolute_scope_name('q_func'))\n",
    "\n",
    "        # target q network evalution\n",
    "        q_tp1 = q_func(obs_tp1_input_get, num_actions, scope=\"target_q_func\")\n",
    "        target_q_func_vars = U.scope_vars(U.absolute_scope_name(\"target_q_func\"))\n",
    "\n",
    "        # q scores for actions which we know were selected in the given state.\n",
    "        q_t_selected = tf.reduce_sum(q_t * tf.one_hot(act_t_ph, num_actions), 1)\n",
    "\n",
    "        # compute estimate of best possible value starting from state at t + 1\n",
    "        if double_q:\n",
    "            q_tp1_using_online_net = q_func(obs_tp1_input_get, num_actions, scope='q_func', reuse=True)\n",
    "            q_tp1_best_using_online_net = tf.arg_max(q_tp1_using_online_net, 1)\n",
    "            q_tp1_best = tf.reduce_sum(q_tp1 * tf.one_hot(q_tp1_best_using_online_net, num_actions), 1)\n",
    "        else:\n",
    "            q_tp1_best = tf.reduce_max(q_tp1, 1)\n",
    "        q_tp1_best_masked = (1.0 - done_mask_ph) * q_tp1_best\n",
    "\n",
    "        # compute RHS of bellman equation\n",
    "        q_t_selected_target = rew_t_ph + gamma * q_tp1_best_masked\n",
    "\n",
    "        # compute the error (potentially clipped)\n",
    "        td_error = q_t_selected - tf.stop_gradient(q_t_selected_target)\n",
    "        errors = U.huber_loss(td_error)\n",
    "        weighted_error = tf.reduce_mean(importance_weights_ph * errors)\n",
    "\n",
    "        # compute optimization op (potentially with gradient clipping)\n",
    "        if grad_norm_clipping is not None:\n",
    "            optimize_expr = U.minimize_and_clip(optimizer,\n",
    "                                                weighted_error,\n",
    "                                                var_list=q_func_vars,\n",
    "                                                clip_val=grad_norm_clipping)\n",
    "        else:\n",
    "            optimize_expr = optimizer.minimize(weighted_error, var_list=q_func_vars)\n",
    "\n",
    "        # update_target_fn will be called periodically to copy Q network to target Q network\n",
    "        update_target_expr = []\n",
    "        for var, var_target in zip(sorted(q_func_vars, key=lambda v: v.name),\n",
    "                                   sorted(target_q_func_vars, key=lambda v: v.name)):\n",
    "            update_target_expr.append(var_target.assign(var))\n",
    "        update_target_expr = tf.group(*update_target_expr)\n",
    "\n",
    "        # Create callable functions\n",
    "        train = U.function(\n",
    "            inputs=[\n",
    "                obs_t_input,\n",
    "                act_t_ph,\n",
    "                rew_t_ph,\n",
    "                obs_tp1_input,\n",
    "                done_mask_ph,\n",
    "                importance_weights_ph\n",
    "            ],\n",
    "            outputs=td_error,\n",
    "            updates=[optimize_expr]\n",
    "        )\n",
    "        update_target = U.function([], [], updates=[update_target_expr])\n",
    "\n",
    "        q_values = U.function([obs_t_input], q_t)\n",
    "\n",
    "    return act_f, train, update_target, {'q_values': q_values}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def co_dqn_learn(\n",
    "    env,\n",
    "    q_func,\n",
    "    lr=1e-3,\n",
    "    max_timesteps=100000,\n",
    "    buffer_size=50000,\n",
    "    train_freq=1,\n",
    "    batch_size=32,\n",
    "    print_freq=1,\n",
    "    checkpoint_freq=10000,\n",
    "    learning_starts=1000,\n",
    "    gamma=1.0,\n",
    "    target_network_update_freq=500,\n",
    "    exploration_fraction=0.1,\n",
    "    exploration_final_eps=0.02,\n",
    "    num_cpu=5,\n",
    "    callback=None,\n",
    "    scope='deepq',\n",
    "    pilot_tol=0,\n",
    "    pilot_is_human=False,\n",
    "    reuse=False):\n",
    "    \n",
    "    # Create all the functions necessary to train the model\n",
    "\n",
    "    sess = U.get_session()\n",
    "    if sess is None:\n",
    "      sess = U.make_session(num_cpu=num_cpu)\n",
    "      sess.__enter__()\n",
    "\n",
    "    def make_obs_ph(name):\n",
    "        return U.BatchInput(env.observation_space.shape, name=name)\n",
    "      \n",
    "    using_control_sharing = pilot_tol > 0\n",
    "    \n",
    "    act, train, update_target, debug = co_build_train(\n",
    "        scope=scope,\n",
    "        make_obs_ph=make_obs_ph,\n",
    "        q_func=q_func,\n",
    "        num_actions=env.action_space.n,\n",
    "        optimizer=tf.train.AdamOptimizer(learning_rate=lr),\n",
    "        gamma=gamma,\n",
    "        grad_norm_clipping=10,\n",
    "        reuse=reuse,\n",
    "        using_control_sharing=using_control_sharing\n",
    "    )\n",
    "    \n",
    "    act_params = {\n",
    "        'make_obs_ph': make_obs_ph,\n",
    "        'q_func': q_func,\n",
    "        'num_actions': env.action_space.n,\n",
    "    }\n",
    "\n",
    "    replay_buffer = ReplayBuffer(buffer_size)\n",
    "\n",
    "    # Initialize the parameters and copy them to the target network.\n",
    "    U.initialize()\n",
    "    update_target()\n",
    "\n",
    "    episode_rewards = [0.0]\n",
    "    episode_outcomes = []\n",
    "    saved_mean_reward = None\n",
    "    obs = env.reset()\n",
    "    prev_t = 0\n",
    "    rollouts = []\n",
    "    \n",
    "    if pilot_is_human:\n",
    "      global human_agent_action\n",
    "      human_agent_action = init_human_action()\n",
    "    \n",
    "    #if not using_control_sharing:\n",
    "    exploration = LinearSchedule(schedule_timesteps=int(exploration_fraction * max_timesteps),\n",
    "                               initial_p=1.0,\n",
    "                               final_p=exploration_final_eps)\n",
    "\n",
    "    with tempfile.TemporaryDirectory() as td:\n",
    "        model_saved = False\n",
    "        model_file = os.path.join(td, 'model')\n",
    "        for t in range(max_timesteps):\n",
    "            act_kwargs = {}\n",
    "            if using_control_sharing:\n",
    "              pilot_action = env.unwrapped.pilot_policy(obs[None, :n_obs_dim])\n",
    "              act_kwargs['pilot_action'] = pilot_action\n",
    "              act_kwargs['pilot_tol'] = pilot_tol if pilot_action != 8 else 0\n",
    "            else:\n",
    "              act_kwargs['update_eps'] = exploration.value(t)\n",
    "              \n",
    "            action = act(obs[None, :], **act_kwargs)[0][0]\n",
    "            if using_control_sharing and np.random.random() < exploration.value(t):\n",
    "              action = random.randint(0, 8)\n",
    "            new_obs, rew, done, info = env.step(action)\n",
    "\n",
    "            if pilot_is_human:\n",
    "              env.render()\n",
    "\n",
    "            # Store transition in the replay buffer.\n",
    "            replay_buffer.add(obs, action, rew, new_obs, float(done))\n",
    "            obs = new_obs\n",
    "\n",
    "            episode_rewards[-1] += rew\n",
    "\n",
    "            if done:\n",
    "                if t > learning_starts:\n",
    "                  for _ in range(t - prev_t):\n",
    "                    obses_t, actions, rewards, obses_tp1, dones = replay_buffer.sample(batch_size)\n",
    "                    weights, batch_idxes = np.ones_like(rewards), None\n",
    "                    td_errors = train(obses_t, actions, rewards, obses_tp1, dones, weights)\n",
    "\n",
    "                obs = env.reset()\n",
    "\n",
    "                episode_outcomes.append(rew)\n",
    "                episode_rewards.append(0.0)\n",
    "\n",
    "                if pilot_is_human:\n",
    "                  global human_agent_action\n",
    "                  human_agent_action = init_human_action()\n",
    "\n",
    "                prev_t = t\n",
    "                    \n",
    "                if pilot_is_human:\n",
    "                  time.sleep(1)\n",
    "\n",
    "            if t > learning_starts and t % target_network_update_freq == 0:\n",
    "                # Update target network periodically.\n",
    "                update_target()\n",
    "\n",
    "            mean_100ep_reward = round(np.mean(episode_rewards[-101:-1]), 1)\n",
    "            mean_100ep_succ = round(np.mean([1 if x==100 else 0 for x in episode_outcomes[-101:-1]]), 2)\n",
    "            mean_100ep_crash = round(np.mean([1 if x==-100 else 0 for x in episode_outcomes[-101:-1]]), 2)\n",
    "            num_episodes = len(episode_rewards)\n",
    "            if done and print_freq is not None and len(episode_rewards) % print_freq == 0:\n",
    "                logger.record_tabular(\"steps\", t)\n",
    "                logger.record_tabular(\"episodes\", num_episodes)\n",
    "                logger.record_tabular(\"mean 100 episode reward\", mean_100ep_reward)\n",
    "                logger.record_tabular(\"mean 100 episode succ\", mean_100ep_succ)\n",
    "                logger.record_tabular(\"mean 100 episode crash\", mean_100ep_crash)\n",
    "                logger.dump_tabular()\n",
    "\n",
    "            if checkpoint_freq is not None and t > learning_starts and num_episodes > 100 and t % checkpoint_freq == 0 and (saved_mean_reward is None or mean_100ep_reward > saved_mean_reward):\n",
    "                if print_freq is not None:\n",
    "                    print('Saving model due to mean reward increase:')\n",
    "                    print(saved_mean_reward, mean_100ep_reward)\n",
    "                U.save_state(model_file)\n",
    "                model_saved = True\n",
    "                saved_mean_reward = mean_100ep_reward\n",
    "\n",
    "        if model_saved:\n",
    "            U.load_state(model_file)\n",
    "\n",
    "    reward_data = {\n",
    "      'rewards': episode_rewards,\n",
    "      'outcomes': episode_outcomes\n",
    "    }\n",
    "          \n",
    "    return ActWrapper(act, act_params), reward_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_co_policy(\n",
    "  env, scope=None, pilot_tol=0, pilot_is_human=False, \n",
    "  n_eps=n_training_episodes, copilot_scope=None, \n",
    "  copilot_q_func=None,\n",
    "  reuse=False, **extras):\n",
    "  \n",
    "  if copilot_scope is not None:\n",
    "    scope = copilot_scope\n",
    "  elif scope is None:\n",
    "    scope = str(uuid.uuid4())\n",
    "  q_func = copilot_q_func if copilot_scope is not None else make_q_func()\n",
    "    \n",
    "  return (scope, q_func), co_dqn_learn(\n",
    "    env,\n",
    "    scope=scope,\n",
    "    q_func=q_func,\n",
    "    max_timesteps=max_ep_len*n_eps,\n",
    "    pilot_tol=pilot_tol,\n",
    "    pilot_is_human=pilot_is_human,\n",
    "    reuse=reuse,\n",
    "    **copilot_dqn_learn_kwargs\n",
    "  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def str_of_config(pilot_tol, pilot_type, embedding_type):\n",
    "  return \"{'pilot_type': '%s', 'pilot_tol': %s, 'embedding_type': '%s'}\" % (pilot_type, pilot_tol, embedding_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "load pretrained copilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "copilot_path = os.path.join(data_dir, 'pretrained_noop_copilot')\n",
    "copilot_scope = '1873ae46-2fb1-4753-a632-25a713204f2d'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "co_env = make_co_env(noop_pilot_policy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "(scope, q_func), (raw_copilot_policy, reward_data) = make_co_policy(\n",
    "  co_env, pilot_tol=1e-3, pilot_is_human=False, n_eps=1,\n",
    "  copilot_scope=copilot_scope,\n",
    "  copilot_q_func=make_q_func(),\n",
    "  reuse=False,\n",
    "  pilot_policy=noop_pilot_policy\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "load_tf_vars(copilot_scope, copilot_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "train copilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "config_kwargs = {\n",
    "  'pilot_tol': 1e-9,\n",
    "  'pilot_policy': noop_pilot_policy\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "co_env = make_co_env(**config_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "(copilot_scope, copilot_q_func), (raw_copilot_policy, reward_data) = make_co_policy(\n",
    "  co_env, **config_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "save trained copilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "copilot_path = os.path.join(data_dir, 'pretrained_noop_copilot')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "copilot_scope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "save_tf_vars(copilot_scope, copilot_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "sanity-check trained copilot with human pilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def copilot_policy(obs):\n",
    "  with tf.variable_scope(copilot_scope, reuse=None):\n",
    "    pilot_action = human_pilot_policy(obs[None, :n_obs_dim])\n",
    "    pilot_tol = 1 if pilot_action in [4, 5, 6, 7] else 0\n",
    "    return raw_copilot_policy._act(\n",
    "      obs, \n",
    "      pilot_tol=pilot_tol, \n",
    "      pilot_action=pilot_action\n",
    "    )[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "co_env = make_co_env(pilot_policy=copilot_policy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "co_env.render()\n",
    "co_env.unwrapped.viewer.window.on_key_press = key_press\n",
    "co_env.unwrapped.viewer.window.on_key_release = key_release"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for _ in range(100):\n",
    "  run_ep(copilot_policy, co_env, render=True)\n",
    "  time.sleep(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.2"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}