packages/robots/diffusion_policy/diffusion_policy_state_pusht_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VrX4VTl5pYNq"
   },
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "L5E-nR6ornyg"
   },
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OknH8Qfqrtc9"
   },
   "outputs": [],
   "source": [
    "#@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, _ = 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]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "vHepJOFBucwg"
   },
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9ZiHF3lzvB6k"
   },
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "X-XRB_g3vsgf"
   },
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4APZkqh336-M"
   },
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "93E9RdnR4D8v"
   },
   "outputs": [],
   "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",
    "    parameters=noise_pred_net.parameters(),\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.parameters())\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 = noise_pred_net\n",
    "ema.copy_to(ema_noise_pred_net.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6F3hUbIuxGdO"
   },
   "outputs": [],
   "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.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OyLjlNQk5nr9"
   },
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}