1.2-lunarlander-human.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",
    "from scipy.stats import binom_test, ttest_1samp\n",
    "\n",
    "from pyglet.window import key as pygkey"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data_dir = os.path.join('data', 'lunarlander-human')"
   ]
  },
  {
   "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 # 2 x 3\n",
    "n_act_true_dim = 2\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"
   ]
  },
  {
   "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": [
    "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",
    "      global human_agent_active\n",
    "      human_agent_action = init_human_action()\n",
    "      human_agent_active = False\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 noop_pilot_policy(obs):\n",
    "  return 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": "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_trajectories = []\n",
    "    episode_actions = []\n",
    "    episode_rewards = []\n",
    "    episode_outcomes = []\n",
    "    saved_mean_reward = None\n",
    "    obs = env.reset()\n",
    "    prev_t = 0\n",
    "    episode_reward = 0\n",
    "    episode_trajectory = []\n",
    "    episode_action = []\n",
    "    \n",
    "    if pilot_is_human:\n",
    "      global human_agent_action\n",
    "      global human_agent_active\n",
    "      human_agent_action = init_human_action()\n",
    "      human_agent_active = False\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",
    "            episode_trajectory.append(obs)\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",
    "            episode_action.append(action)\n",
    "\n",
    "            if pilot_is_human:\n",
    "              env.render()\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_reward += 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(episode_reward)\n",
    "                episode_trajectories.append(episode_trajectory + [new_obs])\n",
    "                episode_actions.append(episode_action)\n",
    "                episode_trajectory = []\n",
    "                episode_action = []\n",
    "                episode_reward = 0\n",
    "\n",
    "                if pilot_is_human:\n",
    "                  global human_agent_action\n",
    "                  global human_agent_active\n",
    "                  human_agent_action = init_human_action()\n",
    "                  human_agent_active = False\n",
    "                  \n",
    "                prev_t = t\n",
    "                    \n",
    "                if pilot_is_human:\n",
    "                  time.sleep(2)\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[-10:]), 1)\n",
    "            mean_100ep_succ = round(np.mean([1 if x==100 else 0 for x in episode_outcomes[-100:]]), 2)\n",
    "            mean_100ep_crash = round(np.mean([1 if x==-100 else 0 for x in episode_outcomes[-100:]]), 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",
    "      'trajectories': episode_trajectories,\n",
    "      'actions': episode_actions\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=None, 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": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "init_human_action = lambda: [0, 1] # noop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "human_agent_action = init_human_action()\n",
    "human_agent_active = False\n",
    "\n",
    "LEFT = pygkey.LEFT\n",
    "RIGHT = pygkey.RIGHT\n",
    "UP = pygkey.UP\n",
    "DOWN = pygkey.DOWN\n",
    "\n",
    "def key_press(key, mod):\n",
    "    global human_agent_action\n",
    "    global human_agent_active\n",
    "    a = int(key)\n",
    "    if a == LEFT:\n",
    "        human_agent_action[1] = 0\n",
    "        human_agent_active = True\n",
    "    elif a == RIGHT:\n",
    "        human_agent_action[1] = 2\n",
    "        human_agent_active = True\n",
    "    elif a == UP:\n",
    "        human_agent_action[0] = 1\n",
    "        human_agent_active = True\n",
    "    elif a == DOWN:\n",
    "        human_agent_action[0] = 0\n",
    "        human_agent_active = True\n",
    "\n",
    "def key_release(key, mod):\n",
    "    global human_agent_action\n",
    "    global human_agent_active\n",
    "    a = int(key)\n",
    "    if a == LEFT or a == RIGHT:\n",
    "        human_agent_action[1] = 1\n",
    "        human_agent_active = False\n",
    "    elif a == UP or a == DOWN:\n",
    "        human_agent_action[0] = 0\n",
    "        human_agent_active = False\n",
    "    \n",
    "def encode_human_action(action):\n",
    "    return action[0]*3+action[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def human_pilot_policy(obs):\n",
    "    global human_agent_action\n",
    "    return encode_human_action(human_agent_action)"
   ]
  },
  {
   "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 = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co_env = make_co_env(\n",
    "  noop_pilot_policy, \n",
    "  build_goal_decoder=None, \n",
    "  using_lander_reward_shaping=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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",
    "  using_lander_reward_shaping=False,\n",
    "  pilot_policy=noop_pilot_policy,\n",
    "  build_goal_decoder=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_tf_vars(copilot_scope, copilot_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "balance and randomize experiment order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for i in range(20):\n",
    "  x, y = '12', '21'\n",
    "  if np.random.random() < 0.5:\n",
    "    x, y = y, x\n",
    "  print('E%d: %s\\nE%d: %s' % (2*i, x, 2*i+1, y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "E0: 21\n",
    "E1: 12\n",
    "E2: 12\n",
    "E3: 21\n",
    "E4: 12\n",
    "E5: 21\n",
    "E6: 21\n",
    "E7: 12\n",
    "E8: 12\n",
    "E9: 21\n",
    "E10: 21\n",
    "E11: 12\n",
    "E12: 21\n",
    "E13: 12\n",
    "E14: 12\n",
    "E15: 21\n",
    "E16: 12\n",
    "E17: 21\n",
    "E18: 21\n",
    "E19: 12\n",
    "E20: 21\n",
    "E21: 12\n",
    "E22: 21\n",
    "E23: 12\n",
    "E24: 12\n",
    "E25: 21\n",
    "E26: 21\n",
    "E27: 12\n",
    "E28: 21\n",
    "E29: 12\n",
    "E30: 12\n",
    "E31: 21\n",
    "E32: 21\n",
    "E33: 12\n",
    "E34: 12\n",
    "E35: 21\n",
    "E36: 21\n",
    "E37: 12\n",
    "E38: 21\n",
    "E39: 12"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "intro solo human pilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pilot_id = 'spike'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_intro_eps = 20"
   ]
  },
  {
   "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
   },
   "outputs": [],
   "source": [
    "intro_rollouts = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "time.sleep(10)\n",
    "for _ in range(n_intro_eps - len(intro_rollouts)):\n",
    "  intro_rollouts.append(run_ep(human_pilot_policy, env, render=True))\n",
    "  time.sleep(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, '%s_pilot_intro.pkl' % pilot_id), 'wb') as f:\n",
    "  pickle.dump({pilot_id: list(zip(*(intro_rollouts)))}, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "evaluate solo human pilot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_eval_eps = 30"
   ]
  },
  {
   "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
   },
   "outputs": [],
   "source": [
    "eval_rollouts = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "time.sleep(10)\n",
    "for _ in range(n_eval_eps - len(eval_rollouts)):\n",
    "  eval_rollouts.append(run_ep(human_pilot_policy, env, render=True))\n",
    "  time.sleep(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, '%s_pilot_eval.pkl' % pilot_id), 'wb') as f:\n",
    "  pickle.dump({pilot_id: list(zip(*(eval_rollouts)))}, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "evaluate copilot with human pilot while fine-tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co_env = make_co_env(\n",
    "  human_pilot_policy,\n",
    "  build_goal_decoder=None,\n",
    "  using_lander_reward_shaping=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "co_env.render()\n",
    "co_env.env.unwrapped.viewer.window.on_key_press = key_press\n",
    "co_env.env.unwrapped.viewer.window.on_key_release = key_release"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_eps = 30\n",
    "pilot_tol = 0.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "time.sleep(10)\n",
    "(scope, q_func), (raw_copilot_policy, reward_data) = make_co_policy(\n",
    "  co_env, pilot_tol=pilot_tol, pilot_is_human=True, n_eps=n_eps,\n",
    "  copilot_scope=copilot_scope,\n",
    "  copilot_q_func=make_q_func(),\n",
    "  reuse=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "co_env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "config_name = str_of_config(pilot_tol, pilot_id, 'rawaction', False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "reward_logs = {config_name: defaultdict(list)}\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": [
    "with open(os.path.join(data_dir, '%s_reward_logs.pkl' % pilot_id), 'wb') as f:\n",
    "  pickle.dump(reward_logs, 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
}