1.1-tabular-ime.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "import pickle\n",
    "import os\n",
    "import types\n",
    "import random\n",
    "import uuid\n",
    "from copy import deepcopy as copy\n",
    "\n",
    "import gym\n",
    "from gym import spaces\n",
    "from gym.envs.classic_control import rendering\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from scipy.misc import logsumexp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": 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": [
    "univ_scramble = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data_dir = os.path.join(\n",
    "  'data', '1.2-tabular-ime', 'univ_scramble' if univ_scramble else 'scramble')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sess = tf.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "create envs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "gw_size = 7\n",
    "n_train_tasks = 49\n",
    "n_act_dim = 4\n",
    "n_obs_dim = gw_size**2 + 1\n",
    "succ_rew_bonus = 1\n",
    "crash_rew_penalty = -1\n",
    "gamma = 0.99\n",
    "max_ep_len = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "is_succ = lambda r: r[-1][2] > succ_rew_bonus / 2\n",
    "is_crash = lambda r: r[-1][2] < crash_rew_penalty / 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "newton_act_labels = [list(range(n_act_dim)) for _ in range(n_obs_dim)]\n",
    "newton_act_labels = np.array(newton_act_labels).astype(int)\n",
    "\n",
    "reverse_inner = lambda x: [x[0]] + x[1:-1][::-1] + [x[-1]]\n",
    "\n",
    "aristotle_act_labels = [list(range(n_act_dim)) for _ in range(n_obs_dim)]\n",
    "for i in range(n_obs_dim):\n",
    "  if univ_scramble:\n",
    "    aristotle_act_labels[i] = list(reversed(aristotle_act_labels[i]))\n",
    "  else:\n",
    "    random.shuffle(aristotle_act_labels[i])\n",
    "aristotle_act_labels = np.array(aristotle_act_labels).astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "all_goals = list(zip(*[x.ravel() for x in np.meshgrid(\n",
    "    np.arange(0, gw_size, 1), np.arange(0, gw_size, 1))]))\n",
    "train_goals = [all_goals[i] for i in np.random.choice(list(range(len(all_goals))), n_train_tasks, replace=False)]\n",
    "train_goals = np.array(train_goals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(train_goals[:, 0], train_goals[:, 1], linewidth=0, color='gray', s=100, marker='*')\n",
    "plt.xlim([-0.5, gw_size-0.5])\n",
    "plt.ylim([-0.5, gw_size-0.5])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'train_goals.pkl'), 'wb') as f:\n",
    "  pickle.dump(train_goals, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'train_goals.pkl'), 'rb') as f:\n",
    "  train_goals = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def make_reward_func(goal):\n",
    "  def pos_from_obs(obs):\n",
    "    x = obs // gw_size\n",
    "    y = obs % gw_size\n",
    "    return np.array([x, y])\n",
    "  \n",
    "  def reward_shaping(obs):\n",
    "    return -np.linalg.norm((pos_from_obs(obs) - goal) / gw_size)\n",
    "\n",
    "  def reward_func(prev_obs, action, obs):\n",
    "    pos = pos_from_obs(obs)\n",
    "    if (pos < 0).any() or (pos >= gw_size).any():\n",
    "      r = crash_rew_penalty\n",
    "    elif (pos == goal).all():\n",
    "      r = succ_rew_bonus\n",
    "    else:\n",
    "      r = 0\n",
    "    r += gamma * reward_shaping(obs) - reward_shaping(prev_obs)\n",
    "    return r\n",
    "  \n",
    "  return reward_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class GridWorldNav(gym.Env):\n",
    "  metadata = {\n",
    "    'render.modes': ['human']\n",
    "  }\n",
    "  \n",
    "  def __init__(\n",
    "      self, \n",
    "      act_labels=None,\n",
    "      max_ep_len=max_ep_len, \n",
    "      reward_func=None,\n",
    "      goal=None\n",
    "    ):\n",
    "    self.observation_space = spaces.Discrete(n_obs_dim)\n",
    "    self.action_space = spaces.Discrete(n_act_dim)\n",
    "    \n",
    "    self.pos = None\n",
    "    self.curr_step = None\n",
    "    self.viewer = None\n",
    "    self.curr_obs = None\n",
    "    self.next_obs = None\n",
    "    \n",
    "    self.succ_rew_bonus = succ_rew_bonus\n",
    "    self.max_ep_len = max_ep_len\n",
    "    self.reward_func = reward_func\n",
    "    self.act_labels = act_labels\n",
    "    self.goal = goal\n",
    "    \n",
    "    self.R = np.zeros((n_obs_dim, n_act_dim, n_obs_dim))\n",
    "    for s in range(n_obs_dim):\n",
    "      for sp in range(n_obs_dim):\n",
    "        self.R[s, :, sp] = self.reward_func(s, None, sp)\n",
    "          \n",
    "    self.T = np.zeros((n_obs_dim, n_act_dim, n_obs_dim))\n",
    "    for s in range(n_obs_dim-1):\n",
    "      x = s // gw_size\n",
    "      y = s % gw_size\n",
    "      self.T[s, self.act_labels[s, 0], x*gw_size+(y-1) if y > 0 else -1] = 1\n",
    "      self.T[s, self.act_labels[s, 1], x*gw_size+(y+1) if y < gw_size-1 else -1] = 1\n",
    "      self.T[s, self.act_labels[s, 2], (x-1)*gw_size+y if x > 0 else -1] = 1\n",
    "      self.T[s, self.act_labels[s, 3], (x+1)*gw_size+y if x < gw_size-1 else -1] = 1\n",
    "    self.T[-1, :, -1] = 1\n",
    "    \n",
    "  def _obs(self):\n",
    "    self.curr_obs = int(self.pos[0]*gw_size + self.pos[1])\n",
    "    if self.curr_obs < 0 or self.curr_obs >= gw_size**2:\n",
    "      self.curr_obs = gw_size**2\n",
    "    return self.curr_obs\n",
    "\n",
    "  def _step(self, action):\n",
    "    if self.next_obs is None:\n",
    "      if action == self.act_labels[self.curr_obs, 0]: # left\n",
    "        self.pos[1] -= 1\n",
    "      elif action == self.act_labels[self.curr_obs, 1]: # right\n",
    "        self.pos[1] += 1\n",
    "      elif action == self.act_labels[self.curr_obs, 2]: # up\n",
    "        self.pos[0] -= 1\n",
    "      elif action == self.act_labels[self.curr_obs, 3]: # down\n",
    "        self.pos[0] += 1\n",
    "      else:\n",
    "        raise ValueError('invalid action')\n",
    "    else:\n",
    "      self.pos = np.array([self.next_obs // gw_size, self.next_obs % gw_size])\n",
    "          \n",
    "    self.curr_step += 1\n",
    "    succ = (self.pos == self.goal).all()\n",
    "    oob = (self.pos < 0).any() or (self.pos >= gw_size).any()\n",
    "    oot = self.curr_step >= self.max_ep_len\n",
    "    \n",
    "    obs = self._obs()\n",
    "    r = self.reward_func(self.prev_obs, action, obs)\n",
    "    done = oot or succ or oob\n",
    "    info = {}\n",
    "    self.prev_obs = obs\n",
    "    \n",
    "    return obs, r, done, info\n",
    "    \n",
    "  def _reset(self):\n",
    "    pos = (np.random.choice(gw_size**2-1) + self.goal[0]*gw_size + self.goal[1]) % (gw_size**2)\n",
    "    self.pos = np.array([pos // gw_size, pos % gw_size])\n",
    "    \n",
    "    self.curr_step = 0\n",
    "    self.prev_obs = self._obs()\n",
    "    self.next_obs = None\n",
    "    return self.prev_obs\n",
    "  \n",
    "  def _render(self, mode='human', close=False):\n",
    "    if close:\n",
    "      if self.viewer is not None:\n",
    "        self.viewer.close()\n",
    "        self.viewer = None\n",
    "      return\n",
    "    \n",
    "    if self.viewer is None:\n",
    "      self.viewer = rendering.SimpleImageViewer()\n",
    "    \n",
    "    fig = plt.figure()\n",
    "    canvas = FigureCanvas(fig)\n",
    "    \n",
    "    plt.scatter([self.goal[0]], [self.goal[1]], color='gray', linewidth=0, alpha=0.75, marker='*')\n",
    "    plt.scatter([self.pos[0]], [self.pos[1]], color='orange', linewidth=0, alpha=0.75)\n",
    "    plt.xlim([-1, gw_size+1])\n",
    "    plt.ylim([-1, gw_size+1])\n",
    "    plt.axis('off')\n",
    "    \n",
    "    agg = canvas.switch_backends(FigureCanvas)\n",
    "    agg.draw()\n",
    "    width, height = fig.get_size_inches() * fig.get_dpi()\n",
    "    self.viewer.imshow(np.fromstring(agg.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3))\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_reward_funcs = [make_reward_func(goal) for goal in train_goals]\n",
    "train_newton_envs = [GridWorldNav(reward_func=r, goal=train_goals[i], act_labels=newton_act_labels) for i, r in enumerate(train_reward_funcs)]\n",
    "train_aristotle_envs = [GridWorldNav(reward_func=r, goal=train_goals[i], act_labels=aristotle_act_labels) for i, r in enumerate(train_reward_funcs)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def run_ep(policy, env, max_ep_len=max_ep_len, render=False, task_idx=None):\n",
    "  obs = env.reset()\n",
    "  done = False\n",
    "  totalr = 0.\n",
    "  prev_obs = obs\n",
    "  rollout = []\n",
    "  for step_idx in range(max_ep_len+1):\n",
    "    if done:\n",
    "      break\n",
    "    action = policy(obs)\n",
    "    obs, r, done, info = env.step(action)\n",
    "    rollout.append((prev_obs, action, r, obs, float(done), task_idx))\n",
    "    prev_obs = obs\n",
    "    if render:\n",
    "      env.render()\n",
    "    totalr += r\n",
    "  return rollout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def make_aristotle_pilot_policy(goal):\n",
    "  gx, gy = goal\n",
    "  def aristotle_pilot_policy(obs):\n",
    "    x = obs // gw_size\n",
    "    y = obs % gw_size\n",
    "    up = gx<x\n",
    "    down = gx>x\n",
    "    left = gy<y\n",
    "    right = gy>y\n",
    "    lr = left or right\n",
    "    ud = up or down\n",
    "    if lr and (not ud or np.random.random() < 0.5):\n",
    "      if left:\n",
    "        return aristotle_act_labels[obs, 0]\n",
    "      elif right:\n",
    "        return aristotle_act_labels[obs, 1]\n",
    "    elif ud:\n",
    "      if up:\n",
    "        return aristotle_act_labels[obs, 2]\n",
    "      elif down:\n",
    "        return aristotle_act_labels[obs, 3]\n",
    "    return aristotle_act_labels[obs, 0]\n",
    "  return aristotle_pilot_policy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "train agent with tabular Q-learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def tabsoftq_iter(R, T, maxiter=1000, verbose=True, Q_init=None, learning_rate=1, ftol=0):\n",
    "  n, m = R.shape[:2]\n",
    "  Q = np.zeros((n, m)) if Q_init is None else copy(Q_init)\n",
    "  prevQ = copy(Q)\n",
    "  if verbose:\n",
    "    diffs = []\n",
    "  for iter_idx in range(maxiter):\n",
    "    V = logsumexp(prevQ, axis=1)\n",
    "    V_broad = V.reshape((1, 1, n))\n",
    "    Q = np.sum(T * (R + gamma * V_broad), axis=2)\n",
    "    Q = (1 - learning_rate) * prevQ + learning_rate * Q\n",
    "    diff = np.mean((Q - prevQ)**2)/(np.std(Q)**2)\n",
    "    if verbose:\n",
    "      diffs.append(diff)\n",
    "    if diff < ftol:\n",
    "      break\n",
    "    prevQ = copy(Q)\n",
    "  if verbose:\n",
    "    plt.xlabel('Number of Iterations')\n",
    "    plt.ylabel('Avg. Squared Bellman Error')\n",
    "    plt.title('Soft Q Iteration')\n",
    "    plt.plot(diffs)\n",
    "    plt.yscale('log')\n",
    "    plt.show()\n",
    "  return Q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def tabsoftq_learn(env):\n",
    "  R = env.unwrapped.R\n",
    "  T = env.unwrapped.T\n",
    "  return tabsoftq_iter(R, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "aristotle_softq_pilot_temp = 1\n",
    "def make_tabsoftq_policy(Q):\n",
    "  def tabsoftq_policy(obs):\n",
    "    return np.argmax(aristotle_softq_pilot_temp*Q[obs, :] + np.random.gumbel(0, 1, n_act_dim))\n",
    "  return tabsoftq_policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "Q = np.stack([tabsoftq_learn(train_aristotle_envs[train_task_idx]) for train_task_idx in range(n_train_tasks)], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for i in range(n_train_tasks):\n",
    "  plt.imshow(np.argmax(Q[i, :-1], axis=1).reshape((gw_size, gw_size)).T)\n",
    "  plt.scatter(train_goals[i, 0], train_goals[i, 1], linewidth=0, color='gray', s=200, marker='*')\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for i in range(n_train_tasks):\n",
    "  plt.imshow(np.max(Q[i, :-1, :], axis=1).reshape((gw_size, gw_size)).T)\n",
    "  plt.scatter(train_goals[i, 0], train_goals[i, 1], linewidth=0, color='gray', s=200, marker='*')\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "aristotle_tabsoftq_pilot_path = os.path.join(data_dir, 'train_aristotle_tabsoftq_pilots.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(aristotle_tabsoftq_pilot_path, 'wb') as f:\n",
    "  pickle.dump(Q, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(aristotle_tabsoftq_pilot_path, 'rb') as f:\n",
    "  Q = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "aristotle_pilot_policies = [make_tabsoftq_policy(Q[i]) for i in range(n_train_tasks)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "sanity-check envs, agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_task_idx = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "run_ep(aristotle_pilot_policies[train_task_idx], train_aristotle_envs[train_task_idx], render=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_aristotle_envs[train_task_idx].close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "run_ep(aristotle_pilot_policies[train_task_idx], train_newton_envs[train_task_idx], render=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_newton_envs[train_task_idx].close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "fit internal dynamics model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n_train_rollouts_per_env = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "demo_rollouts = [[run_ep(aristotle_pilot_policies[train_task_idx], newton_env, render=False, task_idx=train_task_idx)\n",
    "                  for _ in range(n_train_rollouts_per_env)] \n",
    "                 for train_task_idx, newton_env in enumerate(train_newton_envs)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'aristotle_pilot_policy_demo_rollouts.pkl'), 'wb') as f:\n",
    "  pickle.dump(demo_rollouts, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'aristotle_pilot_policy_demo_rollouts.pkl'), 'rb') as f:\n",
    "  demo_rollouts = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def build_mlp(\n",
    "    input_placeholder,\n",
    "    output_size,\n",
    "    scope,\n",
    "    n_layers=2,\n",
    "    size=500,\n",
    "    activation=tf.nn.relu,\n",
    "    output_activation=None,\n",
    "    reuse=False\n",
    "  ):\n",
    "  out = input_placeholder\n",
    "  with tf.variable_scope(scope, reuse=reuse):\n",
    "    for _ in range(n_layers):\n",
    "      out = tf.layers.dense(out, size, activation=activation)\n",
    "    out = tf.layers.dense(out, output_size, activation=output_activation)\n",
    "  return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def vectorize_rollouts(rollouts):\n",
    "  obs = []\n",
    "  actions = []\n",
    "  rewards = []\n",
    "  next_obs = []\n",
    "  dones = []\n",
    "  task_idxes = []\n",
    "  for rollout in rollouts:\n",
    "    more_obs, more_actions, more_rewards, more_next_obs, more_dones, more_task_idxes = list(zip(*rollout))\n",
    "    obs.extend(more_obs)\n",
    "    actions.extend(more_actions)\n",
    "    rewards.extend(more_rewards)\n",
    "    next_obs.extend(more_next_obs)\n",
    "    dones.extend(more_dones)\n",
    "    task_idxes.extend(more_task_idxes)\n",
    "  return np.array(obs), np.array(actions), np.array(rewards), np.array(next_obs), np.array(dones), np.array(task_idxes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "demo_obs = None\n",
    "demo_actions = None\n",
    "demo_next_obs = None\n",
    "demo_task_idxes = None\n",
    "train_demo_example_idxes = None\n",
    "val_demo_batch = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def process_demo_rollouts(demo_rollouts):\n",
    "  global demo_obs\n",
    "  global demo_actions\n",
    "  global demo_next_obs\n",
    "  global demo_task_idxes\n",
    "  global train_demo_example_idxes\n",
    "  global val_demo_batch\n",
    "\n",
    "  vectorized_demo_rollouts = vectorize_rollouts(demo_rollouts)\n",
    "  \n",
    "  demo_obs, demo_actions, demo_rewards, demo_next_obs, demo_done_masks, demo_task_idxes = vectorized_demo_rollouts\n",
    "  demo_example_idxes = list(range(len(demo_obs)))\n",
    "  \n",
    "  random.shuffle(demo_example_idxes)\n",
    "  n_train_demo_examples = int(0.9 * len(demo_example_idxes))\n",
    "  train_demo_example_idxes = demo_example_idxes[:n_train_demo_examples]\n",
    "  val_demo_example_idxes = demo_example_idxes[n_train_demo_examples:]\n",
    "  val_demo_batch = demo_obs[val_demo_example_idxes], demo_actions[val_demo_example_idxes], demo_next_obs[val_demo_example_idxes], demo_task_idxes[val_demo_example_idxes]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "process_demo_rollouts(sum(demo_rollouts, []))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def sample_batch(size):\n",
    "  idxes = random.sample(train_demo_example_idxes, size)\n",
    "  demo_batch = demo_obs[idxes], demo_actions[idxes], demo_next_obs[idxes], demo_task_idxes[idxes]\n",
    "  return demo_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "gamma = 0.99\n",
    "iterations = 1000\n",
    "learning_rate = 1e-3\n",
    "batch_size = 512\n",
    "sq_td_err_penalty = 1e-3\n",
    "\n",
    "q_n_layers = 0\n",
    "q_layer_size = None\n",
    "q_activation = None\n",
    "q_output_activation = None\n",
    "\n",
    "n_layers = 0\n",
    "layer_size = None\n",
    "activation = None\n",
    "output_activation = tf.nn.softmax\n",
    "\n",
    "val_update_freq = 100\n",
    "n_val_eval_rollouts = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "im_scope = str(uuid.uuid4())\n",
    "q_scope = str(uuid.uuid4())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "demo_obs_t_ph = tf.placeholder(tf.int32, [None])\n",
    "demo_act_t_ph = tf.placeholder(tf.int32, [None])\n",
    "demo_task_t_ph = tf.placeholder(tf.int32, [None])\n",
    "demo_batch_size_ph = tf.placeholder(tf.int32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def featurize_obs(obs):\n",
    "  return tf.one_hot(obs, n_obs_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "demo_batch_idxes = tf.reshape(tf.range(0, demo_batch_size_ph, 1), [demo_batch_size_ph, 1])\n",
    "\n",
    "demo_q_t = tf.stack([build_mlp(\n",
    "  featurize_obs(demo_obs_t_ph), n_act_dim, q_scope+'-'+str(train_task_idx), n_layers=q_n_layers, size=q_layer_size,\n",
    "  activation=q_activation, output_activation=q_output_activation\n",
    ") for train_task_idx in range(n_train_tasks)], axis=0)\n",
    "demo_q_t = tf.gather_nd(demo_q_t, tf.concat([tf.expand_dims(demo_task_t_ph, 1), demo_batch_idxes], axis=1))\n",
    "\n",
    "demo_act_idxes = tf.concat([demo_batch_idxes, tf.reshape(demo_act_t_ph, [demo_batch_size_ph, 1])], axis=1)\n",
    "demo_act_val_t = tf.gather_nd(demo_q_t, demo_act_idxes)\n",
    "state_val_t = tf.reduce_logsumexp(demo_q_t, axis=1)\n",
    "act_log_likelihoods = demo_act_val_t - state_val_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "neg_avg_log_likelihood = -tf.reduce_mean(act_log_likelihoods)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "obs_for_obs_tp1_probs = tf.cast(tf.floor(\n",
    "  tf.range(0, n_obs_dim*n_act_dim, 1) / n_act_dim), dtype=tf.int32)\n",
    "\n",
    "act_for_obs_tp1_probs = tf.floormod(tf.range(\n",
    "  0, n_obs_dim*n_act_dim, 1), n_act_dim)\n",
    "\n",
    "obs_tp1_probs_in = tf.one_hot(\n",
    "  obs_for_obs_tp1_probs*n_act_dim+act_for_obs_tp1_probs, n_obs_dim*n_act_dim)\n",
    "\n",
    "obs_tp1_probs = build_mlp(\n",
    "  obs_tp1_probs_in,\n",
    "  n_obs_dim, im_scope, n_layers=n_layers, size=layer_size,\n",
    "  activation=activation, output_activation=output_activation\n",
    ")\n",
    "obs_tp1_probs = tf.reshape(obs_tp1_probs, [n_obs_dim, n_act_dim, n_obs_dim])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "q_tp1 = tf.stack([build_mlp(\n",
    "  featurize_obs(tf.range(0, n_obs_dim, 1)), n_act_dim, q_scope+'-'+str(train_task_idx), n_layers=q_n_layers, size=q_layer_size,\n",
    "  activation=q_activation, output_activation=q_output_activation, reuse=True\n",
    ") for train_task_idx in range(n_train_tasks)], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "v_tp1 = tf.reduce_logsumexp(q_tp1, axis=2)\n",
    "\n",
    "all_rew = tf.convert_to_tensor(np.stack([env.unwrapped.R for env in train_newton_envs], axis=0), dtype=tf.float32)\n",
    "\n",
    "v_tp1_broad = tf.reshape(v_tp1, [n_train_tasks, 1, 1, n_obs_dim])\n",
    "obs_tp1_probs_broad = tf.expand_dims(obs_tp1_probs, 0)\n",
    "\n",
    "exp_v_tp1 = tf.reduce_sum(obs_tp1_probs_broad * v_tp1_broad, axis=3)\n",
    "exp_rew_t = tf.reduce_sum(obs_tp1_probs_broad * all_rew, axis=3)\n",
    "target_t = exp_rew_t + gamma * exp_v_tp1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "q_t = tf.stack([build_mlp(\n",
    "  featurize_obs(tf.range(0, n_obs_dim, 1)), n_act_dim, q_scope+'-'+str(train_task_idx), n_layers=q_n_layers, size=q_layer_size,\n",
    "  activation=q_activation, output_activation=q_output_activation, reuse=True\n",
    ") for train_task_idx in range(n_train_tasks)], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "td_err = q_t - target_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sq_td_err = tf.reduce_mean(td_err**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "loss = neg_avg_log_likelihood + sq_td_err_penalty * sq_td_err"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "update_op = tf.train.AdamOptimizer(learning_rate).minimize(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_val_assisted_env(obs_tp1):\n",
    "  test_goal = np.random.choice(gw_size, 2)\n",
    "  test_reward_func = make_reward_func(test_goal)\n",
    "  test_aristotle_pilot_policy = make_aristotle_pilot_policy(test_goal)\n",
    "  env = GridWorldNav(\n",
    "    reward_func=test_reward_func, goal=test_goal, act_labels=newton_act_labels)\n",
    "  env.unwrapped._step_orig = env.unwrapped._step\n",
    "  def _step(self, action):\n",
    "    transferred_act = dyn_transfer(self.curr_obs, action, obs_tp1=obs_tp1)\n",
    "    obs, r, done, info = self._step_orig(transferred_act)\n",
    "    return obs, r, done, info\n",
    "  env.unwrapped._step = types.MethodType(_step, env.unwrapped)\n",
    "  return test_aristotle_pilot_policy, env\n",
    "\n",
    "def compute_assisted_perf():\n",
    "  obs_tp1 = np.argmax(sess.run(obs_tp1_probs), axis=2)\n",
    "  assisted_rollouts = [run_ep(*make_val_assisted_env(obs_tp1), render=False) for _ in range(n_val_eval_rollouts)]\n",
    "  assisted_rew = np.mean([sum(x[2] for x in r) for r in assisted_rollouts])\n",
    "  assisted_succ = np.mean([1 if is_succ(r) else 0 for r in assisted_rollouts])\n",
    "  assisted_crash = np.mean([1 if is_crash(r) else 0 for r in assisted_rollouts])\n",
    "  assisted_perf = {\n",
    "    'assisted_rew': assisted_rew,\n",
    "    'assisted_succ': assisted_succ,\n",
    "    'assisted_crash': assisted_crash\n",
    "  }\n",
    "  return assisted_perf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def compute_int_dyn_acc():\n",
    "  states = np.repeat(np.arange(0, n_obs_dim, 1), n_act_dim)\n",
    "  actions = np.tile(np.arange(0, n_act_dim, 1), n_obs_dim)\n",
    "  probs = sess.run(obs_tp1_probs)\n",
    "  obs_tp1_true = np.argmax(train_aristotle_envs[0].unwrapped.T, axis=2)[states, actions]\n",
    "  accuracy = np.mean((np.argmax(probs[states, actions, :], axis=1) == obs_tp1_true).astype(int))\n",
    "  kld = -np.mean(np.log(1e-9+probs[states, actions, obs_tp1_true]))\n",
    "  return {'int_dyn_acc': accuracy, 'int_dyn_kld': kld}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.global_variables_initializer().run(session=sess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n_iters = iterations * len(demo_obs) // batch_size\n",
    "train_logs = {\n",
    "  'loss_evals': [],\n",
    "  'nll_evals': [],\n",
    "  'ste_evals': [],\n",
    "  'val_loss_evals': [],\n",
    "  'val_nll_evals': [],\n",
    "  'val_ste_evals': [],\n",
    "  'assisted_rew_evals': [],\n",
    "  'assisted_succ_evals': [],\n",
    "  'assisted_crash_evals': [],\n",
    "  'int_dyn_acc_evals': [],\n",
    "  'int_dyn_kld_evals': []\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def compute_batch_loss(demo_batch, step=False, t=None):\n",
    "  demo_batch_obs_t, demo_batch_act_t, demo_batch_obs_tp1, demo_batch_task_t = demo_batch\n",
    "  \n",
    "  feed_dict = {\n",
    "    demo_obs_t_ph: demo_batch_obs_t,\n",
    "    demo_act_t_ph: demo_batch_act_t,\n",
    "    demo_task_t_ph: demo_batch_task_t,\n",
    "    demo_batch_size_ph: demo_batch_obs_t.shape[0]\n",
    "  }\n",
    "  \n",
    "  [loss_eval, neg_avg_log_likelihood_eval, sq_td_err_eval] = sess.run(\n",
    "    [loss, neg_avg_log_likelihood, sq_td_err], feed_dict=feed_dict)\n",
    "  \n",
    "  if step:\n",
    "    sess.run(update_op, feed_dict=feed_dict)\n",
    "  \n",
    "  d = {\n",
    "    'loss': loss_eval,\n",
    "    'nll': neg_avg_log_likelihood_eval,\n",
    "    'ste': sq_td_err_eval\n",
    "  }\n",
    "  if not step:\n",
    "    d.update(compute_int_dyn_acc())\n",
    "    d.update(compute_assisted_perf())\n",
    "  return d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "val_log = None\n",
    "while len(train_logs['loss_evals']) < n_iters:\n",
    "  demo_batch = sample_batch(batch_size)\n",
    "  \n",
    "  t = len(train_logs['loss_evals'])\n",
    "  train_log = compute_batch_loss(demo_batch, step=True, t=t)\n",
    "  if val_log is None or len(train_logs['loss_evals']) % val_update_freq == 0:\n",
    "    val_log = compute_batch_loss(val_demo_batch, step=False, t=t)\n",
    "  \n",
    "  print('%d %d %f %f %f %f %f %f %f' % (\n",
    "    t, n_iters, train_log['loss'],\n",
    "    train_log['nll'], train_log['ste'], val_log['loss'],\n",
    "    val_log['nll'], val_log['ste'], val_log['assisted_rew'])\n",
    "  )\n",
    "  \n",
    "  for k, v in train_log.items():\n",
    "    train_logs['%s_evals' % k].append(v)\n",
    "  for k, v in val_log.items():\n",
    "    train_logs['%s%s_evals' % ('val_' if k in ['loss', 'nll', 'ste'] else '', k)].append(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for k in ['val_nll_evals', 'val_ste_evals']:\n",
    "  plt.xlabel('Iterations')\n",
    "  plt.ylabel(k.split('_')[1])\n",
    "  plt.plot(train_logs[k])\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.xlabel('Iterations')\n",
    "plt.ylabel('Reward')\n",
    "plt.axhline(y=np.mean(ideal_rew), linestyle='--', color='teal', label='Optimal')\n",
    "plt.axhline(y=np.mean(unassisted_rew), linestyle=':', color='gray', label='Unassisted')\n",
    "plt.plot(train_logs['assisted_rew_evals'], color='orange', label='Assisted')\n",
    "plt.legend(loc='best')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.xlabel('Iterations')\n",
    "plt.ylabel('Success Rate')\n",
    "plt.axhline(y=np.mean(ideal_succ), linestyle='--', color='teal', label='Optimal')\n",
    "plt.axhline(y=np.mean(unassisted_succ), linestyle=':', color='gray', label='Unassisted')\n",
    "plt.plot(train_logs['assisted_succ_evals'], color='orange', label='Assisted')\n",
    "plt.ylim([-0.05, 1.05])\n",
    "plt.legend(loc='best')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.xlabel('Iterations')\n",
    "plt.ylabel('Crash Rate')\n",
    "plt.axhline(y=np.mean(ideal_crash), linestyle='--', color='teal', label='Optimal')\n",
    "plt.axhline(y=np.mean(unassisted_crash), linestyle=':', color='gray', label='Unassisted')\n",
    "plt.plot(train_logs['assisted_crash_evals'], color='orange', label='Assisted')\n",
    "plt.ylim([-0.05, 1.05])\n",
    "plt.legend(loc='best')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.xlabel('Iterations')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.axhline(y=(1/n_act_dim), linestyle='--', color='gray', label='Uniform')\n",
    "plt.plot(train_logs['int_dyn_acc_evals'], color='orange', label='Ours')\n",
    "plt.ylim([-0.05, 1.05])\n",
    "plt.legend(loc='best')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.xlabel('Iterations')\n",
    "plt.ylabel('KL Divergence')\n",
    "plt.axhline(y=-np.log(1/n_act_dim), linestyle='--', color='gray', label='Uniform')\n",
    "plt.plot(train_logs['int_dyn_kld_evals'], color='orange', label='Ours')\n",
    "plt.ylim([-0.05, None])\n",
    "plt.legend(loc='best')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "sweep `n_train_tasks`, repeat each with ten different random seeds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "list_n_train_tasks = [1, 2, 4, 9, 16, 25, 36, 49]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "master_train_logs = [[] for _ in list_n_train_tasks]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "random.shuffle(demo_rollouts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for n_train_tasks_idx, n_train_tasks in enumerate(list_n_train_tasks):\n",
    "  process_demo_rollouts(sum(demo_rollouts[:n_train_tasks], []))\n",
    "\n",
    "  for _ in range(10):\n",
    "    tf.global_variables_initializer().run(session=sess)\n",
    "\n",
    "    n_iters = 20000\n",
    "    train_logs = {\n",
    "      'loss_evals': [],\n",
    "      'nll_evals': [],\n",
    "      'ste_evals': [],\n",
    "      'val_loss_evals': [],\n",
    "      'val_nll_evals': [],\n",
    "      'val_ste_evals': [],\n",
    "      'assisted_rew_evals': [],\n",
    "      'assisted_succ_evals': [],\n",
    "      'assisted_crash_evals': [],\n",
    "      'int_dyn_acc_evals': [],\n",
    "      'int_dyn_kld_evals': []\n",
    "    }\n",
    "\n",
    "    val_log = None\n",
    "    while len(train_logs['loss_evals']) < n_iters:\n",
    "      demo_batch = sample_batch(batch_size)\n",
    "\n",
    "      t = len(train_logs['loss_evals'])\n",
    "      train_log = compute_batch_loss(demo_batch, step=True, t=t)\n",
    "      if val_log is None or len(train_logs['loss_evals']) % val_update_freq == 0:\n",
    "        val_log = compute_batch_loss(val_demo_batch, step=False, t=t)\n",
    "\n",
    "      if t % 1000 == 0:\n",
    "        print('%d %d %f %f %f %f %f %f %f' % (\n",
    "          t, n_iters, train_log['loss'],\n",
    "          train_log['nll'], train_log['ste'], val_log['loss'],\n",
    "          val_log['nll'], val_log['ste'], val_log['assisted_rew'])\n",
    "        )\n",
    "\n",
    "      for k, v in train_log.items():\n",
    "        train_logs['%s_evals' % k].append(v)\n",
    "      for k, v in val_log.items():\n",
    "        train_logs['%s%s_evals' % ('val_' if k in ['loss', 'nll', 'ste'] else '', k)].append(v)\n",
    "\n",
    "    master_train_logs[n_train_tasks_idx].append(train_logs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'master_train_logs.pkl'), 'wb') as f:\n",
    "  pickle.dump((list_n_train_tasks, master_train_logs), f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "internal2real dynamics transfer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def featurize_state(s):\n",
    "  return np.array([s // gw_size, s % gw_size])\n",
    "\n",
    "def state_dist(s, sp):\n",
    "  return np.sum(np.abs(featurize_state(s) - featurize_state(sp)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "obs_tp1 = np.argmax(sess.run(obs_tp1_probs), axis=2)\n",
    "real_dyn = np.argmax(train_newton_envs[0].unwrapped.T, axis=2)\n",
    "def dyn_transfer(state, action, obs_tp1=obs_tp1):\n",
    "  desired_next_state = obs_tp1[state, action]\n",
    "  next_states = [real_dyn[state, a] for a in range(n_act_dim)]\n",
    "  dists = np.array([state_dist(desired_next_state, sp) for sp in next_states])\n",
    "  return np.argmax(-dists)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_assisted_env():\n",
    "  test_goal = np.random.choice(gw_size, 2)\n",
    "  test_reward_func = make_reward_func(test_goal)\n",
    "  test_aristotle_pilot_policy = make_aristotle_pilot_policy(test_goal)\n",
    "  env = GridWorldNav(\n",
    "    reward_func=test_reward_func, goal=test_goal, act_labels=newton_act_labels)\n",
    "  env.unwrapped._step_orig = env.unwrapped._step\n",
    "  def _step(self, action):\n",
    "    transferred_act = dyn_transfer(self.curr_obs, action)\n",
    "    obs, r, done, info = self._step_orig(transferred_act)\n",
    "    return obs, r, done, info\n",
    "  env.unwrapped._step = types.MethodType(_step, env.unwrapped)\n",
    "  return test_aristotle_pilot_policy, env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_env_without_dyn_transfer(act_labels):\n",
    "  test_goal = np.random.choice(gw_size, 2)\n",
    "  test_reward_func = make_reward_func(test_goal)\n",
    "  test_aristotle_pilot_policy = make_aristotle_pilot_policy(test_goal)\n",
    "  unassisted_env = GridWorldNav(act_labels=act_labels, reward_func=test_reward_func, goal=test_goal)\n",
    "  return test_aristotle_pilot_policy, unassisted_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_env_with_rand_pilot():\n",
    "  test_goal = np.random.choice(gw_size, 2)\n",
    "  test_reward_func = make_reward_func(test_goal)\n",
    "  test_aristotle_pilot_policy = make_aristotle_pilot_policy(test_goal)\n",
    "  unassisted_env = GridWorldNav(act_labels=newton_act_labels, reward_func=test_reward_func, goal=test_goal)\n",
    "  return test_aristotle_pilot_policy, unassisted_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "make_unassisted_env = lambda: make_env_without_dyn_transfer(newton_act_labels)\n",
    "make_ideal_env = lambda: make_env_without_dyn_transfer(aristotle_act_labels)\n",
    "make_rand_env = lambda: make_env_with_rand_pilot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n_eval_rollouts = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "assisted_rollouts = [run_ep(*make_assisted_env(), render=False) for _ in range(n_eval_rollouts)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'aristotle_pilot_policy_assisted_rollouts.pkl'), 'wb') as f:\n",
    "  pickle.dump(assisted_rollouts, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'aristotle_pilot_policy_assisted_rollouts.pkl'), 'rb') as f:\n",
    "  assisted_rollouts = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "unassisted_rollouts = [run_ep(*make_unassisted_env(), render=False) for _ in range(n_eval_rollouts)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'aristotle_pilot_policy_unassisted_rollouts.pkl'), 'wb') as f:\n",
    "  pickle.dump(unassisted_rollouts, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'aristotle_pilot_policy_unassisted_rollouts.pkl'), 'rb') as f:\n",
    "  unassisted_rollouts = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ideal_rollouts = [run_ep(*make_ideal_env(), render=False) for _ in range(n_eval_rollouts)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'aristotle_pilot_policy_ideal_rollouts.pkl'), 'wb') as f:\n",
    "  pickle.dump(ideal_rollouts, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'aristotle_pilot_policy_ideal_rollouts.pkl'), 'rb') as f:\n",
    "  ideal_rollouts = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rand_rollouts = [run_ep(*make_rand_env(), render=False) for _ in range(n_eval_rollouts)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'rand_pilot_policy_unassisted_rollouts.pkl'), 'wb') as f:\n",
    "  pickle.dump(rand_rollouts, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'rand_pilot_policy_unassisted_rollouts.pkl'), 'rb') as f:\n",
    "  rand_rollouts = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "unassisted_rew = [sum(x[2] for x in r) for r in unassisted_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ideal_rew = [sum(x[2] for x in r) for r in ideal_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "assisted_rew = [sum(x[2] for x in r) for r in assisted_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rand_rew = [sum(x[2] for x in r) for r in rand_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "np.mean(unassisted_rew)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(assisted_rew)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(ideal_rew)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(rand_rew)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "unassisted_succ = [1 if is_succ(r) else 0 for r in unassisted_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ideal_succ = [1 if is_succ(r) else 0 for r in ideal_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "assisted_succ = [1 if is_succ(r) else 0 for r in assisted_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rand_succ = [1 if is_succ(r) else 0 for r in rand_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(unassisted_succ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.mean(ideal_succ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(assisted_succ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(rand_succ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "unassisted_crash = [1 if is_crash(r) else 0 for r in unassisted_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ideal_crash = [1 if is_crash(r) else 0 for r in ideal_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "assisted_crash = [1 if is_crash(r) else 0 for r in assisted_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rand_crash = [1 if is_crash(r) else 0 for r in rand_rollouts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(unassisted_crash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(assisted_crash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(ideal_crash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(rand_crash)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "viz master logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data_dir_idx = 1\n",
    "data_dir = os.path.join(\n",
    "  'data', '1.%d-tabular-ime' % data_dir_idx, \n",
    "  'univ_scramble' if univ_scramble else 'scramble')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'master_train_logs.pkl'), 'rb') as f:\n",
    "  list_n_train_tasks, master_train_logs = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "other_data_dir = os.path.join(\n",
    "  'data', '1.%d-tabular-ime' % data_dir_idx, \n",
    "  'scramble' if univ_scramble else 'univ_scramble')\n",
    "with open(os.path.join(other_data_dir, 'master_train_logs.pkl'), 'rb') as f:\n",
    "  other_list_n_train_tasks, other_master_train_logs = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "smooth_win = 100\n",
    "def moving_avg(d, n=smooth_win):\n",
    "  s = np.concatenate((np.zeros(1), np.cumsum(d).astype(float)))\n",
    "  return (s[n:] - s[:-n]) / n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "traj_col_means = lambda x: np.nanmean(x, axis=0)\n",
    "traj_col_stderrs = lambda x: np.nanstd(x, axis=0) / np.sqrt(\n",
    "  np.count_nonzero(~np.isnan(x), axis=0))\n",
    "r_mins = lambda x: traj_col_means(x) - traj_col_stderrs(x)\n",
    "r_maxs = lambda x: traj_col_means(x) + traj_col_stderrs(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def plot_fill(R, color, label):\n",
    "  x = range(R.shape[1] - (smooth_win - 1))\n",
    "  y1 = moving_avg(r_mins(R), n=smooth_win)\n",
    "  y2 = moving_avg(r_maxs(R), n=smooth_win)\n",
    "  plt.fill_between(\n",
    "    x, y1, y2, where=y2 >= y1, interpolate=True, facecolor=color, alpha=0.5)\n",
    "  plt.plot(moving_avg(traj_col_means(R), n=smooth_win), color=color, label=label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "local_label = 'Locally Scrambled Actions'\n",
    "global_label = 'Globally Scrambled Actions'\n",
    "label = global_label if univ_scramble else local_label\n",
    "other_label = local_label if univ_scramble else global_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def succ_vs_ntasks_of_logs(master_train_logs):\n",
    "  R = [[] for _ in master_train_logs]\n",
    "  for i, train_logs_set in enumerate(master_train_logs):\n",
    "    succs = []\n",
    "    for train_logs in train_logs_set:\n",
    "      succs.append(np.mean(train_logs['assisted_succ_evals'][-100:]))\n",
    "    R[i] = succs\n",
    "  R = np.array(R).T\n",
    "  return R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "R = succ_vs_ntasks_of_logs(master_train_logs)\n",
    "other_R = succ_vs_ntasks_of_logs(other_master_train_logs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "int_label = 'No Prior' if data_dir_idx == 2 else 'Action Intent Prior'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "old_R = copy(R)\n",
    "old_int_label = copy(int_label)\n",
    "old_list_n_train_tasks = copy(list_n_train_tasks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'font.size': 20})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Number of Training Tasks')\n",
    "plt.ylabel('Internal Dynamics Accuracy')\n",
    "plt.title(global_label if univ_scramble else local_label)\n",
    "\n",
    "plt.errorbar(\n",
    "  x=list_n_train_tasks, y=traj_col_means(R), \n",
    "  yerr=traj_col_stderrs(R), marker='o', label=int_label, \n",
    "  color='orange', capsize=10)\n",
    "plt.errorbar(\n",
    "  x=old_list_n_train_tasks, y=traj_col_means(old_R), \n",
    "  yerr=traj_col_stderrs(old_R), marker='o', label=old_int_label, \n",
    "  color='teal', capsize=10, linestyle=':')\n",
    "\n",
    "plt.axhline(y=(1/n_obs_dim), linestyle='--', color='gray', label='Random')\n",
    "\n",
    "#plt.ylim([-0.05, None])#1.05])\n",
    "#plt.yticks(np.arange(0, 0.8, 0.2))\n",
    "plt.legend(loc='best', framealpha=0.5)\n",
    "plt.savefig(os.path.join(data_dir, 'acc-vs-intent-type-%s.pdf' % ('global' if univ_scramble else 'local')), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Number of Training Tasks')\n",
    "plt.ylabel('Internal Dynamics Accuracy')\n",
    "plt.title(global_label if univ_scramble else local_label)\n",
    "\n",
    "plt.errorbar(\n",
    "  x=list_n_train_tasks, y=traj_col_means(R), \n",
    "  yerr=traj_col_stderrs(R), marker='o', label=int_label, color='orange',\n",
    "  capsize=10)\n",
    "plt.errorbar(\n",
    "  x=old_list_n_train_tasks, y=traj_col_means(old_R), \n",
    "  yerr=traj_col_stderrs(old_R), marker='o',label=old_int_label, color='teal',\n",
    "  capsize=10, linestyle=':')\n",
    "\n",
    "plt.axhline(y=(1/n_obs_dim), linestyle='--', color='gray', label='Random')\n",
    "\n",
    "plt.ylim([-0.05, None])#1.05])\n",
    "plt.yticks(np.arange(0, 0.8, 0.2))\n",
    "plt.legend(loc='lower right')\n",
    "plt.savefig(os.path.join(data_dir, 'acc-vs-intent-type-%s.pdf' % ('global' if univ_scramble else 'local')), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Number of Training Tasks')\n",
    "plt.ylabel('Assisted Success Rate')\n",
    "plt.title(global_label if univ_scramble else local_label)\n",
    "\n",
    "plt.errorbar(\n",
    "  x=list_n_train_tasks, y=traj_col_means(R), \n",
    "  yerr=traj_col_stderrs(R), marker='o', label=int_label, \n",
    "  color='orange', capsize=10)\n",
    "plt.errorbar(\n",
    "  x=old_list_n_train_tasks, y=traj_col_means(old_R), \n",
    "  yerr=traj_col_stderrs(old_R), marker='o', label=old_int_label, \n",
    "  color='teal', capsize=10, linestyle=':')\n",
    "\n",
    "plt.axhline(y=np.mean(rand_succ), linestyle='--', color='gray', label='Random Policy')\n",
    "\n",
    "#plt.ylim([-0.05, None])#1.05])\n",
    "#plt.yticks(np.arange(0, 0.8, 0.2))\n",
    "plt.legend(loc='best', framealpha=0.5)\n",
    "plt.savefig(os.path.join(data_dir, 'succ-vs-intent-type-%s.pdf' % ('global' if univ_scramble else 'local')), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Number of Training Tasks')\n",
    "plt.ylabel('Assisted Success Rate')\n",
    "plt.title(global_label if univ_scramble else local_label)\n",
    "\n",
    "plt.errorbar(\n",
    "  x=list_n_train_tasks, y=traj_col_means(R), \n",
    "  yerr=traj_col_stderrs(R), marker='o', label=int_label, \n",
    "  color='orange', capsize=10)\n",
    "plt.errorbar(\n",
    "  x=old_list_n_train_tasks, y=traj_col_means(old_R), \n",
    "  yerr=traj_col_stderrs(old_R), marker='o', label=old_int_label, \n",
    "  color='teal', capsize=10, linestyle=':')\n",
    "\n",
    "plt.axhline(y=np.mean(rand_succ), linestyle='--', color='gray', label='Random Policy')\n",
    "\n",
    "#plt.ylim([-0.05, None])#1.05])\n",
    "#plt.yticks(np.arange(0, 0.8, 0.2))\n",
    "plt.legend(loc='best', framealpha=0.5)\n",
    "plt.savefig(os.path.join(data_dir, 'succ-vs-intent-type-%s.pdf' % ('global' if univ_scramble else 'local')), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": 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
}