train_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 67,
     "referenced_widgets": [
      "a8e97e26d65147efa2ddb176214e818a",
      "19e913a07b044b2d91ca7b5d2dbbe8f7",
      "1de97479fc3e44caae336cf0b1082d02",
      "b8d4e152c232416f82f5fb30e2478ef8",
      "d66b3f8971b24c81abaefe111e6383a6",
      "11d14b9bf0f74bf7bd1b58e6a26fd49e",
      "e9fde27bd58b499bb386b36314d9efca",
      "5407db80343c44c79350021cd98204f8",
      "68e9f5994766480dac5dc77ce043a2e9",
      "d00c195ac5644bd0aba81dccf56e2df0",
      "54c35b52c4e24543872c8335c935c410"
     ]
    },
    "id": "VrX4VTl5pYNq",
    "outputId": "63f38652-f298-4723-bd5c-e1bbe5f9d961"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ns/mambaforge/envs/robodiff/lib/python3.9/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: libtorch_cuda_cu.so: cannot open shared object file: No such file or directory\n",
      "  warn(f\"Failed to load image Python extension: {e}\")\n",
      "/home/ns/mambaforge/envs/robodiff/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pygame 2.1.2 (SDL 2.0.16, Python 3.9.18)\n",
      "Hello from the pygame community. https://www.pygame.org/contribute.html\n"
     ]
    }
   ],
   "source": [
    "from typing import Tuple, Sequence, Dict, Union, Optional, Callable\n",
    "import numpy as np\n",
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import collections\n",
    "import zarr\n",
    "from diffusers.schedulers.scheduling_ddpm import DDPMScheduler\n",
    "from diffusers.training_utils import EMAModel\n",
    "from diffusers.optimization import get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import gym\n",
    "from gym import spaces\n",
    "import pygame\n",
    "import pymunk\n",
    "import pymunk.pygame_util\n",
    "from pymunk.space_debug_draw_options import SpaceDebugColor\n",
    "from pymunk.vec2d import Vec2d\n",
    "import shapely.geometry as sg\n",
    "import cv2\n",
    "import skimage.transform as st\n",
    "from skvideo.io import vwrite\n",
    "from IPython.display import Video\n",
    "import gdown\n",
    "import os\n",
    "\n",
    "\n",
    "\n",
    "#hack to get the pusht environment to work\n",
    "import sys   \n",
    "sys.path.append('/home/ns/diffusion_policy/')  \n",
    "from diffusion_policy.env.pusht.pusht_image_env import PushTImageEnv\n",
    "import imageio \n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pusht_data_utils import * \n",
    "from noise_predictor_model import ConditionalUnet1D\n",
    "from vision_model import ResidualBlock, ResNetFe\n",
    "from myddpm import MyScheduler, MyDDPM\n",
    "from vision_model import replace_bn_with_gn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OknH8Qfqrtc9",
    "outputId": "c97c843e-0b23-4c0a-c123-cd72c407e163"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "obs['image'].shape: (3, 96, 96) float32, [0,1]\n",
      "obs['agent_pos'].shape: (2,) float32, [0,512]\n",
      "action.shape:  (2,) float32, [0,512]\n"
     ]
    }
   ],
   "source": [
    "seed=1000\n",
    "\n",
    "env = PushTImageEnv() \n",
    "\n",
    "env.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    " \n",
    "obs = env.reset()\n",
    "\n",
    "# 3. 2D positional action space [0,512]\n",
    "action = env.action_space.sample()\n",
    " \n",
    "obs, reward, terminated, info = env.step(action)\n",
    "\n",
    "\n",
    "print(\"obs['image'].shape:\", obs['image'].shape, \"float32, [0,1]\")\n",
    "print(\"obs['agent_pos'].shape:\", obs['agent_pos'].shape, \"float32, [0,512]\")\n",
    "print(\"action.shape: \", action.shape, \"float32, [0,512]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9ZiHF3lzvB6k",
    "outputId": "e6256ff1-ab20-41d4-a9ea-7c035196bf3c"
   },
   "outputs": [],
   "source": [
    "dataset_path = \"data/pusht/pusht_cchi_v7_replay.zarr\"\n",
    "\n",
    "\n",
    "pred_horizon = 16\n",
    "obs_horizon = 2\n",
    "action_horizon = 8\n",
    "\n",
    "dataset = PushTImageDataset(\n",
    "    dataset_path=dataset_path,\n",
    "    pred_horizon=pred_horizon,\n",
    "    obs_horizon=obs_horizon,\n",
    "    action_horizon=action_horizon\n",
    ")\n",
    "\n",
    "stats = dataset.stats\n",
    "\n",
    "dataloader = torch.utils.data.DataLoader(\n",
    "    dataset,\n",
    "    batch_size=64,\n",
    "    num_workers=4,\n",
    "    shuffle=True,\n",
    "    # accelerate cpu-gpu transfer\n",
    "    pin_memory=True,\n",
    "    # don't kill worker process afte each epoch\n",
    "    persistent_workers=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yXq4r744aMh1"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 749,120\n"
     ]
    }
   ],
   "source": [
    "vision_encoder = ResNetFe(ResidualBlock, [2, 2]) \n",
    "vision_encoder = replace_bn_with_gn(vision_encoder)\n",
    "\n",
    "n_params= sum(p.numel() for p in vision_encoder.parameters())\n",
    "print(f\"Number of parameters: {n_params:,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4APZkqh336-M",
    "outputId": "e362d582-ceac-4cdd-e866-c34858593696"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using custom positional encoding\n",
      "Number of parameters: 43,938,882\n"
     ]
    }
   ],
   "source": [
    "vision_feature_dim = 512\n",
    "lowdim_obs_dim = 2\n",
    "# observation feature has 514 dims in total per step\n",
    "obs_dim = vision_feature_dim + lowdim_obs_dim\n",
    "action_dim = 2\n",
    "\n",
    "\n",
    "noise_pred_net = ConditionalUnet1D(\n",
    "    input_dim=action_dim,\n",
    "    global_cond_dim=obs_dim*obs_horizon\n",
    ")\n",
    "\n",
    "\n",
    "# There are two networks in the model\n",
    "nets = nn.ModuleDict({\n",
    "    'vision_encoder': vision_encoder,\n",
    "    'noise_pred_net': noise_pred_net\n",
    "})\n",
    "\n",
    "\n",
    "\n",
    "device = torch.device('cuda')\n",
    "_ = nets.to(device)\n",
    "\n",
    "n_params= sum(p.numel() for p in nets.parameters())\n",
    "print(f\"Number of parameters: {n_params:,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rollout(env, nets,  seed, max_steps=200):\n",
    "    \n",
    "    nets.eval() \n",
    "    env.seed(200+seed)\n",
    "    obs = env.reset()\n",
    "\n",
    "    # keep a queue of last 2 steps of observations\n",
    "    obs_deque = collections.deque(\n",
    "        [obs] * obs_horizon, maxlen=obs_horizon)\n",
    "    # save visualization and rewards\n",
    "    imgs = [env.render(mode='rgb_array')]\n",
    "    rewards = list()\n",
    "    done = False\n",
    "    step_idx = 0\n",
    "    success=False\n",
    "    with tqdm(total=max_steps, desc=\"Eval PushTImageEnv\") as pbar:\n",
    "        while not done:\n",
    "            B = 1\n",
    "            # stack the last obs_horizon number of observations\n",
    "            images = np.stack([x['image'] for x in obs_deque])\n",
    "            agent_poses = np.stack([x['agent_pos'] for x in obs_deque])\n",
    "\n",
    "            # normalize observation\n",
    "            nagent_poses = normalize_data(agent_poses, stats=stats['agent_pos'])\n",
    "            # images are already normalized to [0,1]\n",
    "            nimages = images\n",
    "\n",
    "            # device transfer\n",
    "            nimages = torch.from_numpy(nimages).to(device, dtype=torch.float32)\n",
    "            # (2,3,96,96)\n",
    "            nagent_poses = torch.from_numpy(nagent_poses).to(device, dtype=torch.float32)\n",
    "            # (2,2)\n",
    "\n",
    "            # infer action\n",
    "            with torch.no_grad(): \n",
    "                image_features = nets['vision_encoder'](nimages)\n",
    "                # (2,512) \n",
    "                # concat with low-dim observations\n",
    "                obs_features = torch.cat([image_features, nagent_poses], dim=-1) \n",
    "                # reshape observation to (B,obs_horizon*obs_dim)\n",
    "                obs_cond = obs_features.unsqueeze(0).flatten(start_dim=1)\n",
    "\n",
    "                naction,xts=ddpm.sample_ddpm(1, sample_shape, obs_cond)\n",
    "                \n",
    "\n",
    "            # unnormalize action\n",
    "            naction = naction.detach().to('cpu').numpy()\n",
    "            # (B, pred_horizon, action_dim)\n",
    "            naction = naction[0]\n",
    "            action_pred = unnormalize_data(naction, stats=stats['action'])\n",
    "\n",
    "            # only take action_horizon number of actions\n",
    "            start = obs_horizon - 1\n",
    "            end = start + action_horizon\n",
    "            action = action_pred[start:end,:]\n",
    "            # (action_horizon, action_dim)\n",
    "\n",
    "            # execute action_horizon number of steps\n",
    "            # without replanning\n",
    "            for i in range(len(action)):\n",
    "                obs, reward, done, info = env.step(action[i])\n",
    "                obs_deque.append(obs)\n",
    "                \n",
    "                rewards.append(reward)\n",
    "                imgs.append(env.render(mode='rgb_array'))\n",
    "\n",
    "                # update progress bar\n",
    "                step_idx += 1\n",
    "                pbar.update(1)\n",
    "                pbar.set_postfix(reward=reward)\n",
    "                if step_idx > max_steps:\n",
    "                    done = True\n",
    "                if done:\n",
    "                    success=True\n",
    "                    break\n",
    "\n",
    "    return max(rewards) , success, imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_diffusion_iters = 100\n",
    "noise_scheduler=MyScheduler(T=num_diffusion_iters, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "ddpm=MyDDPM(noise_scheduler, nets['noise_pred_net'], device=device)\n",
    "sample_shape=(pred_horizon, action_dim) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 81,
     "referenced_widgets": [
      "b34bf5a5d603446fb453bbed31c4469c",
      "0873a2dee0e44b0fb3e8445d94c27171",
      "410476fa4ac14dd2b78377f4d779402f",
      "09152dbe3c8543aa809aa592afdc53f1",
      "e2a91fa6d1514913892c0def2e2ecf78",
      "7fd765bcd2a34db49a02839ba1c4f518",
      "d9054c34284045eda3546a21bb5fafc3",
      "a92fc95c6a3d496997da90c87e24ba01",
      "9a98272475d940b1b98438c4d02dcc05",
      "47b6383ed3b64b7e8c7e66d0e7a468d9",
      "1cdde5afac8041bf9b75e0e80bd80630",
      "70b4b657d68648d9be98c099d061cd5a",
      "8c0237e38cfd4143aa423dea26637965",
      "44c486fa8c4241f4a1a245f4e24da769",
      "12ef3c026774405cb91faf18d7fd8afa",
      "99c0577d549c4dd4a61d9ec0cf6254d5",
      "5124950790874225b5cca7556f122f9a",
      "f609505b750747dcaff8a950131743e9",
      "3c0a27b6addb4c6a977824995548fe90",
      "d14f2c65ef9348348250af2df0f32963",
      "eb23edb22d624aee9fd55ab05e2a517e",
      "a073d659b0f347c9b30c343430d7f6aa",
      "7508fbad995648a39d2a55953d000786",
      "9009dc5a656c4befbf513f2574df5a7d",
      "479d2f35498b4fbea77057f3f26fd287",
      "16b5928585984a1b81b809717f125489",
      "9d436c7ee90349a5966f025d095dd0cc",
      "3e503eeff5d94f0e9dc8a3d31190c6a2",
      "4abc5a8b257240b99535a6abe5e00ef6",
      "6c25264289e3411db29f6f8db936b00f",
      "a0ff8cd1b00544caa12cbd809524dd15",
      "1faac530b70b41a58a8c9c0cf44af69c",
      "7cfefcf49a9a426da4e9203ec4debe38",
      "5f96241543b441fcb13c16d4ca476405",
      "e6bbb6ce052045ca895b526c620d8847",
      "5c1c61cd14fb4aa8b335684fbf652a29",
      "a5a2bc5dc0b54a4796e574308bb6c4bc",
      "febe66177c604395a78be993eeff7c6d",
      "6de3445a866442aa9cc7d855cb0d0740",
      "353cb60a2a73417b81ccbb5e52480ed3",
      "69f323a406004822bf1a0bfd79767c9a",
      "88a8f5924a1e4a86805e314dd90c17be",
      "9ada44d21f234712ac5df797c730495a",
      "c777ac64366c40ea8a1f63408879c2ab"
     ]
    },
    "id": "93E9RdnR4D8v",
    "outputId": "6592f2ca-4d83-41a4-c650-564042a44b0d"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cf4a9877998448d99d049758e9178ad9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch:   0%|          | 0/400 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f61f2701cb6e48f7a10468e7c138bdd4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a5c48203bc77444f81dce0caeb8faaed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b5458046c34446e8250548f26f72faf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0ef6a4761aa49e2ac01ec9e95b1be91",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0f51713fd75940a0a44d2fac254e5bcc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "99a1320634e548df98462cff56e1242d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ab22d88607b4a49b5645912e550ed32",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64a186cb71bb45859efcd10f6c8bc32e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b4929d6e65f24e3194d94ccbeb7ff955",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1396382871f44e63afc6560f91a1565b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 100, mean_r: 0.6919100382089165 std_r: 0.3847147647227042\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "96d905453f2642339d198f71c3ed80b8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0d0d019e5edc470fb335643488f9f871",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "37c370837d96442a8477830a8b353524",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "47287435cd5647b4884f2d82d1997360",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "362dbab83cc845aea2ca15ed78b352ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27b6eda8d00a4433aeea7ca6324ca04a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "605fc640facd40588fcb061285839c03",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7ad725688eb4816aeded4da33e53474",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4615f73d3b264791be3a992e58db7a8f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dcac5721000242d299089a4d6e67956c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 200, mean_r: 0.6624416888383055 std_r: 0.3920713531911076\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "47c4ee74518c4e57b84ec144c9a1c362",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78676d0c435b4ab08a8f983d50b751e0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "951892bfae8b46e2896e32354e6dd862",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bef005b22af146d4ab7246d1b03c858e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4001250c1a9749c5b30208654e2b56fd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d920d62b36c345849d8fabbde93d3255",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c6c4716b2b8f44f9a4634ebc6379780f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "914d431615ea48bdae97c2cc80b9b625",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17d9581cc64147b9878eb2ee14fd9ab1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f1958dbfb0bf44deb2acd54a3a65cef0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 300, mean_r: 0.5555869272972284 std_r: 0.4503132616840976\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 400\n",
    "\n",
    " \n",
    "optimizer = torch.optim.AdamW(\n",
    "    params=nets.parameters(),\n",
    "    lr=1e-4, weight_decay=1e-6)\n",
    " \n",
    "lr_scheduler = get_scheduler(\n",
    "    name='cosine',\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=500,\n",
    "    num_training_steps=len(dataloader) * num_epochs\n",
    ")\n",
    "\n",
    "with tqdm(range(num_epochs), desc='Epoch') as tglobal: \n",
    "    for epoch_idx in tglobal:  \n",
    "        nets.train()\n",
    "        epoch_loss = list() \n",
    "        for nbatch in dataloader: \n",
    "            nimage = nbatch['image'][:,:obs_horizon].to(device)\n",
    "            nagent_pos = nbatch['agent_pos'][:,:obs_horizon].to(device)\n",
    "            naction = nbatch['action'].to(device)\n",
    "            B = nagent_pos.shape[0]\n",
    " \n",
    "            image_features = nets['vision_encoder'](\n",
    "                nimage.flatten(end_dim=1))\n",
    "            image_features = image_features.reshape(\n",
    "                *nimage.shape[:2],-1) \n",
    "            obs_features = torch.cat([image_features, nagent_pos], dim=-1)\n",
    "            obs_cond = obs_features.flatten(start_dim=1)\n",
    "          \n",
    "            timesteps = torch.randint(0, noise_scheduler.T,(B,), device=device).long()\n",
    "\n",
    "            noisy_actions , noise= noise_scheduler.get_xt(naction, timesteps)\n",
    "            noise_pred = noise_pred_net(noisy_actions, timesteps, global_cond=obs_cond)\n",
    "\n",
    "            loss = nn.functional.mse_loss(noise_pred, noise)\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "            lr_scheduler.step()\n",
    "\n",
    " \n",
    "            loss_cpu = loss.item()\n",
    "            epoch_loss.append(loss_cpu) \n",
    "        tglobal.set_postfix(loss=np.mean(epoch_loss))\n",
    "\n",
    "        if epoch_idx>0 and epoch_idx %100==0:\n",
    "            rewards = [rollout(env, nets, i*10, 200)[0] for i in range(10)]\n",
    "            mean_r = np.mean(rewards)\n",
    "            std_r = np.std(rewards)\n",
    "            print(f\"epoch: {epoch_idx}, mean_r: {mean_r} std_r: {std_r}\") \n",
    "ema_nets = nets "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d4f1888d4664fe5a0e91e9568e6d353",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02e3054ed8534afe8ef3bcf1a9b0011c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3c014b753c1640e181bfc08aa839c603",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4bbb14c0e1ec4585a5d0dbb9c4084be0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0d257dd012f40148aac28d7b0f388bb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "36e26d3f70fc443db03959a3d8267742",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ffc83f6f00be4e859dc9d063969e7717",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b7abf0f6062d4a36bb45c2657195c126",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a85dcf4e32414e999aa4552c40d93b6c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17eb66f34c344419a91199c3e96f8bc4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean_r: 0.834332207599916 std_r: 0.3166511912504874\n"
     ]
    }
   ],
   "source": [
    "rewards = [rollout(env, nets, i*10, 200)[0] for i in range(10)]\n",
    "mean_r = np.mean(rewards)\n",
    "std_r = np.std(rewards)\n",
    "print(f\"mean_r: {mean_r} std_r: {std_r}\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.9681960031766111,\n",
       " 0.992990162179459,\n",
       " 0.98410073653022,\n",
       " 1.0,\n",
       " 0.4805411708349545,\n",
       " 1.0,\n",
       " 0.0,\n",
       " 0.9946874386600564,\n",
       " 1.0,\n",
       " 0.9228065646178574]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "31903db7ac7347ae86b965bc0e7a032d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ed4313f42a994f3ba94bb89033a1412e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95d3e047c90a4404b200b78a9f21cfbf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a31c939144414794b878bc830aeef94f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8853c6f5aa4d483a910302197a332c1c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cf7013d783134539a8969d7b4d14474f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7eb81d9fd74444869557c5478c225cc4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a822c001a8b4795966b4fb6957cc7ae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eb09c144d40b47d3a9302a49069ab582",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e741210c459a4829a7e626768e03ad7d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "331a1b8948fe4da28ca60295996a4cb7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "415093196c164ced8ff476e3d86eb0d4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d44c9c6dfb3b4c8ebf88fdb8dd04d43b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "48f94810ae2e4ff8918892d0d65d0a3c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e68ca1b67b94706906db99436c55695",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ece74b6adb264a0bab4c4482f0931c59",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3efd860ee67f47cdb920492379ef9140",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a444f887b08488a910a9379f71cc232",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9431364d913148bba1d6e77ce8643a30",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be0707eb068748d98e66b0df3a30b355",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "780a55064149466ea08c418e777d33d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8312c394162e4fa0a2df67d6c245f4b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "838a0dcfb8ca4253afc6805a86492984",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae9d37db294745669121d8e1a092a3c4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b13a00d2a3504999bb95b4dfb4acacba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa945f74fbe44102972ff023b13563dc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdbef237fb674031954f3b4be494e0a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ec91ea877c64d35aca00bedcd51b1b8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3dd2cf223e804b888cd1615453df5dab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8264d3a0a2544c53b16bfc8e97ac446e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ecab6e8d4784b02bba8958baa292ebf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "702058262a804f50928001383a4991dc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "040536b0c96a41a7b4a669f000b4ffc1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8c2e34d5d5ba4f6997c64b43c7f99c49",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "37c5aa9438cb470b8f05ad1b8fd81021",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af5d734984f44ddaa155d9d98cdcadc8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8cdb7bd7982e4ff08101c8fe05591671",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "240d6d87406244e98f4bfe53a44e3b6a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6892e886b1764222b8bd99e6fe4d9e6c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3cb2b3bc74964ce0a255357343c01ab6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ba5afb2a816a406ba42cb27284d73f3b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9263e430e9c64ffc81311d5b52ce4654",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d7874620b3ee40d3b564da4e801d767e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8bd55e207f314543ac70b9c61af0692d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "75f7693e216242eea1a7429e5fff5440",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d3e751cd435f41ed91bb8844c8471650",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e0037eed8de54cf4900fe8c5d495833a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "79d601bdec96440d9fe6de943e2c5444",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d06ac18a0524f2e87b42e4e9571f834",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3317bc12b12a427eb3eada346154f2f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Eval PushTImageEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean_r: 0.7649656332717022 std_r: 0.37320718935672514\n"
     ]
    }
   ],
   "source": [
    "seed=40\n",
    "\n",
    "env = PushTImageEnv() \n",
    "\n",
    "env.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed) \n",
    "\n",
    "rewards = [rollout(env, nets, i*10, 200)[0] for i in range(50)]\n",
    "mean_r = np.mean(rewards)\n",
    "std_r = np.std(rewards)\n",
    "print(f\"mean_r: {mean_r} std_r: {std_r}\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1.0,\n",
       " 1.0,\n",
       " 0.9581792225929594,\n",
       " 1.0,\n",
       " 0.2157577274316913,\n",
       " 0.9773604408702586,\n",
       " 0.0,\n",
       " 0.9596557565779427,\n",
       " 0.9503510825126081,\n",
       " 0.7958965506477205,\n",
       " 1.0,\n",
       " 0.0,\n",
       " 0.999683707311073,\n",
       " 0.5930374558963504,\n",
       " 0.0,\n",
       " 1.0,\n",
       " 0.0,\n",
       " 0.9816286314858859,\n",
       " 1.0,\n",
       " 0.0,\n",
       " 1.0,\n",
       " 0.0,\n",
       " 0.9650778022527365,\n",
       " 0.0,\n",
       " 0.9705912467213306,\n",
       " 0.0,\n",
       " 0.957430005170317,\n",
       " 0.31374807985773745,\n",
       " 0.9991597933143729,\n",
       " 0.5420587940013708,\n",
       " 1.0,\n",
       " 1.0,\n",
       " 1.0,\n",
       " 1.0,\n",
       " 0.6576442347152985,\n",
       " 0.9906368821546335,\n",
       " 0.8785877563284036,\n",
       " 1.0,\n",
       " 0.9829195212488347,\n",
       " 0.9738330642763778,\n",
       " 0.8898527167923332,\n",
       " 1.0,\n",
       " 0.9996975314718549,\n",
       " 0.9997818998820144,\n",
       " 1.0,\n",
       " 0.9970955766400648,\n",
       " 1.0,\n",
       " 0.7441541750818337,\n",
       " 0.9544620083491021,\n",
       " 1.0]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.save(nets.state_dict(), 'pusht_enc_ms_worked.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.18"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0873a2dee0e44b0fb3e8445d94c27171": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7fd765bcd2a34db49a02839ba1c4f518",
      "placeholder": "​",
      "style": "IPY_MODEL_d9054c34284045eda3546a21bb5fafc3",
      "value": "Epoch:   2%"
     }
    },
    "09152dbe3c8543aa809aa592afdc53f1": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_47b6383ed3b64b7e8c7e66d0e7a468d9",
      "placeholder": "​",
      "style": "IPY_MODEL_1cdde5afac8041bf9b75e0e80bd80630",
      "value": " 2/100 [03:14&lt;2:38:57, 97.32s/it, loss=0.0192]"
     }
    },
    "11d14b9bf0f74bf7bd1b58e6a26fd49e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "12ef3c026774405cb91faf18d7fd8afa": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_eb23edb22d624aee9fd55ab05e2a517e",
      "placeholder": "​",
      "style": "IPY_MODEL_a073d659b0f347c9b30c343430d7f6aa",
      "value": " 379/379 [01:37&lt;00:00,  4.40it/s, loss=0.0185]"
     }
    },
    "16b5928585984a1b81b809717f125489": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_1faac530b70b41a58a8c9c0cf44af69c",
      "placeholder": "​",
      "style": "IPY_MODEL_7cfefcf49a9a426da4e9203ec4debe38",
      "value": " 379/379 [01:37&lt;00:00,  4.57it/s, loss=0.0103]"
     }
    },
    "182afbf676cd4c0982b90890bbdbeeef": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "185ce2207cec4ae1a6ffdecadae23894": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "19e913a07b044b2d91ca7b5d2dbbe8f7": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_11d14b9bf0f74bf7bd1b58e6a26fd49e",
      "placeholder": "​",
      "style": "IPY_MODEL_e9fde27bd58b499bb386b36314d9efca",
      "value": ""
     }
    },
    "1b96a78d8b8542eb830e3c446a0dcf42": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "1cdde5afac8041bf9b75e0e80bd80630": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "1d0a57ed8e914d31bcc1c36a56451df7": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_4cc800669dfd45a784864dd4a3881daa",
      "max": 200,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_95dc6c44f97e4e8882a3657bd2fd66fb",
      "value": 200
     }
    },
    "1de97479fc3e44caae336cf0b1082d02": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5407db80343c44c79350021cd98204f8",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_68e9f5994766480dac5dc77ce043a2e9",
      "value": 0
     }
    },
    "1faac530b70b41a58a8c9c0cf44af69c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "2c31ffecd0494e2a8d0a8cefbea5df8a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5d4a78f691ec4555bf179e5ef5d828ac",
      "placeholder": "​",
      "style": "IPY_MODEL_1b96a78d8b8542eb830e3c446a0dcf42",
      "value": "Eval PushTImageEnv: "
     }
    },
    "333b5148e59e47f3998a5cf5df531baf": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "353cb60a2a73417b81ccbb5e52480ed3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "3c0a27b6addb4c6a977824995548fe90": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "3e503eeff5d94f0e9dc8a3d31190c6a2": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "410476fa4ac14dd2b78377f4d779402f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a92fc95c6a3d496997da90c87e24ba01",
      "max": 100,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_9a98272475d940b1b98438c4d02dcc05",
      "value": 2
     }
    },
    "44c486fa8c4241f4a1a245f4e24da769": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_3c0a27b6addb4c6a977824995548fe90",
      "max": 379,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_d14f2c65ef9348348250af2df0f32963",
      "value": 379
     }
    },
    "479d2f35498b4fbea77057f3f26fd287": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6c25264289e3411db29f6f8db936b00f",
      "max": 379,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_a0ff8cd1b00544caa12cbd809524dd15",
      "value": 379
     }
    },
    "47b6383ed3b64b7e8c7e66d0e7a468d9": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "4abc5a8b257240b99535a6abe5e00ef6": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "4cc800669dfd45a784864dd4a3881daa": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5124950790874225b5cca7556f122f9a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5407db80343c44c79350021cd98204f8": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": "20px"
     }
    },
    "54c35b52c4e24543872c8335c935c410": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "5c1c61cd14fb4aa8b335684fbf652a29": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_69f323a406004822bf1a0bfd79767c9a",
      "max": 379,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_88a8f5924a1e4a86805e314dd90c17be",
      "value": 11
     }
    },
    "5d4a78f691ec4555bf179e5ef5d828ac": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5f96241543b441fcb13c16d4ca476405": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_e6bbb6ce052045ca895b526c620d8847",
       "IPY_MODEL_5c1c61cd14fb4aa8b335684fbf652a29",
       "IPY_MODEL_a5a2bc5dc0b54a4796e574308bb6c4bc"
      ],
      "layout": "IPY_MODEL_febe66177c604395a78be993eeff7c6d"
     }
    },
    "68e9f5994766480dac5dc77ce043a2e9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "69f323a406004822bf1a0bfd79767c9a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6c25264289e3411db29f6f8db936b00f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6de3445a866442aa9cc7d855cb0d0740": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "70b4b657d68648d9be98c099d061cd5a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_8c0237e38cfd4143aa423dea26637965",
       "IPY_MODEL_44c486fa8c4241f4a1a245f4e24da769",
       "IPY_MODEL_12ef3c026774405cb91faf18d7fd8afa"
      ],
      "layout": "IPY_MODEL_99c0577d549c4dd4a61d9ec0cf6254d5"
     }
    },
    "7508fbad995648a39d2a55953d000786": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_9009dc5a656c4befbf513f2574df5a7d",
       "IPY_MODEL_479d2f35498b4fbea77057f3f26fd287",
       "IPY_MODEL_16b5928585984a1b81b809717f125489"
      ],
      "layout": "IPY_MODEL_9d436c7ee90349a5966f025d095dd0cc"
     }
    },
    "7cfefcf49a9a426da4e9203ec4debe38": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "7fd765bcd2a34db49a02839ba1c4f518": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "88a8f5924a1e4a86805e314dd90c17be": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "8c0237e38cfd4143aa423dea26637965": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5124950790874225b5cca7556f122f9a",
      "placeholder": "​",
      "style": "IPY_MODEL_f609505b750747dcaff8a950131743e9",
      "value": "Batch: 100%"
     }
    },
    "9009dc5a656c4befbf513f2574df5a7d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_3e503eeff5d94f0e9dc8a3d31190c6a2",
      "placeholder": "​",
      "style": "IPY_MODEL_4abc5a8b257240b99535a6abe5e00ef6",
      "value": "Batch: 100%"
     }
    },
    "95dc6c44f97e4e8882a3657bd2fd66fb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "99c0577d549c4dd4a61d9ec0cf6254d5": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": "hidden",
      "width": null
     }
    },
    "9a98272475d940b1b98438c4d02dcc05": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "9ada44d21f234712ac5df797c730495a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "9d436c7ee90349a5966f025d095dd0cc": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": "hidden",
      "width": null
     }
    },
    "a073d659b0f347c9b30c343430d7f6aa": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "a0ff8cd1b00544caa12cbd809524dd15": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "a5a2bc5dc0b54a4796e574308bb6c4bc": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_9ada44d21f234712ac5df797c730495a",
      "placeholder": "​",
      "style": "IPY_MODEL_c777ac64366c40ea8a1f63408879c2ab",
      "value": " 11/379 [00:03&lt;01:34,  3.91it/s, loss=0.0256]"
     }
    },
    "a8e97e26d65147efa2ddb176214e818a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_19e913a07b044b2d91ca7b5d2dbbe8f7",
       "IPY_MODEL_1de97479fc3e44caae336cf0b1082d02",
       "IPY_MODEL_b8d4e152c232416f82f5fb30e2478ef8"
      ],
      "layout": "IPY_MODEL_d66b3f8971b24c81abaefe111e6383a6"
     }
    },
    "a92fc95c6a3d496997da90c87e24ba01": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "b0a9f64f47d34769a47e165291d2c550": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_2c31ffecd0494e2a8d0a8cefbea5df8a",
       "IPY_MODEL_1d0a57ed8e914d31bcc1c36a56451df7",
       "IPY_MODEL_e1bb86eb510b47d9ac2311926455bf1f"
      ],
      "layout": "IPY_MODEL_182afbf676cd4c0982b90890bbdbeeef"
     }
    },
    "b34bf5a5d603446fb453bbed31c4469c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_0873a2dee0e44b0fb3e8445d94c27171",
       "IPY_MODEL_410476fa4ac14dd2b78377f4d779402f",
       "IPY_MODEL_09152dbe3c8543aa809aa592afdc53f1"
      ],
      "layout": "IPY_MODEL_e2a91fa6d1514913892c0def2e2ecf78"
     }
    },
    "b8d4e152c232416f82f5fb30e2478ef8": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d00c195ac5644bd0aba81dccf56e2df0",
      "placeholder": "​",
      "style": "IPY_MODEL_54c35b52c4e24543872c8335c935c410",
      "value": " 0/0 [00:00&lt;?, ?it/s]"
     }
    },
    "c777ac64366c40ea8a1f63408879c2ab": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "d00c195ac5644bd0aba81dccf56e2df0": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "d14f2c65ef9348348250af2df0f32963": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "d66b3f8971b24c81abaefe111e6383a6": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "d9054c34284045eda3546a21bb5fafc3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "e1bb86eb510b47d9ac2311926455bf1f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_333b5148e59e47f3998a5cf5df531baf",
      "placeholder": "​",
      "style": "IPY_MODEL_185ce2207cec4ae1a6ffdecadae23894",
      "value": " 201/? [00:34&lt;00:00,  6.11it/s, reward=0.828]"
     }
    },
    "e2a91fa6d1514913892c0def2e2ecf78": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "e6bbb6ce052045ca895b526c620d8847": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6de3445a866442aa9cc7d855cb0d0740",
      "placeholder": "​",
      "style": "IPY_MODEL_353cb60a2a73417b81ccbb5e52480ed3",
      "value": "Batch:   3%"
     }
    },
    "e9fde27bd58b499bb386b36314d9efca": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "eb23edb22d624aee9fd55ab05e2a517e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "f609505b750747dcaff8a950131743e9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "febe66177c604395a78be993eeff7c6d": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}