1.0-lunarlander-sim.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",
    "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 import deepq\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 pyglet.window import key as pygkey"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "mpl.rc('savefig', dpi=300)\n",
    "mpl.rc('text', usetex=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data_dir = os.path.join('data', 'lunarlander-sim')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "train synthetic pilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "throttle_mag = 0.75\n",
    "def disc_to_cont(action):\n",
    "    if type(action) == np.ndarray:\n",
    "        return action\n",
    "    # main engine\n",
    "    if action < 3:\n",
    "        m = -throttle_mag\n",
    "    elif action < 6:\n",
    "        m = throttle_mag\n",
    "    else:\n",
    "        raise ValueError\n",
    "    # steering\n",
    "    if action % 3 == 0:\n",
    "        s = -throttle_mag\n",
    "    elif action % 3 == 1:\n",
    "        s = 0\n",
    "    else:\n",
    "        s = throttle_mag\n",
    "    return np.array([m, s])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def mask_helipad(obs, replace=0):\n",
    "  obs = copy(obs)\n",
    "  if len(obs.shape) == 1:\n",
    "    obs[8] = replace\n",
    "  else:\n",
    "    obs[:, 8] = replace\n",
    "  return obs\n",
    "\n",
    "def traj_mask_helipad(traj):\n",
    "  return [mask_helipad(obs) for obs in traj]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_act_dim = 6\n",
    "n_obs_dim = 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def onehot_encode(i, n=n_act_dim):\n",
    "    x = np.zeros(n)\n",
    "    x[i] = 1\n",
    "    return x\n",
    "\n",
    "def onehot_decode(x):\n",
    "    l = np.nonzero(x)[0]\n",
    "    assert len(l) == 1\n",
    "    return l[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_env(using_lander_reward_shaping=False):\n",
    "  env = gym.make('LunarLanderContinuous-v2')\n",
    "  env.action_space = spaces.Discrete(n_act_dim)\n",
    "  env.unwrapped._step_orig = env.unwrapped._step\n",
    "  def _step(self, action):\n",
    "      obs, r, done, info = self._step_orig(disc_to_cont(action))\n",
    "      return obs, r, done, info\n",
    "  env.unwrapped._step = types.MethodType(_step, env.unwrapped)\n",
    "  env.unwrapped.using_lander_reward_shaping = using_lander_reward_shaping\n",
    "  return env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make_env(using_lander_reward_shaping=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "max_ep_len = 1000\n",
    "n_training_episodes = 500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "make_q_func = lambda: deepq.models.mlp([64, 64])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pilot_dqn_learn_kwargs = {\n",
    "  'lr': 1e-3,\n",
    "  'exploration_fraction': 0.1,\n",
    "  'exploration_final_eps': 0.02,\n",
    "  'target_network_update_freq': 1500,\n",
    "  'print_freq': 100,\n",
    "  'num_cpu': 5,\n",
    "  'gamma': 0.99\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "full_pilot_scope = 'full_pilot'\n",
    "full_pilot_q_func = make_q_func()\n",
    "load_pretrained_full_pilot = True"
   ]
  },
  {
   "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": {},
   "outputs": [],
   "source": [
    "raw_full_pilot_policy, full_pilot_reward_data = deepq.learn(\n",
    "  env,\n",
    "  q_func=full_pilot_q_func,\n",
    "  max_timesteps=max_timesteps,\n",
    "  scope=full_pilot_scope,\n",
    "  **pilot_dqn_learn_kwargs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'full_pilot_reward_data.pkl'), 'wb') as f:\n",
    "  pickle.dump(full_pilot_reward_data, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'full_pilot_reward_data.pkl'), 'rb') as f:\n",
    "  full_pilot_reward_data = pickle.load(f)"
   ]
  },
  {
   "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 = None\n",
    "    actions = []\n",
    "    for step_idx in range(max_ep_len+1):\n",
    "        if done:\n",
    "            trajectory = info['trajectory']\n",
    "            break\n",
    "        action = policy(obs[None, :])\n",
    "        obs, r, done, info = env.step(action)\n",
    "        actions.append(action)\n",
    "        if render:\n",
    "          env.render()\n",
    "        totalr += r\n",
    "    outcome = r if r % 100 == 0 else 0\n",
    "    return totalr, outcome, trajectory, actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def full_pilot_policy(obs):\n",
    "  with tf.variable_scope(full_pilot_scope, reuse=None):\n",
    "    return raw_full_pilot_policy._act(obs)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class LaggyPilotPolicy(object):\n",
    "  def __init__(self):\n",
    "    self.last_laggy_pilot_act = None\n",
    "    \n",
    "  def __call__(self, obs, lag_prob=0.8):\n",
    "    if self.last_laggy_pilot_act is None or np.random.random() >= lag_prob:\n",
    "      action = full_pilot_policy(obs)\n",
    "      self.last_laggy_pilot_act = action\n",
    "    return self.last_laggy_pilot_act"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "laggy_pilot_policy = LaggyPilotPolicy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def noisy_pilot_policy(obs, noise_prob=0.15):\n",
    "  action = full_pilot_policy(obs)\n",
    "  if np.random.random() < noise_prob:\n",
    "    action = (action + 3) % 6\n",
    "  if np.random.random() < noise_prob:\n",
    "    action = action//3*3 + (action + np.random.randint(1, 3)) % 3\n",
    "  return action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def noop_pilot_policy(obs):\n",
    "  return 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def sensor_pilot_policy(obs, thresh=0.1):\n",
    "  d = obs[0, 8] - obs[0, 0] # horizontal dist to helipad\n",
    "  if d < -thresh:\n",
    "    return 0\n",
    "  elif d > thresh:\n",
    "    return 2\n",
    "  else:\n",
    "    return 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# begin debug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "run_ep(full_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": [
    "# end debug"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "full_pilot_path = os.path.join(data_dir, 'full_pilot.tf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "save_tf_vars(full_pilot_scope, full_pilot_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "load_tf_vars(full_pilot_scope, full_pilot_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "evaluate synthetic pilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pilot_names = ['full', 'laggy', 'noisy', 'noop', 'sensor']"
   ]
  },
  {
   "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": [
    "pilot_evals = [list(zip(*[run_ep(eval('%s_pilot_policy' % pilot_name), env, render=False) for _ in range(n_eval_eps)])) for pilot_name in pilot_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'pilot_evals.pkl'), 'wb') as f:\n",
    "  pickle.dump(dict(zip(pilot_names, pilot_evals)), f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mean_rewards = [np.mean(pilot_eval[0]) for pilot_eval in pilot_evals]\n",
    "outcome_distrns = [Counter(pilot_eval[1]) for pilot_eval in pilot_evals]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print('\\n'.join([str(x) for x in zip(pilot_names, mean_rewards, outcome_distrns)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_videos = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for pilot_name in pilot_names:\n",
    "  for i in range(n_videos):\n",
    "    wrapped_env = wrappers.Monitor(env, os.path.join(data_dir, 'videos', '%s_pilot.%d' % (pilot_name, i)), force=True)\n",
    "    run_ep(eval('%s_pilot_policy' % pilot_name), wrapped_env, render=False)\n",
    "    wrapped_env.close()\n",
    "    env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "train supervised goal decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pilot_name = 'full'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pilot_policy = eval('%s_pilot_policy' % pilot_name)\n",
    "n_rollouts = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "rollouts = [run_ep(pilot_policy, env, render=False)[2:] for _ in range(n_rollouts)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, '%s_pilot_policy_rollouts.pkl' % pilot_name), 'wb') as f:\n",
    "  pickle.dump(rollouts, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, '%s_pilot_policy_rollouts.pkl' % pilot_name), 'rb') as f:\n",
    "  rollouts = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_val_rollouts = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rollouts, val_rollouts = rollouts[:-n_val_rollouts], rollouts[-n_val_rollouts:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def combined_rollout(states, actions):\n",
    "  return np.array([np.concatenate((\n",
    "    np.array(obs),\n",
    "    onehot_encode(action))) for obs, action in zip(\n",
    "      states[:-1] if len(states) == len(actions) + 1 else states, actions)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def format_rollouts(rollouts):\n",
    "  X_dat = np.zeros((len(rollouts), max_ep_len, n_obs_dim + n_act_dim))\n",
    "  Y_dat = np.zeros((len(rollouts), max_ep_len))\n",
    "  M_dat = np.zeros((len(rollouts), max_ep_len))\n",
    "  for i, (states, actions) in enumerate(rollouts):\n",
    "    Y_dat[i, :] = states[0][-1]\n",
    "    X_dat[i, :len(actions), :] = traj_mask_helipad(combined_rollout(states, actions))\n",
    "    M_dat[i, :len(actions)] = 1\n",
    "  return X_dat, Y_dat, M_dat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "X_dat, Y_dat, M_dat = format_rollouts(rollouts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "val_X_dat, val_Y_dat, val_M_dat = format_rollouts(val_rollouts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "example_idxes = list(range(X_dat.shape[0]))\n",
    "def next_batch(batch_size):\n",
    "  batch_idxes = random.sample(example_idxes, batch_size)\n",
    "  return X_dat[batch_idxes, :, :], Y_dat[batch_idxes, :], M_dat[batch_idxes, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Training Parameters\n",
    "learning_rate = 1e-2\n",
    "training_steps = 1000\n",
    "batch_size = 128\n",
    "display_step = training_steps // 10\n",
    "\n",
    "# Network Parameters\n",
    "num_input = X_dat.shape[2]\n",
    "timesteps = X_dat.shape[1] # timesteps\n",
    "num_hidden = 32 # hidden layer num of features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "gd_scope = 'gd_scope'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with tf.variable_scope(gd_scope, reuse=False):\n",
    "  # tf Graph input\n",
    "  X = tf.placeholder(\"float\", [None, timesteps, num_input])\n",
    "  Y = tf.placeholder(\"float\", [None, timesteps])\n",
    "  M = tf.placeholder(\"float\", [None, timesteps]) # mask for variable length sequences\n",
    "  INIT_STATE_A = tf.placeholder(\"float\", [None, num_hidden])\n",
    "  INIT_STATE_B = tf.placeholder(\"float\", [None, num_hidden])\n",
    "\n",
    "  weights = {\n",
    "      'out': tf.Variable(tf.random_normal([num_hidden, 1]))\n",
    "  }\n",
    "  biases = {\n",
    "      'out': tf.Variable(tf.random_normal([1]))\n",
    "  }\n",
    "  \n",
    "  unstacked_X = tf.unstack(X, timesteps, 1)\n",
    "\n",
    "  lstm_cell = rnn.BasicLSTMCell(num_hidden, forget_bias=1.0)\n",
    "\n",
    "  state = (INIT_STATE_A, INIT_STATE_B)\n",
    "  rnn_outputs = []\n",
    "  rnn_states = []\n",
    "  for input_ in unstacked_X:\n",
    "    output, state = lstm_cell(input_, state)\n",
    "    rnn_outputs.append(output)\n",
    "    rnn_states.append(state)\n",
    "\n",
    "  prediction = tf.reshape(\n",
    "    tf.concat([tf.matmul(output, weights['out']) + biases['out'] for output in rnn_outputs], axis=1), \n",
    "    shape=[tf.shape(X)[0], timesteps])\n",
    "  \n",
    "  predictions = [tf.matmul(output, weights['out']) + biases['out'] for output in rnn_outputs]\n",
    "\n",
    "  loss_op = tf.reduce_sum((prediction - Y)**2 * M) / tf.reduce_sum(M)\n",
    "\n",
    "  optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n",
    "  train_op = optimizer.minimize(loss_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sess = U.get_session()\n",
    "if sess is None:\n",
    "  sess = U.make_session(num_cpu=5)\n",
    "  sess.__enter__()\n",
    "  \n",
    "sess.run(tf.variables_initializer([v for v in tf.global_variables() if v.name.startswith(gd_scope + '/')]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with tf.variable_scope(gd_scope, reuse=False):\n",
    "  for step in range(1, training_steps+1):\n",
    "      batch_x, batch_y, batch_mask = next_batch(batch_size)\n",
    "      sess.run(train_op, feed_dict={X: batch_x, Y: batch_y, M: batch_mask, \n",
    "                                   INIT_STATE_A: np.zeros((batch_size, num_hidden)),\n",
    "                                   INIT_STATE_B: np.zeros((batch_size, num_hidden))})\n",
    "      if step % display_step == 0 or step == 1:\n",
    "          loss = sess.run(loss_op, feed_dict={X: X_dat,\n",
    "                                             Y: Y_dat,\n",
    "                                             M: M_dat,\n",
    "                                             INIT_STATE_A: np.zeros((X_dat.shape[0], num_hidden)),\n",
    "                                             INIT_STATE_B: np.zeros((X_dat.shape[0], num_hidden))})\n",
    "          val_loss = sess.run(loss_op, feed_dict={X: val_X_dat,\n",
    "                                                 Y: val_Y_dat,\n",
    "                                                 M: val_M_dat,\n",
    "                                                 INIT_STATE_A: np.zeros((val_X_dat.shape[0], num_hidden)),\n",
    "                                                 INIT_STATE_B: np.zeros((val_X_dat.shape[0], num_hidden))})\n",
    "          print(\"Step \" + str(step) + \", Training Loss= \" + \\\n",
    "                \"{:.4f}\".format(loss), \", Validation Loss= \" + \"{:.4f}\".format(val_loss))\n",
    "\n",
    "  print(\"Optimization Finished!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train_supervised_goal_decoder(gd_scope, rollouts):\n",
    "  X_dat, Y_dat, M_dat = format_rollouts(rollouts)\n",
    "  \n",
    "  example_idxes = list(range(X_dat.shape[0]))\n",
    "  def next_batch(batch_size):\n",
    "    batch_idxes = random.sample(example_idxes, batch_size)\n",
    "    return X_dat[batch_idxes, :, :], Y_dat[batch_idxes, :], M_dat[batch_idxes, :]\n",
    "  \n",
    "  # Training Parameters\n",
    "  learning_rate = 1e-2\n",
    "  training_steps = 1000\n",
    "  batch_size = 128\n",
    "  display_step = training_steps // 10\n",
    "\n",
    "  # Network Parameters\n",
    "  num_input = X_dat.shape[2]\n",
    "  timesteps = X_dat.shape[1] # timesteps\n",
    "  num_hidden = 32 # hidden layer num of features\n",
    "    \n",
    "  sess = U.get_session()\n",
    "  if sess is None:\n",
    "    sess = U.make_session(num_cpu=5)\n",
    "    sess.__enter__()\n",
    "\n",
    "  sess.run(tf.variables_initializer([v for v in tf.global_variables() if v.name.startswith(gd_scope + '/')]))\n",
    "  \n",
    "  with tf.variable_scope(gd_scope, reuse=False):\n",
    "    for step in range(1, training_steps+1):\n",
    "        batch_x, batch_y, batch_mask = next_batch(batch_size)\n",
    "        sess.run(train_op, feed_dict={X: batch_x, Y: batch_y, M: batch_mask, \n",
    "                                     INIT_STATE_A: np.zeros((batch_size, num_hidden)),\n",
    "                                     INIT_STATE_B: np.zeros((batch_size, num_hidden))})\n",
    "        if step % display_step == 0 or step == 1:\n",
    "            loss = sess.run(loss_op, feed_dict={X: X_dat,\n",
    "                                               Y: Y_dat,\n",
    "                                               M: M_dat,\n",
    "                                               INIT_STATE_A: np.zeros((X_dat.shape[0], num_hidden)),\n",
    "                                               INIT_STATE_B: np.zeros((X_dat.shape[0], num_hidden))})\n",
    "            print(\"Step \" + str(step) + \", Training Loss={:.4f}\".format(loss))\n",
    "\n",
    "    print(\"Optimization Finished!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def build_retrain_goal_decoder(pilot_name):\n",
    "  with open(os.path.join(data_dir, '%s_pilot_policy_rollouts.pkl' % pilot_name), 'rb') as f:\n",
    "    off_pol_rollouts = pickle.load(f)\n",
    "  def retrain_goal_decoder(on_pol_rollouts):\n",
    "    train_supervised_goal_decoder(gd_scope, off_pol_rollouts + on_pol_rollouts)\n",
    "  return retrain_goal_decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "gd_path = os.path.join(data_dir, '%s_pilot_goal_decoder.tf' % pilot_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "save_tf_vars(gd_scope, gd_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "load_tf_vars(gd_scope, gd_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def decode_goal(trajectory, init_state=None, only_final=False):\n",
    "  traj_X = np.zeros((1, max_ep_len, n_obs_dim + n_act_dim))\n",
    "  traj_X[0, :len(trajectory), :] = np.array(trajectory)\n",
    "  with tf.variable_scope(gd_scope, reuse=False):\n",
    "    feed_dict = {X: traj_X}\n",
    "    if init_state is not None:\n",
    "      feed_dict[INIT_STATE_A] = init_state[0]\n",
    "      feed_dict[INIT_STATE_B] = init_state[1]\n",
    "    else:\n",
    "      feed_dict[INIT_STATE_A] = np.zeros((1, num_hidden))\n",
    "      feed_dict[INIT_STATE_B] = np.zeros((1, num_hidden))\n",
    "    if only_final:\n",
    "      g, s = sess.run(\n",
    "        [predictions[len(trajectory)-1], rnn_states[len(trajectory)-1]], \n",
    "        feed_dict=feed_dict\n",
    "      )\n",
    "      return g[0, 0], s\n",
    "    else:\n",
    "      g, s = sess.run(\n",
    "        [predictions, rnn_states[len(trajectory)-1]], \n",
    "        feed_dict=feed_dict\n",
    "      )\n",
    "      return [x[0, 0] for x in g], s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def build_build_goal_decoder(pilot_name):\n",
    "  def build_goal_decoder():\n",
    "    load_tf_vars(gd_scope, os.path.join(data_dir, '%s_pilot_goal_decoder.tf' % pilot_name))\n",
    "    return decode_goal\n",
    "  return build_goal_decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "build model-based goal decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "goals = np.arange(-0.8, 1, 0.05)\n",
    "n_goals = len(goals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sess = U.get_session()\n",
    "if sess is None:\n",
    "  sess = U.make_session(num_cpu=5)\n",
    "  sess.__enter__()\n",
    "    \n",
    "with tf.variable_scope(full_pilot_scope, reuse=None):\n",
    "  Q_obs = tf.get_variable(\"Q_obs\", (n_goals, n_obs_dim))\n",
    "\n",
    "sess.run(tf.variables_initializer([Q_obs]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with tf.variable_scope(full_pilot_scope, reuse=True):\n",
    "  Q_values = full_pilot_q_func(Q_obs, n_act_dim, scope=\"q_func\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def compute_map_est_goal(s, a, log_prior, temp=1000):\n",
    "  states = []\n",
    "  for g in goals:\n",
    "    state = copy(s)\n",
    "    state[8] = g\n",
    "    states.append(state)\n",
    "  with tf.variable_scope(full_pilot_scope, reuse=True):\n",
    "    Q = sess.run(\n",
    "        Q_values,\n",
    "        feed_dict={Q_obs: np.array(states)}\n",
    "    )\n",
    "    \n",
    "  Q *= temp\n",
    "  \n",
    "  action = onehot_decode(a)\n",
    "  log_cond_likelihood = Q[:, action] - logsumexp(Q, axis=1)\n",
    "  log_marginal_likelihood = logsumexp(log_cond_likelihood) - np.log(n_goals)\n",
    "  log_likelihood = log_cond_likelihood - log_marginal_likelihood\n",
    "  log_posterior = log_likelihood + log_prior\n",
    "  map_est_goal = goals[max(range(n_goals), key=lambda i: log_posterior[i])]\n",
    "  return log_posterior, map_est_goal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "zero_goal_idx = len(goals)//2-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def mb_decode_goal(trajectory, init_state=None, only_final=False):\n",
    "  if init_state is None:\n",
    "    prior = np.ones(n_goals) / n_goals\n",
    "    prior[zero_goal_idx] *= 2\n",
    "    prior = prior / prior.sum()\n",
    "    log_prior = np.log(prior)\n",
    "    map_est_goals = []\n",
    "  else:\n",
    "    log_prior, map_est_goals = init_state\n",
    "    trajectory = trajectory[-1:]\n",
    "  for t in trajectory:\n",
    "    s = np.array(t[:-n_act_dim])\n",
    "    a = np.array(t[-n_act_dim:])\n",
    "    log_posterior, map_est_goal = compute_map_est_goal(s, a, log_prior)\n",
    "    map_est_goals.append(map_est_goal)\n",
    "    log_prior = log_posterior\n",
    "  return (map_est_goals[-1] if only_final else map_est_goals), (log_posterior, map_est_goals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "decode_goal = mb_decode_goal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def build_build_goal_decoder(pilot_name):\n",
    "  def build_goal_decoder():\n",
    "    return decode_goal\n",
    "  return build_goal_decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# begin debug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rollout = rollouts[925]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "goal = rollout[0][0][-1]\n",
    "traj = traj_mask_helipad(combined_rollout(*rollout))\n",
    "pred_goal, _ = decode_goal(traj)\n",
    "mb_pred_goal, _ = mb_decode_goal(traj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.xlabel('Step')\n",
    "plt.ylabel('X-Coordinate')\n",
    "plt.axhline(y=goal, label='True Goal', linestyle='--', linewidth=5, color='gray', alpha=0.5)\n",
    "plt.plot(pred_goal[:len(rollout[0])], label='Predicted Goal (SL)', color='orange')\n",
    "plt.plot(mb_pred_goal[:len(rollout[0])], label='Predicted Goal (BI)', color='teal')\n",
    "plt.legend(loc='best')\n",
    "plt.ylim([-1, 1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#rollout = rollouts[986]\n",
    "for rollout in rollouts[800:850]:\n",
    "  goal = rollout[0][0][-1]\n",
    "  traj = traj_mask_helipad(combined_rollout(*rollout))\n",
    "  pred_goal, _ = decode_goal(traj)\n",
    "  mb_pred_goal, _ = mb_decode_goal(traj)\n",
    "  plt.xlabel('Step')\n",
    "  plt.ylabel('X-Coordinate')\n",
    "  plt.axhline(y=goal, label='True Goal', linestyle='--', linewidth=5, color='gray', alpha=0.5)\n",
    "  plt.plot(pred_goal[:len(rollout[0])], label='Predicted Goal (SL)', color='orange')\n",
    "  plt.plot(mb_pred_goal[:len(rollout[0])], label='Predicted Goal (BI)', color='teal')\n",
    "  plt.legend(loc='best')\n",
    "  plt.ylim([-1, 1])\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "y_trues = []\n",
    "y_preds = []\n",
    "for rollout in rollouts:\n",
    "  goal = rollout[0][0][-1]\n",
    "  traj = traj_mask_helipad(combined_rollout(*rollout))\n",
    "  pred_goal, final_states = mb_decode_goal(traj)\n",
    "  y_trues.extend([goal] * len(pred_goal))\n",
    "  y_preds.extend(pred_goal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "y_trues = np.array(y_trues)\n",
    "y_preds = np.array(y_preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean((y_trues - y_preds)**2), np.mean((y_trues - 0)**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mb_pred_goal, final_states = mb_decode_goal(traj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.ylabel('Timestep')\n",
    "plt.xlabel('Horizontal Location')\n",
    "plt.title('Sample Episode from Optimal Synthetic Pilot')\n",
    "plt.axvline(x=goal, label='True Goal', linestyle='--', linewidth=1, color='green')\n",
    "plt.plot(list(reversed(pred_goal[:len(rollout[0])])), range(len(pred_goal[:len(rollout[0])])), label='Inferred Goal (Supervised Learning)', color='teal')\n",
    "plt.plot(list(reversed(mb_pred_goal[:len(rollout[0])])), range(len(mb_pred_goal[:len(rollout[0])])), label='Inferred Goal (Bayesian Inference)', color='gray')\n",
    "plt.yticks([0, 100, 200, 300, 400], ['400', '300', '200', '100', '0'])\n",
    "plt.legend(loc='best')\n",
    "plt.xlim([-1, 1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# end debug"
   ]
  },
  {
   "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.models.mlp([64, 64])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "copilot_dqn_learn_kwargs = {\n",
    "  'lr': 1e-3,\n",
    "  'exploration_fraction': 0.1,\n",
    "  'exploration_final_eps': 0.02,\n",
    "  'target_network_update_freq': 1500,\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, build_goal_decoder=None, using_lander_reward_shaping=False, **extras):\n",
    "    env = gym.make('LunarLanderContinuous-v2')\n",
    "    env.unwrapped.using_lander_reward_shaping = using_lander_reward_shaping\n",
    "    env.action_space = spaces.Discrete(n_act_dim)\n",
    "    env.unwrapped.pilot_policy = pilot_policy\n",
    "    if build_goal_decoder is None:\n",
    "      obs_box = env.observation_space\n",
    "      env.observation_space = spaces.Box(np.concatenate((obs_box.low, np.zeros(n_act_dim))), \n",
    "                                         np.concatenate((obs_box.high, np.ones(n_act_dim))))\n",
    "    \n",
    "    env.unwrapped._step_orig = env.unwrapped._step\n",
    "    if build_goal_decoder is None:\n",
    "      def _step(self, action):\n",
    "        obs, r, done, info = self._step_orig(disc_to_cont(action))\n",
    "        obs = np.concatenate((obs, onehot_encode(self.pilot_policy(obs[None, :]))))\n",
    "        return obs, r, done, info\n",
    "    else:\n",
    "      goal_decoder = build_goal_decoder()\n",
    "      def _step(self, action):\n",
    "        obs, r, done, info = self._step_orig(disc_to_cont(action))\n",
    "        self.actions.append(self.pilot_policy(obs[None, :]))\n",
    "        traj = traj_mask_helipad(combined_rollout(self.trajectory[-1:], self.actions[-1:]))\n",
    "        goal, self.init_state = goal_decoder(traj, init_state=self.init_state, only_final=True)\n",
    "        obs = mask_helipad(obs, replace=goal)\n",
    "        return obs, r, done, info\n",
    "    env.unwrapped._step = types.MethodType(_step, env.unwrapped)\n",
    "    \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",
    "      # if necessary, switch steering and keep main\n",
    "      mixed_actions = 3 * (pi_actions // 3) + (opt_actions % 3)\n",
    "      mixed_act_idxes = tf.concat([batch_idxes, tf.reshape(mixed_actions, [batch_size, 1])], axis=1)\n",
    "      mixed_act_q_values = tf.gather_nd(q_values, mixed_act_idxes)\n",
    "      mixed_actions = tf.where(pi_act_q_values >= (1 - pilot_tol_ph) * opt_q_values, pi_actions, mixed_actions)\n",
    "\n",
    "      # if necessary, keep steering and switch main\n",
    "      mixed_act_idxes = tf.concat([batch_idxes, tf.reshape(mixed_actions, [batch_size, 1])], axis=1)\n",
    "      mixed_act_q_values = tf.gather_nd(q_values, mixed_act_idxes)\n",
    "      steer_mixed_actions = 3 * (opt_actions // 3) + (pi_actions % 3)\n",
    "      mixed_actions = tf.where(mixed_act_q_values >= (1 - pilot_tol_ph) * opt_q_values, mixed_actions, steer_mixed_actions)\n",
    "\n",
    "      # if necessary, switch steering and main\n",
    "      mixed_act_idxes = tf.concat([batch_idxes, tf.reshape(mixed_actions, [batch_size, 1])], axis=1)\n",
    "      mixed_act_q_values = tf.gather_nd(q_values, mixed_act_idxes)\n",
    "      actions = tf.where(mixed_act_q_values >= (1 - pilot_tol_ph) * opt_q_values, mixed_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": {},
   "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",
    "    using_supervised_goal_decoder=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",
    "            masked_obs = obs if using_supervised_goal_decoder else mask_helipad(obs)\n",
    "\n",
    "            act_kwargs = {}\n",
    "            if using_control_sharing:\n",
    "              act_kwargs['pilot_action'] = env.unwrapped.pilot_policy(obs[None, :n_obs_dim])\n",
    "              act_kwargs['pilot_tol'] = pilot_tol if not pilot_is_human or (pilot_is_human and human_agent_active) else 0\n",
    "            else:\n",
    "              act_kwargs['update_eps'] = exploration.value(t)\n",
    "              \n",
    "            action = act(masked_obs[None, :], **act_kwargs)[0][0]\n",
    "            new_obs, rew, done, info = env.step(action)\n",
    "\n",
    "            if pilot_is_human:\n",
    "              env.render()\n",
    "              time.sleep(sim_delay_for_human)\n",
    "\n",
    "            # Store transition in the replay buffer.\n",
    "            masked_new_obs = new_obs if using_supervised_goal_decoder else mask_helipad(new_obs)\n",
    "            replay_buffer.add(masked_obs, action, rew, masked_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, build_goal_decoder=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",
    "    using_supervised_goal_decoder=(build_goal_decoder is not None),\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, using_lander_reward_shaping):\n",
    "  return \"{'pilot_type': '%s', 'pilot_tol': %s, 'embedding_type': '%s', 'using_lander_reward_shaping': %s}\" % (pilot_type, pilot_tol, embedding_type, str(using_lander_reward_shaping))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "train and evaluate copilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_reps = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pilot_ids = ['sensor']\n",
    "pilot_policies = [eval('%s_pilot_policy' % pilot_name) for pilot_name in pilot_ids]\n",
    "embedding_type = 'rawaction'\n",
    "using_lander_reward_shaping = True\n",
    "pilot_tols = [0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "configs = []\n",
    "for pilot_id, pilot_policy in zip(pilot_ids, pilot_policies):\n",
    "  if embedding_type != 'rawaction':\n",
    "    build_goal_decoder = build_build_goal_decoder(pilot_id)\n",
    "  else:\n",
    "    build_goal_decoder = None \n",
    "  for pilot_tol in pilot_tols:\n",
    "    configs.append((\n",
    "      str_of_config(pilot_tol, pilot_id, embedding_type, using_lander_reward_shaping), \n",
    "      {\n",
    "        'pilot_tol': pilot_tol,\n",
    "        'build_goal_decoder': build_goal_decoder,\n",
    "        'pilot_policy': pilot_policy,\n",
    "        'using_lander_reward_shaping': using_lander_reward_shaping,\n",
    "        'reuse': False\n",
    "      }))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "reward_logs = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for config_name, config_kwargs in configs:\n",
    "  print(config_name)\n",
    "  reward_logs[config_name] = defaultdict(list)\n",
    "  co_env = make_co_env(**config_kwargs)\n",
    "  for i in range(n_reps):\n",
    "    (copilot_scope, copilot_q_func), (raw_copilot_policy, reward_data) = make_co_policy(\n",
    "      co_env, **config_kwargs)\n",
    "    for k, v in reward_data.items():\n",
    "      reward_logs[config_name][k].append(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "reward_log_file = 'reward_logs.pkl'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, reward_log_file), 'wb') as f:\n",
    "  pickle.dump(reward_logs, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Train and test on different pilots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pilot_tol_of_id = {\n",
    "  'noop': 0,\n",
    "  'laggy': 0.7,\n",
    "  'noisy': 0.4,\n",
    "  'sensor': 0\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "training_pilot_ids = list(pilot_tol_of_id.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "copilot_of_training_pilot = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "copilot_path_of_training_pilot = lambda training_pilot_id: os.path.join(data_dir, 'pretrained_%s_copilot')\n",
    "copilot_scope_of_training_pilot = lambda training_pilot_id: ('pretrained_%s_copilot_scope' % training_pilot_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for training_pilot_id, pilot_tol in pilot_tol_of_id.items():\n",
    "  pilot_policy = eval('%s_pilot_policy' % training_pilot_id)\n",
    "  copilot_scope = copilot_scope_of_training_pilot(training_pilot_id)\n",
    "  config_kwargs = {\n",
    "    'pilot_policy': pilot_policy,\n",
    "    'pilot_tol': pilot_tol,\n",
    "    'copilot_scope': copilot_scope,\n",
    "    'copilot_q_func': make_q_func()\n",
    "  }\n",
    "  co_env = make_co_env(**config_kwargs)\n",
    "  (copilot_scope, copilot_q_func), (raw_copilot_policy, reward_data) = make_co_policy(co_env, **config_kwargs)\n",
    "  \n",
    "  copilot_of_training_pilot[training_pilot_id] = (copilot_scope, raw_copilot_policy)\n",
    "  copilot_path = copilot_path_of_training_pilot(training_pilot_id)\n",
    "  save_tf_vars(copilot_scope, copilot_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_copilot_policy(training_pilot_id, eval_pilot_policy, pilot_tol):\n",
    "  copilot_scope, raw_copilot_policy = copilot_of_training_pilot[training_pilot_id]\n",
    "  def copilot_policy(obs):\n",
    "    with tf.variable_scope(copilot_scope, reuse=None):\n",
    "      masked_obs = mask_helipad(obs)[0]\n",
    "      pilot_action = eval_pilot_policy(masked_obs[None, :n_obs_dim])\n",
    "      \n",
    "      if masked_obs.size == n_obs_dim:\n",
    "        feed_obs = np.concatenate((masked_obs, onehot_encode(pilot_action)))\n",
    "      else:\n",
    "        feed_obs = masked_obs\n",
    "\n",
    "      return raw_copilot_policy._act(\n",
    "        feed_obs[None, :], \n",
    "        pilot_tol=pilot_tol, \n",
    "        pilot_action=pilot_action\n",
    "      )[0][0]\n",
    "  return copilot_policy"
   ]
  },
  {
   "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": [
    "cross_evals = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for training_pilot_id, training_pilot_tol in pilot_tol_of_id.items():\n",
    "  # load pretrained copilot\n",
    "  copilot_scope = copilot_scope_of_training_pilot(training_pilot_id)\n",
    "  training_pilot_policy = eval('%s_pilot_policy' % training_pilot_id)\n",
    "  config_kwargs = {\n",
    "    'pilot_policy': training_pilot_policy,\n",
    "    'pilot_tol': training_pilot_tol,\n",
    "    'copilot_scope': copilot_scope,\n",
    "    'copilot_q_func': make_q_func(),\n",
    "    'reuse': True\n",
    "  }\n",
    "  co_env = make_co_env(**config_kwargs)\n",
    "  make_co_policy(co_env, **config_kwargs)\n",
    "  copilot_path = copilot_path_of_training_pilot(training_pilot_id)\n",
    "  load_tf_vars(copilot_scope, copilot_path)\n",
    "  # evaluate copilot with different pilots\n",
    "  for eval_pilot_id, eval_pilot_tol in pilot_tol_of_id.items():\n",
    "    eval_pilot_policy = eval('%s_pilot_policy' % eval_pilot_id)\n",
    "    copilot_policy = make_copilot_policy(training_pilot_id, eval_pilot_policy, eval_pilot_tol)\n",
    "    co_env = make_co_env(pilot_policy=eval_pilot_policy)\n",
    "    cross_evals[(training_pilot_id, eval_pilot_id)] = [run_ep(copilot_policy, co_env, render=False)[:2] for _ in range(n_eval_eps)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'cross_evals.pkl'), 'wb') as f:\n",
    "  pickle.dump(cross_evals, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "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
}