demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo of BISCUIT-NF on the Embodied AI iTHOR dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates the application and evaluation of BISCUIT-NF on the Embodied AI iTHOR dataset. We show how BISCUIT-NF allows for simulating interventions in latent space, and how it models a causal representation of the environment. Further, we reproduce the visualizations of the interaction maps in the paper. The notebook can be run locally or on Google Colab.\n",
    "\n",
    "> **Note:** This notebook is intended to be run on a GPU. If you are running this notebook on Google Colab, make sure to enable GPU acceleration by going to `Runtime -> Change runtime type -> Hardware accelerator -> GPU`. Any GPU, such as a T4, P4, or P100, should work."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "In order to run the notebook, you need to install the dependencies. If you are running the notebook on Google Colab, you can skip this step. Otherwise, install the conda environment `biscuit` following the instructions in the [README](../README.md). \n",
    "\n",
    "Further, to run the notebook on Google Colab, we need to clone the BISCUIT repository to the workspace. This is done in the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# If you are running locally, the cloning is skipped\n",
    "if not os.path.isdir('models/'):\n",
    "    print(\"Cloning the repository...\")\n",
    "    !git clone https://github.com/phlippe/BISCUIT\n",
    "    sys.path.append('BISCUIT/')\n",
    "    # On Google Colab, we need to install the dependencies\n",
    "    !pip install --quiet pytorch-lightning>=2.0 optuna ai2thor gdown\n",
    "    !pip install --quiet --upgrade gdown"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then import our needed packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Any, Dict, Optional\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import hsv_to_rgb\n",
    "from matplotlib.patches import Patch\n",
    "from tqdm.auto import tqdm\n",
    "from glob import glob\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "from models.biscuit_nf import BISCUITNF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, we will make use of the GPU if available:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we need to download the pretrained model and the small test sequence on which we want to perform the evaluation and visualization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Paths where to download the model and data to (relative to the current directory)\n",
    "# Overall download size is about 50MB\n",
    "CHECKPOINT_FILE = 'pretrained_models/BISCUITNF_iTHOR/BISCUITNF_40l_64hid.ckpt'\n",
    "DATA_FOLDER = 'demo_data/ithor/'\n",
    "\n",
    "# Download the model\n",
    "if not os.path.exists(CHECKPOINT_FILE):\n",
    "    print(\"Downloading the model...\")\n",
    "    !gdown 1GKFeEcoFFqLY3uMvrefS70THy6EcogqK\n",
    "    !tar -xzvf pretrained_models.tar.gz\n",
    "    !rm pretrained_models.tar.gz\n",
    "\n",
    "# Download the data\n",
    "if not os.path.exists(DATA_FOLDER):\n",
    "    print(\"Downloading the data...\")\n",
    "    !gdown 1PEqpxrHdx8-RoApAPeeu8OQqK3qmqKqH\n",
    "    !tar -xzvf demo_data.tar.gz\n",
    "    !rm demo_data.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading data\n",
    "\n",
    "As a first step to verify that the data was downloaded correctly, we can load the test sequence and visualize it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple function for plotting images\n",
    "def show_img(imgs: Any,\n",
    "             figure_title: str = None,\n",
    "             titles: str = None):\n",
    "    if not isinstance(imgs, list):\n",
    "        imgs = [imgs]\n",
    "    fig, axes = plt.subplots(1, len(imgs), figsize=(len(imgs) * 3.5, 3.5))\n",
    "    if len(imgs) == 1:\n",
    "        axes = [axes]\n",
    "    for i, ax in enumerate(axes):\n",
    "        if len(imgs[i].shape) == 3 and imgs[i].shape[0] in [3,4]:\n",
    "            imgs[i] = imgs[i].permute(1, 2, 0)\n",
    "        if isinstance(imgs[i], torch.Tensor):\n",
    "            imgs[i] = imgs[i].detach().cpu().numpy()\n",
    "        ax.imshow(imgs[i])\n",
    "        ax.axis('off')\n",
    "        if titles is not None:\n",
    "            ax.set_title(titles[i], weight='bold')\n",
    "    if figure_title is not None:\n",
    "        fig.suptitle(figure_title, weight='bold', size=16)\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the frames and autoencoder encodings from the data\n",
    "seq_file = sorted(glob(DATA_FOLDER + '*.npz'))[0]\n",
    "encoding_file = seq_file.replace('.npz', '_encodings.npz')\n",
    "\n",
    "frames = torch.from_numpy(np.load(seq_file)['frames']).float().to(device) / 255.0\n",
    "encodings = torch.from_numpy(np.load(encoding_file)['encodings']).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_img([frames[90], frames[91], frames[92]], figure_title='Example Sequence')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see a sequence of three images, showing the kitchen environment of iTHOR. At the first time step, the agent is turning off the microwave, and at the second time step, the agent is turning off the front-right stove burner."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading model\n",
    "\n",
    "Next, we load the pretrained model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BISCUITNF.load_from_checkpoint(CHECKPOINT_FILE)\n",
    "model.to(device)\n",
    "model.freeze()\n",
    "_ = model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "During the visualizations, we will use the model to perform interventions in latent space. For this, we need to find an alignment between the latent space and the causal variables. We use our correlation evaluation for this and assign each latent variable to the causal variable with the highest correlation. We do this below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CAUSAL_VARIABLE_NAMES = [\n",
    "    'Cabinet',\n",
    "    'Egg',\n",
    "    'Microwave - Active',\n",
    "    'Microwave - Open',\n",
    "    'Plate',\n",
    "    'Stove Front-Right',\n",
    "    'Stove Front-Left',\n",
    "    'Stove Back-Left',\n",
    "    'Stove Back-Right',\n",
    "    'Toaster'\n",
    "]\n",
    "causal_to_latents = {\n",
    "    k: torch.where(model.last_target_assignment[:,idx] == 1)[0] for idx, k in enumerate(CAUSAL_VARIABLE_NAMES)\n",
    "}\n",
    "print('-' * 50)\n",
    "print('Latent Dimensions for Causal Variables:')\n",
    "print('-' * 50)\n",
    "for k in causal_to_latents:\n",
    "    print(f'{k}: {causal_to_latents[k].cpu().numpy()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Commonly, the Egg and the Plate with the potato have the most latent variables assigned to them since they are the most versatile objects in the environment. Note that not all latent variables necessarily contain information, and those without information are arbitrarily assigned to any causal variable here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performing interventions in the latent space\n",
    "\n",
    "With the model and dataset loaded, we are ready to analyse the learned causal representation. First, we reproduce the triplet prediction results, i.e. the interventions in latent space, from the paper. Note that the provided pretrained model slightly differs from the model used in the paper, and thus the results may differ slightly as well.\n",
    "\n",
    "We first define a helper function to perform the intervention in latent space. The function takes (1) the input image on which we want to perform the interventions, (2) the image from which we want to take the intervention values, and (3) which latent/causal variables we want to intervene on. The first step is to encode both images into the learned causal representation. On the representation of the first input image, we replace the values of the latent variables corresponding to the specified causal variables in (3) with the values of the second input image. Afterwards, we simply reverse the flow to map the new latent vector back into the autoencoder representation space, and decode the image. The function then visualizes the image with the intervention applied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def triplet_prediction(\n",
    "        image1_idx: int,\n",
    "        image2_idx: int,\n",
    "        vars_from_image2: List[str]):\n",
    "    assert image1_idx >= 0 and image2_idx >= 0 and image1_idx < len(frames) and image2_idx < len(frames), \\\n",
    "        f'Invalid image indices {image1_idx}, {image2_idx}. Must be in [0, {len(frames)}).'\n",
    "    image1_frame = frames[image1_idx].permute(1, 2, 0).cpu().numpy()\n",
    "    image2_frame = frames[image2_idx].permute(1, 2, 0).cpu().numpy()\n",
    "    image_encodings = torch.stack([encodings[image1_idx], encodings[image2_idx]], dim=0)\n",
    "    image_latents, _ = model.flow.forward(image_encodings)\n",
    "    image1_latents = image_latents[0]\n",
    "    image2_latents = image_latents[1]\n",
    "    image3_latents = image1_latents.clone()\n",
    "    for var in vars_from_image2:\n",
    "        image3_latents[causal_to_latents[var]] = image2_latents[causal_to_latents[var]]\n",
    "    image3_encodings = model.flow.reverse(image3_latents[None])\n",
    "    image3_frame = model.autoencoder.decoder(image3_encodings)[0]\n",
    "    image3_frame = (image3_frame + 1.0) / 2.0\n",
    "    image3_frame = image3_frame.permute(1, 2, 0).cpu().numpy()\n",
    "    show_img([image1_frame, image2_frame, image3_frame],\n",
    "             figure_title=f'Manipulating {\", \".join(vars_from_image2)}',\n",
    "             titles=['Input Image 1', 'Input Image 2', 'Generated Output'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using this function, we can apply various interventions. For example, we can intervene on the Microwave to turn it off:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triplet_prediction(image1_idx=90, \n",
    "                   image2_idx=99, \n",
    "                   vars_from_image2=['Microwave - Active'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how the microwave is turned off in the generated output while all other variables are identical to the original input image. Similarly, we can turn on the front-left stove burner, even simultaneously with the microwave:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triplet_prediction(image1_idx=90, \n",
    "                   image2_idx=99, \n",
    "                   vars_from_image2=['Stove Front-Left', 'Microwave - Active'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notably, the Egg does not change its state although the stove is burning. This is because we only change the state of the stove, but force the state of the Egg to be identical to the original input image. This shows the strong disentanglement BISCUIT has learned, and its ability to perform novel interventions unseen during training.\n",
    "\n",
    "Various other interventions are possible, of which we show a few below. Feel free to try out your own interventions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triplet_prediction(image1_idx=90, \n",
    "                   image2_idx=99, \n",
    "                   vars_from_image2=['Stove Front-Left'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triplet_prediction(image1_idx=30, \n",
    "                   image2_idx=99, \n",
    "                   vars_from_image2=['Microwave - Open', 'Cabinet'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triplet_prediction(image1_idx=99, \n",
    "                   image2_idx=42, \n",
    "                   vars_from_image2=['Microwave - Open', 'Stove Front-Right'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the results of the paper, we discussed how the Egg and Plate are the most difficult causal variables to model in the environment. Although BISCUIT was not able to perfectly disentangle them, its causal representation often allows for interventions on the Egg and Plate nonetheless. One example of such is shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triplet_prediction(image1_idx=30, \n",
    "                   image2_idx=90, \n",
    "                   vars_from_image2=['Microwave - Open', 'Egg'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing identified interaction variables\n",
    "\n",
    "Next, we visualize the interaction variables identified by BISCUIT. For this, we take as input an image, and evaluate the predicted interaction variables for various action positions on the image. In iTHOR, an action corresponds to a position in the image, and thus we visualize the interaction variables by evaluating them for a grid of action positions. We first define a helper function to evaluate the interaction variables for a given image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def create_interaction_grid(image_idx: int, \n",
    "                            batch_size: int = 1024):\n",
    "    assert image_idx >= 0 and image_idx < len(frames), \\\n",
    "        f'Invalid image index {image_idx}. Must be in [0, {len(frames)}).'\n",
    "    # Encode the image on which the action should be applied\n",
    "    prev_img_encoding = encodings[image_idx]\n",
    "    prev_img_latents = model.encode(prev_img_encoding[None])\n",
    "    # Create a grid of actions\n",
    "    resolution = frames.shape[-1]\n",
    "    x, y = torch.meshgrid(torch.linspace(-1, 1, steps=resolution, device=device),\n",
    "                          torch.linspace(-1, 1, steps=resolution, device=device),\n",
    "                          indexing='xy')\n",
    "    actions = torch.stack([x, y], dim=-1).flatten(0, 1)\n",
    "    # Predict the interaction variables for each action\n",
    "    outs = []\n",
    "    for i in tqdm(range(0, actions.shape[0], batch_size), leave=False):\n",
    "        batch_actions = actions[i:i+batch_size]\n",
    "        batch_prev_state = prev_img_latents.expand(batch_actions.shape[0], -1)\n",
    "        outs.append(model.prior_t1.get_interaction_quantization(batch_actions, \n",
    "                                                                prev_state=batch_prev_state,\n",
    "                                                                soft=True))\n",
    "    pred_intv = torch.cat(outs, dim=0)\n",
    "    pred_intv = pred_intv.unflatten(0, (x.shape[0], y.shape[0]))\n",
    "    return pred_intv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each interaction variable is predicted in the range [-1,1]. We visualize them by binarizing them, e.g. setting all values above 0 to 1 and all values below 0 to 0. We then visualize the binarized interaction variables for each causal variable on the image. For causal variables that have multiple latent variables, we only visualize one to reduce the clutter in the image. We do this below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_interaction_grid(image_idx: int, \n",
    "                               grid: torch.Tensor, \n",
    "                               latent_sub_index: Dict[str, int] = None, \n",
    "                               threshold: float = -0.2):\n",
    "    # Select the latents to visualize\n",
    "    if latent_sub_index is None:\n",
    "        latent_sub_index = {}\n",
    "    latents = [causal_to_latents[k][latent_sub_index.get(k, 0)].item() \n",
    "               for k in CAUSAL_VARIABLE_NAMES]\n",
    "    # Binarize interaction variables\n",
    "    binarized_grid = (grid.transpose(1, 0, 2) > threshold).astype(np.float32)\n",
    "    # Create the image\n",
    "    num_vars = len(latents)\n",
    "    hues = [hsv_to_rgb([i/num_vars*0.9, 1.0, 1.0]) for i in range(num_vars)]\n",
    "    img = np.zeros((binarized_grid.shape[0], binarized_grid.shape[1], 3), dtype=np.float32)\n",
    "    for i, l in enumerate(latents):\n",
    "        if binarized_grid[:, :, l].sum() > np.prod(binarized_grid.shape[:2])/2:\n",
    "            binarized_grid[:, :, l] = 1 - binarized_grid[:, :, l]\n",
    "        img += binarized_grid[:, :, l:l+1] * hues[i][None, None, :]\n",
    "    # Normalize image\n",
    "    counts_sum = binarized_grid[:,:,latents].sum(axis=-1, keepdims=True)\n",
    "    img = img / np.maximum(counts_sum, 1)\n",
    "    img += (counts_sum == 0) * np.array([[[0.9, 0.9, 0.9]]])\n",
    "    # Overlay with original image\n",
    "    orig_img = frames[image_idx].permute(1, 2, 0).cpu().numpy()\n",
    "    new_img = 0.5 * img + 0.5 * orig_img\n",
    "    # Plot\n",
    "    _, axes = plt.subplots(1, 3, figsize=(3 * 3.5, 3.5))\n",
    "    axes[0].imshow(orig_img)\n",
    "    axes[0].set_title('Original Image')\n",
    "    axes[1].imshow(img)\n",
    "    axes[1].set_title('Interaction Grid')\n",
    "    axes[2].imshow(new_img)\n",
    "    axes[2].set_title('Overlay')\n",
    "    axes[2].legend(handles=[Patch(color=hues[i], label=CAUSAL_VARIABLE_NAMES[i]) for i in range(num_vars)], \n",
    "                   bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0., title='$\\\\bf{Causal\\ Variables}$')\n",
    "    for ax in axes:\n",
    "        ax.axis('off')\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we use this function to visualize the interaction variables for some example images. Note that the `latent_sub_index` parameter specifies which latent variable to use for causal variables with multiple latent variables. This may require adjustment for new models. Further, for the Plate and Egg, different dimensions show slightly different behaviors. Feel free to try out your own images!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_sub_index = {\n",
    "    'Cabinet': 1,\n",
    "    'Egg': 0,\n",
    "    'Microwave - Active': 0,\n",
    "    'Microwave - Open': 0,\n",
    "    'Plate': 0,\n",
    "    'Stove Front-Right': 1,\n",
    "    'Stove Front-Left': 0,\n",
    "    'Stove Back-Left': 0,\n",
    "    'Stove Back-Right': 0,\n",
    "    'Toaster': 0\n",
    "}\n",
    "\n",
    "for image_idx in [25, 70, 90]:\n",
    "    grid = create_interaction_grid(image_idx=image_idx)\n",
    "    visualize_interaction_grid(image_idx=image_idx, \n",
    "                               grid=grid.cpu().numpy(),\n",
    "                               latent_sub_index=latent_sub_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The interaction variables both dependent on the position and the image. For example, the interaction variable for the Egg switches to the front-left stove knob in the last image, since turning on the stove will cause the egg to get cooked. Similarly, the interaction variable for the Plate changes position depending on where the Plate is in the current frame. This shows that BISCUIT has learned to identify the interaction variables for each causal variable, and that these interaction variables are dependent on the image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulating Actions in the Environment\n",
    "\n",
    "Finally, we show how BISCUIT can be used to simulate actions in the environment. For this, we first define a helper function to simulate an action in the environment. The function takes (1) the input image on which we want to perform the action, and (2) the action we want to perform. The function then encodes the input image into the latent space, and uses the learned transition prior to sample a next state based on the provided action. Finally, it decodes the latent vector back into the image space and visualizes the result. To be a bit more robust to potential reconstruction errors, we allow to pass the latents directly for long-term simulations. We do this below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def next_step_prediction(\n",
    "        image: torch.Tensor,\n",
    "        action: torch.Tensor,\n",
    "        latents: Optional[torch.Tensor] = None,\n",
    "        plot_images: bool = True):\n",
    "    if latents is None:\n",
    "        input_image = (image * 2.0) - 1.0\n",
    "        latents = model.autoencoder.encoder(input_image[None])\n",
    "        latents, _ = model.flow.forward(latents)\n",
    "    new_latents, _ = model.prior_t1.sample(latents, action[None], num_samples=1)\n",
    "    new_latents = new_latents.squeeze(1)\n",
    "    new_encodings = model.flow.reverse(new_latents)\n",
    "    new_image = model.autoencoder.decoder(new_encodings)[0]\n",
    "    new_image = (new_image + 1.0) / 2.0\n",
    "    if plot_images:\n",
    "        new_image_frame = new_image.permute(1, 2, 0).cpu().numpy()\n",
    "        old_image_frame = image.permute(1, 2, 0).cpu().numpy()\n",
    "        clicked_image_frame = np.copy(old_image_frame)\n",
    "        if action.ndim == 1:\n",
    "            action = action[None]\n",
    "        for i in range(action.shape[0]):\n",
    "            pixel_y = int((action[i, 0].item() + 1.0) / 2.0 * image.shape[-2])\n",
    "            pixel_x = int((action[i, 1].item() + 1.0) / 2.0 * image.shape[-1])\n",
    "            clicked_image_frame[max(0, pixel_y-5):pixel_y+6, \n",
    "                                max(0, pixel_x-5):pixel_x+6, \n",
    "                                :] = np.array([1.0, 0.0, 0.0])\n",
    "        show_img([old_image_frame, clicked_image_frame, new_image_frame],\n",
    "                figure_title=f'Performing action {(action if action.ndim == 1 else action[0]).squeeze().cpu().numpy()}',\n",
    "                titles=['Previous Frame', 'Click Location', 'New Sample'])\n",
    "    return new_image, new_latents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a first example, we can simulate four actions: \n",
    "* the agent turning on the front-left stove burner;\n",
    "* the agent opening the cabinet;\n",
    "* the agent taking the Plate from the counter;\n",
    "* the agent opening the microwave.\n",
    "\n",
    "Since the sampling is stochastic and the model is not perfect, the results may not be perfectly stable for many steps and differ across seeds. Feel free to try out your own actions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latents = None\n",
    "image = frames[0]\n",
    "pl.seed_everything(42)\n",
    "\n",
    "action_sequence = [\n",
    "    (0.82, 0.12),\n",
    "    (-0.5, -0.5),\n",
    "    (0.7, -0.3),\n",
    "    (-0.5, 0.2),\n",
    "]\n",
    "\n",
    "for action in action_sequence:\n",
    "    image, latents = next_step_prediction(image=image, \n",
    "                                          action=torch.tensor(action, device=device), \n",
    "                                          latents=latents,\n",
    "                                          plot_images=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note in the example above how the stove takes several frames to fully fire up. This shows how BISCUIT learned to model the dynamics of the environment, and how it can be used to simulate actions in the environment. \n",
    "\n",
    "### Multiple Actions at the Same Time\n",
    "\n",
    "Since each latent variable is sampled independently conditioned on the previous time step, we can also simulate multiple actions at the same time. For example, we can simulate the agent turning on the microwave while picking up the plate at the same time. We do this by passing different actions to different latent variables, therefore simulating as if the agent is performing both actions at the same time. We do this below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latents = None\n",
    "image = frames[0]\n",
    "pl.seed_everything(42)\n",
    "\n",
    "default_action = (0.0, 0.0)\n",
    "causal_var_actions = {\n",
    "    'Microwave - Active': (-0.5, 0.8),\n",
    "    'Plate': (0.65, -0.43)\n",
    "}\n",
    "\n",
    "action = torch.tensor(default_action, device=device)\n",
    "action = action[None].repeat(model.hparams.num_latents, 1)\n",
    "for k in causal_var_actions:\n",
    "    for latent_idx in causal_to_latents[k]:\n",
    "        action[latent_idx] = torch.tensor(causal_var_actions[k], device=device)\n",
    "\n",
    "image, latents = next_step_prediction(image=image,\n",
    "                                      action=action,\n",
    "                                      latents=latents,\n",
    "                                      plot_images=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Such a combination of actions is not possible in the simulator and would be very difficult to achieve from entangled representations. This shows the benefits of learned causal representations with BISCUIT, and how they can be used to simulate unseen combinations of actions in the environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interactive Demo\n",
    "\n",
    "If you are running this demo on Google Colab, we can visualize the prediction by an interactive demo. Feel free to click anywhere on the image below, and we will automatically show the prediction of the model of the next step (the prediction and reloading of the image can take a few seconds on Colab). To restart the demo from the beginning or choose a different starting image, simply re-run the cell. Note that since the model itself is not perfect, the model can sometimes face difficulty in predicting actions on the Egg and Plate, especially with multi-modal outcomes. Feel free to try out your own actions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from google.colab import output\n",
    "except ImportError:\n",
    "    print(\"Not running in Google Colab. Skipping the widget manager.\")\n",
    "else:\n",
    "    !pip install --quiet ipympl\n",
    "    output.enable_custom_widget_manager()\n",
    "\n",
    "    %matplotlib ipympl\n",
    "\n",
    "    latents = None\n",
    "    image = frames[0]\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.imshow(image.permute(1, 2, 0).cpu().numpy())\n",
    "    load_text = ax.text(image.shape[-1]//2, \n",
    "                        image.shape[-2]//2, \n",
    "                        'Loading...', \n",
    "                        fontsize='x-large',\n",
    "                        weight='bold',\n",
    "                        va='center',\n",
    "                        ha='center',\n",
    "                        backgroundcolor=(1.0, 0.8, 0.8))\n",
    "    load_text.set_visible(False)\n",
    "    ax.axis('off')\n",
    "\n",
    "    def onclick(event):\n",
    "        global image, latents\n",
    "        load_text.set_visible(True)\n",
    "        fig.canvas.draw()\n",
    "        ix, iy = event.xdata, event.ydata\n",
    "        ix = (ix / image.shape[-1] - 0.5) * 2.0\n",
    "        iy = (iy / image.shape[-2] - 0.5) * 2.0\n",
    "        print(f'Clicked at x={ix:.2f}, y={iy:.2f}')\n",
    "        image, latents = next_step_prediction(image=image,\n",
    "                                              action=torch.tensor([iy, ix], \n",
    "                                                                  dtype=torch.float32,\n",
    "                                                                  device=device),\n",
    "                                              latents=latents,\n",
    "                                              plot_images=False)\n",
    "        ax.imshow(image.permute(1, 2, 0).cpu().numpy())\n",
    "        load_text.set_visible(False)\n",
    "        fig.canvas.draw()\n",
    "        \n",
    "    cid = fig.canvas.mpl_connect('button_press_event', onclick)"
   ]
  }
 ],
 "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.8.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "fa7e3afb224b4c936d9b4ea809be26926770ba5c2d1e177536460cba3b3336ad"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}