{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "01de6b24",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "44e67151",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from accelerate import Accelerator\n",
    "\n",
    "from promptrl.utils import Logger\n",
    "from promptrl.task import make_env\n",
    "from scripts.load_args import CaptionArgs, VP2AuxArgs\n",
    "from scripts.load import load_checkpoint_agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3feddc7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "plt.rcParams['figure.dpi'] = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5ff22d9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.set_device(15)\n",
    "accelerator = Accelerator()\n",
    "logger = Logger(None, use_wandb=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b39cd47a",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = VP2AuxArgs()\n",
    "args.load_checkpoint = '../' + args.load_checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "44865f6c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2022-09-24 16:13:42,789 INFO]: Loading language model...\n",
      "[2022-09-24 16:13:47,162 INFO]: Loading prompt model...\n"
     ]
    }
   ],
   "source": [
    "agent = load_checkpoint_agent(accelerator, logger, args)\n",
    "agent.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e7aaa27d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': <function _get_alf_env_loader.<locals>._initializer at 0x7f07815439d0>, 'ood': <function _get_alf_env_loader.<locals>._initializer at 0x7f0781543940>}\n"
     ]
    }
   ],
   "source": [
    "envs = make_env('alf-pick-place', 'img', agent.tokenizer)\n",
    "print(envs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "294a2844",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initialize AlfredThorEnv...\n",
      "Overall we have 46 games...\n",
      "Evaluating with 46 games\n",
      "Overall we have 46 games...\n",
      "Evaluating with 46 games\n",
      "Found path: /home/anthony_liu/.ai2thor/releases/thor-201909061227-Linux64/thor-201909061227-Linux64\n",
      "Mono path[0] = '/home/anthony_liu/.ai2thor/releases/thor-201909061227-Linux64/thor-201909061227-Linux64_Data/Managed'\n",
      "Mono config path = '/home/anthony_liu/.ai2thor/releases/thor-201909061227-Linux64/thor-201909061227-Linux64_Data/Mono/etc'\n",
      "Unable to preload the following plugins:\n",
      "\tScreenSelector.so\n",
      "Display 0 'NVIDIA VGX  32\"': 1024x768 (primary device).\n",
      "Display 1 'NVIDIA VGX  32\"': 1024x768 (secondary device).\n",
      "Display 2 'NVIDIA VGX  32\"': 1024x768 (secondary device).\n",
      "Display 3 'NVIDIA VGX  32\"': 1024x768 (secondary device).\n",
      "Display 4 'NVIDIA VGX  32\"': 1024x768 (secondary device).\n",
      "Display 5 'NVIDIA VGX  32\"': 1024x768 (secondary device).\n",
      "Display 6 'NVIDIA VGX  32\"': 1024x768 (secondary device).\n",
      "Display 7 'NVIDIA VGX  32\"': 1024x768 (secondary device).\n",
      "Logging to /home/anthony_liu/.config/unity3d/Allen Institute for Artificial Intelligence/AI2-Thor/Player.log\n",
      "ThorEnv started.\n"
     ]
    }
   ],
   "source": [
    "env_id = envs['id']()\n",
    "#env_ood = envs['ood']()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9cdc7ad1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _process_first_obs(obs, infos):\n",
    "    _, _, goal = infos['feedback'][0].split('\\n\\n')\n",
    "    return obs, goal\n",
    "\n",
    "def _process_obs(obs, info):\n",
    "    return obs\n",
    "\n",
    "def _log_eval_obs(obs, infos):\n",
    "    return infos[\"feedback\"][0]\n",
    "\n",
    "def _decode_action(action):\n",
    "    gen_action_s = agent.tokenizer.decode(action[0], skip_special_tokens=True)\n",
    "    gen_action_clean = gen_action_s.split('[', 2)[1]\n",
    "    gen_action_clean = gen_action_clean.replace('SEP]', '').strip()\n",
    "    return gen_action_clean\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0174916d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_grid(obs, max_width=4):\n",
    "    N = obs.shape[0]\n",
    "    if 1 < N <= max_width:\n",
    "        fig, ax = plt.subplots(1, N)\n",
    "        for j in range(N):\n",
    "            ax[j].set_axis_off()\n",
    "            ax[j].imshow(obs[j])\n",
    "    elif N > max_width:\n",
    "        fig, ax = plt.subplots(1 + N // max_width, max_width)\n",
    "        for h in range(1 + N // max_width):\n",
    "            for w in range(max_width):\n",
    "                i = h * max_width + w\n",
    "                ax[h, w].set_axis_off()\n",
    "                ax[h, w].set_aspect('equal')\n",
    "                if i >= N:\n",
    "                    continue\n",
    "                ax[h, w].imshow(obs[i])\n",
    "        fig.subplots_adjust(wspace=0, hspace=0)\n",
    "    else:\n",
    "        fig, ax = plt.subplots(1, N)\n",
    "        ax.set_axis_off()\n",
    "        ax.imshow(obs[0])\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "95afa573",
   "metadata": {},
   "outputs": [],
   "source": [
    "caption_goal = 'Your task is to: caption the following observation'\n",
    "caption_goal_ids = agent.tokenizer(caption_goal, return_tensors='pt').input_ids.squeeze(0).to(accelerator.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5de47b98",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_loop(env, agent):\n",
    "    obs, infos = env.reset()\n",
    "    agent.reset()\n",
    "    init_obs, goal = _process_first_obs(obs, infos)\n",
    "    goal_ids = agent.tokenizer([goal], return_tensors='pt').input_ids.to(accelerator.device)\n",
    "    agent.observe_first(init_obs, goal_ids, infos)\n",
    "    finished = False\n",
    "    sample_obs = [init_obs[0]]\n",
    "    sample_actions = []\n",
    "    oracle_obs = []\n",
    "    while not finished:\n",
    "        # generate action\n",
    "        gen_action = agent.sample()\n",
    "        gen_action_clean = _decode_action(gen_action)\n",
    "        show_grid(sample_obs[-1])\n",
    "        caption_ids = agent.direct_sample([sample_obs[-1]], [caption_goal_ids], {}, task_id=1, max_length=50)\n",
    "        caption = _decode_action(caption_ids)\n",
    "        print('Predicted caption:', caption)\n",
    "        print('Predicted action:', gen_action_clean)\n",
    "        # env step\n",
    "        obs, scores, dones, infos = env.step([gen_action_clean])\n",
    "        finished = dones[0]\n",
    "        sample_actions.append(gen_action_clean)\n",
    "        lang_obs = _log_eval_obs(obs, infos)\n",
    "        oracle_obs.append(lang_obs)\n",
    "        success = float(infos['won'][0])\n",
    "\n",
    "        # update agent\n",
    "        obs = _process_obs(obs, infos)\n",
    "        sample_obs.append(obs[0])\n",
    "        agent.observe(obs, goal_ids, gen_action_clean, infos)\n",
    "\n",
    "    return success, goal, sample_obs, sample_actions, oracle_obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "686b42d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Resetting ThorEnv\n",
      "Task: put some cellphone on bed\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 68240 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 20 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the bed, you see a laptop, a book, a cd, a mug, a creditcard, a cellphone, a cd, and a creditcard.\n",
      "Predicted action: look\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83000 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the shelf, you see a mug, a book, a mug, and a keychain.\n",
      "Predicted action: go to shelf\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83000 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the shelf, you see a mug, a book, a mug, and a keychain.\n",
      "Predicted action: go to shelf\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83000 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the shelf, you see a mug, a book, a mug, and a keychain.\n",
      "Predicted action: go to shelf\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83000 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the shelf, you see a mug, a book, a mug, and a keychain.\n",
      "Predicted action: go to shelf\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83000 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the shelf, you see a mug, a book, a mug, and a keychain.\n",
      "Predicted action: go to shelf\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 78292 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the shelf, you see a mug.\n",
      "Predicted action: go to shelf\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 78292 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the shelf, you see a mug.\n",
      "Predicted action: go to shelf\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 80908 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the shelf, you see a mug.\n",
      "Predicted action: go to drawer\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 79080 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: The drawer is open. In it, you see a desklamp, a alarmclock, a cd, a creditcard, a pen, and a creditcard.\n",
      "Predicted action: open drawer\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 80032 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: The drawer is open. In it, you see a desklamp, a alarmclock, a cd, a creditcard, and a pen.\n",
      "Predicted action: take cellphone from drawer\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 78884 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: The drawer is open. In it, you see a desklamp, a alarmclock, a cd, a creditcard, and a pen.\n",
      "Predicted action: close drawer\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 78444 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: The drawer is open. In it, you see a desklamp, a alarmclock, a cd, a creditcard, and a pen.\n",
      "Predicted action: go to bed\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 101372 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted caption: On the bed, you see a pillow, a laptop, a laptop, and a pillow.\n",
      "Predicted action: put cellphone in/on bed\n"
     ]
    }
   ],
   "source": [
    "success, goal, sample_obs, sample_actions, oracle_obs = eval_loop(env_id, agent)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
