diffusion_policy_state_pusht_demo.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3 (ipykernel)"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU",
  "gpuClass": "standard",
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "89ebea621ae847039742f3edc488b6ff": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "model_module_version": "1.5.0",
     "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_a6da49e27cd245ce847cf94cab24dce5",
       "IPY_MODEL_575e02d16dce4c3885f5918b77194512",
       "IPY_MODEL_ba06215466fe4676a375517132ca996c"
      ],
      "layout": "IPY_MODEL_6d782f4fe81341f19b1cccb46ce88677"
     }
    },
    "a6da49e27cd245ce847cf94cab24dce5": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "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_1d3c1f93be2a4bfc8d762392b5d2d93c",
      "placeholder": "​",
      "style": "IPY_MODEL_6a492b09c81549e59ee509ab7bf5a38d",
      "value": "Eval PushTStateEnv:  77%"
     }
    },
    "575e02d16dce4c3885f5918b77194512": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "model_module_version": "1.5.0",
     "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": "danger",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_e96f4ab68adc4ecab2d868f6d0ae7e40",
      "max": 200,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_78557990a7334246bf4700ba2d91fb29",
      "value": 154
     }
    },
    "ba06215466fe4676a375517132ca996c": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "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_b182573737c64d238123139066829297",
      "placeholder": "​",
      "style": "IPY_MODEL_0a36f7e44224497d9b6c3a54ce227268",
      "value": " 154/200 [00:28&lt;00:07,  5.79it/s, reward=1]"
     }
    },
    "6d782f4fe81341f19b1cccb46ce88677": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "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
     }
    },
    "1d3c1f93be2a4bfc8d762392b5d2d93c": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "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
     }
    },
    "6a492b09c81549e59ee509ab7bf5a38d": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "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": ""
     }
    },
    "e96f4ab68adc4ecab2d868f6d0ae7e40": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "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
     }
    },
    "78557990a7334246bf4700ba2d91fb29": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "model_module_version": "1.5.0",
     "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": ""
     }
    },
    "b182573737c64d238123139066829297": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "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
     }
    },
    "0a36f7e44224497d9b6c3a54ce227268": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "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": ""
     }
    }
   }
  }
 },
 "cells": [
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Installing pip packages**\n",
    "#@markdown - Diffusion Model: [PyTorch](https://pytorch.org) & [HuggingFace diffusers](https://huggingface.co/docs/diffusers/index)\n",
    "#@markdown - Dataset Loading: [Zarr](https://zarr.readthedocs.io/en/stable/) & numcodecs\n",
    "#@markdown - Push-T Env: gym, pygame, pymunk & shapely\n",
    "!python --version\n",
    "!pip3 install torch==1.13.1 torchvision==0.14.1 diffusers==0.18.2 \\\n",
    "scikit-image==0.19.3 scikit-video==1.1.11 zarr==2.12.0 numcodecs==0.10.2 \\\n",
    "pygame==2.1.2 pymunk==6.2.1 gym==0.26.2 shapely==1.8.4 \\\n",
    "&> /dev/null # mute output"
   ],
   "metadata": {
    "id": "2QwO2gAgiJS2",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "80ebffdf-e858-4c8c-ad42-fccffa89dbe4"
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Python 3.10.12\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Imports**\n",
    "# diffusion policy import\n",
    "from typing import Tuple, Sequence, Dict, Union, Optional\n",
    "import numpy as np\n",
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\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",
    "# env import\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"
   ],
   "metadata": {
    "id": "VrX4VTl5pYNq",
    "ExecuteTime": {
     "end_time": "2023-12-08T00:34:50.569055Z",
     "start_time": "2023-12-08T00:34:49.173566Z"
    }
   },
   "execution_count": 1,
   "outputs": [
    {
     "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"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Environment**\n",
    "#@markdown Defines a PyMunk-based Push-T environment `PushTEnv`.\n",
    "#@markdown\n",
    "#@markdown **Goal**: push the gray T-block into the green area.\n",
    "#@markdown\n",
    "#@markdown Adapted from [Implicit Behavior Cloning](https://implicitbc.github.io/)\n",
    "\n",
    "\n",
    "positive_y_is_up: bool = False\n",
    "\"\"\"Make increasing values of y point upwards.\n",
    "\n",
    "When True::\n",
    "\n",
    "    y\n",
    "    ^\n",
    "    |      . (3, 3)\n",
    "    |\n",
    "    |   . (2, 2)\n",
    "    |\n",
    "    +------ > x\n",
    "\n",
    "When False::\n",
    "\n",
    "    +------ > x\n",
    "    |\n",
    "    |   . (2, 2)\n",
    "    |\n",
    "    |      . (3, 3)\n",
    "    v\n",
    "    y\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "def to_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]:\n",
    "    \"\"\"Convenience method to convert pymunk coordinates to pygame surface\n",
    "    local coordinates.\n",
    "\n",
    "    Note that in case positive_y_is_up is False, this function wont actually do\n",
    "    anything except converting the point to integers.\n",
    "    \"\"\"\n",
    "    if positive_y_is_up:\n",
    "        return round(p[0]), surface.get_height() - round(p[1])\n",
    "    else:\n",
    "        return round(p[0]), round(p[1])\n",
    "\n",
    "\n",
    "def light_color(color: SpaceDebugColor):\n",
    "    color = np.minimum(1.2 * np.float32([color.r, color.g, color.b, color.a]), np.float32([255]))\n",
    "    color = SpaceDebugColor(r=color[0], g=color[1], b=color[2], a=color[3])\n",
    "    return color\n",
    "\n",
    "class DrawOptions(pymunk.SpaceDebugDrawOptions):\n",
    "    def __init__(self, surface: pygame.Surface) -> None:\n",
    "        \"\"\"Draw a pymunk.Space on a pygame.Surface object.\n",
    "\n",
    "        Typical usage::\n",
    "\n",
    "        >>> import pymunk\n",
    "        >>> surface = pygame.Surface((10,10))\n",
    "        >>> space = pymunk.Space()\n",
    "        >>> options = pymunk.pygame_util.DrawOptions(surface)\n",
    "        >>> space.debug_draw(options)\n",
    "\n",
    "        You can control the color of a shape by setting shape.color to the color\n",
    "        you want it drawn in::\n",
    "\n",
    "        >>> c = pymunk.Circle(None, 10)\n",
    "        >>> c.color = pygame.Color(\"pink\")\n",
    "\n",
    "        See pygame_util.demo.py for a full example\n",
    "\n",
    "        Since pygame uses a coordiante system where y points down (in contrast\n",
    "        to many other cases), you either have to make the physics simulation\n",
    "        with Pymunk also behave in that way, or flip everything when you draw.\n",
    "\n",
    "        The easiest is probably to just make the simulation behave the same\n",
    "        way as Pygame does. In that way all coordinates used are in the same\n",
    "        orientation and easy to reason about::\n",
    "\n",
    "        >>> space = pymunk.Space()\n",
    "        >>> space.gravity = (0, -1000)\n",
    "        >>> body = pymunk.Body()\n",
    "        >>> body.position = (0, 0) # will be positioned in the top left corner\n",
    "        >>> space.debug_draw(options)\n",
    "\n",
    "        To flip the drawing its possible to set the module property\n",
    "        :py:data:`positive_y_is_up` to True. Then the pygame drawing will flip\n",
    "        the simulation upside down before drawing::\n",
    "\n",
    "        >>> positive_y_is_up = True\n",
    "        >>> body = pymunk.Body()\n",
    "        >>> body.position = (0, 0)\n",
    "        >>> # Body will be position in bottom left corner\n",
    "\n",
    "        :Parameters:\n",
    "                surface : pygame.Surface\n",
    "                    Surface that the objects will be drawn on\n",
    "        \"\"\"\n",
    "        self.surface = surface\n",
    "        super(DrawOptions, self).__init__()\n",
    "\n",
    "    def draw_circle(\n",
    "        self,\n",
    "        pos: Vec2d,\n",
    "        angle: float,\n",
    "        radius: float,\n",
    "        outline_color: SpaceDebugColor,\n",
    "        fill_color: SpaceDebugColor,\n",
    "    ) -> None:\n",
    "        p = to_pygame(pos, self.surface)\n",
    "\n",
    "        pygame.draw.circle(self.surface, fill_color.as_int(), p, round(radius), 0)\n",
    "        pygame.draw.circle(self.surface, light_color(fill_color).as_int(), p, round(radius-4), 0)\n",
    "\n",
    "        circle_edge = pos + Vec2d(radius, 0).rotated(angle)\n",
    "        p2 = to_pygame(circle_edge, self.surface)\n",
    "        line_r = 2 if radius > 20 else 1\n",
    "        # pygame.draw.lines(self.surface, outline_color.as_int(), False, [p, p2], line_r)\n",
    "\n",
    "    def draw_segment(self, a: Vec2d, b: Vec2d, color: SpaceDebugColor) -> None:\n",
    "        p1 = to_pygame(a, self.surface)\n",
    "        p2 = to_pygame(b, self.surface)\n",
    "\n",
    "        pygame.draw.aalines(self.surface, color.as_int(), False, [p1, p2])\n",
    "\n",
    "    def draw_fat_segment(\n",
    "        self,\n",
    "        a: Tuple[float, float],\n",
    "        b: Tuple[float, float],\n",
    "        radius: float,\n",
    "        outline_color: SpaceDebugColor,\n",
    "        fill_color: SpaceDebugColor,\n",
    "    ) -> None:\n",
    "        p1 = to_pygame(a, self.surface)\n",
    "        p2 = to_pygame(b, self.surface)\n",
    "\n",
    "        r = round(max(1, radius * 2))\n",
    "        pygame.draw.lines(self.surface, fill_color.as_int(), False, [p1, p2], r)\n",
    "        if r > 2:\n",
    "            orthog = [abs(p2[1] - p1[1]), abs(p2[0] - p1[0])]\n",
    "            if orthog[0] == 0 and orthog[1] == 0:\n",
    "                return\n",
    "            scale = radius / (orthog[0] * orthog[0] + orthog[1] * orthog[1]) ** 0.5\n",
    "            orthog[0] = round(orthog[0] * scale)\n",
    "            orthog[1] = round(orthog[1] * scale)\n",
    "            points = [\n",
    "                (p1[0] - orthog[0], p1[1] - orthog[1]),\n",
    "                (p1[0] + orthog[0], p1[1] + orthog[1]),\n",
    "                (p2[0] + orthog[0], p2[1] + orthog[1]),\n",
    "                (p2[0] - orthog[0], p2[1] - orthog[1]),\n",
    "            ]\n",
    "            pygame.draw.polygon(self.surface, fill_color.as_int(), points)\n",
    "            pygame.draw.circle(\n",
    "                self.surface,\n",
    "                fill_color.as_int(),\n",
    "                (round(p1[0]), round(p1[1])),\n",
    "                round(radius),\n",
    "            )\n",
    "            pygame.draw.circle(\n",
    "                self.surface,\n",
    "                fill_color.as_int(),\n",
    "                (round(p2[0]), round(p2[1])),\n",
    "                round(radius),\n",
    "            )\n",
    "\n",
    "    def draw_polygon(\n",
    "        self,\n",
    "        verts: Sequence[Tuple[float, float]],\n",
    "        radius: float,\n",
    "        outline_color: SpaceDebugColor,\n",
    "        fill_color: SpaceDebugColor,\n",
    "    ) -> None:\n",
    "        ps = [to_pygame(v, self.surface) for v in verts]\n",
    "        ps += [ps[0]]\n",
    "\n",
    "        radius = 2\n",
    "        pygame.draw.polygon(self.surface, light_color(fill_color).as_int(), ps)\n",
    "\n",
    "        if radius > 0:\n",
    "            for i in range(len(verts)):\n",
    "                a = verts[i]\n",
    "                b = verts[(i + 1) % len(verts)]\n",
    "                self.draw_fat_segment(a, b, radius, fill_color, fill_color)\n",
    "\n",
    "    def draw_dot(\n",
    "        self, size: float, pos: Tuple[float, float], color: SpaceDebugColor\n",
    "    ) -> None:\n",
    "        p = to_pygame(pos, self.surface)\n",
    "        pygame.draw.circle(self.surface, color.as_int(), p, round(size), 0)\n",
    "\n",
    "\n",
    "def pymunk_to_shapely(body, shapes):\n",
    "    geoms = list()\n",
    "    for shape in shapes:\n",
    "        if isinstance(shape, pymunk.shapes.Poly):\n",
    "            verts = [body.local_to_world(v) for v in shape.get_vertices()]\n",
    "            verts += [verts[0]]\n",
    "            geoms.append(sg.Polygon(verts))\n",
    "        else:\n",
    "            raise RuntimeError(f'Unsupported shape type {type(shape)}')\n",
    "    geom = sg.MultiPolygon(geoms)\n",
    "    return geom\n",
    "\n",
    "# env\n",
    "class PushTEnv(gym.Env):\n",
    "    metadata = {\"render.modes\": [\"human\", \"rgb_array\"], \"video.frames_per_second\": 10}\n",
    "    reward_range = (0., 1.)\n",
    "\n",
    "    def __init__(self,\n",
    "            legacy=False,\n",
    "            block_cog=None, damping=None,\n",
    "            render_action=True,\n",
    "            render_size=96,\n",
    "            reset_to_state=None\n",
    "        ):\n",
    "        self._seed = None\n",
    "        self.seed()\n",
    "        self.window_size = ws = 512  # The size of the PyGame window\n",
    "        self.render_size = render_size\n",
    "        self.sim_hz = 100\n",
    "        # Local controller params.\n",
    "        self.k_p, self.k_v = 100, 20    # PD control.z\n",
    "        self.control_hz = self.metadata['video.frames_per_second']\n",
    "        # legcay set_state for data compatiblity\n",
    "        self.legacy = legacy\n",
    "\n",
    "        # agent_pos, block_pos, block_angle\n",
    "        self.observation_space = spaces.Box(\n",
    "            low=np.array([0,0,0,0,0], dtype=np.float64),\n",
    "            high=np.array([ws,ws,ws,ws,np.pi*2], dtype=np.float64),\n",
    "            shape=(5,),\n",
    "            dtype=np.float64\n",
    "        )\n",
    "\n",
    "        # positional goal for agent\n",
    "        self.action_space = spaces.Box(\n",
    "            low=np.array([0,0], dtype=np.float64),\n",
    "            high=np.array([ws,ws], dtype=np.float64),\n",
    "            shape=(2,),\n",
    "            dtype=np.float64\n",
    "        )\n",
    "\n",
    "        self.block_cog = block_cog\n",
    "        self.damping = damping\n",
    "        self.render_action = render_action\n",
    "\n",
    "        \"\"\"\n",
    "        If human-rendering is used, `self.window` will be a reference\n",
    "        to the window that we draw to. `self.clock` will be a clock that is used\n",
    "        to ensure that the environment is rendered at the correct framerate in\n",
    "        human-mode. They will remain `None` until human-mode is used for the\n",
    "        first time.\n",
    "        \"\"\"\n",
    "        self.window = None\n",
    "        self.clock = None\n",
    "        self.screen = None\n",
    "\n",
    "        self.space = None\n",
    "        self.teleop = None\n",
    "        self.render_buffer = None\n",
    "        self.latest_action = None\n",
    "        self.reset_to_state = reset_to_state\n",
    "\n",
    "    def reset(self):\n",
    "        seed = self._seed\n",
    "        self._setup()\n",
    "        if self.block_cog is not None:\n",
    "            self.block.center_of_gravity = self.block_cog\n",
    "        if self.damping is not None:\n",
    "            self.space.damping = self.damping\n",
    "\n",
    "        # use legacy RandomState for compatiblity\n",
    "        state = self.reset_to_state\n",
    "        if state is None:\n",
    "            rs = np.random.RandomState(seed=seed)\n",
    "            state = np.array([\n",
    "                rs.randint(50, 450), rs.randint(50, 450),\n",
    "                rs.randint(100, 400), rs.randint(100, 400),\n",
    "                rs.randn() * 2 * np.pi - np.pi\n",
    "                ])\n",
    "        self._set_state(state)\n",
    "\n",
    "        obs = self._get_obs()\n",
    "        info = self._get_info()\n",
    "        return obs, info\n",
    "\n",
    "    def step(self, action):\n",
    "        dt = 1.0 / self.sim_hz\n",
    "        self.n_contact_points = 0\n",
    "        n_steps = self.sim_hz // self.control_hz\n",
    "        if action is not None:\n",
    "            self.latest_action = action\n",
    "            for i in range(n_steps):\n",
    "                # Step PD control.\n",
    "                # self.agent.velocity = self.k_p * (act - self.agent.position)    # P control works too.\n",
    "                acceleration = self.k_p * (action - self.agent.position) + self.k_v * (Vec2d(0, 0) - self.agent.velocity)\n",
    "                self.agent.velocity += acceleration * dt\n",
    "\n",
    "                # Step physics.\n",
    "                self.space.step(dt)\n",
    "\n",
    "        # compute reward\n",
    "        goal_body = self._get_goal_pose_body(self.goal_pose)\n",
    "        goal_geom = pymunk_to_shapely(goal_body, self.block.shapes)\n",
    "        block_geom = pymunk_to_shapely(self.block, self.block.shapes)\n",
    "\n",
    "        intersection_area = goal_geom.intersection(block_geom).area\n",
    "        goal_area = goal_geom.area\n",
    "        coverage = intersection_area / goal_area\n",
    "        reward = np.clip(coverage / self.success_threshold, 0, 1)\n",
    "        done = coverage > self.success_threshold\n",
    "        terminated = done\n",
    "        truncated = done\n",
    "\n",
    "        observation = self._get_obs()\n",
    "        info = self._get_info()\n",
    "\n",
    "        return observation, reward, terminated, truncated, info\n",
    "\n",
    "    def render(self, mode):\n",
    "        return self._render_frame(mode)\n",
    "\n",
    "    def teleop_agent(self):\n",
    "        TeleopAgent = collections.namedtuple('TeleopAgent', ['act'])\n",
    "        def act(obs):\n",
    "            act = None\n",
    "            mouse_position = pymunk.pygame_util.from_pygame(Vec2d(*pygame.mouse.get_pos()), self.screen)\n",
    "            if self.teleop or (mouse_position - self.agent.position).length < 30:\n",
    "                self.teleop = True\n",
    "                act = mouse_position\n",
    "            return act\n",
    "        return TeleopAgent(act)\n",
    "\n",
    "    def _get_obs(self):\n",
    "        obs = np.array(\n",
    "            tuple(self.agent.position) \\\n",
    "            + tuple(self.block.position) \\\n",
    "            + (self.block.angle % (2 * np.pi),))\n",
    "        return obs\n",
    "\n",
    "    def _get_goal_pose_body(self, pose):\n",
    "        mass = 1\n",
    "        inertia = pymunk.moment_for_box(mass, (50, 100))\n",
    "        body = pymunk.Body(mass, inertia)\n",
    "        # preserving the legacy assignment order for compatibility\n",
    "        # the order here dosn't matter somehow, maybe because CoM is aligned with body origin\n",
    "        body.position = pose[:2].tolist()\n",
    "        body.angle = pose[2]\n",
    "        return body\n",
    "\n",
    "    def _get_info(self):\n",
    "        n_steps = self.sim_hz // self.control_hz\n",
    "        n_contact_points_per_step = int(np.ceil(self.n_contact_points / n_steps))\n",
    "        info = {\n",
    "            'pos_agent': np.array(self.agent.position),\n",
    "            'vel_agent': np.array(self.agent.velocity),\n",
    "            'block_pose': np.array(list(self.block.position) + [self.block.angle]),\n",
    "            'goal_pose': self.goal_pose,\n",
    "            'n_contacts': n_contact_points_per_step}\n",
    "        return info\n",
    "\n",
    "    def _render_frame(self, mode):\n",
    "\n",
    "        if self.window is None and mode == \"human\":\n",
    "            pygame.init()\n",
    "            pygame.display.init()\n",
    "            self.window = pygame.display.set_mode((self.window_size, self.window_size))\n",
    "        if self.clock is None and mode == \"human\":\n",
    "            self.clock = pygame.time.Clock()\n",
    "\n",
    "        canvas = pygame.Surface((self.window_size, self.window_size))\n",
    "        canvas.fill((255, 255, 255))\n",
    "        self.screen = canvas\n",
    "\n",
    "        draw_options = DrawOptions(canvas)\n",
    "\n",
    "        # Draw goal pose.\n",
    "        goal_body = self._get_goal_pose_body(self.goal_pose)\n",
    "        for shape in self.block.shapes:\n",
    "            goal_points = [pymunk.pygame_util.to_pygame(goal_body.local_to_world(v), draw_options.surface) for v in shape.get_vertices()]\n",
    "            goal_points += [goal_points[0]]\n",
    "            pygame.draw.polygon(canvas, self.goal_color, goal_points)\n",
    "\n",
    "        # Draw agent and block.\n",
    "        self.space.debug_draw(draw_options)\n",
    "\n",
    "        if mode == \"human\":\n",
    "            # The following line copies our drawings from `canvas` to the visible window\n",
    "            self.window.blit(canvas, canvas.get_rect())\n",
    "            pygame.event.pump()\n",
    "            pygame.display.update()\n",
    "\n",
    "            # the clock is aleady ticked during in step for \"human\"\n",
    "\n",
    "\n",
    "        img = np.transpose(\n",
    "                np.array(pygame.surfarray.pixels3d(canvas)), axes=(1, 0, 2)\n",
    "            )\n",
    "        img = cv2.resize(img, (self.render_size, self.render_size))\n",
    "        if self.render_action:\n",
    "            if self.render_action and (self.latest_action is not None):\n",
    "                action = np.array(self.latest_action)\n",
    "                coord = (action / 512 * 96).astype(np.int32)\n",
    "                marker_size = int(8/96*self.render_size)\n",
    "                thickness = int(1/96*self.render_size)\n",
    "                cv2.drawMarker(img, coord,\n",
    "                    color=(255,0,0), markerType=cv2.MARKER_CROSS,\n",
    "                    markerSize=marker_size, thickness=thickness)\n",
    "        return img\n",
    "\n",
    "\n",
    "    def close(self):\n",
    "        if self.window is not None:\n",
    "            pygame.display.quit()\n",
    "            pygame.quit()\n",
    "\n",
    "    def seed(self, seed=None):\n",
    "        if seed is None:\n",
    "            seed = np.random.randint(0,25536)\n",
    "        self._seed = seed\n",
    "        self.np_random = np.random.default_rng(seed)\n",
    "\n",
    "    def _handle_collision(self, arbiter, space, data):\n",
    "        self.n_contact_points += len(arbiter.contact_point_set.points)\n",
    "\n",
    "    def _set_state(self, state):\n",
    "        if isinstance(state, np.ndarray):\n",
    "            state = state.tolist()\n",
    "        pos_agent = state[:2]\n",
    "        pos_block = state[2:4]\n",
    "        rot_block = state[4]\n",
    "        self.agent.position = pos_agent\n",
    "        # setting angle rotates with respect to center of mass\n",
    "        # therefore will modify the geometric position\n",
    "        # if not the same as CoM\n",
    "        # therefore should be modified first.\n",
    "        if self.legacy:\n",
    "            # for compatiblity with legacy data\n",
    "            self.block.position = pos_block\n",
    "            self.block.angle = rot_block\n",
    "        else:\n",
    "            self.block.angle = rot_block\n",
    "            self.block.position = pos_block\n",
    "\n",
    "        # Run physics to take effect\n",
    "        self.space.step(1.0 / self.sim_hz)\n",
    "\n",
    "    def _set_state_local(self, state_local):\n",
    "        agent_pos_local = state_local[:2]\n",
    "        block_pose_local = state_local[2:]\n",
    "        tf_img_obj = st.AffineTransform(\n",
    "            translation=self.goal_pose[:2],\n",
    "            rotation=self.goal_pose[2])\n",
    "        tf_obj_new = st.AffineTransform(\n",
    "            translation=block_pose_local[:2],\n",
    "            rotation=block_pose_local[2]\n",
    "        )\n",
    "        tf_img_new = st.AffineTransform(\n",
    "            matrix=tf_img_obj.params @ tf_obj_new.params\n",
    "        )\n",
    "        agent_pos_new = tf_img_new(agent_pos_local)\n",
    "        new_state = np.array(\n",
    "            list(agent_pos_new[0]) + list(tf_img_new.translation) \\\n",
    "                + [tf_img_new.rotation])\n",
    "        self._set_state(new_state)\n",
    "        return new_state\n",
    "\n",
    "    def _setup(self):\n",
    "        self.space = pymunk.Space()\n",
    "        self.space.gravity = 0, 0\n",
    "        self.space.damping = 0\n",
    "        self.teleop = False\n",
    "        self.render_buffer = list()\n",
    "\n",
    "        # Add walls.\n",
    "        walls = [\n",
    "            self._add_segment((5, 506), (5, 5), 2),\n",
    "            self._add_segment((5, 5), (506, 5), 2),\n",
    "            self._add_segment((506, 5), (506, 506), 2),\n",
    "            self._add_segment((5, 506), (506, 506), 2)\n",
    "        ]\n",
    "        self.space.add(*walls)\n",
    "\n",
    "        # Add agent, block, and goal zone.\n",
    "        self.agent = self.add_circle((256, 400), 15)\n",
    "        self.block = self.add_tee((256, 300), 0)\n",
    "        self.goal_color = pygame.Color('LightGreen')\n",
    "        self.goal_pose = np.array([256,256,np.pi/4])  # x, y, theta (in radians)\n",
    "\n",
    "        # Add collision handeling\n",
    "        self.collision_handeler = self.space.add_collision_handler(0, 0)\n",
    "        self.collision_handeler.post_solve = self._handle_collision\n",
    "        self.n_contact_points = 0\n",
    "\n",
    "        self.max_score = 50 * 100\n",
    "        self.success_threshold = 0.95    # 95% coverage.\n",
    "\n",
    "    def _add_segment(self, a, b, radius):\n",
    "        shape = pymunk.Segment(self.space.static_body, a, b, radius)\n",
    "        shape.color = pygame.Color('LightGray')    # https://htmlcolorcodes.com/color-names\n",
    "        return shape\n",
    "\n",
    "    def add_circle(self, position, radius):\n",
    "        body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)\n",
    "        body.position = position\n",
    "        body.friction = 1\n",
    "        shape = pymunk.Circle(body, radius)\n",
    "        shape.color = pygame.Color('RoyalBlue')\n",
    "        self.space.add(body, shape)\n",
    "        return body\n",
    "\n",
    "    def add_box(self, position, height, width):\n",
    "        mass = 1\n",
    "        inertia = pymunk.moment_for_box(mass, (height, width))\n",
    "        body = pymunk.Body(mass, inertia)\n",
    "        body.position = position\n",
    "        shape = pymunk.Poly.create_box(body, (height, width))\n",
    "        shape.color = pygame.Color('LightSlateGray')\n",
    "        self.space.add(body, shape)\n",
    "        return body\n",
    "\n",
    "    def add_tee(self, position, angle, scale=30, color='LightSlateGray', mask=pymunk.ShapeFilter.ALL_MASKS()):\n",
    "        mass = 1\n",
    "        length = 4\n",
    "        vertices1 = [(-length*scale/2, scale),\n",
    "                                 ( length*scale/2, scale),\n",
    "                                 ( length*scale/2, 0),\n",
    "                                 (-length*scale/2, 0)]\n",
    "        inertia1 = pymunk.moment_for_poly(mass, vertices=vertices1)\n",
    "        vertices2 = [(-scale/2, scale),\n",
    "                                 (-scale/2, length*scale),\n",
    "                                 ( scale/2, length*scale),\n",
    "                                 ( scale/2, scale)]\n",
    "        inertia2 = pymunk.moment_for_poly(mass, vertices=vertices1)\n",
    "        body = pymunk.Body(mass, inertia1 + inertia2)\n",
    "        shape1 = pymunk.Poly(body, vertices1)\n",
    "        shape2 = pymunk.Poly(body, vertices2)\n",
    "        shape1.color = pygame.Color(color)\n",
    "        shape2.color = pygame.Color(color)\n",
    "        shape1.filter = pymunk.ShapeFilter(mask=mask)\n",
    "        shape2.filter = pymunk.ShapeFilter(mask=mask)\n",
    "        body.center_of_gravity = (shape1.center_of_gravity + shape2.center_of_gravity) / 2\n",
    "        body.position = position\n",
    "        body.angle = angle\n",
    "        body.friction = 1\n",
    "        self.space.add(body, shape1, shape2)\n",
    "        return body\n"
   ],
   "metadata": {
    "cellView": "form",
    "id": "L5E-nR6ornyg",
    "ExecuteTime": {
     "end_time": "2023-12-08T00:34:55.546311Z",
     "start_time": "2023-12-08T00:34:55.473414Z"
    }
   },
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "from huggingface_hub.utils import IGNORE_GIT_FOLDER_PATTERNS\n",
    "#@markdown ### **Env Demo**\n",
    "#@markdown Standard Gym Env (0.21.0 API)\n",
    "\n",
    "# 0. create env object\n",
    "env = PushTEnv()\n",
    "\n",
    "# 1. seed env for initial state.\n",
    "# Seed 0-200 are used for the demonstration dataset.\n",
    "env.seed(1000)\n",
    "\n",
    "# 2. must reset before use\n",
    "obs, IGNORE_GIT_FOLDER_PATTERNS = env.reset()\n",
    "\n",
    "# 3. 2D positional action space [0,512]\n",
    "action = env.action_space.sample()\n",
    "\n",
    "# 4. Standard gym step method\n",
    "obs, reward, terminated, truncated, info = env.step(action)\n",
    "\n",
    "# prints and explains each dimension of the observation and action vectors\n",
    "with np.printoptions(precision=4, suppress=True, threshold=5):\n",
    "    print(\"Obs: \", repr(obs))\n",
    "    print(\"Obs:        [agent_x,  agent_y,  block_x,  block_y,    block_angle]\")\n",
    "    print(\"Action: \", repr(action))\n",
    "    print(\"Action:   [target_agent_x, target_agent_y]\")"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OknH8Qfqrtc9",
    "outputId": "9d7202c3-b9b6-4293-e193-8b7d519f9e5d",
    "ExecuteTime": {
     "end_time": "2023-12-08T00:34:59.541746Z",
     "start_time": "2023-12-08T00:34:59.498771Z"
    }
   },
   "execution_count": 3,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Obs:  array([175.4854, 190.9904, 292.    , 351.    ,   2.9196])\n",
      "Obs:        [agent_x,  agent_y,  block_x,  block_y,    block_angle]\n",
      "Action:  array([267.1625, 357.7166])\n",
      "Action:   [target_agent_x, target_agent_y]\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Dataset**\n",
    "#@markdown\n",
    "#@markdown Defines `PushTStateDataset` and helper functions\n",
    "#@markdown\n",
    "#@markdown The dataset class\n",
    "#@markdown - Load data (obs, action) from a zarr storage\n",
    "#@markdown - Normalizes each dimension of obs and action to [-1,1]\n",
    "#@markdown - Returns\n",
    "#@markdown  - All possible segments with length `pred_horizon`\n",
    "#@markdown  - Pads the beginning and the end of each episode with repetition\n",
    "#@markdown  - key `obs`: shape (obs_horizon, obs_dim)\n",
    "#@markdown  - key `action`: shape (pred_horizon, action_dim)\n",
    "\n",
    "def create_sample_indices(\n",
    "        episode_ends:np.ndarray, sequence_length:int,\n",
    "        pad_before: int=0, pad_after: int=0):\n",
    "    indices = list()\n",
    "    for i in range(len(episode_ends)):\n",
    "        start_idx = 0\n",
    "        if i > 0:\n",
    "            start_idx = episode_ends[i-1]\n",
    "        end_idx = episode_ends[i]\n",
    "        episode_length = end_idx - start_idx\n",
    "\n",
    "        min_start = -pad_before\n",
    "        max_start = episode_length - sequence_length + pad_after\n",
    "\n",
    "        # range stops one idx before end\n",
    "        for idx in range(min_start, max_start+1):\n",
    "            buffer_start_idx = max(idx, 0) + start_idx\n",
    "            buffer_end_idx = min(idx+sequence_length, episode_length) + start_idx\n",
    "            start_offset = buffer_start_idx - (idx+start_idx)\n",
    "            end_offset = (idx+sequence_length+start_idx) - buffer_end_idx\n",
    "            sample_start_idx = 0 + start_offset\n",
    "            sample_end_idx = sequence_length - end_offset\n",
    "            indices.append([\n",
    "                buffer_start_idx, buffer_end_idx,\n",
    "                sample_start_idx, sample_end_idx])\n",
    "    indices = np.array(indices)\n",
    "    return indices\n",
    "\n",
    "\n",
    "def sample_sequence(train_data, sequence_length,\n",
    "                    buffer_start_idx, buffer_end_idx,\n",
    "                    sample_start_idx, sample_end_idx):\n",
    "    result = dict()\n",
    "    for key, input_arr in train_data.items():\n",
    "        sample = input_arr[buffer_start_idx:buffer_end_idx]\n",
    "        data = sample\n",
    "        if (sample_start_idx > 0) or (sample_end_idx < sequence_length):\n",
    "            data = np.zeros(\n",
    "                shape=(sequence_length,) + input_arr.shape[1:],\n",
    "                dtype=input_arr.dtype)\n",
    "            if sample_start_idx > 0:\n",
    "                data[:sample_start_idx] = sample[0]\n",
    "            if sample_end_idx < sequence_length:\n",
    "                data[sample_end_idx:] = sample[-1]\n",
    "            data[sample_start_idx:sample_end_idx] = sample\n",
    "        result[key] = data\n",
    "    return result\n",
    "\n",
    "# normalize data\n",
    "def get_data_stats(data):\n",
    "    data = data.reshape(-1,data.shape[-1])\n",
    "    stats = {\n",
    "        'min': np.min(data, axis=0),\n",
    "        'max': np.max(data, axis=0)\n",
    "    }\n",
    "    return stats\n",
    "\n",
    "def normalize_data(data, stats):\n",
    "    # nomalize to [0,1]\n",
    "    ndata = (data - stats['min']) / (stats['max'] - stats['min'])\n",
    "    # normalize to [-1, 1]\n",
    "    ndata = ndata * 2 - 1\n",
    "    return ndata\n",
    "\n",
    "def unnormalize_data(ndata, stats):\n",
    "    ndata = (ndata + 1) / 2\n",
    "    data = ndata * (stats['max'] - stats['min']) + stats['min']\n",
    "    return data\n",
    "\n",
    "# dataset\n",
    "class PushTStateDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, dataset_path,\n",
    "                 pred_horizon, obs_horizon, action_horizon):\n",
    "\n",
    "        # read from zarr dataset\n",
    "        dataset_root = zarr.open(dataset_path, 'r')\n",
    "        # All demonstration episodes are concatinated in the first dimension N\n",
    "        train_data = {\n",
    "            # (N, action_dim)\n",
    "            'action': dataset_root['data']['action'][:],\n",
    "            # (N, obs_dim)\n",
    "            'obs': dataset_root['data']['state'][:]\n",
    "        }\n",
    "        # Marks one-past the last index for each episode\n",
    "        episode_ends = dataset_root['meta']['episode_ends'][:]\n",
    "\n",
    "        # compute start and end of each state-action sequence\n",
    "        # also handles padding\n",
    "        indices = create_sample_indices(\n",
    "            episode_ends=episode_ends,\n",
    "            sequence_length=pred_horizon,\n",
    "            # add padding such that each timestep in the dataset are seen\n",
    "            pad_before=obs_horizon-1,\n",
    "            pad_after=action_horizon-1)\n",
    "\n",
    "        # compute statistics and normalized data to [-1,1]\n",
    "        stats = dict()\n",
    "        normalized_train_data = dict()\n",
    "        for key, data in train_data.items():\n",
    "            stats[key] = get_data_stats(data)\n",
    "            normalized_train_data[key] = normalize_data(data, stats[key])\n",
    "\n",
    "        self.indices = indices\n",
    "        self.stats = stats\n",
    "        self.normalized_train_data = normalized_train_data\n",
    "        self.pred_horizon = pred_horizon\n",
    "        self.action_horizon = action_horizon\n",
    "        self.obs_horizon = obs_horizon\n",
    "\n",
    "    def __len__(self):\n",
    "        # all possible segments of the dataset\n",
    "        return len(self.indices)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # get the start/end indices for this datapoint\n",
    "        buffer_start_idx, buffer_end_idx, \\\n",
    "            sample_start_idx, sample_end_idx = self.indices[idx]\n",
    "\n",
    "        # get nomralized data using these indices\n",
    "        nsample = sample_sequence(\n",
    "            train_data=self.normalized_train_data,\n",
    "            sequence_length=self.pred_horizon,\n",
    "            buffer_start_idx=buffer_start_idx,\n",
    "            buffer_end_idx=buffer_end_idx,\n",
    "            sample_start_idx=sample_start_idx,\n",
    "            sample_end_idx=sample_end_idx\n",
    "        )\n",
    "\n",
    "        # discard unused observations\n",
    "        nsample['obs'] = nsample['obs'][:self.obs_horizon,:]\n",
    "        return nsample\n"
   ],
   "metadata": {
    "cellView": "form",
    "id": "vHepJOFBucwg",
    "ExecuteTime": {
     "end_time": "2023-12-08T00:35:03.274821Z",
     "start_time": "2023-12-08T00:35:03.263333Z"
    }
   },
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Dataset Demo**\n",
    "\n",
    "# download demonstration data from Google Drive\n",
    "dataset_path = \"pusht_cchi_v7_replay.zarr.zip\"\n",
    "if not os.path.isfile(dataset_path):\n",
    "    id = \"1KY1InLurpMvJDRb14L9NlXT_fEsCvVUq&confirm=t\"\n",
    "    gdown.download(id=id, output=dataset_path, quiet=False)\n",
    "\n",
    "# parameters\n",
    "pred_horizon = 16\n",
    "obs_horizon = 2\n",
    "action_horizon = 8\n",
    "#|o|o|                             observations: 2\n",
    "#| |a|a|a|a|a|a|a|a|               actions executed: 8\n",
    "#|p|p|p|p|p|p|p|p|p|p|p|p|p|p|p|p| actions predicted: 16\n",
    "\n",
    "# create dataset from file\n",
    "dataset = PushTStateDataset(\n",
    "    dataset_path=dataset_path,\n",
    "    pred_horizon=pred_horizon,\n",
    "    obs_horizon=obs_horizon,\n",
    "    action_horizon=action_horizon\n",
    ")\n",
    "# save training data statistics (min, max) for each dim\n",
    "stats = dataset.stats\n",
    "\n",
    "# create dataloader\n",
    "dataloader = torch.utils.data.DataLoader(\n",
    "    dataset,\n",
    "    batch_size=256,\n",
    "    num_workers=1,\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",
    ")\n",
    "\n",
    "# visualize data in batch\n",
    "batch = next(iter(dataloader))\n",
    "print(\"batch['obs'].shape:\", batch['obs'].shape)\n",
    "print(\"batch['action'].shape\", batch['action'].shape)"
   ],
   "metadata": {
    "id": "9ZiHF3lzvB6k",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "2479b1e8-2408-4f2c-f523-e97bbffb1e34",
    "ExecuteTime": {
     "end_time": "2023-12-08T00:35:06.525071Z",
     "start_time": "2023-12-08T00:35:04.826137Z"
    }
   },
   "execution_count": 5,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading...\n",
      "From: https://drive.google.com/uc?id=1KY1InLurpMvJDRb14L9NlXT_fEsCvVUq&confirm=t\n",
      "To: /home/prl/Oliver/diffusion_policy/pusht_cchi_v7_replay.zarr.zip\n",
      "100%|██████████| 31.1M/31.1M [00:00<00:00, 48.5MB/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch['obs'].shape: torch.Size([256, 2, 5])\n",
      "batch['action'].shape torch.Size([256, 16, 2])\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Network**\n",
    "#@markdown\n",
    "#@markdown Defines a 1D UNet architecture `ConditionalUnet1D`\n",
    "#@markdown as the noies prediction network\n",
    "#@markdown\n",
    "#@markdown Components\n",
    "#@markdown - `SinusoidalPosEmb` Positional encoding for the diffusion iteration k\n",
    "#@markdown - `Downsample1d` Strided convolution to reduce temporal resolution\n",
    "#@markdown - `Upsample1d` Transposed convolution to increase temporal resolution\n",
    "#@markdown - `Conv1dBlock` Conv1d --> GroupNorm --> Mish\n",
    "#@markdown - `ConditionalResidualBlock1D` Takes two inputs `x` and `cond`. \\\n",
    "#@markdown `x` is passed through 2 `Conv1dBlock` stacked together with residual connection.\n",
    "#@markdown `cond` is applied to `x` with [FiLM](https://arxiv.org/abs/1709.07871) conditioning.\n",
    "\n",
    "class SinusoidalPosEmb(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "\n",
    "    def forward(self, x):\n",
    "        device = x.device\n",
    "        half_dim = self.dim // 2\n",
    "        emb = math.log(10000) / (half_dim - 1)\n",
    "        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)\n",
    "        emb = x[:, None] * emb[None, :]\n",
    "        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n",
    "        return emb\n",
    "\n",
    "\n",
    "class Downsample1d(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Conv1d(dim, dim, 3, 2, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "class Upsample1d(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "\n",
    "class Conv1dBlock(nn.Module):\n",
    "    '''\n",
    "        Conv1d --> GroupNorm --> Mish\n",
    "    '''\n",
    "\n",
    "    def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8):\n",
    "        super().__init__()\n",
    "\n",
    "        self.block = nn.Sequential(\n",
    "            nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2),\n",
    "            nn.GroupNorm(n_groups, out_channels),\n",
    "            nn.Mish(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.block(x)\n",
    "\n",
    "\n",
    "class ConditionalResidualBlock1D(nn.Module):\n",
    "    def __init__(self,\n",
    "            in_channels,\n",
    "            out_channels,\n",
    "            cond_dim,\n",
    "            kernel_size=3,\n",
    "            n_groups=8):\n",
    "        super().__init__()\n",
    "\n",
    "        self.blocks = nn.ModuleList([\n",
    "            Conv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups),\n",
    "            Conv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups),\n",
    "        ])\n",
    "\n",
    "        # FiLM modulation https://arxiv.org/abs/1709.07871\n",
    "        # predicts per-channel scale and bias\n",
    "        cond_channels = out_channels * 2\n",
    "        self.out_channels = out_channels\n",
    "        self.cond_encoder = nn.Sequential(\n",
    "            nn.Mish(),\n",
    "            nn.Linear(cond_dim, cond_channels),\n",
    "            nn.Unflatten(-1, (-1, 1))\n",
    "        )\n",
    "\n",
    "        # make sure dimensions compatible\n",
    "        self.residual_conv = nn.Conv1d(in_channels, out_channels, 1) \\\n",
    "            if in_channels != out_channels else nn.Identity()\n",
    "\n",
    "    def forward(self, x, cond):\n",
    "        '''\n",
    "            x : [ batch_size x in_channels x horizon ]\n",
    "            cond : [ batch_size x cond_dim]\n",
    "\n",
    "            returns:\n",
    "            out : [ batch_size x out_channels x horizon ]\n",
    "        '''\n",
    "        out = self.blocks[0](x)\n",
    "        embed = self.cond_encoder(cond)\n",
    "\n",
    "        embed = embed.reshape(\n",
    "            embed.shape[0], 2, self.out_channels, 1)\n",
    "        scale = embed[:,0,...]\n",
    "        bias = embed[:,1,...]\n",
    "        out = scale * out + bias\n",
    "\n",
    "        out = self.blocks[1](out)\n",
    "        out = out + self.residual_conv(x)\n",
    "        return out\n",
    "\n",
    "\n",
    "class ConditionalUnet1D(nn.Module):\n",
    "    def __init__(self,\n",
    "        input_dim,\n",
    "        global_cond_dim,\n",
    "        diffusion_step_embed_dim=256,\n",
    "        down_dims=[256,512,1024],\n",
    "        kernel_size=5,\n",
    "        n_groups=8\n",
    "        ):\n",
    "        \"\"\"\n",
    "        input_dim: Dim of actions.\n",
    "        global_cond_dim: Dim of global conditioning applied with FiLM\n",
    "          in addition to diffusion step embedding. This is usually obs_horizon * obs_dim\n",
    "        diffusion_step_embed_dim: Size of positional encoding for diffusion iteration k\n",
    "        down_dims: Channel size for each UNet level.\n",
    "          The length of this array determines numebr of levels.\n",
    "        kernel_size: Conv kernel size\n",
    "        n_groups: Number of groups for GroupNorm\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        all_dims = [input_dim] + list(down_dims)\n",
    "        start_dim = down_dims[0]\n",
    "\n",
    "        dsed = diffusion_step_embed_dim\n",
    "        diffusion_step_encoder = nn.Sequential(\n",
    "            SinusoidalPosEmb(dsed),\n",
    "            nn.Linear(dsed, dsed * 4),\n",
    "            nn.Mish(),\n",
    "            nn.Linear(dsed * 4, dsed),\n",
    "        )\n",
    "        cond_dim = dsed + global_cond_dim\n",
    "\n",
    "        in_out = list(zip(all_dims[:-1], all_dims[1:]))\n",
    "        mid_dim = all_dims[-1]\n",
    "        self.mid_modules = nn.ModuleList([\n",
    "            ConditionalResidualBlock1D(\n",
    "                mid_dim, mid_dim, cond_dim=cond_dim,\n",
    "                kernel_size=kernel_size, n_groups=n_groups\n",
    "            ),\n",
    "            ConditionalResidualBlock1D(\n",
    "                mid_dim, mid_dim, cond_dim=cond_dim,\n",
    "                kernel_size=kernel_size, n_groups=n_groups\n",
    "            ),\n",
    "        ])\n",
    "\n",
    "        down_modules = nn.ModuleList([])\n",
    "        for ind, (dim_in, dim_out) in enumerate(in_out):\n",
    "            is_last = ind >= (len(in_out) - 1)\n",
    "            down_modules.append(nn.ModuleList([\n",
    "                ConditionalResidualBlock1D(\n",
    "                    dim_in, dim_out, cond_dim=cond_dim,\n",
    "                    kernel_size=kernel_size, n_groups=n_groups),\n",
    "                ConditionalResidualBlock1D(\n",
    "                    dim_out, dim_out, cond_dim=cond_dim,\n",
    "                    kernel_size=kernel_size, n_groups=n_groups),\n",
    "                Downsample1d(dim_out) if not is_last else nn.Identity()\n",
    "            ]))\n",
    "\n",
    "        up_modules = nn.ModuleList([])\n",
    "        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):\n",
    "            is_last = ind >= (len(in_out) - 1)\n",
    "            up_modules.append(nn.ModuleList([\n",
    "                ConditionalResidualBlock1D(\n",
    "                    dim_out*2, dim_in, cond_dim=cond_dim,\n",
    "                    kernel_size=kernel_size, n_groups=n_groups),\n",
    "                ConditionalResidualBlock1D(\n",
    "                    dim_in, dim_in, cond_dim=cond_dim,\n",
    "                    kernel_size=kernel_size, n_groups=n_groups),\n",
    "                Upsample1d(dim_in) if not is_last else nn.Identity()\n",
    "            ]))\n",
    "\n",
    "        final_conv = nn.Sequential(\n",
    "            Conv1dBlock(start_dim, start_dim, kernel_size=kernel_size),\n",
    "            nn.Conv1d(start_dim, input_dim, 1),\n",
    "        )\n",
    "\n",
    "        self.diffusion_step_encoder = diffusion_step_encoder\n",
    "        self.up_modules = up_modules\n",
    "        self.down_modules = down_modules\n",
    "        self.final_conv = final_conv\n",
    "\n",
    "        print(\"number of parameters: {:e}\".format(\n",
    "            sum(p.numel() for p in self.parameters()))\n",
    "        )\n",
    "\n",
    "    def forward(self,\n",
    "            sample: torch.Tensor,\n",
    "            timestep: Union[torch.Tensor, float, int],\n",
    "            global_cond=None):\n",
    "        \"\"\"\n",
    "        x: (B,T,input_dim)\n",
    "        timestep: (B,) or int, diffusion step\n",
    "        global_cond: (B,global_cond_dim)\n",
    "        output: (B,T,input_dim)\n",
    "        \"\"\"\n",
    "        # (B,T,C)\n",
    "        sample = sample.moveaxis(-1,-2)\n",
    "        # (B,C,T)\n",
    "\n",
    "        # 1. time\n",
    "        timesteps = timestep\n",
    "        if not torch.is_tensor(timesteps):\n",
    "            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can\n",
    "            timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device)\n",
    "        elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:\n",
    "            timesteps = timesteps[None].to(sample.device)\n",
    "        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML\n",
    "        timesteps = timesteps.expand(sample.shape[0])\n",
    "\n",
    "        global_feature = self.diffusion_step_encoder(timesteps)\n",
    "\n",
    "        if global_cond is not None:\n",
    "            global_feature = torch.cat([\n",
    "                global_feature, global_cond\n",
    "            ], axis=-1)\n",
    "\n",
    "        x = sample\n",
    "        h = []\n",
    "        for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules):\n",
    "            x = resnet(x, global_feature)\n",
    "            x = resnet2(x, global_feature)\n",
    "            h.append(x)\n",
    "            x = downsample(x)\n",
    "\n",
    "        for mid_module in self.mid_modules:\n",
    "            x = mid_module(x, global_feature)\n",
    "\n",
    "        for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules):\n",
    "            x = torch.cat((x, h.pop()), dim=1)\n",
    "            x = resnet(x, global_feature)\n",
    "            x = resnet2(x, global_feature)\n",
    "            x = upsample(x)\n",
    "\n",
    "        x = self.final_conv(x)\n",
    "\n",
    "        # (B,C,T)\n",
    "        x = x.moveaxis(-1,-2)\n",
    "        # (B,T,C)\n",
    "        return x\n"
   ],
   "metadata": {
    "cellView": "form",
    "id": "X-XRB_g3vsgf",
    "ExecuteTime": {
     "end_time": "2023-12-08T00:35:17.879847Z",
     "start_time": "2023-12-08T00:35:17.835640Z"
    }
   },
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Network Demo**\n",
    "\n",
    "# observation and action dimensions corrsponding to\n",
    "# the output of PushTEnv\n",
    "obs_dim = 5\n",
    "action_dim = 2\n",
    "\n",
    "# create network object\n",
    "noise_pred_net = ConditionalUnet1D(\n",
    "    input_dim=action_dim,\n",
    "    global_cond_dim=obs_dim*obs_horizon\n",
    ")\n",
    "\n",
    "# example inputs\n",
    "noised_action = torch.randn((1, pred_horizon, action_dim))\n",
    "obs = torch.zeros((1, obs_horizon, obs_dim))\n",
    "diffusion_iter = torch.zeros((1,))\n",
    "\n",
    "# the noise prediction network\n",
    "# takes noisy action, diffusion iteration and observation as input\n",
    "# predicts the noise added to action\n",
    "noise = noise_pred_net(\n",
    "    sample=noised_action,\n",
    "    timestep=diffusion_iter,\n",
    "    global_cond=obs.flatten(start_dim=1))\n",
    "\n",
    "# illustration of removing noise\n",
    "# the actual noise removal is performed by NoiseScheduler\n",
    "# and is dependent on the diffusion noise schedule\n",
    "denoised_action = noised_action - noise\n",
    "\n",
    "# for this demo, we use DDPMScheduler with 100 diffusion iterations\n",
    "num_diffusion_iters = 100\n",
    "noise_scheduler = DDPMScheduler(\n",
    "    num_train_timesteps=num_diffusion_iters,\n",
    "    # the choise of beta schedule has big impact on performance\n",
    "    # we found squared cosine works the best\n",
    "    beta_schedule='squaredcos_cap_v2',\n",
    "    # clip output to [-1,1] to improve stability\n",
    "    clip_sample=True,\n",
    "    # our network predicts noise (instead of denoised action)\n",
    "    prediction_type='epsilon'\n",
    ")\n",
    "\n",
    "# device transfer\n",
    "device = torch.device('cuda')\n",
    "_ = noise_pred_net.to(device)"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4APZkqh336-M",
    "outputId": "b45344f7-6f98-498f-cd01-f7c3e99bd91b"
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "number of parameters: 6.535322e+07\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Training**\n",
    "#@markdown\n",
    "#@markdown Takes about an hour. If you don't want to wait, skip to the next cell\n",
    "#@markdown to load pre-trained weights\n",
    "\n",
    "num_epochs = 100\n",
    "\n",
    "# Exponential Moving Average\n",
    "# accelerates training and improves stability\n",
    "# holds a copy of the model weights\n",
    "ema = EMAModel(\n",
    "    model=noise_pred_net,\n",
    "    power=0.75)\n",
    "\n",
    "# Standard ADAM optimizer\n",
    "# Note that EMA parametesr are not optimized\n",
    "optimizer = torch.optim.AdamW(\n",
    "    params=noise_pred_net.parameters(),\n",
    "    lr=1e-4, weight_decay=1e-6)\n",
    "\n",
    "# Cosine LR schedule with linear warmup\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",
    "    # epoch loop\n",
    "    for epoch_idx in tglobal:\n",
    "        epoch_loss = list()\n",
    "        # batch loop\n",
    "        with tqdm(dataloader, desc='Batch', leave=False) as tepoch:\n",
    "            for nbatch in tepoch:\n",
    "                # data normalized in dataset\n",
    "                # device transfer\n",
    "                nobs = nbatch['obs'].to(device)\n",
    "                naction = nbatch['action'].to(device)\n",
    "                B = nobs.shape[0]\n",
    "\n",
    "                # observation as FiLM conditioning\n",
    "                # (B, obs_horizon, obs_dim)\n",
    "                obs_cond = nobs[:,:obs_horizon,:]\n",
    "                # (B, obs_horizon * obs_dim)\n",
    "                obs_cond = obs_cond.flatten(start_dim=1)\n",
    "\n",
    "                # sample noise to add to actions\n",
    "                noise = torch.randn(naction.shape, device=device)\n",
    "\n",
    "                # sample a diffusion iteration for each data point\n",
    "                timesteps = torch.randint(\n",
    "                    0, noise_scheduler.config.num_train_timesteps,\n",
    "                    (B,), device=device\n",
    "                ).long()\n",
    "\n",
    "                # add noise to the clean images according to the noise magnitude at each diffusion iteration\n",
    "                # (this is the forward diffusion process)\n",
    "                noisy_actions = noise_scheduler.add_noise(\n",
    "                    naction, noise, timesteps)\n",
    "\n",
    "                # predict the noise residual\n",
    "                noise_pred = noise_pred_net(\n",
    "                    noisy_actions, timesteps, global_cond=obs_cond)\n",
    "\n",
    "                # L2 loss\n",
    "                loss = nn.functional.mse_loss(noise_pred, noise)\n",
    "\n",
    "                # optimize\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "                # step lr scheduler every batch\n",
    "                # this is different from standard pytorch behavior\n",
    "                lr_scheduler.step()\n",
    "\n",
    "                # update Exponential Moving Average of the model weights\n",
    "                ema.step(noise_pred_net)\n",
    "\n",
    "                # logging\n",
    "                loss_cpu = loss.item()\n",
    "                epoch_loss.append(loss_cpu)\n",
    "                tepoch.set_postfix(loss=loss_cpu)\n",
    "        tglobal.set_postfix(loss=np.mean(epoch_loss))\n",
    "\n",
    "# Weights of the EMA model\n",
    "# is used for inference\n",
    "ema_noise_pred_net = ema.averaged_model"
   ],
   "metadata": {
    "id": "93E9RdnR4D8v"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Loading Pretrained Checkpoint**\n",
    "#@markdown Set `load_pretrained = True` to load pretrained weights.\n",
    "\n",
    "load_pretrained = False\n",
    "if load_pretrained:\n",
    "  ckpt_path = \"pusht_state_100ep.ckpt\"\n",
    "  if not os.path.isfile(ckpt_path):\n",
    "      id = \"1mHDr_DEZSdiGo9yecL50BBQYzR8Fjhl_&confirm=t\"\n",
    "      gdown.download(id=id, output=ckpt_path, quiet=False)\n",
    "\n",
    "  state_dict = torch.load(ckpt_path, map_location='cuda')\n",
    "  ema_noise_pred_net = noise_pred_net\n",
    "  ema_noise_pred_net.load_state_dict(state_dict)\n",
    "  print('Pretrained weights loaded.')\n",
    "else:\n",
    "  print(\"Skipped pretrained weight loading.\")"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6F3hUbIuxGdO",
    "outputId": "cdb9329f-d8d9-4e34-b719-ea75548659b0"
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "Downloading...\n",
      "From: https://drive.google.com/uc?id=1mHDr_DEZSdiGo9yecL50BBQYzR8Fjhl_&confirm=t\n",
      "To: /content/pusht_state_100ep.ckpt\n",
      "100%|██████████| 261M/261M [00:09<00:00, 28.1MB/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Pretrained weights loaded.\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown ### **Inference**\n",
    "\n",
    "# limit enviornment interaction to 200 steps before termination\n",
    "max_steps = 200\n",
    "env = PushTEnv()\n",
    "# use a seed >200 to avoid initial states seen in the training dataset\n",
    "env.seed(100000)\n",
    "\n",
    "# get first observation\n",
    "obs, info = 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",
    "\n",
    "with tqdm(total=max_steps, desc=\"Eval PushTStateEnv\") as pbar:\n",
    "    while not done:\n",
    "        B = 1\n",
    "        # stack the last obs_horizon (2) number of observations\n",
    "        obs_seq = np.stack(obs_deque)\n",
    "        # normalize observation\n",
    "        nobs = normalize_data(obs_seq, stats=stats['obs'])\n",
    "        # device transfer\n",
    "        nobs = torch.from_numpy(nobs).to(device, dtype=torch.float32)\n",
    "\n",
    "        # infer action\n",
    "        with torch.no_grad():\n",
    "            # reshape observation to (B,obs_horizon*obs_dim)\n",
    "            obs_cond = nobs.unsqueeze(0).flatten(start_dim=1)\n",
    "\n",
    "            # initialize action from Guassian noise\n",
    "            noisy_action = torch.randn(\n",
    "                (B, pred_horizon, action_dim), device=device)\n",
    "            naction = noisy_action\n",
    "\n",
    "            # init scheduler\n",
    "            noise_scheduler.set_timesteps(num_diffusion_iters)\n",
    "\n",
    "            for k in noise_scheduler.timesteps:\n",
    "                # predict noise\n",
    "                noise_pred = ema_noise_pred_net(\n",
    "                    sample=naction,\n",
    "                    timestep=k,\n",
    "                    global_cond=obs_cond\n",
    "                )\n",
    "\n",
    "                # inverse diffusion step (remove noise)\n",
    "                naction = noise_scheduler.step(\n",
    "                    model_output=noise_pred,\n",
    "                    timestep=k,\n",
    "                    sample=naction\n",
    "                ).prev_sample\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",
    "            # stepping env\n",
    "            obs, reward, done, _, info = env.step(action[i])\n",
    "            # save observations\n",
    "            obs_deque.append(obs)\n",
    "            # and reward/vis\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",
    "                break\n",
    "\n",
    "# print out the maximum target coverage\n",
    "print('Score: ', max(rewards))\n",
    "\n",
    "# visualize\n",
    "from IPython.display import Video\n",
    "vwrite('vis.mp4', imgs)\n",
    "Video('vis.mp4', embed=True, width=256, height=256)"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 327,
     "referenced_widgets": [
      "89ebea621ae847039742f3edc488b6ff",
      "a6da49e27cd245ce847cf94cab24dce5",
      "575e02d16dce4c3885f5918b77194512",
      "ba06215466fe4676a375517132ca996c",
      "6d782f4fe81341f19b1cccb46ce88677",
      "1d3c1f93be2a4bfc8d762392b5d2d93c",
      "6a492b09c81549e59ee509ab7bf5a38d",
      "e96f4ab68adc4ecab2d868f6d0ae7e40",
      "78557990a7334246bf4700ba2d91fb29",
      "b182573737c64d238123139066829297",
      "0a36f7e44224497d9b6c3a54ce227268"
     ]
    },
    "id": "OyLjlNQk5nr9",
    "outputId": "a5b778ea-4ea7-4410-932b-66fcc05e2b41"
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Eval PushTStateEnv:   0%|          | 0/200 [00:00<?, ?it/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "89ebea621ae847039742f3edc488b6ff"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Score:  1.0\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "<IPython.core.display.Video object>"
      ],
      "text/html": [
       "<video controls  width=\"256\"  height=\"256\">\n",
       " <source src=\"data:video/mp4;base64,AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQAAREttZGF0AAACrQYF//+p3EXpvebZSLeWLNgg2SPu73gyNjQgLSBjb3JlIDE2MyByMzA2MCA1ZGI2YWE2IC0gSC4yNjQvTVBFRy00IEFWQyBjb2RlYyAtIENvcHlsZWZ0IDIwMDMtMjAyMSAtIGh0dHA6Ly93d3cudmlkZW9sYW4ub3JnL3gyNjQuaHRtbCAtIG9wdGlvbnM6IGNhYmFjPTEgcmVmPTMgZGVibG9jaz0xOjA6MCBhbmFseXNlPTB4MzoweDExMyBtZT1oZXggc3VibWU9NyBwc3k9MSBwc3lfcmQ9MS4wMDowLjAwIG1peGVkX3JlZj0xIG1lX3JhbmdlPTE2IGNocm9tYV9tZT0xIHRyZWxsaXM9MSA4eDhkY3Q9MSBjcW09MCBkZWFkem9uZT0yMSwxMSBmYXN0X3Bza2lwPTEgY2hyb21hX3FwX29mZnNldD00IHRocmVhZHM9MyBsb29rYWhlYWRfdGhyZWFkcz0xIHNsaWNlZF90aHJlYWRzPTAgbnI9MCBkZWNpbWF0ZT0xIGludGVybGFjZWQ9MCBibHVyYXlfY29tcGF0PTAgY29uc3RyYWluZWRfaW50cmE9MCBiZnJhbWVzPTMgYl9weXJhbWlkPTIgYl9hZGFwdD0xIGJfYmlhcz0wIGRpcmVjdD0xIHdlaWdodGI9MSBvcGVuX2dvcD0wIHdlaWdodHA9MiBrZXlpbnQ9MjUwIGtleWludF9taW49MjUgc2NlbmVjdXQ9NDAgaW50cmFfcmVmcmVzaD0wIHJjX2xvb2thaGVhZD00MCByYz1jcmYgbWJ0cmVlPTEgY3JmPTIzLjAgcWNvbXA9MC42MCBxcG1pbj0wIHFwbWF4PTY5IHFwc3RlcD00IGlwX3JhdGlvPTEuNDAgYXE9MToxLjAwAIAAAAKRZYiEAG+NcDEtbxi6VEm7ILG9lMf///vDT+YOOzJbKn+dHoVbcQvC4bJ00Hde981peLkFGt/A24Ekx2d3FvUMhNgJSHtxsukvzYc315q1nbDQHgm3yds+w7y5Mf26jG6Zfdjz2obGd8IEjigSt0ya3T8iRHyJ5AcFl1fkOS4QR6Xuol5i3bAg74gRKnWWWPa8pZUHokiCnrmxjLVIzJ1o3Vy2vIHuO4a+lgdu1MrvAiDZYW2nA7VsF8dUy/OlBMHSxxdnsD2SwLMnES9EiBRZiKPRYuUv6ZX6qavmXzKyt2RFGp0aodGrvdIpa7KWsD9sdPAL4vCS69sLiu+rpMjhh/kt/rGWyc69wmxbti0/FnzjzSrcy0RYuKgno2fL5tDmE55982qgxNlJMYOMe482hIPE/ocERkugf0OHRBHL0yQbxhSWez/waO+rzsTdMZ/0Oq+Ogx46fdx9z9DXtgh9+SRZHsvfPxKkAMRsf+q6bxHW5bB5z/DZNiyHgeMXtmHrkKkP2scJYqqSlfCOlYT4was/Sy2Fo5p6M6w05QYQrsSNVbP/LroTW583KMrZNypNbP1FBlY1D7klNiJ8VGdHkxeqInt/Yjtku2nKnLfe0ws05BnABkrEzajhBWx5SZmEE9K4x/xoTL6wFXm3PJArviM+d3vvk+zpQWl6evEwPBAhZau+g3TXDlYPUmVu4+zYnrGLd6ycE+YMefsHUDtyWcJ6i/ggYfxDSGjnvrYND7UO/4gEvzmSP5npk1+TonOM4xe7K9hinfO39t/V5ARleqHgn6eQ9I7b1aX9Sj0bWZZHnJqs8dLQB2/0HLOCIe814WQ8WtvOM4bX6zRkQbpCQgLn7yOJX9RWv3dcJA8h+fS1AAAAhEGaJGxG//bKI+XPoBb9ACOq/CG1RHbFjFMZSyH3/45vGHLGTBOSUHIckmj6yFTK55d+39gp5tenKZ0+TIp69VbAU0MEr63OYoPbYIeZ+0hU6MxvecDc5O9acfEr5mb9Hxb13fN2KOyn7ueqrFb9B7Mcp+ntOfoab/kz/fXBPkhqkpR7GAAAAHVBnkJ4k/+2AV73ld2/78T6R0AmAM4pMFCMov11Jmo+tamLJnL+GFPjKKY9k/0mUyGY9i//Nd+w5cBPwnFLZb4nNBfrpfIWjaqWIbuYSJkLg1GlpP+Tb3zlxM1EV0QfFE/31pFlZ3dUz+UVJf2zjgRxHTs8ZIEAAABUAZ5hdER/wQY+Htvz9l4HN0mKSSLJygls+Lg0uMzO+ugyCy48+dgykyXKhPA1YMeYFIrylk88DYYEnc6/XqxCS5B5YizTs9B6aZZe9XEgnZIAHqHgAAAAKwGeY2pEf7kVrN5UgkOSSJdvC6uj+mrkSCfv+nVwyKGnNOTtiBV4kVMow8EAAACFQZpoSahBaJlMCN/0LDoiPjwcQvAesAL+9gtbOnGU0zReJsewWMj5JA9f2XkYg5gNw7NzEfDAzyM/5Gm3diqGS4k45BKJATBYCP3ULbtvxYfPYKUTfCdz0Ht0EqbVglc3AWYfMsPRMCAqXfvoB1nnqwjOZgYbD2+UZrv+kirETibpk+kzZwAAAElBnoZFESyfrnND8e1EAlX40xY6DK1Yn+qR0LNkeIAkXdILzr4nVyVcp51nY7sr6xstO9/SYD7aQOmieH5Vq4RIYjdyoid2qASBAAAAMAGepXREf7FDugwA/nRyfZsY5binmyjg3Zkam04CAJqA37zrHK676rH9cSutzXLyQQAAAC0BnqdqRH+2SpdfbZb0f9mM1F1n59jOcJJMnage7/i6E2L5hIH0/dfxsiqfmSAAAACDQZqrSahBbJlMCN/5UNXQAtcS8/73dqCO5nBfqSHNknLndgrL0Ch5suIw772xcU9URFwOnyLTixK5AOUG5tcAVLGsZ8/xt2nuWn9hvzMHxi9sNLrGUk3vI/izEVRf2wfJ6uoBCcLfducX+1SUU5WZd7Mg6sf/neGN2O64WM8NtW/xB+wAAAAuQZ7JRRUs36FT3TBR4J+JBHtj5m1rKkWrdISrAEMYPCLmXpTkr/w3cAZ8fFfPgQAAAEMBnupqRH+jJma8sM0kMB4YgWUuLzOQQmfj+/qF/IPIZV1ry7KNx6u8jw/GVvaHji/pNO+Xe2ljM11UsLlf2KZ3ewtIAAAAWUGa7EmoQWyZTAjf+LrbWIFGAF+04f5eHKfOOU1jybbvVbbrGcfUTPCg5dbqMxX/O8BBo3+r41VavE1NyZuHQkhZ0Ma/+LBcvfHnAeSrpm7qO/5MGjcVGcVUAAAAT0GbDknhClJlMFFSxv/4uttVs26AIQp66dVYdfR49fcXI/kalRF47N48xssWrUKkK0AO56ie8R7L8+Y1LM5zztzClEsSC5JU7s3rSEIyUlEAAAA/AZ8takR/ozRgbSzyZZh09jccj2ji06zY5CZXGw7WgEWL13AUN8pV8oOYlZ1T6LFRNhqtRME6ZVe5tjrZXiWxAAAAkkGbMUnhDomUwI3/+Y1khuDwRPWAS8kFrjYKpZ38aRiS/A0yrDwcayFBVCPMqAPW9P6lwqAgABoVzrKtJFJ2K7M0rOG7z/rYpB5wK8RLGWNff32i6Lfp/4Z8gmLpNHq7iVOnpnZJ+Nc0x8z+Jm80uzJSvLMg6NNlFEnEnmXuVD2ZhR+5kmnzu7THZEnZ9mMduJyBAAAAZ0GfT0UVPN+WfK2Axgf6tW4EiDSrn0KWgKR9iGNkL+UIEnjF+mVhTRUKuvN3mC0Hl4teb27amkowsK7/CDkQ5aL9V/0SNB6uWSTG/IL+e8ZcYoUN1N/cYFYspJgLSp1zmIqkUZ8zGCAAAAA3AZ9wakR/YA8cRdCtb7lUMMSibgFmDr9/TpTW8H5ORuVrC1uA1kykgeLh6XqoMq2PeU5jxxgqYAAAAJtBm3NJqEFomUwU8b/6fw5gA1bd/UkX/+5LX+P9d1NKtE84RZcvMqhbopWGv0PSJXmhuHs1ti5oGwMsS09a8qWRwFF1/TMCc9TPkVt5uQmSKYy/6GlXKdM4RcE0ATHYq9hWPkWznxsDlGPFowqVUhuWRorZ4MvG1tAHkHf0ne6dibAV/bujK7mYet3gfj+D38ywMkhH7aYLqDUqwQAAADcBn5JqRH9kIjFJylYJaPGUvhlvGhYavLAshX29IcxG18UWcPfPjhm8pdEwqrv3T3M/UcpUVTJAAAAAYkGblEnhClJlMCN/+nzL/LHQAbfP4Q2e8R+I2ICdm6ILmkdmAkNnxM333lVEaGzHjMJidbOci09H8zxyFcU4aAOo0/vCuePmIybXoSqGSDlI6oKdupSbofQVi2s4tZWa2kqQAAAASkGbtUnhDomUwI3/+nzL+xu6ANcvXqLyt1oMMrgrjT7HS2aICfWzlcflEz2KS0T59D/VwRzIf8Coyp0Ggq6iSqlP4eGNpEcbG85BAAAAwkGb10nhDyZTBRE8b/pNhbML254LnHr12C23dAPXEje+fM9rieH0CIDldF30OZFIOjsw4jGR7kiUCszljjlLsztSRC8+QqJA9+pjQwq9VNvj3rEnAYhRosCQfNsp48BwbE7ZAGFsWzO2er34mmzGJlT04AgD9bMJYUpdDkHj2kJMVlbzlhHB8khOL2G/ikEU7S9mlAIfvxburLtfu60Hp9K7dM71jsZm/14F9X6T7gHdBDjsnjilGEXqPbMsE73rZuLgAAAAQgGf9mpEf5iS/xWP0Tl01rcB9p4exluCKxwfrxOaTN5KqdafeZuYk7y6CuVqFhIxX9Ht5fQ7LwRBKKUR52DWrXATwQAAAKZBm/hJ4Q8mUwI3//qfQ5xoP57jgZp4AvLSbV6BJYwL5nH/XeUz9kxuiJzoTpcuCYJKX5fPcWC7xD6Rq+wnSq6EWGo7pDw49S17ieTvzD9eZQa7cKgnOHwZWC6v+NMgFL9KKcTV20aSRaVq2wUaX/pV9FL+KGqZgVbd8iszYsz8q5nHIlTd6laN+qn3S2rt97/+BKxp0x2VEQRheDiiGjgGKc/pnZITAAAAoEGaGknhDyZTBRE8b/qT+CmDYq0A6FimX2GfxK8GyJV2TcIlHwA6uQwuPwxjygcR6EXOnMAn30q9FqkYEYo613iMYpG+gfvwxOJXXDwN2oRxgdWszjQPAIwIbmPP7cwKoIEcOiSIMIYP0muhNLB3WmsDC4lfegXgFQaCubUitYkXvw1qt38xNARu5LQpgW4CzYg7+CAq+dy0yC3X4mjB79AAAABEAZ45akR/bHVOygn59HTGJPeCE7eK8JYKfxw4IJZsBnfk/t7bUXnOYimlLi1bMzy8r/JNlVD4+hcoMHGsA4K9j3emjkUAAADSQZo8SeEPJlMFPG/6nupWSuAMYPna5FYBjHJSPx6aUL7tMFcsmdbyjcKxx+GufWOJpu6UI83oJ+saLd3QIKyvPafL17Oed/oacXmsVbnt31UIDn3KzfSy6dCgGLnRb+JJHKoDZiH3f90nfrpFo6e7hNfdUoj3blHPXZle1Znu5CzRHhBq4ikqS2rOSnIDptRE5Zvk7beQpajh9oWHoZLBJ3QKOrj4ZIU/MjmCRcX+wI901SgpWi+nIMLOt1jD6Q01g5LIIh8BzQ0Qe77/V+KcH28uAAAAWwGeW2pEf3Ku4u+jJfriu0ilvVGxChUl1tGV4abd8h38Blx77MAVHCZULGyANeGOQXGkJxsaKwdPgVgO9hdK4wyjj9DhsCl0uJRP6n+ETx3wIbKuJ6EcOWWQMBkAAACbQZpeSeEPJlMFPG/6l6HAGMmb4le5m0Bmzz2CRwhvJigwOfph/x+ICX2mfeDwSlSMuC1HpPLgJB8WO8KaUM2UEnqhh2ITX5BIDB+wedCQAgACGmU28NeiYkY8fdZ03qSkWFy8As+6FkELBKulpxkLP5WuISbwovI7QwOBHNInkZPoKI5W+ftkmIQZH6PrSqXPH6+B25LsgvipwIUAAABAAZ59akR/ZA/Zll5pqPvkg3xNhW77OeE2/MNripAaGo4mNvZEs+3dUWHtlxugKmNAdqo69B53AkpLk7CwOZsR/gAAAKxBmmBJ4Q8mUwU8b/qbTklFAYL+YaSIXn44IiK0Wg/e8T+qE+bHjAj29oBzln40bWpCi39MzVHhdNyGcGnAhIamQYYh24RlXAXIELfJpjc37i4t1wy/G880aKfHKpCMRFZyHXGOjcKCL9m6f/dcMFRAuXqoR+DnKxwB4htst12wGXtnx5yxNs3gR7JDQ/Cco2aCqKq7xfXfmfwM534njS9m0qp2S6EctKijNP/kAAAASgGen2pEf2wArN7EbKlrQS582e83rfkNvEbPKFn/mVYfZ1bVT4Tp9JMen9GYIqcvFT/o7Ey8R+7Ia36vacFGcULDIqZ6dC101TrFAAAAjkGagknhDyZTBTxv+pP4KYEk2oDDI6SNbxeij6pCAAe9cP4/ONgvSYux+AZBOYSg+f0v9fKvF9xOg/I6BPfoGLOAV9VEJF3o61CMR+D7KYV6r8TGr76CbaeqCw7c8PJEUn2dGXkjT2ofSwybkWSKqQv0RY6bwv4grwcRkWZnlLrmz85qg7Z83bcKp+VTO4sAAABgAZ6hakR/bPO/JGDgi+p9MLxMmZh+zfwkKj1jSXluJqlxLberJYGppZieJ7N0ElN8B4ahYZcVWlt0Ps90xenaUhRMdlWJioNLhCIN6wdd3NFgZ0T8/XvkGnx+IL15ydshAAAAnEGapEnhDyZTBTxv+p3UVQ7SASe5OteX3JtHOlBcmaC8BtsK/yy/ieYTQSzkNxHBY+nc3XrbsXY0WuwLZ1q/lCP5bvFLUKlKrDGFhb8kJKeppKCNDlPuEg+NMjuTrbv7Y7EUZRjsCjPc78RnBk1pDWkq6UnlQvnns8VsQCIPjnJPSn30S6ZptD5RpEmZKKdXizBcR+6IJYJGjCPYOAAAAE0BnsNqRH9sCr6whSUtFxmo70NMpoSk862dxLIxJeB61oPBOa5qAqESbPnXkoIEqSOwcaz8F9F4fZk26RCVeH5b1NXQTczxw0LTXl9+MwAAAMFBmsVJ4Q8mUwI3//qibKR81rWwJg2OSHXppgL7UlLgp1Puc5PlSEeULyc4gqiLwIdhL8PDJ8zWsDP5vf/crnX1yN+QrtRkNOS21WOCGNqirrJqrPaOR3WFxkH3T/WZSOUyesZOHpSV8TRm8+9aQ6L0y2OFZeXsdSbssaJ0Ni81AXrzGx4mljxAB7h9iirYC+eBI2UPgiap2VZ9/Y4pjkFdNSPTLNXCdIF+Et81rxpiGQDQ8FXpPxqLRwsJSId5KjiRAAAA60Ga6EnhDyZTAjf/+pQr2MDxxVQR0sCvAhfwAOMLtLwvaSTn2H0vf+97OaOlVSO9XxjuAKVe9k7zYxAAxqoznEaWiI7ugVWdwVT8lMRIWEkbF17lwT7YzEF2dqSYMA5lJMqcCgXw8Xb82B8Ub7d2lQd+XwFeIzeLyfjMjhJktaNg9EFrwomGgxr9g0M9DnMLMR1LGeDSBwzoyKEXAD0s3Qh6xh96VBWiGq7WRYIYLfOmKO9wDYSoQv4ht6qS0UHZXTN4cXJG9PCtZF7VLfkmKeoP6z9/EtvB2ZLGfWPG9K20+gfpD/4BAFU7Ti0AAABQQZ8GRRE832VtteDFAgL0xWIf+CDgnttHtiXqYOIv7NM3eBuis8BPyuNs7mDYp+s7FT/3q+Q8Pp5PrCW/WBy3D/jGcWmlJeXoWkUiydD9NCkAAABXAZ8nakR/dBbAkYKR3qQi7/sdLGg4WMjAE90HI27kErSfarGqqUdAOXA+3K9r/jJvgNxGXt/02b5vaN+DFOjVZIje8VIdZbDJnwssDZyqz03q0lQ7ITdAAAAA+UGbKkmoQWiZTBTxv/rgHCHgFGbSGPViJL9TuWGNO454jDpa+H2zIIACbHS2QL1tJTkJGdsTmdw4P3XOobvzw/wihA7dabJntCAkRTn/2COwkBaDUavOHgewS682CZBG0VhBVWEmB2uOHjNTia67kjCtw6pv6gOenyUgtbdGLvE6bGLNjMhPZmctGiTPcnU8/LYKDRfhp1jg0Wi3iwUFerF9EYtUXdJ+CVMylcwhlewnFMoesbG0kpe5Q4mkkQZts/6RdAGk4m7O64vYn9HIyi345CzEjV8DetmrPwrRPGKcwpTZc2CD5NB/j/iY9uJ7Acrjse0kagBgMAAAAE8Bn0lqRH9t2BmUUyrfJfayEPA2j4f5rd93Dgehy2BTS6zq93/ZSv3WatVs9BR18NhUj0m/Xz5/nwTQSPfAQ7xtY89cI4cmNv8aBP42XX0hAAAA5kGbS0nhClJlMCN/+pPnqeD4hUSqykicb53LTuCfgJJ5sdEAN04ZxMsnwUiaNE1glre5ZKVbXCLvA3SbjM12q6q+g8qDoLMwalAk6HB+QHJz34Kwoon30XbFmS8FihqpUMhN/CJxn7vbXdFg0UDg+F8yJ1dNlaj+o8pq1yEAmG+zJoZQYStfxs9NT+ns4tTbAL9aex60cc/Lj/Kox7lO1FTi23+c3qeQogsumOXypCQCc/MskI+8Euvt9z0NiNjw8K5wQIHgunHGAUDhoK1Lnv/l2KFhEUDy+mgG5EuJD6XkAfkywf8gAAABCUGbbUnhDomUwU0TG//64Bwiw34osn04rJvtU5PxcRX2NposU2kfbqC4cNJAYrSOqVoHlAPNWHXPG83nm5DWYqN6q55Q3PFHVFuNx+luX4G4kZnEV6NxOK3eL3FD+dpXv1OoynYhANfwm+hKqrKHw8XZBkKXFutB8t5C07ToIZJQGQN9KBPRJDSx6IiEBoTweyMk8+dhWvd9hu2eOt8VYNPLybUwF+8rLvfGfeNlVoZ3XEZ1itgNP7tZrexUUTnbZswra98A9vSBllM6W4r5lWl/cUAWVM/jCYg0WwP9fiIEJIyYGMEEGNffbgbCxJHNM82YpeuF8Fo4NFFAnJr4EJ4tazzCeD87o4AAAAAyAZ+MakR/SV9tZJ4GB47aHQyCiYU8hAn0FrtWzh6xct/GMCLAK32GKXxN7lxMHHyvweEAAABoQZuOSeEPJlMCN//6l/VuVDm9cqWxB7gC/+1R9qBgWg4Fncnwxt9nIwaQNAgbD2JC0xyKqTQms/q5ukY1BkTb5a6qhFfBaERv/LsYkAZhBFAnNssBz3HawYY31C71TH4aJZAYGXO90XcAAABQQZuxSeEPJlMCN//6XIsHYzw+gjwYp//qLcf9l+cHa1L3xo66Zx/95n5gsJVfsExzeGDieas2BFvGB+AmOdHz/Ik10yUDuuXU34Ux4PeRyOEAAAAlQZ/PRRE830M0GSmKls/Ab8M0yKa5fsZixtfOm+W/ATxUZX7NEAAAACcBn/BqRH9HFNhOfQcUeVbFwLjn3JEGq98dmg5Cq3AqGJtyeMhY8PgAAABzQZvySahBaJlMCN/6XL9FgmGToWC2GlEmYkJDYXZqo/B5HVv1mAKXZFd9+Xx2JRqUBB/Pf9uUDDsJ7kbXZW/X15rbsx04Ax3FCxdrmeuNt9x1PC/aIr5FMG+l9SFa/65UQCQYlcp2DFIjQC3abqPJX/rFQQAAAFZBmhRJ4QpSZTBREsb/+ly89QmRALXJJ7TFSc60X+hnXZKUCQ5+EdcFvv0wJhSE7oWRGr+OKfGlW8pY8Q2o/Y6bjrdlyrsYEcwPZ0X/FIJc2B57BPdP6AAAABgBnjNqRH9J25Fb23/w5rKBSrT27nRbsdAAAABFQZo1SeEOiZTAjf/6XL2IBFN7kKBzeVyR50CmrK4eWlEwo73oY5NBynF1VP9kEZ3pVY6zZjKQbdi6GNX7cbn71/VN9aqBAAAAI0GaVknhDyZTAjf/+ljDRVewBl8k5cFklkH/FNoZ1q4iKmOAAAAAIEGad0nhDyZTAjf/+li9OfwDRMr7jJXCjgp3GlG/6MHhAAAAXEGamUnhDyZTBRE8b/pcvYgEWqVvKqcBcshYb5hHpMSF3GzwGVbKlin89gFv0jaJauY18LjpHOb42r++ukOERt90XS7tUzEu8l4qf0hQt7m4RZHTTHfvP4RkH2VJAAAACgGeuGpEf0hOuUQAAABZQZq8SeEPJlMCN//6XIy9hnbUIU4U0Y8jSP9QABK9TsVDBrqMZDUzZoYPan6YZJYxlxFl7g+zv6YGiRtYqIQOxpRdYG4OxQpdPzELlhtI3P40VnX7EpnY9MMAAAA0QZ7aRRE830JUgrHA9mIf36j+H1QFnz+EhbqP4QfZ+NyGLcO5x32/rvXPJCqlgDh6lfFwQAAAABUBnvtqRH9IUoViB12TS4SYJRH+AIEAAABwQZr+SahBaJlMFPG/+uAcNs3AouQ24AoFAoeoVCJ1pKMOnvwXJx30Br5iWaq0xNLC6AZmKamsy6LKSXnRjFqRi7k38QTP6Kn1vgT0FJux47v8zJ8564m0xS/tvR/sSxOJzFhj0bLeQ0tIpqB/8P//gQAAABcBnx1qRH90Ftjy1XOCLP2t5421v7zE6AAAAOhBmwJJ4QpSZTAjf/pe+TaskAh+ETn/DJ+qF8jVPtceeRYpE222sMDw9tOlKWDpmw2M73N/UYy6PFUTypfyfu58ZEtID/rvuIUDI+8eFuTWFsppb8xhNXCB+jxJLSujWCKIRzaSAniuaGnrVBPr590XqJLwBTcrwjq5pMK8KzkUu4zpYpiAQZniMn8gfbt1EAVmhLy04BT3Q5l1kfY5e2atW7ivo4uDjB44352x9hsAi+KHFigVjjirfAWMeRpYG4dAX/+lnNioWZGKO4rQwGok3Hn1Tyex+zgummyCv4QKsCDp+Tl/OgxgAAAAhEGfIEU0TJ88kDOSn27WeR4G3ZPBF3kEPLA5t2HMHVBv79u66OKGXvhaAL4E6WsHV/6VfRUzzLAH9pR03QULc0Po5tYYb70t/+DrvvTT+gaipytuS47lchOIU6DKxCAz8Kmhg0U4wkMptuvz2cXpodWZ9sah+2qw37JT42AD6dJjQf0OQQAAAHIBn190RH9Hy8/JxqGeCZHITqti98S9VMoSlIBYb6yAitZqQjbxLED/8shEIJBT3MZJfmg/bTmIdaU72hRONjC3Q3l7SYedA/6onm4A3ezKLzFD+pvsPwlxJjnhfl2RKlnop6MtCaB9l11NelVfjOpSxLEAAABFAZ9BakR/RejboOil6n1plmyWCdsK9SmYgB0EvYE/bILe8rquJyeYSSrlUq/u5t7rTciQuKnH1ri5+oYw3jvT9j5Zbe6BAAAAkUGbREmoQWiZTBTxv/pbWGygAdHX/O7+UpEjQBBEQNHOaTyeUUWOjUySh8DVgh7Mn1pubeujzke53fyUY++wqtEBMdatZpcFMo1jcP2i0/Yl3mE8GV51SjKxXq39w8EyJSUVhD6MXFepWqdpNMp+QDMoEOEeN8CAnSYXVQtQVOt7X4hvtkRVCcBjCys4TWy5OUAAAAAvAZ9jakR/RbtZ/ZeBSX1qXbA9iLKy2vnyiRX0LiAARgEY4AZJGwiXdU61io7EqyEAAACTQZtoSeEKUmUwI3/6fvfABYzTJ1m+9b9HUKXTBGS/bNGURTqhJJzk4+/rZLL7yyVWbfKechD1ESF3F+0DM+rdikZHs60jmqLemJlp6D/V4M+EiO/tabXsUVjficFognps16Ph6NB+s2MIjeQYrVdb7U9xz5pmVVG59QjgDzOjFBPF7pZkl8d9SBFV2HPI9xn+zmKDAAAAWEGfhkU0TJ9WxY5O3ld2/sLGP8SO3LyX2vrgIHoJAUPwXj7Bh2X5aUR9ztzO4PGSTT9L+d93LYRN6Yub6DHfvC8MsDcOwaUNamDJe6KT6Fyrk8arNO1AdPEAAABFAZ+ldER/YkVP20u+K+fYVCRNeNgQuKvivn1U9254KKFQNulhkX/7OqDnp9nMac/QNU/PV3/ZycV5D8B+TztD8n/NzUmRAAAAPQGfp2pEf2UuKWHArS8cB3Kt78wNyh+PkrO5EtZLxgBb1sqHWrv3qDHmkZPbQIN2rvkzDb+G9p/PR3u7brkAAACDQZuqSahBaJlMFPG/+nzMJuQkrgA0H9/224F6n4c5ujri4iSjS6YAEq7/aB9LXbZqshVfxK5V/2zN23LPjq/v5qfYy/BQBPrNOWD7ec8p8TzxsekA8aJEihWOrs2vptKBzLAwe6RP0HLqwZqVLqIMpjh3oPLvveLCZTnfMwAnkI7EfpQAAABJAZ/JakR/ZVTUOPDbLyCL9W3zg43y5x8Mk79O/AaMsNMj4mlVacWjRzGq4z1+uoX4/SQ8legLORP1bMcUEU0icCkIvQPfRukZUQAAAElBm8tJ4QpSZTAjf/qBai9gAhBysmfRmhkGvFCO04oCARkmcy/ubaZ7x7Xu4c6yKXxtQ+MrsdeijALFmYXWIk0wblE8gBCbujcIAAAAZUGb7EnhDomUwI3/+uAcbyKcSVIBYzTJ0sbBfKGAemChtvkGAQ2koMd0LIoxhXtbsppCuSruc7LIpEI2rWpeMSp8KpHRtHRMz96GEHr8RqAUjB9WF43GXmEppviYv4AmI8gQnUOMAAAALkGaDUnhDyZTAjf/+uAcbyKOGYBBk29//nSkcRelmVIJ7my7T1oxka727IXrqCEAAABbQZouSeEPJlMCN//64Bxt2Zn12hfqjpGgD4NidkhP/dgx2k7qNTkWWy3ZRxHFAaYlNlouEymfJSH+I4I4vRhdX6/NzOsYSue2FkYinIKlTVh/lnoZ11szU5bq6QAAAD5BmlBJ4Q8mUwURPG/6U/W//2p7Ku0+u0L9RcQUEdHK6QTDW57UaXgyFDnKTRozgDnW3AcDBkZcdj4X4RD9gQAAABcBnm9qRH+YkuOsuUAFDth1+XyaMX/LZAAAAGlBmnNJ4Q8mUwI3//p+P+n0+AG3z6HwvCHrA/tl4vRAyTVNcj1poA1reku28PcM2fOR4niyzAKruFoOxN8ytUf2w5zjX6yGN9SYw/4CPTkaUAjHyJfe1jo6kvye+qgwp2ugu1r2Kwy2nYAAAAA+QZ6RRRE8313SMo6vBvXKF4SxK9zNsnnfd2DaKAnSFhBIsd5hbf9H1ABF8WTjl45LConTnOzHoeTksE/CfDkAAAAgAZ6yakR/ZWzCUJS8EED7hzkEOjZnWJ39gC0LZDBzp+AAAADHQZq1SahBaJlMFPG/+psfgyMQJpdI52yItPbTM4x6f9t9VIMmQAWXZThGWLSS0LOxWuXgw/4EOxQPv35WAv4/IA0o4QhDGiYTyRVKoHTyYXRf4omFSGlX85GmZ9ccSjPqJOSfHwdZNXaWi4XFp67Zm2vxfdsxvn58pC2SmUeOgLoFNRQ+cIhDpZ0FxlhGQZZ4QKYhBpafm34r3uC0gLBm/V3J0bB0rLMvEskQmnZ7W/5QLRZtwOCWHPG5XSB/HX+h3HnDvx+mgAAAADwBntRqRH9xFUSQrqg0TbsDIwb1uUeSj//hU0ufeUGPQawPKWX3R0+u8VzZWV+FiByCMSvbzRJDuRc1DlkAAAEJQZrZSeEKUmUwI3/64Bwhlbgy2K8Z3gm/A7eda7y6jRCqYwMNz+FAqtTd/G6tXL1lwT0A1MbJurFtpxNNYye158DGdFFxdgjYRUtmwjCVHA9gAIvGZzgCQdf6/NZ4GkP4SJNxaBE7o0nd3vzvRdSXsUzZTA8188sREd+lbR8BU0hua2RQMhd6hXnlJxeaEDFTT/1OLyMA2MFJ2jQ35leCy052UyrJRsBSQ1H3tOheQNNcBqfvyHM607SdRpu0Ahta7SP8TFvbuMTEeHBExY5uWGCexxHwY44hhzxFTh4J0ks6bdfVbJFlPzIeamgqKpWLsBhVc9W5lf6BMeSFa0cO7S+P2gGw75Kq0AAAAIhBnvdFNEyfY/2/UvcaVVcJsYSOh8+zFHI6BqG35STIuiBt8KvA9gn7aLELcXU4Gv07EQsuJeDc96vZHddZa+GIKq4GhgkpKbTnyWbFuRoy63GsOoQg4txnde2lQfXPD6Nz22HDaCS23/wDn3m1W1gG+S4SpG2QeJBBOmdpfsrsWjqySDrCos3BAAAAZQGfFnREf2/llvGf7sgTrdYcwPvXq3J8r9qj1YkUENV8tEzytCkbDnwu23cWYLol1Ql6wiyJs8F0S4UP+qYaynZoFMVegzmZEsO6KoIBPmK3/WV0vFnMH/eE7kPUPguh0RLy57qhAAAAIAGfGGpEf244wiRKH7FYNLwmxtaT+VqqCJH/SF1oFBggAAAAr0GbG0moQWiZTBTxv/qXocBLdA/qP6lTUdCjisKGh/zuy6IeoPt2baYXOwBRDlndQEFF5ZsVxuoGJyFwJe1c1CmbMQ4dMG2Xu5WEnBhzqgLQ1sOpyedH9Bwf8hxnYswSyc3lyhG6psgLCHX6w85qjvjosyvPq8P5HdMT2ReyFk/cB3Fj7ZV4X4YzePkmib2TTcIphgS6VnM1tWZh0tbsvk/BixX4RztjvcPBoHCzV8EAAAA2AZ86akR/ZB9ul6WDjJakHPvivVIYKbU4JXY281j6RxhwTEGpCGteqgtUqZ20U6zZRu9BEc+8AAAAYUGbPUnhClJlMFLG//p+fFa+gAvGZrLsoPR692fCo3BQnhzliwBzdjqE+pDssttQU3ILP9NSeEOqENFwhMUrkli194vxrQNeN0jk2ksPLijRd7Xj+ow/ttFCv+RRkvk4/oEAAAA0AZ9cakR/YJMHGIZpQv/mUUSfYQdcZwOQxyGodMxwmhyz5Wez1qzsdy+G88Q/v8hgePed8wAAAGZBm15J4Q6JlMCN//qZDWyLLRITsATLUDBZOZ99BCH8vj7ZTE5moGOV1muRwk8NVjmW7Fp5lf2kfQpsdj/i5q39ei7bxk4hjk4pAXblpxU0/vHphNDFQjqbqnICuBn+iH7MfL2FfL0AAABRQZtgSeEPJlMFFTxv+nzGZuRJR1lqmCoAdBQgNU+Uoqau38ZK3oym9UNIGZA7jx7BTSQReyhkn0fIX0WBjWIEsNgR+RiyY02iVj3t0uUxuvi4AAAAMQGfn2pEf6FLId1ishSqivQSksTW+TvEv/Mgd71xJQRch8U2+NbuD4lTES0BKe2BCVkAAABLQZuCSeEPJlMFPG/6mRDqY5GY97t4An5dnYfPDethK/2mftEOHn1aA+RlE7FHtKJ5Z/ledsdWmYEN+7KQoj2SUYh26z8K0iZUgJuAAAAAKwGfoWpEf2R0BSo22fL+xBEo9qY+jZwkoEjVLKMCkT3vClCAVS6OMjaBAJkAAABlQZukSeEPJlMFPG/6mYc4AY4/500NYqs//OLSGs9kQL11LU9HgJQtPLe0iR1eIotda7Law088RLsWH784Ddti83De63j7o1EnDeGRyWIL3TfBD5lVcv57yA7Jm6266L7wz8A2OKAAAAA2AZ/DakR/aajYSdGkr/6qfQY2CvitRSwrRfN5zuDWHUEvGTytavqbXj3eWInTMTkAadBPzR2FAAAARUGbxknhDyZTBTxv+pkZdAPMkFrlNMxXg2vKTBHlsUHK3/wgiVTBw4cD/KqGHp0uL0X8hWBQWEIARNZsmaVVZamx8wmOKQAAACEBn+VqRH9yqmKkFzu0Gqxiv/nuwlSFn0cu9L1N/dTz170AAAB+QZvoSeEPJlMFPG/6mYc4AY3+WnUjGYGvCLHGxHvs44RMP//YCSon9kSn8IcHdCNtYx+kq/wHdSa+N7oqrQHZvMRNJGNMiTZQZvq76FibPvDqIJUdFJcC0ZlebLkPXIKpYakx+4ub3nGcBa6DZvOaFLCu6JJ4YNhgsu1jv8vxAAAANAGeB2pEf3QWso+BJ0ZdX5MQAFqFN+SN3aId+NG6sFOLjNvRk3MMOgXZVIc4mK98gv/x1+EAAAA4QZoJSeEPJlMCN//6l43QCkbILSMcsiWNkR2QTECS147rjgrWS/4ZF+lNKP8bFexzlP76thYBscUAAADzQZotSeEPJlMCN//639lAQojjz+EOUMitiaO0r/gc04/dSBGjnd8UMxDJ+36Z+hcAaYic/brqPBid0s1ySqgg1CmJwQ8K/0x5IMThDjqZfZ5BC+cPFqcNhCoX/ndJ5x1kMmluOk+ZPKVGcM7l7eQIlosXxFs6MFO7H+BfqrS6cRbSU+Y7FUXyafbO49/ex+CooTwt8S35r2TOjnc6ccrbJ0JSLNExqOvS+Lln4dKWnyi9JKV7nm+OLuN7rL+lSUiMZmAqpFCRwh66LPpgiV2Rp9hBaTRgm6A0kN4Cld3TJesyv6BrwasVCwb5ZnDh9SpNQvgpAAAAhUGeS0URPJ9iXUxKnVJOUcFcNdiSEVOZTqHSleLCzgHl8alVh3/fSqcJfzg8LuGE2VXLTaKmZ4TL3ZOY0Spn9FIXbOHpMiNNPNxEH2tiTWRvbNb+DzSRQNosWA4KnoCp/D5NODcH95gJZt3uRTTihp74u6YquFRW8SZMErwakra+PguaWFcAAABLAZ5qdER/a11k7yRbPkbKlmtOAiJhLWnkEthaXfmkyPDwT8OJYiQs837Yx5ERYhArYcOVM5P74KVS0hPgb92pGrHZqA9vptu3OT+QAAAAWwGebGpEf29lM69bdR7noGHPCr7B1w4lVYNhl+g7A9vIcNrXY9BojmiJC913DGVYzEB9NVjn5BnZ0QT2kAgV5lbV4MhUxZPzIpLED+vh6vnKN8G0dsaGKALYEBEAAACoQZpvSahBaJlMFPG/+sy28HTQmQ6Y9QcP7w45285DsQPQw1NsSIYQSxCajmeiqZ4nE95PkhG1J3Df7HKjSl+AJYMm0vZk3UJDopt4giZF6AaYahuIlVAijJsWRcuvjXRVyaWKPkfOeKFAqWGeZdy+f+0pzP4sqmSjzOqAvt/NdL0HDVXIExNFA1WuVRd+s2HPbSVzSAjoNfJ62AeQFeDvfirE0xkkw0GxAAAASQGejmpEf2+/w0K+b8hJPz/Bl419Or3I0O2P2/dxjeFR2vIdgQPD8qCBPvNjQC1bbTh9mtQ0EIeg8w+20BHysG1LxolFOvMFPJEAAADaQZqSSeEKUmUwI3/6y+UB+FaWpB63j8NE1ktfgG6MkP8yxVk5ZebU3YyDsrx9QPw/vCePwABygaMSyta1EAlmFPcbYvQetbi7A8AIDDoXq9Uy/BT4lu1cm789aEtRR3HI2kZtiSJifRFxWCkMHzdTpn7cMqNwqeqrtoMdTaWdPVfJCuYMEjc2IiJMf4qrmU6QsXb+68VoDO1GOWWf4GPftdDZqBJXt7Bem8UIH55oNRpkBjOEOKcrhtpzuphkSG1SHL78Kx1U2S53mdWyXts8SXmiyVzeOm4gBnAAAABzQZ6wRTRM327Xjpxkoz+OBAEIzZN+j5TFVlKdKyIbgyQrMAFaQhom8HxY3u2i6l4BQqoXBsqkWqtx+6EiGallCT//9OO8XEzBOGN0lm0uHOuS5oyWbXt9F4NJE3dHTCc951z85RTYPSXrQvwieNhxb4/vYAAAAG4BntFqRH9twy2BPwWAObg/+GNiEoIjx9bY6qxCW2rw4ckRDMMxNKyrHX9f+PaG7WlQdNkr/TRBHfapnJdtrN3pMxgvhhd5ec8ckHHcpAh9tkpYSJBFS5BUHgGXEFP2Y46gOgOTTlpxQ61XIHLESwAAAS9BmtZJqEFomUwI3/rZ4DsV9/7/L0RvT7TbRAZpysmCZ/Bt60wTLhrxr1lOGs1LRbM3lEFfDXizg7xwZLtIbfYxV1SV2HwPEs4xJsxB/miCnusXvr58iXnPlFI9WaswW1/nY/wFSfUKHfM2NOmtrCQ/UVFLLk7VrXrV2ziDsRUGCAKhZnwCNXB6q/fZDmnQ0NDDfOD34H3KvxgMvGWHQ9vKsHtBKQrLlYao9Uswb1wAhkuMu8/DGZjdoUMzowS3mb8hCRMWWxLln5RHbOfTZGcZiMF4AD85YaOtf7HdTqxPRZGQIqSOsvI9m+MxjKAKNziuf/EtgFYfZMOlNPYcrzsbscVw8ZGT87Ey3u28TJksha6qZaGWni7rKjgoBmodYhDVpKZAjfV8i7GX3/X0d+oAAABnQZ70RREsn2TuGFQx8ARCkQVOZ4gQtLywL4ydk+6nPRzmSRH4n3nt6G6/SkpvZl+fZSMZNEIWZr9JsMDrdnP+0sHT0L5k2+RoMIbnSiNDfe1Sfo/TrB2G+KbYhyga39gecTr2+DebNgAAAFwBnxN0RH9yiu4JGqww7STIacK47/fv7h+a//caOvXRkKjR+ABmubX2CZtD/BIzMglupA20jwrl2Q3AO/zZDx7pmCu/SI7prjszeD/iHwLIHjS/T9vDRx0ovBkOuQAAAF8BnxVqRH9xMECM9RQLiDiuY7/Pt43+FaDkalW8BXF7ie00O9C5NzBDluR27aH85CrlMmNRDbsBGkaVZncWgzlCt12K6B3Vz1m2p3vZr7BMzYZM1ZiJFk6tFJVCiIsx6AAAAQ5BmxlJqEFsmUwI3/rL1J9k9Z6cOBPIHcCccgD30qsvT5145WjgUj5kBDwnzBVJAewJXptQTUz3bH75rMYCmagu5cz+8ZyIEzvKNjvZMGqrbpeh/FJpJjdTOnzAvAtWnOT4VUvVuAzpJdHcATzTZYqOkugeRXUjiQH8MHqQ7rPNEtXJ/1Fp7Ve/oScwvITr2rsapCWcyJRFGRp+2aUuFC/MH6w9Z59FmXdcKSVzhx1ktRA8WCroFm3FgWXW/mjfFupXXLRENZ91O9gfJN1JmssSdhCev10RIq88NTT38yRBP/WPhK13QeXOIrhBSy+RoP7htuals807rLoW/OIHPWXbRzv4n91to04Cb4+h46sAAABJQZ83RRUs32mlCrPHjO0z5qHbNqPZeLe/TukRJMTUDc7vr2JLwtUwmYunJF+6LJeF0rOSzIYAWwhTyi+Lsl8OnnG48JyGPycCgQAAAGkBn1hqRH9uBzVYPTopV4o3XJGHrnvnqaIEhuYJtxIkOjLmkE1FOcMH+XxzeFScl+iVnaUcG2zDxgfbKf8cGYBGdtunIoQ/58LJihkj1PJihK4BMPCFcl9b0CxHzYn6rSeDI10DCbiMh5QAAADoQZtbSahBbJlMFExv+l6ZsV1V7iDQD8cPvtyTVGJKMes8fpe5qeNBxPyxTKqjaMxB1/5BpkWYODlvklKjpGU/QnYekKE43Ds+69nfQ6+qw5+YV+GtdPVounSiJ9mtb9Tos8dV5KOxhJH9M/83E5mHTNA09s93YRKXbpPKwLzAqfLvn9hz3yDUVFXNOi8rxdtU4LY9SE2N0T9JNkNdM2HnMYiM5VUUgMCXaTF704pSWII/Lo5cAGBbpdKIhRrfsBAsXOqTQYKPfr9WuX4V9VZJ4zMgzYBSzsVWMrLqOIFadQa+GBl+/+J3QQAAAEoBn3pqRH9RzKLygkU49/cdiFxsXJe0fW817ZCFON4o+8lJkMYbS+upms8+98JNIx28mlCxDFUZNlcNYTB/JYDxf5Z6UX1v8Ha4uAAAAHVBm31J4QpSZTBSxf/6X0cs21AYyGHo+UOBDr9qL0xamkGGfMPSGJ1/X2DqKlHkuYsWqzsUeg99IIBs4f6PHpzOItT/iHi3HEWOgDOXfhimS4byFol6oS8pKQOx4B3P8u6ynurw0gCRj9ZxeWI36xpff39HAy0AAAApAZ+cakR/UBBo8QREbe5YTX9+bsuBqLS0zRZUEtw1gCfmExV7IAbj72EAAACUQZueSeEOiZTAi//6Y1G6Uf6IZqCuIxNrUqcCVPVQfVAT/M4zpGtQoz+gNEyPd+tUdv9c2kLJ5JBTdiL+5ELOH1FzfD0+dWjIznZD/fjzVjPsWt9LYU2O7C8M5XRLLUbN/vfa/H7z28pwvO7HOg6DdRGbPQhdM+xSOv0ooGHfT8LuucAcwZVY/1JjkgOn1nFH/KlSmAAAAH5Bm79J4Q8mUwIv//peWeYl9lBEr/JFRruv/EjXxZ2OADVEUnXBJYA24gaNLkysDRgt3d2Zie5C7uXnei5S7Cei2GtoxGnckxQ/ZqBmjXnwuVY8vULtjCpbjHRP48A4cXpCjPUI7JO6cZFVi2WqNDLG6LGLwhcLz3f91c2F/zAAAAClQZvASeEPJlMCL//6l46cB+Vc3I47Qm9HODAys8/yxtUpHHrLlDRe/mt+iP7GbC867ZGaerUag+alsJ+ZIm4mMH374Pcf38ZOAKag488ZobJOjZHejvT+2jDptYKfKIboYkDH3q62EUqxGb2i8MmT9iOsqmE4yUMJFXWZb1LOzYvaY8yox1EHfwYt8MsPJD2gRh0YuxTMtcZMg3/VCPORw0jv1cdrAAAAqUGb5EnhDyZTAi//+ydCIJJlSRT0/YsBNgS1Q7FljNxoKGbkTtzpyroiWO3kxPMXGf3FkBqr5UR2TlQ/4nb9Xm6jcjslFWIxgJNTTriSsv/2ZYBSgWUQhjuFH/p50mfebeWp+5j3soYroJQenbSqGcUownaYFpA6qvz22GShNVouoabtE2fg92wk9zucQAS3/fZFrC+7uycIR6I5fdGcUUZOLwbu6rZn50AAAABnQZ4CRRE8n7rXXgxHv9MJuo9JVFPVuIRV6K5G2RhB++jIgD0hx0Zm+pR7L0Z3Oc3uV7Q2sCeIFaRd5KEh+su8E92ZWU+tZzQvlvikSabJaxy7/hHTATkWByemnN1TQ5wnw8/lgBMkgQAAACwBniF0RH91KypuXcQtJQSdFweOvNr+gHGKaBBER+RAuMdsy0xIfEivnsQoPAAAAEkBniNqRH95/fvXE/D4qhrVSche0smNKjrsPyy7+B9QYEwkPfjXTUptU4YpWY7LfiapUUajHe6IZztSuRVYxIEnB8SdWpr0UqTNAAAAQEGaJkmoQWiZTBTxP/RXw2KG+AhSbLQwzQGSsiVBueAfi28uCuTcmK8xd0hvj20YXU9COfqtqufdU9awty9XuoEAAAAxAZ5FakR/esg/I9YARKRkChKJLVjPX94lC/SzVYidGHYM2bA/a7Je1qWaoxGCwTkMwQAAAEpBmkdJ4QpSZTAif/RZyXhSAoWlwupLokvWsvhMa7iAYslipwssIobiaAZFOg2ErrPF5+H4vhxixgLlj9jlblpqlT3DIeSBYry0AQAAADVBmmhJ4Q6JlMCJ//SW1q+u8PQFApatBPTc/3iAsmvmked0YMseUifV5B5XYHb/wEQm+yDJwAAAAI1BmolJ4Q8mUwIn//SmhAZ8uAPlGvj8W5J+5nbHWQrALalH6FIT3oIhRYyOU8MfMiWOpXv0gIPc2waRKiwwWh1GzG04qf7sADzfw20AMFMzvv9YIdASOzP0bpe+WXXOEyg8tPxvQIUxKfJA8uZIakhCIHfbHKEXDH/i0In07uoeqVu6XTfjyjyW7ekWcRAAAAA7QZqqSeEPJlMCJ//0fW0C9f1WOkQDAQJYlGvGbx71qFwY06fD6dGbXJ3or+UXDaGwkz1+A/JKr8iiSRUAAABFQZrMSeEPJlMFETxP9Kdhn3wbNjbqXmWujVANXUA0yrO1ay6pkbrjiw5Uy7yYGIJj0nuqVv3S++B03w9rHxFD0HSC6joCAAAAJwGe62pEf3eHXlhrRBQRWTnn2YyiWIh9H/c8s0+Sy0oWNY22Px1ymAAAAHhBmvBJ4Q8mUwJ/5GHS/4snXAQk/lwYAS1/ZSPchzU4rNF0PEaBvMfK4MAKev5XzMy24FVr5xQmvXYIBgpVLEtDkMMgAywFykKlD4c38fW1z/gzmTrpHsiZwntXrLkuUgQdX2TZKpEO8L6i8hN09/3+zDOAkxRD/FMAAABPQZ8ORRE8n0e6YV+bSwREIpwG6q0LL/4eqYftL4pN0Lg1BJzI+26SuyO+pJRU7RBR046mHI9KwngefKWWYjsv/q5SOrIaUBxp/K3XoNJCcQAAADUBny10RH9SSsmdmv033+fdsEsiW7E9WEqBm0K+1Q6KRxo/2MvB9ozxJU7sTidgINQ+EcwWgQAAADgBny9qRH9ReuP5/J/rp/f3Jcmvyoy1PENRmUVXVrqjtezrlvp+1WWPCqPncgzBmPhonsrkJaMiMAAAAHZBmzNJqEFomUwJf4dT50t/AJzUJ2/qYHdnE6tjvpTz9WmxhJfZR33TZs8gR67FG5vYO+Di8EIEgknCew0Ib1KNDlUNnU58vcHB8TjWBSpjzugT4EZKyx6vfCS2s262ogMuFD8ARl/p3XKc9Lv+6ZN2LPlNoeCrAAAASkGfUUURLN9LNgZ5Rbx/RYWH/EV2kX148q3H5qeKQd31mXOhE3/NbvJQbkWNvxXNm6Kaj7D8bnRZCVKKvybqoD5kR5XMa1S3OS9BAAAAOAGfcmpEf1Hl+i3J1AJFzB2lQztgp6kdWi859yp7RM8Mt6bJ+NOod+C975LTe3gU4s1PCvAh8jA+AAAAwUGbdUmoQWyZTBRN/wBDgQXHaSpBAtnZAyxHz3SyTIyZ7cI60Pvbpf9cz9H5SaLtGG2+2o0tWo+Lb+n8EtceVST6EjjQQo6tzBJtNIvJv50Kag8PwHxcAR7ijc65ZPR56/rXUybyqQwDFTOiT48RJlBOuOzsyNXFyXze6LboHPQaeUMG0EWCSq9HS7WI1Xi6RiZCrGII3UUkhMi7fsMoCScH5fP+BVbXWB2IOb7lQXIanSGXB986nCGo7dA5Boz5NBgAAAA7AZ+UakR/UolkdhGyudAWE5bhegV3ZIATVTuVKPYxtVyJmpjh+yuee+grjOqio4EUTm3EWf7aXUVkzoEAAACgQZuYSeEKUmUwK/8Brfrn+KqVjtpA7gPZnXZm8Cr2k0YQvD2qwxfXBvfjv2Vt6w7z9HxvOKN92vP7W0TtPntoqTHepLdOnq1/t7DoRAavckw0KNig6nDtsSLTlAyTw9VShLH9AkFdNk25QKWTIxpeyqHheadtzPtfIHaRDyTi2/o5wbeoZMYKv/mJp1j9qV/dZAA6wCbRId9qiDKY/PCQwAAAAGJBn7ZFNEzfSaJrwG9yFE53Bzw22EbC5Va/dh5nq8xcanv7Ikoe8EOt5amAOLWNPyYB8TCqlIXQHFeh8PypwYwPv7TJ54YMpWZcKQ/jtToSVo0wDmH9y8MJ/WsKbKKOTrpGwQAAAEABn9dqRH9QEKouPgX1Bi5fI14f1IMR7/QNmn30MPvRWuZvJD3b6yW4Z6DFm/HawBXwShNngDy6XL08djNeNNH9AAAAX0Gb2kmoQWiZTBTxHwTjrykZW0ZpoShpL3KinEklwYSHJLYvMN4dJgMQ+RZcyRWIPPZXcBI3OcuU5QciTatdHsL8RMAUhns0k6a6BCUljrfasDMx1DFfTAv5MUbWBh9AAAAALwGf+WpEf1FujWgC8sP8L9qwJju/cPEfi8CO9F4BVHsYxPh87w9iIjT6Ork9/vKBAAAJxG1vb3YAAABsbXZoZAAAAAAAAAAAAAAAAAAAA+gAABg4AAEAAAEAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAIAAAjudHJhawAAAFx0a2hkAAAAAwAAAAAAAAAAAAAAAQAAAAAAABg4AAAAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAQAAAAABgAAAAYAAAAAAAJGVkdHMAAAAcZWxzdAAAAAAAAAABAAAYOAAABAAAAQAAAAAIZm1kaWEAAAAgbWRoZAAAAAAAAAAAAAAAAAAAMgAAATYAVcQAAAAAAC1oZGxyAAAAAAAAAAB2aWRlAAAAAAAAAAAAAAAAVmlkZW9IYW5kbGVyAAAACBFtaW5mAAAAFHZtaGQAAAABAAAAAAAAAAAAAAAkZGluZgAAABxkcmVmAAAAAAAAAAEAAAAMdXJsIAAAAAEAAAfRc3RibAAAAK1zdHNkAAAAAAAAAAEAAACdYXZjMQAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAABgAGAASAAAAEgAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABj//wAAADNhdmNDAfQACv/hABZn9AAKkZsoxtCAAAADAIAAABkHiRLLAQAGaOvjxEhE//j4AAAAABRidHJ0AAAAAAAAWBQAAFgUAAAAGHN0dHMAAAAAAAAAAQAAAJsAAAIAAAAAFHN0c3MAAAAAAAAAAQAAAAEAAARAY3R0cwAAAAAAAACGAAAAAQAABAAAAAABAAAKAAAAAAEAAAQAAAAAAQAAAAAAAAABAAACAAAAAAEAAAoAAAAAAQAABAAAAAABAAAAAAAAAAEAAAIAAAAAAQAACAAAAAACAAACAAAAAAEAAAQAAAAAAQAABgAAAAABAAACAAAAAAEAAAgAAAAAAgAAAgAAAAABAAAGAAAAAAEAAAIAAAAAAgAABAAAAAABAAAGAAAAAAEAAAIAAAAAAQAABAAAAAABAAAGAAAAAAEAAAIAAAAAAQAABgAAAAABAAACAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAGAAAAAAEAAAIAAAAAAQAABgAAAAABAAACAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAEAAAAAAEAAAgAAAAAAgAAAgAAAAABAAAGAAAAAAEAAAIAAAAAAQAABAAAAAABAAAGAAAAAAEAAAIAAAAAAQAABAAAAAABAAAIAAAAAAIAAAIAAAAAAQAABAAAAAABAAAGAAAAAAEAAAIAAAAAAwAABAAAAAABAAAGAAAAAAEAAAIAAAAAAQAACAAAAAACAAACAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAKAAAAAAEAAAQAAAAAAQAAAAAAAAABAAACAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAKAAAAAAEAAAQAAAAAAQAAAAAAAAABAAACAAAAAAEAAAYAAAAAAQAAAgAAAAAEAAAEAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAIAAAAAAIAAAIAAAAAAQAABgAAAAABAAACAAAAAAEAAAoAAAAAAQAABAAAAAABAAAAAAAAAAEAAAIAAAAAAQAABgAAAAABAAACAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAEAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAGAAAAAAEAAAIAAAAAAQAABgAAAAABAAACAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAGAAAAAAEAAAIAAAAAAQAABAAAAAABAAAKAAAAAAEAAAQAAAAAAQAAAAAAAAABAAACAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAIAAAAAAIAAAIAAAAAAQAACgAAAAABAAAEAAAAAAEAAAAAAAAAAQAAAgAAAAABAAAIAAAAAAIAAAIAAAAAAQAABgAAAAABAAACAAAAAAEAAAYAAAAAAQAAAgAAAAADAAAEAAAAAAEAAAoAAAAAAQAABAAAAAABAAAAAAAAAAEAAAIAAAAAAQAABgAAAAABAAACAAAAAAQAAAQAAAAAAQAABgAAAAABAAACAAAAAAEAAAoAAAAAAQAABAAAAAABAAAAAAAAAAEAAAIAAAAAAQAACAAAAAACAAACAAAAAAEAAAYAAAAAAQAAAgAAAAABAAAIAAAAAAIAAAIAAAAAAQAABgAAAAABAAACAAAAABxzdHNjAAAAAAAAAAEAAAABAAAAmwAAAAEAAAKAc3RzegAAAAAAAAAAAAAAmwAABUYAAACIAAAAeQAAAFgAAAAvAAAAiQAAAE0AAAA0AAAAMQAAAIcAAAAyAAAARwAAAF0AAABTAAAAQwAAAJYAAABrAAAAOwAAAJ8AAAA7AAAAZgAAAE4AAADGAAAARgAAAKoAAACkAAAASAAAANYAAABfAAAAnwAAAEQAAACwAAAATgAAAJIAAABkAAAAoAAAAFEAAADFAAAA7wAAAFQAAABbAAAA/QAAAFMAAADqAAABDQAAADYAAABsAAAAVAAAACkAAAArAAAAdwAAAFoAAAAcAAAASQAAACcAAAAkAAAAYAAAAA4AAABdAAAAOAAAABkAAAB0AAAAGwAAAOwAAACIAAAAdgAAAEkAAACVAAAAMwAAAJcAAABcAAAASQAAAEEAAACHAAAATQAAAE0AAABpAAAAMgAAAF8AAABCAAAAGwAAAG0AAABCAAAAJAAAAMsAAABAAAABDQAAAIwAAABpAAAAJAAAALMAAAA6AAAAZQAAADgAAABqAAAAVQAAADUAAABPAAAALwAAAGkAAAA6AAAASQAAACUAAACCAAAAOAAAADwAAAD3AAAAiQAAAE8AAABfAAAArAAAAE0AAADeAAAAdwAAAHIAAAEzAAAAawAAAGAAAABjAAABEgAAAE0AAABtAAAA7AAAAE4AAAB5AAAALQAAAJgAAACCAAAAqQAAAK0AAABrAAAAMAAAAE0AAABEAAAANQAAAE4AAAA5AAAAkQAAAD8AAABJAAAAKwAAAHwAAABTAAAAOQAAADwAAAB6AAAATgAAADwAAADFAAAAPwAAAKQAAABmAAAARAAAAGMAAAAzAAAAFHN0Y28AAAAAAAAAAQAAADAAAABidWR0YQAAAFptZXRhAAAAAAAAACFoZGxyAAAAAAAAAABtZGlyYXBwbAAAAAAAAAAAAAAAAC1pbHN0AAAAJal0b28AAAAdZGF0YQAAAAEAAAAATGF2ZjU4Ljc2LjEwMA==\" type=\"video/mp4\">\n",
       " Your browser does not support the video tag.\n",
       " </video>"
      ]
     },
     "metadata": {},
     "execution_count": 12
    }
   ]
  }
 ]
}