1.3-lunarlander-human-analysis.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",
    "\n",
    "from scipy.stats import binom_test, ttest_1samp"
   ]
  },
  {
   "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-human')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pilot_ids = ['spike', 'jet', 'faye', 'vicious', 'ed', 'ein', 'julia', 'punch', 'judy', 'lin', 'grencia', 'laughingbull']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "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": [
    "NOOP = 1\n",
    "def compute_in_events_for_comb_traj(traj, acts):\n",
    "  ts = [t for t, obs in enumerate(traj[:-1]) if onehot_decode(obs[-6:]) != NOOP]\n",
    "  return list(np.array(ts) / len(traj))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "eval_rollouts = []\n",
    "comb_ins = []\n",
    "comb_acts = []\n",
    "comb_rews = []\n",
    "comb_outcomes = []\n",
    "comb_states = []\n",
    "for pilot_id in pilot_ids:\n",
    "  with open(os.path.join(data_dir, '%s_pilot_eval.pkl' % pilot_id), 'rb') as f:\n",
    "    eval_rollouts.extend(list(zip(*(pickle.load(f)[pilot_id]))))\n",
    "  with open(os.path.join(data_dir, '%s_reward_logs.pkl' % pilot_id), 'rb') as f:\n",
    "    reward_logs = pickle.load(f)\n",
    "    reward_data = list(reward_logs.values())[0]\n",
    "    comb_acts.extend(reward_data['actions'][0])\n",
    "    comb_rews.extend(reward_data['rewards'][0])\n",
    "    comb_outcomes.extend(reward_data['outcomes'][0])\n",
    "    comb_ins.extend([x for traj, acts in zip(reward_data['trajectories'][0], reward_data['actions'][0]) for x in compute_in_events_for_comb_traj(traj, acts)])\n",
    "    comb_states.extend(reward_data['trajectories'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "comb_outcomes = [x if x in [-100, 100] else 0 for x in comb_outcomes]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def n_act_changes(seq):\n",
    "  return np.mean([1 if x != y else 0 for x, y in zip(seq[:-1], seq[1:])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def compute_act_durations(actions):\n",
    "  durations = []\n",
    "  last_dur = 1\n",
    "  for prev, nxt in zip(actions[:-1], actions[1:]):\n",
    "    if prev != nxt:\n",
    "      durations.append(last_dur)\n",
    "      last_dur = 1\n",
    "    else:\n",
    "      last_dur += 1\n",
    "  return durations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def compute_speeds(states):\n",
    "  return [np.linalg.norm(x[2:4]) for x in states]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_changes = []\n",
    "rews = []\n",
    "outcomes = []\n",
    "ins = []\n",
    "act_durations = []\n",
    "speeds = []\n",
    "for rollout in eval_rollouts:\n",
    "  n_changes.append(n_act_changes(rollout[3]))\n",
    "  act_durations.append(np.mean(compute_act_durations(rollout[3])))\n",
    "  speeds.append(np.mean(compute_speeds(rollout[2])))\n",
    "  rews.append(rollout[0])\n",
    "  outcomes.append(rollout[1])\n",
    "  ts = [t for t, a in enumerate(rollout[3]) if a != NOOP]\n",
    "  ins.extend(list(np.array(ts) / len(rollout[3])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "outcomes = [x if x in [-100, 100] else 0 for x in outcomes]\n",
    "comb_n_changes = [n_act_changes(x) for x in comb_acts]\n",
    "comb_act_durations = [np.mean(compute_act_durations(x)) for x in comb_acts]\n",
    "comb_speeds = [np.mean(compute_speeds(x)) for x in comb_states]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'font.size': 16})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Fraction of Actions that Differ from Previous Action')\n",
    "plt.ylabel('Reward')\n",
    "plt.scatter(n_changes, rews, label='Solo Human Pilot', alpha=0.25, color='gray')\n",
    "plt.scatter(comb_n_changes, comb_rews, label='Human Pilot + Copilot', alpha=0.25, color='orange')\n",
    "plt.legend(loc='best')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "all_n_changes = np.array(n_changes + comb_n_changes)\n",
    "all_outcomes = np.array(outcomes + comb_outcomes)\n",
    "n_changes = np.array(n_changes)\n",
    "comb_n_changes = np.array(comb_n_changes)\n",
    "outcomes = np.array(outcomes)\n",
    "comb_outcomes = np.array(comb_outcomes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Fraction of Actions that Differ from Previous Action')\n",
    "plt.ylabel('Frequency (Normalized Number of Trajectories)')\n",
    "plt.title('Lunar Lander - Assisted Trajectories')\n",
    "plt.hist(comb_n_changes[comb_outcomes == 100], alpha=0.5, linewidth=0, label='Success', color='green', normed=True)\n",
    "plt.hist(comb_n_changes[comb_outcomes == -100], alpha=0.5, linewidth=0, label='Crash', color='red', normed=True)\n",
    "plt.legend(loc='best')\n",
    "plt.savefig(os.path.join(data_dir, 'lander-act-switch-vs-succ-assisted.pdf'), bbox_inches='tight')\n",
    "plt.savefig(os.path.join(data_dir, 'lander-act-switch-vs-succ-assisted.png'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Fraction of Actions that Differ from Previous Action')\n",
    "plt.ylabel('Frequency (Normalized Number of Trajectories)')\n",
    "#plt.title('All Trajectories')\n",
    "plt.hist(all_n_changes[all_outcomes == 100], alpha=0.5, linewidth=0, label='Success', color='green', normed=True)\n",
    "plt.hist(all_n_changes[all_outcomes == -100], alpha=0.5, linewidth=0, label='Crash', color='red', normed=True)\n",
    "plt.legend(loc='best')\n",
    "plt.title('Lunar Lander - All Trajectories')\n",
    "plt.savefig(os.path.join(data_dir, 'lander-act-switch-vs-succ.pdf'), bbox_inches='tight')\n",
    "plt.savefig(os.path.join(data_dir, 'lander-act-switch-vs-succ.png'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Fraction of Actions that Differ from Previous Action')\n",
    "plt.ylabel('Frequency (Normalized Number of Trajectories)')\n",
    "#plt.title('All Trajectories')\n",
    "plt.hist(n_changes, alpha=0.5, linewidth=0, label='Solo Human Pilot', color='gray', normed=True)\n",
    "plt.hist(comb_n_changes, alpha=0.5, linewidth=0, label='Human Pilot + Copilot', color='orange', normed=True)\n",
    "plt.legend(loc='best')\n",
    "plt.title('Lunar Lander')\n",
    "plt.savefig(os.path.join(data_dir, 'lander-act-switch-vs-pilot.pdf'), bbox_inches='tight')\n",
    "plt.savefig(os.path.join(data_dir, 'lander-act-switch-vs-pilot.png'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bin_heights, bin_edges = np.histogram(comb_ins, bins=20, normed=True)\n",
    "bin_centers = 0.5 * (bin_edges[1:] + bin_edges[:-1])\n",
    "plt.plot(bin_centers, bin_heights, label='Human Pilot + Copilot', color='orange')\n",
    "\n",
    "bin_heights, bin_edges = np.histogram(ins, bins=20, normed=True)\n",
    "bin_centers = 0.5 * (bin_edges[1:] + bin_edges[:-1])\n",
    "plt.plot(bin_centers, bin_heights, label='Solo Human Pilot', color='gray')\n",
    "\n",
    "plt.title('Lunar Lander')\n",
    "plt.xlabel('Time (Fraction of Trajectory)')\n",
    "plt.ylabel('User Input Frequency (Normalized)')\n",
    "plt.legend(loc='best')\n",
    "plt.savefig(os.path.join(data_dir, 'lander-input-freq-vs-time.pdf'), bbox_inches='tight')\n",
    "plt.savefig(os.path.join(data_dir, 'lander-input-freq-vs-time.png'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Actions Per Minute (APM)')\n",
    "plt.ylabel('Number of Episodes')\n",
    "plt.hist(60 / (np.array(act_durations) * 0.2), bins=20, alpha=0.5, label='Solo Human Pilot', color='gray')#, normed=True)\n",
    "plt.hist(60 / (np.array(comb_act_durations) * 0.2), bins=20, alpha=0.5, label='Human Pilot + Copilot', color='orange')#, normed=True)\n",
    "plt.legend(loc='best')\n",
    "plt.title('Lunar Lander User Study')\n",
    "plt.xlim([None, 250])\n",
    "plt.savefig(os.path.join(data_dir, 'lander-apm.pdf'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "apm = int(60 / (np.mean(act_durations) * 0.2))\n",
    "comb_apm = int(60 / (np.mean(comb_act_durations) * 0.2))\n",
    "apm, comb_apm, comb_apm / apm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.mean(speeds), np.mean(comb_speeds), (1 - np.mean(comb_speeds) / np.mean(speeds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "stats_of_pilot = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for pilot_id in pilot_ids:\n",
    "  with open(os.path.join(data_dir, '%s_pilot_eval.pkl' % pilot_id), 'rb') as f:\n",
    "    eval_rollouts = list(zip(*(pickle.load(f)[pilot_id])))\n",
    "  with open(os.path.join(data_dir, '%s_reward_logs.pkl' % pilot_id), 'rb') as f:\n",
    "    reward_logs = pickle.load(f)\n",
    "    reward_data = list(reward_logs.values())[0]\n",
    "  unassisted_succ = [1 if x[1] == 100 else 0 for x in eval_rollouts]\n",
    "  assisted_succ = [1 if x == 100 else 0 for x in reward_data['outcomes'][0]]\n",
    "  unassisted_fail = [1 if x[1] == -100 else 0 for x in eval_rollouts]\n",
    "  assisted_fail = [1 if x == -100 else 0 for x in reward_data['outcomes'][0]]\n",
    "  stats_of_pilot[pilot_id] = (\n",
    "    np.mean(unassisted_succ), np.std(unassisted_succ) / np.sqrt(len(unassisted_succ)), len(unassisted_succ),\n",
    "    np.mean(assisted_succ), np.std(assisted_succ) / np.sqrt(len(assisted_succ)), len(assisted_succ),\n",
    "    np.mean(unassisted_fail), np.std(unassisted_fail) / np.sqrt(len(unassisted_fail)), len(unassisted_fail),\n",
    "    np.mean(assisted_fail), np.std(assisted_fail) / np.sqrt(len(assisted_fail)), len(assisted_fail)\n",
    "  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "out = []\n",
    "for pilot_id in pilot_ids:\n",
    "  s = stats_of_pilot[pilot_id]\n",
    "  ctrl_succ = s[0]\n",
    "  ctrl_fail = s[6]\n",
    "  treat_succ = s[3]\n",
    "  treat_fail = s[9]\n",
    "  out.append([pilot_id, 0, ctrl_succ, ctrl_fail])\n",
    "  out.append([pilot_id, 1, treat_succ, treat_fail])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# for JMP\n",
    "with open(os.path.join(data_dir, 'lander_hyp_test.csv'), 'w') as f:\n",
    "  f.write('userid,assisted,successrate,crashrate\\n')\n",
    "  f.write('\\n'.join([','.join([str(z) for z in x]) for x in out]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "solo_pi_succs = []\n",
    "solo_pi_crashes = []\n",
    "comb_succs = []\n",
    "comb_crashes = []\n",
    "for pilot_id in pilot_ids:\n",
    "  s = stats_of_pilot[pilot_id]\n",
    "  ctrl_succ = s[0]\n",
    "  ctrl_fail = s[6]\n",
    "  treat_succ = s[3]\n",
    "  treat_fail = s[9]\n",
    "  solo_pi_succs.append(ctrl_succ)\n",
    "  solo_pi_crashes.append(ctrl_fail)\n",
    "  comb_succs.append(treat_succ)\n",
    "  comb_crashes.append(treat_fail)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'font.size': 16})\n",
    "plt.xlabel('Crash Rate')\n",
    "plt.ylabel('Success Rate')\n",
    "plt.title(r'Lunar Lander User Study ($n = %d$)' % len(solo_pi_crashes))\n",
    "plt.scatter(solo_pi_crashes, solo_pi_succs, label='Solo Human Pilot', color='gray', s=100, marker='o')\n",
    "plt.scatter(comb_crashes, comb_succs, label='Human Pilot + Copilot', color='orange', s=100, marker='^')\n",
    "plt.scatter([0.156], [0.026], label='Solo Copilot', color='teal', s=100, marker='s')\n",
    "plt.legend(loc='upper right')\n",
    "plt.ylim([-0.05, 1.05])\n",
    "plt.xlim([-0.05, 1.05])\n",
    "plt.savefig(os.path.join(data_dir, 'lander-user-study-fig.pdf'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for pilot_id in pilot_ids:\n",
    "  ctrl_p, ctrl_std, ctrl_n, treat_p, treat_std, treat_n = stats_of_pilot[pilot_id][:6]\n",
    "  print(\n",
    "    binom_test(ctrl_p * ctrl_n, ctrl_n, treat_p),\n",
    "    binom_test(treat_p * treat_n, treat_n, ctrl_p)\n",
    "  )\n",
    "  ctrl_p, ctrl_std, ctrl_n, treat_p, treat_std, treat_n = stats_of_pilot[pilot_id][6:]\n",
    "  print(\n",
    "    binom_test(ctrl_p * ctrl_n, ctrl_n, treat_p),\n",
    "    binom_test(treat_p * treat_n, treat_n, ctrl_p)\n",
    "  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "unassisted_succ = []\n",
    "assisted_succ = []\n",
    "unassisted_fail = []\n",
    "assisted_fail = []\n",
    "unassisted_rew = []\n",
    "assisted_rew = []\n",
    "for pilot_id in pilot_ids:\n",
    "  with open(os.path.join(data_dir, '%s_pilot_eval.pkl' % pilot_id), 'rb') as f:\n",
    "    eval_rollouts = list(zip(*(pickle.load(f)[pilot_id])))\n",
    "  with open(os.path.join(data_dir, '%s_reward_logs.pkl' % pilot_id), 'rb') as f:\n",
    "    reward_logs = pickle.load(f)\n",
    "    reward_data = list(reward_logs.values())[0]\n",
    "  unassisted_succ.extend([1 if x[1] == 100 else 0 for x in eval_rollouts])\n",
    "  assisted_succ.extend([1 if x == 100 else 0 for x in reward_data['outcomes'][0]])\n",
    "  unassisted_fail.extend([1 if x[1] == -100 else 0 for x in eval_rollouts])\n",
    "  assisted_fail.extend([1 if x == -100 else 0 for x in reward_data['outcomes'][0]])\n",
    "  unassisted_rew.extend([x[0] for x in eval_rollouts])\n",
    "  assisted_rew.extend(reward_data['rewards'][0])\n",
    "stats_of_pilots = (\n",
    "  np.mean(unassisted_rew), np.std(unassisted_rew) / np.sqrt(len(unassisted_rew)),\n",
    "  np.mean(assisted_rew), np.std(assisted_rew) / np.sqrt(len(assisted_rew)),\n",
    "  np.mean(unassisted_succ), np.std(unassisted_succ) / np.sqrt(len(unassisted_succ)),\n",
    "  np.mean(assisted_succ), np.std(assisted_succ) / np.sqrt(len(assisted_succ)), \n",
    "  np.mean(unassisted_fail), np.std(unassisted_fail) / np.sqrt(len(unassisted_fail)), \n",
    "  np.mean(assisted_fail), np.std(assisted_fail) / np.sqrt(len(assisted_fail)), \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(3):\n",
    "  print('%f \\pm %f & %f \\pm %f \\\\\\\\' % (stats_of_pilots[4*i], stats_of_pilots[4*i+1], stats_of_pilots[4*i+2], stats_of_pilots[4*i+3]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binom_test(np.mean(assisted_succ), len(assisted_succ), 0.026)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binom_test(np.mean(assisted_fail), len(assisted_fail), 0.156)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('%f \\pm %f & %f & %f & %f \\pm %f & %f & %f \\\\\\\\' % (\n",
    "  stats_of_pilots[0], stats_of_pilots[1], stats_of_pilots[4], stats_of_pilots[8],\n",
    "  stats_of_pilots[2], stats_of_pilots[3], stats_of_pilots[6], stats_of_pilots[10],\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "solo_outcomes = []\n",
    "solo_traj = []\n",
    "assisted_outcomes = []\n",
    "assisted_traj = []\n",
    "for pilot_id in pilot_ids:\n",
    "  with open(os.path.join(data_dir, '%s_pilot_eval.pkl' % pilot_id), 'rb') as f:\n",
    "    eval_rollouts = pickle.load(f)[pilot_id]\n",
    "  with open(os.path.join(data_dir, '%s_reward_logs.pkl' % pilot_id), 'rb') as f:\n",
    "    reward_logs = pickle.load(f)\n",
    "    reward_data = list(reward_logs.values())[0]\n",
    "  solo_outcomes.extend(eval_rollouts[1])\n",
    "  solo_traj.extend(eval_rollouts[2])\n",
    "  assisted_outcomes.extend(reward_data['outcomes'][0])\n",
    "  assisted_traj.extend(reward_data['trajectories'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Counter([t[0][8] for t in solo_traj])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "goals = [round(float(x), 1) for x in np.arange(-0.8, 1, 0.2)]\n",
    "SUCCESS = 100\n",
    "CRASH = -100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def plot_trajectories(outcomes, trajectories, title, file_name=None, G=-0.8, show_goal=True):\n",
    "  mpl.rcParams.update({'font.size': 16})\n",
    "  plt.title(title)\n",
    "\n",
    "  for outcome, trajectory in zip(outcomes, trajectories):\n",
    "    x, y, vx, vy, a, av, lc, rc, g = list(zip(*trajectory[::5]))[:9]\n",
    "    if g[0] != G:\n",
    "      continue\n",
    "    if outcome == SUCCESS:\n",
    "      cmap = mpl.cm.YlGn\n",
    "    elif outcome == CRASH:\n",
    "      cmap = mpl.cm.YlOrRd\n",
    "    else:\n",
    "      cmap = mpl.cm.gray\n",
    "    plt.scatter(x, y, c=range(len(x)), cmap=cmap, alpha=0.75, linewidth=0)\n",
    "    if show_goal:\n",
    "      plt.scatter([g[0]], [0], marker='*', color='yellow', edgecolor='black', linewidth=1, s=300, alpha=0.5)\n",
    "    \n",
    "  plt.xlim([-1, 1])\n",
    "  plt.ylim([-0.1, 1.1])\n",
    "  plt.xticks([])\n",
    "  plt.yticks([])\n",
    "  plt.axis('off')\n",
    "  if file_name is not None:\n",
    "    plt.savefig(os.path.join(data_dir, file_name), bbox_inches='tight')\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_trajectories(solo_outcomes, solo_traj, 'Solo Human Pilot', 'human-pilot-solo-traj.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_trajectories(assisted_outcomes, assisted_traj, 'Human Pilot + Copilot', 'human-pilot-assisted-traj.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "survey_of_pilot = {\n",
    "  'vicious': [7, 5, 6, 7, 4, 7, 6, 7, 3, 7, 3, 7, 7],\n",
    "  'jet': [2, 5, 5, 3, 3, 6, 5, 3, 1, 1, 3, 6, 5],\n",
    "  'grencia': [6, 6, 6, 5, 5, 4, 7, 6, 5, 5, 5, 6, 6],\n",
    "  'ed': [1, 7, 6, 4, 6, 6, 5, 4, 2, 3, 4, 7, 7],\n",
    "  'julia': [2, 7, 5, 4, 3, 6, 6, 4, 5, 6, 3, 6, 6],\n",
    "  'faye': [5, 6, 5, 5, 4, 5, 5, 6, 6, 6, 7, 7, 5],\n",
    "  'lin': [3, 6, 7, 4, 3, 4, 7, 5, 4, 7, 6, 7, 7],\n",
    "  'judy': [5, 2, 7, 5, 6, 7, 7, 1, 6, 7, 3, 7, 7],\n",
    "  'spike': [4, 5, 7, 7, 5, 3, 7, 7, 5, 4, 5, 7, 7],\n",
    "  'laughingbull': [5, 2, 4, 2, 3, 3, 6, 4, 2, 4, 6, 7, 3],\n",
    "  'ein': [5, 7, 6, 5, 5, 4, 6, 3, 2, 2, 5, 7, 7],\n",
    "  'punch': [7, 7, 7, 6, 3, 6, 7, 4, 6, 6, 6, 7, 7]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "survey_questions = [\n",
    "  'I improved over time',\n",
    "  'The game was too difficult',\n",
    "  'The copilot was generally helpful in completing the task',\n",
    "  'I understood what the copilot was trying to do',\n",
    "  'I could anticipate the copilot’s behavior',\n",
    "  'The copilot improved over time',\n",
    "  'The copilot was helpful in avoiding crashing',\n",
    "  'The copilot was helpful in avoiding flying out of bounds',\n",
    "  'The copilot was helpful in landing quickly before time ran out',\n",
    "  'The copilot was helpful in landing between the flags',\n",
    "  'The copilot was too aggressive and didn’t give me enough control',\n",
    "  'I play differently with the copilot than without the copilot',\n",
    "  'I prefer playing with the copilot'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "survey_elts = [[] for _ in range(len(list(survey_of_pilot.values())[0]))]\n",
    "for k, v in survey_of_pilot.items():\n",
    "  for j, x in enumerate(v):\n",
    "    survey_elts[j].append(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "survey_ps = [ttest_1samp(x, 4)[1] for x in survey_elts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "survey_ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "survey_counts = np.zeros((len(survey_of_pilot), len(list(survey_of_pilot.values())[0]), 7))\n",
    "for i, (k, v) in enumerate(survey_of_pilot.items()):\n",
    "  for j, x in enumerate(v):\n",
    "    survey_counts[i, j, (8-x)-1] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "survey_counts = survey_counts.sum(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "survey_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "question_avgs = (survey_counts * np.tile(np.arange(1, 8, 1), survey_counts.shape[0]).reshape(survey_counts.shape))\n",
    "question_avgs = question_avgs.sum(axis=1) / len(pilot_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question_avgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def fmt_pval(p):\n",
    "  if p < 0.001:\n",
    "    return '<0.001'\n",
    "  elif p < 0.01:\n",
    "    return '<0.01'\n",
    "  elif p < 0.05:\n",
    "    return '<0.05'\n",
    "  else:\n",
    "    return '>0.05'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, x in enumerate(survey_questions):\n",
    "  if survey_ps[i] < 0.05 and question_avgs[i] > 4:\n",
    "    print('& \\\\textbf{%s} & $\\\\mathbf{%s}$ & $\\\\mathbf{%0.2f}$ & %d & %d & %d & %d & %d & %d & %d \\\\\\\\' % (x, fmt_pval(survey_ps[i]), round(question_avgs[i], 2), *list(survey_counts[i])))\n",
    "  else:\n",
    "    print('& %s & $%s$ & %0.2f & %d & %d & %d & %d & %d & %d & %d \\\\\\\\' % (x, fmt_pval(survey_ps[i]), round(question_avgs[i], 2), *list(survey_counts[i])))"
   ]
  },
  {
   "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
}