{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "01de6b24",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "c57c3654",
   "metadata": {},
   "outputs": [],
   "source": [
    "import PIL\n",
    "import IPython.display\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "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_vhome_envs\n",
    "from scripts.load_args import VHomeCaptionArgs, VHomeVP2Args\n",
    "from scripts.load import load_checkpoint_agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3feddc7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "plt.rcParams['figure.dpi'] = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": 12,
   "id": "b39cd47a",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = VHomeVP2Args()\n",
    "args.load_checkpoint = '../' + args.load_checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "44865f6c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2022-11-02 19:19:19,427 INFO]: Loading language model...\n",
      "[2022-11-02 19:19:24,612 INFO]: Loading prompt model...\n"
     ]
    }
   ],
   "source": [
    "agent = load_checkpoint_agent(accelerator, logger, args)\n",
    "agent.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e7aaa27d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': <function make_vhome_loader.<locals>._initializer at 0x7f98b0f24200>}\n"
     ]
    }
   ],
   "source": [
    "envs = make_vhome_envs('virtualhome-id', 0)\n",
    "print(envs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "294a2844",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading inital scene graph from /z/anthony/prompt-subtasks/promptrl/envs/virtualhome/LID/data/test_init_env/InDistributation.p\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2022-11-02 19:21:33,334 WARNING]: Retrying (Retry(total=4, connect=4, read=5, redirect=None, status=None)) after connection broken by 'NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7f9479270890>: Failed to establish a new connection: [Errno 111] Connection refused')': /\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['/z/anthony/prompt-subtasks/promptrl/envs/virtualhome/LID/virtualhome/simulation/unity_simulator/v2.2.5/linux_exec.v2.2.5_beta.x86_64', '-batchmode', '-http-port=8904', '-logFile /z/anthony/prompt-subtasks/nb/Player_8904.log']\n",
      "Getting connection...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2022-11-02 19:21:37,338 WARNING]: Retrying (Retry(total=3, connect=3, read=5, redirect=None, status=None)) after connection broken by 'NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7f9479272550>: Failed to establish a new connection: [Errno 111] Connection refused')': /\n",
      "ALSA lib confmisc.c:767:(parse_card) cannot find card '0'\n",
      "ALSA lib conf.c:4528:(_snd_config_evaluate) function snd_func_card_driver returned error: No such file or directory\n",
      "ALSA lib confmisc.c:392:(snd_func_concat) error evaluating strings\n",
      "ALSA lib conf.c:4528:(_snd_config_evaluate) function snd_func_concat returned error: No such file or directory\n",
      "ALSA lib confmisc.c:1246:(snd_func_refer) error evaluating name\n",
      "ALSA lib conf.c:4528:(_snd_config_evaluate) function snd_func_refer returned error: No such file or directory\n",
      "ALSA lib conf.c:5007:(snd_config_expand) Evaluate error: No such file or directory\n",
      "ALSA lib pcm.c:2495:(snd_pcm_open_noupdate) Unknown PCM default\n",
      "ALSA lib confmisc.c:767:(parse_card) cannot find card '0'\n",
      "ALSA lib conf.c:4528:(_snd_config_evaluate) function snd_func_card_driver returned error: No such file or directory\n",
      "ALSA lib confmisc.c:392:(snd_func_concat) error evaluating strings\n",
      "ALSA lib conf.c:4528:(_snd_config_evaluate) function snd_func_concat returned error: No such file or directory\n",
      "ALSA lib confmisc.c:1246:(snd_func_refer) error evaluating name\n",
      "ALSA lib conf.c:4528:(_snd_config_evaluate) function snd_func_refer returned error: No such file or directory\n",
      "ALSA lib conf.c:5007:(snd_config_expand) Evaluate error: No such file or directory\n",
      "ALSA lib pcm.c:2495:(snd_pcm_open_noupdate) Unknown PCM default\n",
      "[2022-11-02 19:21:45,346 WARNING]: Retrying (Retry(total=2, connect=2, read=5, redirect=None, status=None)) after connection broken by 'NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7f9479272c90>: Failed to establish a new connection: [Errno 111] Connection refused')': /\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VHome Env id loaded. 199 tasks. (1999 total)\n"
     ]
    }
   ],
   "source": [
    "env_id = envs['id']()\n",
    "#env_ood = envs['ood']()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9cdc7ad1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _process_first_obs(obs, infos):\n",
    "    raw_goal = infos['goal'][0]\n",
    "    assert isinstance(raw_goal, list)\n",
    "    #assert len(raw_goal) > 0\n",
    "    if len(raw_goal) == 1:\n",
    "        goal = f'Your task is to: {raw_goal[0]}.'\n",
    "    else:\n",
    "        goal = 'Your task is to: ' + ', '.join(raw_goal[:-1]) + ', and ' + raw_goal[-1] + '.'\n",
    "    return obs, goal, []\n",
    "\n",
    "def _process_obs(obs, info):\n",
    "    return obs\n",
    "\n",
    "def _log_eval_obs(obs, infos):\n",
    "    visible_interactable = infos['interactable_objects'][0] & set(infos['vis_objects_raw'][0][1].values())\n",
    "    visible_interactable.discard('character')\n",
    "    objs_list = ['a ' + o for o in sorted(visible_interactable)]\n",
    "    caption = 'In front of you, you see '\n",
    "    if len(objs_list) == 0:\n",
    "        caption += 'nothing.'\n",
    "    elif len(objs_list) == 1:\n",
    "        caption += objs_list[0] + '.'\n",
    "    else:\n",
    "        caption += ', '.join(objs_list[:-1]) + ', and ' + objs_list[-1] + '.'\n",
    "    return caption\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": 21,
   "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": 49,
   "id": "eac38098",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_grid_img(images_old, nrows=None, maxncol=10):\n",
    "    if nrows is None:\n",
    "        nrows = int(np.ceil(len(images_old) / maxncol))\n",
    "    images = [x for x in images_old]\n",
    "    h, w, _ = images[0].shape\n",
    "    ncols = int((len(images)+nrows-1)/nrows)\n",
    "    missing = ncols - (len(images)%ncols)\n",
    "    for m in range(missing):\n",
    "        images.append(np.zeros((h, w, 3)).astype(np.uint8))\n",
    "    img_final = []\n",
    "    for it_r in range(nrows):\n",
    "        init_ind = it_r * ncols \n",
    "        end_ind = init_ind + ncols\n",
    "        images_take = [images[it] for it in range(init_ind, end_ind)]\n",
    "        img_final.append(np.concatenate(images_take, 1))\n",
    "    img_final = np.concatenate(img_final, 0)\n",
    "    img_final = PIL.Image.fromarray(img_final[:,:,::-1])\n",
    "    return img_final"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "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": 33,
   "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",
    "    print(goal)\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=0, 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": 59,
   "id": "57ffb6cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "env_id.part_task_id = 1# set partition task id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "686b42d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your task is to: put 2 milk on the kitchentable.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 236264 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to kitchen\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 189828 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to kitchencabinet\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 136684 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: open kitchencabinet\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 138736 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: grab milk\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 138736 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to kitchen\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 148104 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to kitchentable\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 150156 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: put milk on kitchentable\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 150156 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to kitchencabinet\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 132092 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: open kitchencabinet\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 138148 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: grab milk\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 137808 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: put milk on kitchentable\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 138036 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to kitchencabinet\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 125672 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: open kitchencabinet\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 122412 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: grab milk\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 122412 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: put milk on kitchentable\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 122384 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 125248 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: open dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 133944 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: grab plate\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 133944 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: put plate on kitchentable\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 133944 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to livingroom\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to plate\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: grab plate\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: put plate on kitchentable\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on switch\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: walk to switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126420 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126452 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted action: switch on dishwasher\n"
     ]
    }
   ],
   "source": [
    "success, goal, sample_obs, sample_actions, oracle_obs = eval_loop(env_id, agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "34686fd5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1278660 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1344x2016>"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "display_grid_img([o[0] for o in sample_obs], maxncol=6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fbb8e93",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
