view_demos.ipynb

{
 "metadata": {
  "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.6.13"
  },
  "orig_nbformat": 4,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.6.13 64-bit ('babyai_wandb': conda)"
  },
  "interpreter": {
   "hash": "6cd421c9139f6051ec65a59284aaabcb657bde1ea04070b3d624f966b5bd0c2d"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import gym\n",
    "from babyai.levels.levelgen import *"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "register_levels('bonus_levels', globals())"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "\n",
    "env = gym.make('BabyAI-UnlockRGB-v0')\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "from array2gif import write_gif"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "import babyai.utils as utils"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "import numpy"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "d_path = \"/data/graceduansu/demos/BabyAI-UnlockRGB-v0\""
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "source": [
    "demos_path = utils.get_demos_path(d_path, None, None, valid=False)\n",
    "train_demos = utils.load_demos(demos_path)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "batch = [train_demos[i] for i in range(len(train_demos))]"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "source": [
    "demo_list = utils.demos.transform_demos(batch)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "source": [
    "print(demo_list[0])"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [4, 1, 2],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [5, 1, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]]], dtype=uint8), 'direction': 0, 'mission': 'open the green door'}, <Actions.forward: 2>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [4, 1, 2],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [5, 1, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]]], dtype=uint8), 'direction': 0, 'mission': 'open the green door'}, <Actions.right: 1>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 1, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]]], dtype=uint8), 'direction': 1, 'mission': 'open the green door'}, <Actions.forward: 2>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 1, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]]], dtype=uint8), 'direction': 1, 'mission': 'open the green door'}, <Actions.left: 0>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [4, 1, 2],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [5, 1, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]]], dtype=uint8), 'direction': 0, 'mission': 'open the green door'}, <Actions.forward: 2>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [4, 1, 2],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [5, 1, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]]], dtype=uint8), 'direction': 0, 'mission': 'open the green door'}, <Actions.pickup: 3>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [4, 1, 2],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 1, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]]], dtype=uint8), 'direction': 0, 'mission': 'open the green door'}, <Actions.forward: 2>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [4, 1, 2],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [5, 1, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]]], dtype=uint8), 'direction': 0, 'mission': 'open the green door'}, <Actions.left: 0>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 1, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [4, 1, 2],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]]], dtype=uint8), 'direction': 3, 'mission': 'open the green door'}, <Actions.forward: 2>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0],\n",
      "        [5, 1, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [4, 1, 2],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]]], dtype=uint8), 'direction': 3, 'mission': 'open the green door'}, <Actions.forward: 2>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [5, 1, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [4, 1, 2]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]]], dtype=uint8), 'direction': 3, 'mission': 'open the green door'}, <Actions.right: 1>, False), ({'image': array([[[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [2, 5, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [4, 1, 2],\n",
      "        [5, 1, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0]],\n",
      "\n",
      "       [[0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [0, 0, 0],\n",
      "        [2, 5, 0],\n",
      "        [1, 0, 0]]], dtype=uint8), 'direction': 0, 'mission': 'open the green door'}, <Actions.toggle: 5>, True)]\n"
     ]
    }
   ],
   "metadata": {
    "tags": [
     "outputPrepend"
    ]
   }
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "source": [
    "import matplotlib.pyplot as plt"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "source": [
    "def get_demo_gifs(demo_list, location):\n",
    "\n",
    "    for d in range(len(demo_list)):\n",
    "        frames = []\n",
    "        last_img = None\n",
    "        for i in range(len(demo_list[d])):\n",
    "            ex_obs = demo_list[d][i][0]['image']\n",
    "            last_img = env.get_obs_render(\n",
    "                    ex_obs,\n",
    "                    tile_size=8\n",
    "                )\n",
    "            dir = demo_list[d][i][0]['direction']\n",
    "\n",
    "            if dir == 0:\n",
    "                # pointing right\n",
    "                last_img = numpy.rot90(last_img, axes=(1,0))\n",
    "            elif dir == 1:\n",
    "                # pointing down\n",
    "                last_img = numpy.rot90(last_img, k=2, axes=(1,0))\n",
    "            elif dir == 2:\n",
    "                # pointing left\n",
    "                last_img = numpy.rot90(last_img, axes=(0,1))\n",
    "            # else it's pointing up - leave alone\n",
    "\n",
    "            frames.append(numpy.moveaxis(last_img, 2, 0))\n",
    "\n",
    "        for i in range(10):\n",
    "            frames.append(numpy.moveaxis(last_img, 2, 0))\n",
    "\n",
    "        mission = demo_list[d][1][0]['mission']\n",
    "        print(\"Saving gif \" + str(d) + \"... \", end=\"\")\n",
    "        write_gif(numpy.array(frames), location + \"/demo\" + str(d) + \"_\" + mission + \"_\"  + \".gif\", fps=1/0.3)\n",
    "        print(\"Done.\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "source": [
    "# gif for demo index 0\n",
    "# for i in range(len(demo_list[0])):\n",
    "#     env.render('human')\n",
    "#     frames.append(numpy.moveaxis(env.render(\"rgb_array\"), 2, 0))\n",
    "#     action = demo_list[0][i][1]\n",
    "#     obs, reward, done, _ = env.step(action)\n",
    "\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## unlock"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "d_path = '/data/graceduansu/demos/BabyAI-Unlock-v0'"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "source": [
    "demos_path = utils.get_demos_path(d_path, None, None, valid=False)\n",
    "train_demos = utils.load_demos(demos_path)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "source": [
    "len(train_demos)"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "1000000"
      ]
     },
     "metadata": {},
     "execution_count": 43
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "source": [
    "batch = [train_demos[i] for i in range(0, 100)]"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "source": [
    "demo_list = utils.demos.transform_demos(batch)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "source": [
    "get_demo_gifs(demo_list, '/data/graceduansu/demos/1UnlockRGB')"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Saving gif 0... Done.\n",
      "Saving gif 1... Done.\n",
      "Saving gif 2... Done.\n"
     ]
    }
   ],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ]
}