SayCan-Robot-Pick-Place.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CfYKd3ZPB6s_"
   },
   "source": [
    "### Copyright 2022 Google LLC. SPDX-License-Identifier: Apache-2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bLzMRizDB2l5"
   },
   "source": [
    "Copyright 2022 Google LLC. SPDX-License-Identifier: Apache-2.0\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
    "\n",
    "https://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6GrV-nN-CTpl"
   },
   "source": [
    "# SayCan on a Robot Pick and Place Tabletop Environment\n",
    "\n",
    "[SayCan](https://say-can.github.io/) is an algorithm that grounds large language models with robotic affordances for long-horizon planning. Given a set of low-level robotic skills (e.g., \"put the green block in the red bowl\") and a high-level instruction (e.g., \"stack all the blocks\"), it scores what a language model believes will help forward the high-level instruction and scores what a robotic affordance model believes is possible. Together these give a task that is useful and possible and the robot executes the command.\n",
    "\n",
    "<img src=\"https://github.com/say-can/say-can.github.io/blob/main/img/saycan.png?raw=true\" height=\"320px\">\n",
    "\n",
    "This colab runs an example of SayCan for a pick and place robot on a table top.\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/say-can/say-can.github.io/main/img/open_source_tabletop.png\" height=\"320px\">\n",
    "\n",
    "Models used: [GPT-3](https://arxiv.org/abs/2005.14165) (InstructGPT), [CLIP](https://arxiv.org/abs/2103.00020) (ViT-B/32), [ViLD](https://arxiv.org/abs/2104.13921), and [CLIPort](https://cliport.github.io/) variant ([Transporter Nets](https://transporternets.github.io/))\n",
    "\n",
    "### **Quick Start:**\n",
    "\n",
    "**Step 1.** Register for an [OpenAI API key](https://openai.com/blog/openai-api/) to use GPT-3 (there's a free trial) and enter it below\n",
    "\n",
    "**Step 2.** Menu > Change runtime type > Hardware accelerator > \"GPU\"\n",
    "\n",
    "**Step 3.** Menu > Runtime > Run all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d2UdCCcpwbl0"
   },
   "source": [
    "## Setup and Installation\n",
    "Installs required packages (CLIP, PyTorch, etc.) and imports necessary libraries for vision, robotics, and machine learning tasks.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F2mFnHhPSTd_"
   },
   "outputs": [],
   "source": [
    "#@markdown Setup and Installation\n",
    "\n",
    "# Install required packages\n",
    "%pip install ftfy regex tqdm fvcore imageio==2.4.1 imageio-ffmpeg==0.4.5\n",
    "%pip install git+https://github.com/openai/CLIP.git\n",
    "%pip install -U --no-cache-dir gdown --pre\n",
    "%pip install pybullet moviepy flax==0.5.3 openai easydict imageio-ffmpeg\n",
    "%pip install tensorflow==2.7.0  # Uncomment if error: UNIMPLEMENTED: DNN library is not found.\n",
    "\n",
    "# Import required libraries\n",
    "import collections\n",
    "import datetime\n",
    "import os\n",
    "os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'\n",
    "import random\n",
    "import threading\n",
    "import time\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# Computer Vision and Image Processing\n",
    "import cv2\n",
    "import imageio\n",
    "from PIL import Image\n",
    "from moviepy.editor import ImageSequenceClip\n",
    "\n",
    "# Machine Learning Libraries\n",
    "import clip\n",
    "import flax\n",
    "from flax import linen as nn\n",
    "# from flax.training import checkpoints\n",
    "# from flax.metrics import tensorboard\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import torch\n",
    "import tensorflow.compat.v1 as tf\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "# Scientific Computing and Data Processing\n",
    "import numpy as np\n",
    "import optax\n",
    "from heapq import nlargest\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Utilities and Others\n",
    "from easydict import EasyDict\n",
    "import openai\n",
    "import pickle\n",
    "import pybullet\n",
    "\n",
    "import pybullet_data\n",
    "from tqdm import tqdm\n",
    "import IPython\n",
    "\n",
    "# Download required assets if they don't exist\n",
    "if not os.path.exists('ur5e/ur5e.urdf'):\n",
    "    !gdown --id 1Cc_fDSBL6QiDvNT4dpfAEbhbALSVoWcc\n",
    "    !gdown --id 1yOMEm-Zp_DL3nItG9RozPeJAmeOldekX\n",
    "    !gdown --id 1GsqNLhEl9dd4Mc3BM0dX3MibOI1FVWNM\n",
    "    !unzip ur5e.zip\n",
    "    !unzip robotiq_2f_85.zip\n",
    "    !unzip bowl.zip\n",
    "\n",
    "# Download ViLD pretrained model weights\n",
    "!gsutil cp -r gs://cloud-tpu-checkpoints/detection/projects/vild/colab/image_path_v2 ./\n",
    "\n",
    "# Setup TensorBoard\n",
    "%load_ext tensorboard\n",
    "\n",
    "\n",
    "# Display GPU information\n",
    "!nvidia-smi\n",
    "\n",
    "# Check JAX backend\n",
    "from jax.lib import xla_bridge\n",
    "print(\"JAX backend:\", xla_bridge.get_backend().platform)\n",
    "# Configure OpenAI API\n",
    "\n",
    "load_dotenv()  # Load environment variables from .env file\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "if openai_api_key is None:\n",
    "    raise ValueError(\"Please set OPENAI_API_KEY in .env file\")\n",
    "openai.api_key = openai_api_key  # Set the API key for the openai package\n",
    "ENGINE = \"gpt-3.5-turbo-instruct\"\n",
    "# Note for scoring model, due to limitations of the GPT-3 api, each option \n",
    "# requires a separate call and can be expensive. Recommend iterating with ada."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6YErUBSo3RT2"
   },
   "source": [
    "## Setup Environment\n",
    "Define PyBullet-based environment with a UR5e and 2-finger gripper\n",
    "\n",
    "Defines core configuration for the robot environment including:\n",
    "- Available objects (blocks, bowls)\n",
    "- Color definitions\n",
    "- Valid placement targets\n",
    "- Workspace boundaries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D6gTAdeX39yd"
   },
   "outputs": [],
   "source": [
    "#@markdown Global constants: pick and place objects, colors, workspace bounds\n",
    "\n",
    "PICK_TARGETS = {\n",
    "  \"blue block\": None,\n",
    "  \"red block\": None,\n",
    "  \"green block\": None,\n",
    "  \"yellow block\": None,\n",
    "}\n",
    "\n",
    "COLORS = {\n",
    "    \"blue\":   (78/255,  121/255, 167/255, 255/255),\n",
    "    \"red\":    (255/255,  87/255,  89/255, 255/255),\n",
    "    \"green\":  (89/255,  169/255,  79/255, 255/255),\n",
    "    \"yellow\": (237/255, 201/255,  72/255, 255/255),\n",
    "}\n",
    "\n",
    "PLACE_TARGETS = {\n",
    "  \"blue block\": None,\n",
    "  \"red block\": None,\n",
    "  \"green block\": None,\n",
    "  \"yellow block\": None,\n",
    "\n",
    "  \"blue bowl\": None,\n",
    "  \"red bowl\": None,\n",
    "  \"green bowl\": None,\n",
    "  \"yellow bowl\": None,\n",
    "\n",
    "  \"top left corner\":     (-0.3 + 0.05, -0.2 - 0.05, 0),\n",
    "  \"top right corner\":    (0.3 - 0.05,  -0.2 - 0.05, 0),\n",
    "  \"middle\":              (0,           -0.5,        0),\n",
    "  \"bottom left corner\":  (-0.3 + 0.05, -0.8 + 0.05, 0),\n",
    "  \"bottom right corner\": (0.3 - 0.05,  -0.8 + 0.05, 0),\n",
    "}\n",
    "\n",
    "PIXEL_SIZE = 0.00267857\n",
    "BOUNDS = np.float32([[-0.3, 0.3], [-0.8, -0.2], [0, 0.15]])  # X Y Z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Robot Gripper Definition \n",
    "Implements the Robotiq2F85 gripper class for the UR5e robot arm, handling:\n",
    "- Gripper movement control\n",
    "- Object grasping mechanics\n",
    "- Contact detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0bj91lBq6WeK"
   },
   "outputs": [],
   "source": [
    "#@markdown Gripper (Robotiq 2F85) code\n",
    "\n",
    "class Robotiq2F85:\n",
    "  \"\"\"Gripper handling for Robotiq 2F85.\"\"\"\n",
    "\n",
    "  def __init__(self, robot, tool):\n",
    "    self.robot = robot\n",
    "    self.tool = tool\n",
    "    pos = [0.1339999999999999, -0.49199999999872496, 0.5]\n",
    "    rot = pybullet.getQuaternionFromEuler([np.pi, 0, np.pi])\n",
    "    urdf = \"robotiq_2f_85/robotiq_2f_85.urdf\"\n",
    "    self.body = pybullet.loadURDF(urdf, pos, rot)\n",
    "    self.n_joints = pybullet.getNumJoints(self.body)\n",
    "    self.activated = False\n",
    "\n",
    "    # Connect gripper base to robot tool.\n",
    "    pybullet.createConstraint(self.robot, tool, self.body, 0, jointType=pybullet.JOINT_FIXED, jointAxis=[0, 0, 0], parentFramePosition=[0, 0, 0], childFramePosition=[0, 0, -0.07], childFrameOrientation=pybullet.getQuaternionFromEuler([0, 0, np.pi / 2]))\n",
    "\n",
    "    # Set friction coefficients for gripper fingers.\n",
    "    for i in range(pybullet.getNumJoints(self.body)):\n",
    "      pybullet.changeDynamics(self.body, i, lateralFriction=10.0, spinningFriction=1.0, rollingFriction=1.0, frictionAnchor=True)\n",
    "\n",
    "    # Start thread to handle additional gripper constraints.\n",
    "    self.motor_joint = 1\n",
    "    self.constraints_thread = threading.Thread(target=self.step)\n",
    "    self.constraints_thread.daemon = True\n",
    "    self.constraints_thread.start()\n",
    "\n",
    "  # Control joint positions by enforcing hard contraints on gripper behavior.\n",
    "  # Set one joint as the open/close motor joint (other joints should mimic).\n",
    "  def step(self):\n",
    "    while True:\n",
    "      try:\n",
    "        currj = [pybullet.getJointState(self.body, i)[0] for i in range(self.n_joints)]\n",
    "        indj = [6, 3, 8, 5, 10]\n",
    "        targj = [currj[1], -currj[1], -currj[1], currj[1], currj[1]]\n",
    "        pybullet.setJointMotorControlArray(self.body, indj, pybullet.POSITION_CONTROL, targj, positionGains=np.ones(5))\n",
    "      except:\n",
    "        return\n",
    "      time.sleep(0.001)\n",
    "\n",
    "  # Close gripper fingers.\n",
    "  def activate(self):\n",
    "    pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=1, force=10)\n",
    "    self.activated = True\n",
    "\n",
    "  # Open gripper fingers.\n",
    "  def release(self):\n",
    "    pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=-1, force=10)\n",
    "    self.activated = False\n",
    "\n",
    "  # If activated and object in gripper: check object contact.\n",
    "  # If activated and nothing in gripper: check gripper contact.\n",
    "  # If released: check proximity to surface (disabled).\n",
    "  def detect_contact(self):\n",
    "    obj, _, ray_frac = self.check_proximity()\n",
    "    if self.activated:\n",
    "      empty = self.grasp_width() < 0.01\n",
    "      cbody = self.body if empty else obj\n",
    "      if obj == self.body or obj == 0:\n",
    "        return False\n",
    "      return self.external_contact(cbody)\n",
    "  #   else:\n",
    "  #     return ray_frac < 0.14 or self.external_contact()\n",
    "\n",
    "  # Return if body is in contact with something other than gripper\n",
    "  def external_contact(self, body=None):\n",
    "    if body is None:\n",
    "      body = self.body\n",
    "    pts = pybullet.getContactPoints(bodyA=body)\n",
    "    pts = [pt for pt in pts if pt[2] != self.body]\n",
    "    return len(pts) > 0  # pylint: disable=g-explicit-length-test\n",
    "\n",
    "  def check_grasp(self):\n",
    "    while self.moving():\n",
    "      time.sleep(0.001)\n",
    "    success = self.grasp_width() > 0.01\n",
    "    return success\n",
    "\n",
    "  def grasp_width(self):\n",
    "    lpad = np.array(pybullet.getLinkState(self.body, 4)[0])\n",
    "    rpad = np.array(pybullet.getLinkState(self.body, 9)[0])\n",
    "    dist = np.linalg.norm(lpad - rpad) - 0.047813\n",
    "    return dist\n",
    "\n",
    "  def check_proximity(self):\n",
    "    ee_pos = np.array(pybullet.getLinkState(self.robot, self.tool)[0])\n",
    "    tool_pos = np.array(pybullet.getLinkState(self.body, 0)[0])\n",
    "    vec = (tool_pos - ee_pos) / np.linalg.norm((tool_pos - ee_pos))\n",
    "    ee_targ = ee_pos + vec\n",
    "    ray_data = pybullet.rayTest(ee_pos, ee_targ)[0]\n",
    "    obj, link, ray_frac = ray_data[0], ray_data[1], ray_data[2]\n",
    "    return obj, link, ray_frac"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Class\n",
    "Main PyBullet environment implementation including:\n",
    "- Robot arm control\n",
    "- Object spawning and physics\n",
    "- Camera rendering\n",
    "- State management\n",
    "- Action execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ELFr8iE28hVX"
   },
   "outputs": [],
   "source": [
    "#@markdown Gym-style environment code\n",
    "\n",
    "class PickPlaceEnv():\n",
    "\n",
    "  def __init__(self):\n",
    "    self.dt = 1/480\n",
    "    self.sim_step = 0\n",
    "\n",
    "    # Configure and start PyBullet.\n",
    "    # python3 -m pybullet_utils.runServer\n",
    "    # pybullet.connect(pybullet.SHARED_MEMORY)  # pybullet.GUI for local GUI.\n",
    "    pybullet.connect(pybullet.DIRECT)  # pybullet.GUI for local GUI.\n",
    "    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_GUI, 0)\n",
    "    pybullet.setPhysicsEngineParameter(enableFileCaching=0)\n",
    "    assets_path = os.path.dirname(os.path.abspath(\"\"))\n",
    "    pybullet.setAdditionalSearchPath(assets_path)\n",
    "    pybullet.setAdditionalSearchPath(pybullet_data.getDataPath())\n",
    "    pybullet.setTimeStep(self.dt)\n",
    "\n",
    "    self.home_joints = (np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 3 * np.pi / 2, 0)  # Joint angles: (J0, J1, J2, J3, J4, J5).\n",
    "    self.home_ee_euler = (np.pi, 0, np.pi)  # (RX, RY, RZ) rotation in Euler angles.\n",
    "    self.ee_link_id = 9  # Link ID of UR5 end effector.\n",
    "    self.tip_link_id = 10  # Link ID of gripper finger tips.\n",
    "    self.gripper = None\n",
    "\n",
    "  def reset(self, config):\n",
    "    pybullet.resetSimulation(pybullet.RESET_USE_DEFORMABLE_WORLD)\n",
    "    pybullet.setGravity(0, 0, -9.8)\n",
    "    self.cache_video = []\n",
    "\n",
    "    # Temporarily disable rendering to load URDFs faster.\n",
    "    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 0)\n",
    "\n",
    "    # Add robot.\n",
    "    pybullet.loadURDF(\"plane.urdf\", [0, 0, -0.001])\n",
    "    self.robot_id = pybullet.loadURDF(\"ur5e/ur5e.urdf\", [0, 0, 0], flags=pybullet.URDF_USE_MATERIAL_COLORS_FROM_MTL)\n",
    "    self.ghost_id = pybullet.loadURDF(\"ur5e/ur5e.urdf\", [0, 0, -10])  # For forward kinematics.\n",
    "    self.joint_ids = [pybullet.getJointInfo(self.robot_id, i) for i in range(pybullet.getNumJoints(self.robot_id))]\n",
    "    self.joint_ids = [j[0] for j in self.joint_ids if j[2] == pybullet.JOINT_REVOLUTE]\n",
    "\n",
    "    # Move robot to home configuration.\n",
    "    for i in range(len(self.joint_ids)):\n",
    "      pybullet.resetJointState(self.robot_id, self.joint_ids[i], self.home_joints[i])\n",
    "\n",
    "    # Add gripper.\n",
    "    if self.gripper is not None:\n",
    "      while self.gripper.constraints_thread.is_alive():\n",
    "        self.constraints_thread_active = False\n",
    "    self.gripper = Robotiq2F85(self.robot_id, self.ee_link_id)\n",
    "    self.gripper.release()\n",
    "\n",
    "    # Add workspace.\n",
    "    plane_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])\n",
    "    plane_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])\n",
    "    plane_id = pybullet.createMultiBody(0, plane_shape, plane_visual, basePosition=[0, -0.5, 0])\n",
    "    pybullet.changeVisualShape(plane_id, -1, rgbaColor=[0.2, 0.2, 0.2, 1.0])\n",
    "\n",
    "    # Load objects according to config.\n",
    "    self.config = config\n",
    "    self.obj_name_to_id = {}\n",
    "    obj_names = list(self.config[\"pick\"]) + list(self.config[\"place\"])\n",
    "    obj_xyz = np.zeros((0, 3))\n",
    "    for obj_name in obj_names:\n",
    "      if (\"block\" in obj_name) or (\"bowl\" in obj_name):\n",
    "\n",
    "        # Get random position 15cm+ from other objects.\n",
    "        while True:\n",
    "          rand_x = np.random.uniform(BOUNDS[0, 0] + 0.1, BOUNDS[0, 1] - 0.1)\n",
    "          rand_y = np.random.uniform(BOUNDS[1, 0] + 0.1, BOUNDS[1, 1] - 0.1)\n",
    "          rand_xyz = np.float32([rand_x, rand_y, 0.03]).reshape(1, 3)\n",
    "          if len(obj_xyz) == 0:\n",
    "            obj_xyz = np.concatenate((obj_xyz, rand_xyz), axis=0)\n",
    "            break\n",
    "          else:\n",
    "            nn_dist = np.min(np.linalg.norm(obj_xyz - rand_xyz, axis=1)).squeeze()\n",
    "            if nn_dist > 0.15:\n",
    "              obj_xyz = np.concatenate((obj_xyz, rand_xyz), axis=0)\n",
    "              break\n",
    "        \n",
    "        object_color = COLORS[obj_name.split(\" \")[0]]\n",
    "        object_type = obj_name.split(\" \")[1]\n",
    "        object_position = rand_xyz.squeeze()\n",
    "        if object_type == \"block\":\n",
    "          object_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])\n",
    "          object_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])\n",
    "          object_id = pybullet.createMultiBody(0.01, object_shape, object_visual, basePosition=object_position)\n",
    "        elif object_type == \"bowl\":\n",
    "          object_position[2] = 0\n",
    "          object_id = pybullet.loadURDF(\"bowl/bowl.urdf\", object_position, useFixedBase=1)\n",
    "        pybullet.changeVisualShape(object_id, -1, rgbaColor=object_color)\n",
    "        self.obj_name_to_id[obj_name] = object_id\n",
    "\n",
    "    # Re-enable rendering.\n",
    "    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 1)\n",
    "\n",
    "    for _ in range(200):\n",
    "      pybullet.stepSimulation()\n",
    "    return self.get_observation()\n",
    "\n",
    "  def servoj(self, joints):\n",
    "    \"\"\"Move to target joint positions with position control.\"\"\"\n",
    "    pybullet.setJointMotorControlArray(\n",
    "      bodyIndex=self.robot_id,\n",
    "      jointIndices=self.joint_ids,\n",
    "      controlMode=pybullet.POSITION_CONTROL,\n",
    "      targetPositions=joints,\n",
    "      positionGains=[0.01]*6)\n",
    "  \n",
    "  def movep(self, position):\n",
    "    \"\"\"Move to target end effector position.\"\"\"\n",
    "    joints = pybullet.calculateInverseKinematics(\n",
    "        bodyUniqueId=self.robot_id,\n",
    "        endEffectorLinkIndex=self.tip_link_id,\n",
    "        targetPosition=position,\n",
    "        targetOrientation=pybullet.getQuaternionFromEuler(self.home_ee_euler),\n",
    "        maxNumIterations=100)\n",
    "    self.servoj(joints)\n",
    "\n",
    "  def step(self, action=None):\n",
    "    \"\"\"Do pick and place motion primitive.\"\"\"\n",
    "    pick_xyz, place_xyz = action[\"pick\"].copy(), action[\"place\"].copy()\n",
    "\n",
    "    # Set fixed primitive z-heights.\n",
    "    hover_xyz = pick_xyz.copy() + np.float32([0, 0, 0.2])\n",
    "    pick_xyz[2] = 0.03\n",
    "    place_xyz[2] = 0.15\n",
    "\n",
    "    # Move to object.\n",
    "    ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
    "    while np.linalg.norm(hover_xyz - ee_xyz) > 0.01:\n",
    "      self.movep(hover_xyz)\n",
    "      self.step_sim_and_render()\n",
    "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
    "    while np.linalg.norm(pick_xyz - ee_xyz) > 0.01:\n",
    "      self.movep(pick_xyz)\n",
    "      self.step_sim_and_render()\n",
    "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
    "\n",
    "    # Pick up object.\n",
    "    self.gripper.activate()\n",
    "    for _ in range(240):\n",
    "      self.step_sim_and_render()\n",
    "    while np.linalg.norm(hover_xyz - ee_xyz) > 0.01:\n",
    "      self.movep(hover_xyz)\n",
    "      self.step_sim_and_render()\n",
    "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
    "    \n",
    "    # Move to place location.\n",
    "    while np.linalg.norm(place_xyz - ee_xyz) > 0.01:\n",
    "      self.movep(place_xyz)\n",
    "      self.step_sim_and_render()\n",
    "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
    "\n",
    "    # Place down object.\n",
    "    while (not self.gripper.detect_contact()) and (place_xyz[2] > 0.03):\n",
    "      place_xyz[2] -= 0.001\n",
    "      self.movep(place_xyz)\n",
    "      for _ in range(3):\n",
    "        self.step_sim_and_render()\n",
    "    self.gripper.release()\n",
    "    for _ in range(240):\n",
    "      self.step_sim_and_render()\n",
    "    place_xyz[2] = 0.2\n",
    "    ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
    "    while np.linalg.norm(place_xyz - ee_xyz) > 0.01:\n",
    "      self.movep(place_xyz)\n",
    "      self.step_sim_and_render()\n",
    "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
    "    place_xyz = np.float32([0, -0.5, 0.2])\n",
    "    while np.linalg.norm(place_xyz - ee_xyz) > 0.01:\n",
    "      self.movep(place_xyz)\n",
    "      self.step_sim_and_render()\n",
    "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
    "\n",
    "    observation = self.get_observation()\n",
    "    reward = self.get_reward()\n",
    "    done = False\n",
    "    info = {}\n",
    "    return observation, reward, done, info\n",
    "\n",
    "  def set_alpha_transparency(self, alpha: float) -> None:\n",
    "    for id in range(20):\n",
    "      visual_shape_data = pybullet.getVisualShapeData(id)\n",
    "      for i in range(len(visual_shape_data)):\n",
    "        object_id, link_index, _, _, _, _, _, rgba_color = visual_shape_data[i]\n",
    "        rgba_color = list(rgba_color[0:3]) +  [alpha]\n",
    "        pybullet.changeVisualShape(\n",
    "            self.robot_id, linkIndex=i, rgbaColor=rgba_color)      \n",
    "        pybullet.changeVisualShape(\n",
    "            self.gripper.body, linkIndex=i, rgbaColor=rgba_color)\n",
    "\n",
    "  def step_sim_and_render(self):\n",
    "    pybullet.stepSimulation()\n",
    "    self.sim_step += 1\n",
    "\n",
    "    # Render current image at 8 FPS.\n",
    "    if self.sim_step % 60 == 0:\n",
    "      self.cache_video.append(self.get_camera_image())\n",
    "\n",
    "  def get_camera_image(self):\n",
    "    image_size = (240, 240)\n",
    "    intrinsics = (120., 0, 120., 0, 120., 120., 0, 0, 1)\n",
    "    color, _, _, _, _ = env.render_image(image_size, intrinsics)\n",
    "    return color\n",
    "\n",
    "  def get_camera_image_top(self, \n",
    "                           image_size=(240, 240), \n",
    "                           intrinsics=(2000., 0, 2000., 0, 2000., 2000., 0, 0, 1),\n",
    "                           position=(0, -0.5, 5),\n",
    "                           orientation=(0, np.pi, -np.pi / 2),\n",
    "                           zrange=(0.01, 1.),\n",
    "                           set_alpha=True):\n",
    "    set_alpha and self.set_alpha_transparency(0)\n",
    "    color, _, _, _, _ = env.render_image_top(image_size, \n",
    "                                             intrinsics,\n",
    "                                             position,\n",
    "                                             orientation,\n",
    "                                             zrange)\n",
    "    set_alpha and self.set_alpha_transparency(1)\n",
    "    return color\n",
    "\n",
    "  def get_reward(self):\n",
    "    return 0  # TODO: check did the robot follow text instructions?\n",
    "\n",
    "  def get_observation(self):\n",
    "    observation = {}\n",
    "\n",
    "    # Render current image.\n",
    "    color, depth, position, orientation, intrinsics = self.render_image()\n",
    "\n",
    "    # Get heightmaps and colormaps.\n",
    "    points = self.get_pointcloud(depth, intrinsics)\n",
    "    position = np.float32(position).reshape(3, 1)\n",
    "    rotation = pybullet.getMatrixFromQuaternion(orientation)\n",
    "    rotation = np.float32(rotation).reshape(3, 3)\n",
    "    transform = np.eye(4)\n",
    "    transform[:3, :] = np.hstack((rotation, position))\n",
    "    points = self.transform_pointcloud(points, transform)\n",
    "    heightmap, colormap, xyzmap = self.get_heightmap(points, color, BOUNDS, PIXEL_SIZE)\n",
    "\n",
    "    observation[\"image\"] = colormap\n",
    "    observation[\"xyzmap\"] = xyzmap\n",
    "    observation[\"pick\"] = list(self.config[\"pick\"])\n",
    "    observation[\"place\"] = list(self.config[\"place\"])\n",
    "    return observation\n",
    "\n",
    "  def render_image(self, image_size=(720, 720), intrinsics=(360., 0, 360., 0, 360., 360., 0, 0, 1)):\n",
    "\n",
    "    # Camera parameters.\n",
    "    position = (0, -0.85, 0.4)\n",
    "    orientation = (np.pi / 4 + np.pi / 48, np.pi, np.pi)\n",
    "    orientation = pybullet.getQuaternionFromEuler(orientation)\n",
    "    zrange = (0.01, 10.)\n",
    "    noise=True\n",
    "\n",
    "    # OpenGL camera settings.\n",
    "    lookdir = np.float32([0, 0, 1]).reshape(3, 1)\n",
    "    updir = np.float32([0, -1, 0]).reshape(3, 1)\n",
    "    rotation = pybullet.getMatrixFromQuaternion(orientation)\n",
    "    rotm = np.float32(rotation).reshape(3, 3)\n",
    "    lookdir = (rotm @ lookdir).reshape(-1)\n",
    "    updir = (rotm @ updir).reshape(-1)\n",
    "    lookat = position + lookdir\n",
    "    focal_len = intrinsics[0]\n",
    "    znear, zfar = (0.01, 10.)\n",
    "    viewm = pybullet.computeViewMatrix(position, lookat, updir)\n",
    "    fovh = (image_size[0] / 2) / focal_len\n",
    "    fovh = 180 * np.arctan(fovh) * 2 / np.pi\n",
    "\n",
    "    # Notes: 1) FOV is vertical FOV 2) aspect must be float\n",
    "    aspect_ratio = image_size[1] / image_size[0]\n",
    "    projm = pybullet.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)\n",
    "\n",
    "    # Render with OpenGL camera settings.\n",
    "    _, _, color, depth, segm = pybullet.getCameraImage(\n",
    "        width=image_size[1],\n",
    "        height=image_size[0],\n",
    "        viewMatrix=viewm,\n",
    "        projectionMatrix=projm,\n",
    "        shadow=1,\n",
    "        flags=pybullet.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,\n",
    "        renderer=pybullet.ER_BULLET_HARDWARE_OPENGL)\n",
    "\n",
    "    # Get color image.\n",
    "    color_image_size = (image_size[0], image_size[1], 4)\n",
    "    color = np.array(color, dtype=np.uint8).reshape(color_image_size)\n",
    "    color = color[:, :, :3]  # remove alpha channel\n",
    "    if noise:\n",
    "      color = np.int32(color)\n",
    "      color += np.int32(np.random.normal(0, 3, color.shape))\n",
    "      color = np.uint8(np.clip(color, 0, 255))\n",
    "\n",
    "    # Get depth image.\n",
    "    depth_image_size = (image_size[0], image_size[1])\n",
    "    zbuffer = np.float32(depth).reshape(depth_image_size)\n",
    "    depth = (zfar + znear - (2 * zbuffer - 1) * (zfar - znear))\n",
    "    depth = (2 * znear * zfar) / depth\n",
    "    if noise:\n",
    "      depth += np.random.normal(0, 0.003, depth.shape)\n",
    "\n",
    "    intrinsics = np.float32(intrinsics).reshape(3, 3)\n",
    "    return color, depth, position, orientation, intrinsics\n",
    "\n",
    "  def render_image_top(self, \n",
    "                       image_size=(240, 240), \n",
    "                       intrinsics=(2000., 0, 2000., 0, 2000., 2000., 0, 0, 1),\n",
    "                       position=(0, -0.5, 5),\n",
    "                       orientation=(0, np.pi, -np.pi / 2),\n",
    "                       zrange=(0.01, 1.)):\n",
    "\n",
    "    # Camera parameters.\n",
    "    orientation = pybullet.getQuaternionFromEuler(orientation)\n",
    "    noise=True\n",
    "\n",
    "    # OpenGL camera settings.\n",
    "    lookdir = np.float32([0, 0, 1]).reshape(3, 1)\n",
    "    updir = np.float32([0, -1, 0]).reshape(3, 1)\n",
    "    rotation = pybullet.getMatrixFromQuaternion(orientation)\n",
    "    rotm = np.float32(rotation).reshape(3, 3)\n",
    "    lookdir = (rotm @ lookdir).reshape(-1)\n",
    "    updir = (rotm @ updir).reshape(-1)\n",
    "    lookat = position + lookdir\n",
    "    focal_len = intrinsics[0]\n",
    "    znear, zfar = (0.01, 10.)\n",
    "    viewm = pybullet.computeViewMatrix(position, lookat, updir)\n",
    "    fovh = (image_size[0] / 2) / focal_len\n",
    "    fovh = 180 * np.arctan(fovh) * 2 / np.pi\n",
    "\n",
    "    # Notes: 1) FOV is vertical FOV 2) aspect must be float\n",
    "    aspect_ratio = image_size[1] / image_size[0]\n",
    "    projm = pybullet.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)\n",
    "\n",
    "    # Render with OpenGL camera settings.\n",
    "    _, _, color, depth, segm = pybullet.getCameraImage(\n",
    "        width=image_size[1],\n",
    "        height=image_size[0],\n",
    "        viewMatrix=viewm,\n",
    "        projectionMatrix=projm,\n",
    "        shadow=1,\n",
    "        flags=pybullet.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,\n",
    "        renderer=pybullet.ER_BULLET_HARDWARE_OPENGL)\n",
    "\n",
    "    # Get color image.\n",
    "    color_image_size = (image_size[0], image_size[1], 4)\n",
    "    color = np.array(color, dtype=np.uint8).reshape(color_image_size)\n",
    "    color = color[:, :, :3]  # remove alpha channel\n",
    "    if noise:\n",
    "      color = np.int32(color)\n",
    "      color += np.int32(np.random.normal(0, 3, color.shape))\n",
    "      color = np.uint8(np.clip(color, 0, 255))\n",
    "\n",
    "    # Get depth image.\n",
    "    depth_image_size = (image_size[0], image_size[1])\n",
    "    zbuffer = np.float32(depth).reshape(depth_image_size)\n",
    "    depth = (zfar + znear - (2 * zbuffer - 1) * (zfar - znear))\n",
    "    depth = (2 * znear * zfar) / depth\n",
    "    if noise:\n",
    "      depth += np.random.normal(0, 0.003, depth.shape)\n",
    "\n",
    "    intrinsics = np.float32(intrinsics).reshape(3, 3)\n",
    "    return color, depth, position, orientation, intrinsics\n",
    "\n",
    "  def get_pointcloud(self, depth, intrinsics):\n",
    "    \"\"\"Get 3D pointcloud from perspective depth image.\n",
    "    Args:\n",
    "      depth: HxW float array of perspective depth in meters.\n",
    "      intrinsics: 3x3 float array of camera intrinsics matrix.\n",
    "    Returns:\n",
    "      points: HxWx3 float array of 3D points in camera coordinates.\n",
    "    \"\"\"\n",
    "    height, width = depth.shape\n",
    "    xlin = np.linspace(0, width - 1, width)\n",
    "    ylin = np.linspace(0, height - 1, height)\n",
    "    px, py = np.meshgrid(xlin, ylin)\n",
    "    px = (px - intrinsics[0, 2]) * (depth / intrinsics[0, 0])\n",
    "    py = (py - intrinsics[1, 2]) * (depth / intrinsics[1, 1])\n",
    "    points = np.float32([px, py, depth]).transpose(1, 2, 0)\n",
    "    return points\n",
    "\n",
    "  def transform_pointcloud(self, points, transform):\n",
    "    \"\"\"Apply rigid transformation to 3D pointcloud.\n",
    "    Args:\n",
    "      points: HxWx3 float array of 3D points in camera coordinates.\n",
    "      transform: 4x4 float array representing a rigid transformation matrix.\n",
    "    Returns:\n",
    "      points: HxWx3 float array of transformed 3D points.\n",
    "    \"\"\"\n",
    "    padding = ((0, 0), (0, 0), (0, 1))\n",
    "    homogen_points = np.pad(points.copy(), padding,\n",
    "                            \"constant\", constant_values=1)\n",
    "    for i in range(3):\n",
    "      points[Ellipsis, i] = np.sum(transform[i, :] * homogen_points, axis=-1)\n",
    "    return points\n",
    "\n",
    "  def get_heightmap(self, points, colors, bounds, pixel_size):\n",
    "    \"\"\"Get top-down (z-axis) orthographic heightmap image from 3D pointcloud.\n",
    "    Args:\n",
    "      points: HxWx3 float array of 3D points in world coordinates.\n",
    "      colors: HxWx3 uint8 array of values in range 0-255 aligned with points.\n",
    "      bounds: 3x2 float array of values (rows: X,Y,Z; columns: min,max) defining\n",
    "        region in 3D space to generate heightmap in world coordinates.\n",
    "      pixel_size: float defining size of each pixel in meters.\n",
    "    Returns:\n",
    "      heightmap: HxW float array of height (from lower z-bound) in meters.\n",
    "      colormap: HxWx3 uint8 array of backprojected color aligned with heightmap.\n",
    "      xyzmap: HxWx3 float array of XYZ points in world coordinates.\n",
    "    \"\"\"\n",
    "    width = int(np.round((bounds[0, 1] - bounds[0, 0]) / pixel_size))\n",
    "    height = int(np.round((bounds[1, 1] - bounds[1, 0]) / pixel_size))\n",
    "    heightmap = np.zeros((height, width), dtype=np.float32)\n",
    "    colormap = np.zeros((height, width, colors.shape[-1]), dtype=np.uint8)\n",
    "    xyzmap = np.zeros((height, width, 3), dtype=np.float32)\n",
    "\n",
    "    # Filter out 3D points that are outside of the predefined bounds.\n",
    "    ix = (points[Ellipsis, 0] >= bounds[0, 0]) & (points[Ellipsis, 0] < bounds[0, 1])\n",
    "    iy = (points[Ellipsis, 1] >= bounds[1, 0]) & (points[Ellipsis, 1] < bounds[1, 1])\n",
    "    iz = (points[Ellipsis, 2] >= bounds[2, 0]) & (points[Ellipsis, 2] < bounds[2, 1])\n",
    "    valid = ix & iy & iz\n",
    "    points = points[valid]\n",
    "    colors = colors[valid]\n",
    "\n",
    "    # Sort 3D points by z-value, which works with array assignment to simulate\n",
    "    # z-buffering for rendering the heightmap image.\n",
    "    iz = np.argsort(points[:, -1])\n",
    "    points, colors = points[iz], colors[iz]\n",
    "    px = np.int32(np.floor((points[:, 0] - bounds[0, 0]) / pixel_size))\n",
    "    py = np.int32(np.floor((points[:, 1] - bounds[1, 0]) / pixel_size))\n",
    "    px = np.clip(px, 0, width - 1)\n",
    "    py = np.clip(py, 0, height - 1)\n",
    "    heightmap[py, px] = points[:, 2] - bounds[2, 0]\n",
    "    for c in range(colors.shape[-1]):\n",
    "      colormap[py, px, c] = colors[:, c]\n",
    "      xyzmap[py, px, c] = points[:, c]\n",
    "    colormap = colormap[::-1, :, :]  # Flip up-down.\n",
    "    xv, yv = np.meshgrid(np.linspace(BOUNDS[0, 0], BOUNDS[0, 1], height),\n",
    "                         np.linspace(BOUNDS[1, 0], BOUNDS[1, 1], width))\n",
    "    xyzmap[:, :, 0] = xv\n",
    "    xyzmap[:, :, 1] = yv\n",
    "    xyzmap = xyzmap[::-1, :, :]  # Flip up-down.\n",
    "    heightmap = heightmap[::-1, :]  # Flip up-down.\n",
    "    return heightmap, colormap, xyzmap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Initialization\n",
    "Creates environment instance and displays initial camera views:\n",
    "- Side perspective\n",
    "- Top-down view\n",
    "- Neural network input view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QqJTWp1qWRoq"
   },
   "outputs": [],
   "source": [
    "#@markdown Initialize environment \n",
    "\n",
    "if 'env' in locals():\n",
    "  # Safely exit gripper threading before re-initializing environment.\n",
    "  env.gripper.running = False\n",
    "  while env.gripper.constraints_thread.isAlive():\n",
    "    time.sleep(0.01)\n",
    "env = PickPlaceEnv()\n",
    "#@markdown Render images.\n",
    "\n",
    "# Define and reset environment.\n",
    "config = {'pick':  ['yellow block', 'green block', 'blue block'],\n",
    "          'place': ['yellow bowl', 'green bowl', 'blue bowl']}\n",
    "\n",
    "np.random.seed(42)\n",
    "obs = env.reset(config)\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "img = env.get_camera_image()\n",
    "plt.title('Perspective side-view')\n",
    "plt.imshow(img)\n",
    "plt.subplot(1, 2, 2)\n",
    "img = env.get_camera_image_top()\n",
    "img = np.flipud(img.transpose(1, 0, 2))\n",
    "plt.title('Orthographic top-view')\n",
    "plt.imshow(img)\n",
    "plt.show()\n",
    "\n",
    "# Note: orthographic cameras do not exist. But we can approximate them by\n",
    "# projecting a 3D point cloud from an RGB-D camera, then unprojecting that onto\n",
    "# an orthographic plane. Orthographic views are useful for spatial action maps.\n",
    "plt.title('Unprojected orthographic top-view')\n",
    "plt.imshow(obs['image'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Uyc6-l8IgQPS"
   },
   "source": [
    "## **Demo:** ViLD\n",
    "Run zero-shot open-vocabulary object detection with [ViLD](https://arxiv.org/abs/2104.13921) to generate a list of objects as a scene description for a large language model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ViLD Setup\n",
    "Configures Vision-Language Detection (ViLD) model for:\n",
    "- Object detection\n",
    "- Scene understanding\n",
    "- Visual grounding of language commands"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DBkMaGRng1Vd"
   },
   "outputs": [],
   "source": [
    "# Define and reset environment.\n",
    "config = {'pick':  ['yellow block', 'green block', 'blue block'],\n",
    "          'place': ['yellow bowl', 'green bowl', 'blue bowl']}\n",
    "\n",
    "np.random.seed(42)\n",
    "obs = env.reset(config)\n",
    "img = env.get_camera_image_top()\n",
    "img = np.flipud(img.transpose(1, 0, 2))\n",
    "plt.title('ViLD Input Image')\n",
    "plt.imshow(img)\n",
    "plt.show()\n",
    "imageio.imwrite('tmp.jpg', img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CLIP Model Configuration\n",
    "Sets up CLIP model for:\n",
    "- Text-to-image understanding\n",
    "- Visual feature extraction\n",
    "- Language understanding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ckfJ6LwjEiIQ"
   },
   "outputs": [],
   "source": [
    "#@markdown Load CLIP model with memory management\n",
    "\n",
    "# Clear CUDA cache before loading model\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "# Define device globally\n",
    "DEVICE = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.set_device(DEVICE)\n",
    "\n",
    "def load_clip_model(max_retries=3, initial_batch_size=32):\n",
    "    \"\"\"Load CLIP model with automatic memory management\"\"\"\n",
    "    current_device = DEVICE  # Use global device setting\n",
    "    \n",
    "    for attempt in range(max_retries):\n",
    "        try:\n",
    "            # Clear cache before loading\n",
    "            if torch.cuda.is_available():\n",
    "                torch.cuda.empty_cache()\n",
    "            \n",
    "            # Load model\n",
    "            clip_model, clip_preprocess = clip.load(\"ViT-B/32\", device=current_device)\n",
    "            clip_model.eval()\n",
    "            \n",
    "            # Print model info\n",
    "            print(\"Model parameters:\", f\"{np.sum([int(np.prod(p.shape)) for p in clip_model.parameters()]):,}\")\n",
    "            print(\"Input resolution:\", clip_model.visual.input_resolution)\n",
    "            print(\"Context length:\", clip_model.context_length)\n",
    "            print(\"Vocab size:\", clip_model.vocab_size)\n",
    "            \n",
    "            return clip_model, clip_preprocess, initial_batch_size\n",
    "\n",
    "        except RuntimeError as e:\n",
    "            print(f\"CUDA out of memory on attempt {attempt + 1}. Error: {e}\")\n",
    "            if attempt < max_retries - 1:\n",
    "                # Reduce batch size for next attempt\n",
    "                initial_batch_size = initial_batch_size // 2\n",
    "                print(f\"Reducing batch size to {initial_batch_size} and retrying...\")\n",
    "                if torch.cuda.is_available():\n",
    "                    torch.cuda.empty_cache()\n",
    "                time.sleep(1)  # Give GPU a moment to free memory\n",
    "            else:\n",
    "                print(\"Failed all attempts to load model. Trying CPU fallback...\")\n",
    "                current_device = \"cpu\"\n",
    "                clip_model, clip_preprocess = clip.load(\"ViT-B/32\", device=current_device)\n",
    "                clip_model.eval()\n",
    "                return clip_model, clip_preprocess, 16  # Smaller batch size for CPU\n",
    "\n",
    "# Load the model with automatic memory management\n",
    "try:\n",
    "    clip_model, clip_preprocess, batch_size = load_clip_model()\n",
    "    print(f\"Successfully loaded CLIP model on {DEVICE} with batch size {batch_size}\")\n",
    "except Exception as e:\n",
    "    print(f\"Failed to load CLIP model: {e}\")\n",
    "    raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "l-qWw2WDh_dp"
   },
   "outputs": [],
   "source": [
    "#@markdown Define ViLD hyperparameters.\n",
    "FLAGS = {\n",
    "    'prompt_engineering': True,\n",
    "    'this_is': True,\n",
    "    'temperature': 100.0,\n",
    "    'use_softmax': False,\n",
    "}\n",
    "FLAGS = EasyDict(FLAGS)\n",
    "\n",
    "\n",
    "# # Global matplotlib settings\n",
    "# SMALL_SIZE = 16#10\n",
    "# MEDIUM_SIZE = 18#12\n",
    "# BIGGER_SIZE = 20#14\n",
    "\n",
    "# plt.rc('font', size=MEDIUM_SIZE)         # controls default text sizes\n",
    "# plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title\n",
    "# plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels\n",
    "# plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels\n",
    "# plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels\n",
    "# plt.rc('legend', fontsize=MEDIUM_SIZE)   # legend fontsize\n",
    "# plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title\n",
    "\n",
    "\n",
    "# Parameters for drawing figure.\n",
    "display_input_size = (10, 10)\n",
    "overall_fig_size = (18, 24)\n",
    "\n",
    "line_thickness = 1\n",
    "fig_size_w = 35\n",
    "# fig_size_h = min(max(5, int(len(category_names) / 2.5) ), 10)\n",
    "mask_color =   'red'\n",
    "alpha = 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. ViLD Prompt Engineering Cell\n",
    "- **Purpose**: Create language templates for object detection\n",
    "- **Key Functions**: \n",
    " - `article()`: Determines \"a\" vs \"an\"\n",
    " - `processed_name()`: Cleans object names\n",
    " - `build_text_embedding()`: Creates CLIP text embeddings\n",
    "- **Templates**: Multiple prompt formats like:\n",
    " - \"a photo of {}\"\n",
    " - \"This is a {}\"\n",
    " - Various quality descriptions (good/bad/bright/dark)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "TOtCs2B9iDMJ"
   },
   "outputs": [],
   "source": [
    "#@markdown ViLD prompt engineering.\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "def article(name):\n",
    "  return \"an\" if name[0] in \"aeiou\" else \"a\"\n",
    "\n",
    "def processed_name(name, rm_dot=False):\n",
    "  # _ for lvis\n",
    "  # / for obj365\n",
    "  res = name.replace(\"_\", \" \").replace(\"/\", \" or \").lower()\n",
    "  if rm_dot:\n",
    "    res = res.rstrip(\".\")\n",
    "  return res\n",
    "\n",
    "single_template = [\n",
    "    \"a photo of {article} {}.\"\n",
    "]\n",
    "\n",
    "# multiple_templates = [\n",
    "#     \"There is {article} {} in the scene.\",\n",
    "#     \"a painting of a {}.\",\n",
    "# ]\n",
    "\n",
    "multiple_templates = [\n",
    "    'There is {article} {} in the scene.',\n",
    "    'There is the {} in the scene.',\n",
    "    'a photo of {article} {} in the scene.',\n",
    "    'a photo of the {} in the scene.',\n",
    "    'a photo of one {} in the scene.',\n",
    "\n",
    "\n",
    "    'itap of {article} {}.',\n",
    "    'itap of my {}.',  # itap: I took a picture of\n",
    "    'itap of the {}.',\n",
    "    'a photo of {article} {}.',\n",
    "    'a photo of my {}.',\n",
    "    'a photo of the {}.',\n",
    "    'a photo of one {}.',\n",
    "    'a photo of many {}.',\n",
    "\n",
    "    'a good photo of {article} {}.',\n",
    "    'a good photo of the {}.',\n",
    "    'a bad photo of {article} {}.',\n",
    "    'a bad photo of the {}.',\n",
    "    'a photo of a nice {}.',\n",
    "    'a photo of the nice {}.',\n",
    "    'a photo of a cool {}.',\n",
    "    'a photo of the cool {}.',\n",
    "    'a photo of a weird {}.',\n",
    "    'a photo of the weird {}.',\n",
    "\n",
    "    'a photo of a small {}.',\n",
    "    'a photo of the small {}.',\n",
    "    'a photo of a large {}.',\n",
    "    'a photo of the large {}.',\n",
    "\n",
    "    'a photo of a clean {}.',\n",
    "    'a photo of the clean {}.',\n",
    "    'a photo of a dirty {}.',\n",
    "    'a photo of the dirty {}.',\n",
    "\n",
    "    'a bright photo of {article} {}.',\n",
    "    'a bright photo of the {}.',\n",
    "    'a dark photo of {article} {}.',\n",
    "    'a dark photo of the {}.',\n",
    "\n",
    "    'a photo of a hard to see {}.',\n",
    "    'a photo of the hard to see {}.',\n",
    "    'a low resolution photo of {article} {}.',\n",
    "    'a low resolution photo of the {}.',\n",
    "    'a cropped photo of {article} {}.',\n",
    "    'a cropped photo of the {}.',\n",
    "    'a close-up photo of {article} {}.',\n",
    "    'a close-up photo of the {}.',\n",
    "    'a jpeg corrupted photo of {article} {}.',\n",
    "    'a jpeg corrupted photo of the {}.',\n",
    "    'a blurry photo of {article} {}.',\n",
    "    'a blurry photo of the {}.',\n",
    "    'a pixelated photo of {article} {}.',\n",
    "    'a pixelated photo of the {}.',\n",
    "\n",
    "    'a black and white photo of the {}.',\n",
    "    'a black and white photo of {article} {}.',\n",
    "\n",
    "    'a plastic {}.',\n",
    "    'the plastic {}.',\n",
    "\n",
    "    'a toy {}.',\n",
    "    'the toy {}.',\n",
    "    'a plushie {}.',\n",
    "    'the plushie {}.',\n",
    "    'a cartoon {}.',\n",
    "    'the cartoon {}.',\n",
    "\n",
    "    'an embroidered {}.',\n",
    "    'the embroidered {}.',\n",
    "\n",
    "    'a painting of the {}.',\n",
    "    'a painting of a {}.',\n",
    "]\n",
    "\n",
    "def build_text_embedding(categories):\n",
    "    if FLAGS.prompt_engineering:\n",
    "        templates = multiple_templates\n",
    "    else:\n",
    "        templates = single_template\n",
    "\n",
    "    with torch.no_grad():\n",
    "        all_text_embeddings = []\n",
    "        print(\"Building text embeddings...\")\n",
    "        \n",
    "        # Process in smaller batches\n",
    "        batch_size = 32  # Adjust this based on your GPU memory\n",
    "        for i in range(0, len(categories), batch_size):\n",
    "            batch_categories = categories[i:i + batch_size]\n",
    "            batch_embeddings = []\n",
    "            \n",
    "            for category in tqdm(batch_categories):\n",
    "                texts = [\n",
    "                    template.format(processed_name(category[\"name\"], rm_dot=True),\n",
    "                                  article=article(category[\"name\"]))\n",
    "                    for template in templates]\n",
    "                if FLAGS.this_is:\n",
    "                    texts = [\n",
    "                        \"This is \" + text if text.startswith(\"a\") or text.startswith(\"the\") else text \n",
    "                        for text in texts\n",
    "                    ]\n",
    "                texts = clip.tokenize(texts).to(device)\n",
    "                \n",
    "                try:\n",
    "                    text_embeddings = clip_model.encode_text(texts)\n",
    "                    text_embeddings /= text_embeddings.norm(dim=-1, keepdim=True)\n",
    "                    text_embedding = text_embeddings.mean(dim=0)\n",
    "                    text_embedding /= text_embedding.norm()\n",
    "                    batch_embeddings.append(text_embedding)\n",
    "                except RuntimeError as e:\n",
    "                    print(f\"CUDA out of memory. Error: {e}\")\n",
    "                    torch.cuda.empty_cache()\n",
    "                    # Retry with smaller batch\n",
    "                    texts = texts[:len(texts)//2]\n",
    "                    text_embeddings = clip_model.encode_text(texts)\n",
    "                    text_embeddings /= text_embeddings.norm(dim=-1, keepdim=True)\n",
    "                    text_embedding = text_embeddings.mean(dim=0)\n",
    "                    text_embedding /= text_embedding.norm()\n",
    "                    batch_embeddings.append(text_embedding)\n",
    "                \n",
    "                # Clear cache periodically\n",
    "                if len(batch_embeddings) % 8 == 0:\n",
    "                    torch.cuda.empty_cache()\n",
    "            \n",
    "            all_text_embeddings.extend(batch_embeddings)\n",
    "            torch.cuda.empty_cache()\n",
    "        \n",
    "        all_text_embeddings = torch.stack(all_text_embeddings, dim=1)\n",
    "        all_text_embeddings = all_text_embeddings.to(device)\n",
    "    \n",
    "    return all_text_embeddings.cpu().numpy().T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "FK5lnSQEiM4r"
   },
   "outputs": [],
   "source": [
    "#@markdown Load ViLD model.\n",
    "\n",
    "gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.2)\n",
    "session = tf.Session(graph=tf.Graph(), config=tf.ConfigProto(gpu_options=gpu_options))\n",
    "saved_model_dir = \"./image_path_v2\"\n",
    "_ = tf.saved_model.loader.load(session, [\"serve\"], saved_model_dir)\n",
    "\n",
    "numbered_categories = [{\"name\": str(idx), \"id\": idx,} for idx in range(50)]\n",
    "numbered_category_indices = {cat[\"id\"]: cat for cat in numbered_categories}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "gazf5PBkkzlK"
   },
   "outputs": [],
   "source": [
    "#@markdown Non-maximum suppression (NMS).\n",
    "def nms(dets, scores, thresh, max_dets=1000):\n",
    "  \"\"\"Non-maximum suppression.\n",
    "  Args:\n",
    "    dets: [N, 4]\n",
    "    scores: [N,]\n",
    "    thresh: iou threshold. Float\n",
    "    max_dets: int.\n",
    "  \"\"\"\n",
    "  y1 = dets[:, 0]\n",
    "  x1 = dets[:, 1]\n",
    "  y2 = dets[:, 2]\n",
    "  x2 = dets[:, 3]\n",
    "\n",
    "  areas = (x2 - x1) * (y2 - y1)\n",
    "  order = scores.argsort()[::-1]\n",
    "\n",
    "  keep = []\n",
    "  while order.size > 0 and len(keep) < max_dets:\n",
    "    i = order[0]\n",
    "    keep.append(i)\n",
    "\n",
    "    xx1 = np.maximum(x1[i], x1[order[1:]])\n",
    "    yy1 = np.maximum(y1[i], y1[order[1:]])\n",
    "    xx2 = np.minimum(x2[i], x2[order[1:]])\n",
    "    yy2 = np.minimum(y2[i], y2[order[1:]])\n",
    "\n",
    "    w = np.maximum(0.0, xx2 - xx1)\n",
    "    h = np.maximum(0.0, yy2 - yy1)\n",
    "    intersection = w * h\n",
    "    overlap = intersection / (areas[i] + areas[order[1:]] - intersection + 1e-12)\n",
    "\n",
    "    inds = np.where(overlap <= thresh)[0]\n",
    "    order = order[inds + 1]\n",
    "  return keep"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. ViLD Result Visualization Cell\n",
    "- **Purpose**: Visualization tools for detection results\n",
    "- **Key Functions**:\n",
    " - `draw_bounding_box_on_image()`\n",
    " - `draw_mask_on_image_array()`\n",
    " - `visualize_boxes_and_labels_on_image_array()`\n",
    "- **Features**:\n",
    " - Bounding box drawing\n",
    " - Instance mask visualization \n",
    " - Label and score display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "JO63v4Xuk1N6"
   },
   "outputs": [],
   "source": [
    "#@markdown ViLD Result Visualization\n",
    "import PIL.ImageColor as ImageColor\n",
    "import PIL.ImageDraw as ImageDraw\n",
    "import PIL.ImageFont as ImageFont\n",
    "\n",
    "STANDARD_COLORS = [\"White\"]\n",
    "# STANDARD_COLORS = [\n",
    "#     \"AliceBlue\", \"Chartreuse\", \"Aqua\", \"Aquamarine\", \"Azure\", \"Beige\", \"Bisque\",\n",
    "#     \"BlanchedAlmond\", \"BlueViolet\", \"BurlyWood\", \"CadetBlue\", \"AntiqueWhite\",\n",
    "#     \"Chocolate\", \"Coral\", \"CornflowerBlue\", \"Cornsilk\", \"Cyan\",\n",
    "#     \"DarkCyan\", \"DarkGoldenRod\", \"DarkGrey\", \"DarkKhaki\", \"DarkOrange\",\n",
    "#     \"DarkOrchid\", \"DarkSalmon\", \"DarkSeaGreen\", \"DarkTurquoise\", \"DarkViolet\",\n",
    "#     \"DeepPink\", \"DeepSkyBlue\", \"DodgerBlue\", \"FloralWhite\",\n",
    "#     \"ForestGreen\", \"Fuchsia\", \"Gainsboro\", \"GhostWhite\", \"Gold\", \"GoldenRod\",\n",
    "#     \"Salmon\", \"Tan\", \"HoneyDew\", \"HotPink\", \"Ivory\", \"Khaki\",\n",
    "#     \"Lavender\", \"LavenderBlush\", \"LawnGreen\", \"LemonChiffon\", \"LightBlue\",\n",
    "#     \"LightCoral\", \"LightCyan\", \"LightGoldenRodYellow\", \"LightGray\", \"LightGrey\",\n",
    "#     \"LightGreen\", \"LightPink\", \"LightSalmon\", \"LightSeaGreen\", \"LightSkyBlue\",\n",
    "#     \"LightSlateGray\", \"LightSlateGrey\", \"LightSteelBlue\", \"LightYellow\", \"Lime\",\n",
    "#     \"LimeGreen\", \"Linen\", \"Magenta\", \"MediumAquaMarine\", \"MediumOrchid\",\n",
    "#     \"MediumPurple\", \"MediumSeaGreen\", \"MediumSlateBlue\", \"MediumSpringGreen\",\n",
    "#     \"MediumTurquoise\", \"MediumVioletRed\", \"MintCream\", \"MistyRose\", \"Moccasin\",\n",
    "#     \"NavajoWhite\", \"OldLace\", \"Olive\", \"OliveDrab\", \"Orange\",\n",
    "#     \"Orchid\", \"PaleGoldenRod\", \"PaleGreen\", \"PaleTurquoise\", \"PaleVioletRed\",\n",
    "#     \"PapayaWhip\", \"PeachPuff\", \"Peru\", \"Pink\", \"Plum\", \"PowderBlue\", \"Purple\",\n",
    "#     \"RosyBrown\", \"RoyalBlue\", \"SaddleBrown\", \"Green\", \"SandyBrown\",\n",
    "#     \"SeaGreen\", \"SeaShell\", \"Sienna\", \"Silver\", \"SkyBlue\", \"SlateBlue\",\n",
    "#     \"SlateGray\", \"SlateGrey\", \"Snow\", \"SpringGreen\", \"SteelBlue\", \"GreenYellow\",\n",
    "#     \"Teal\", \"Thistle\", \"Tomato\", \"Turquoise\", \"Violet\", \"Wheat\", \"White\",\n",
    "#     \"WhiteSmoke\", \"Yellow\", \"YellowGreen\"\n",
    "# ]\n",
    "\n",
    "def draw_bounding_box_on_image(image,\n",
    "                               ymin,\n",
    "                               xmin,\n",
    "                               ymax,\n",
    "                               xmax,\n",
    "                               color=\"red\",\n",
    "                               thickness=4,\n",
    "                               display_str_list=(),\n",
    "                               use_normalized_coordinates=True):\n",
    "  \"\"\"Adds a bounding box to an image.\n",
    "\n",
    "  Bounding box coordinates can be specified in either absolute (pixel) or\n",
    "  normalized coordinates by setting the use_normalized_coordinates argument.\n",
    "\n",
    "  Each string in display_str_list is displayed on a separate line above the\n",
    "  bounding box in black text on a rectangle filled with the input \"color\".\n",
    "  If the top of the bounding box extends to the edge of the image, the strings\n",
    "  are displayed below the bounding box.\n",
    "\n",
    "  Args:\n",
    "    image: a PIL.Image object.\n",
    "    ymin: ymin of bounding box.\n",
    "    xmin: xmin of bounding box.\n",
    "    ymax: ymax of bounding box.\n",
    "    xmax: xmax of bounding box.\n",
    "    color: color to draw bounding box. Default is red.\n",
    "    thickness: line thickness. Default value is 4.\n",
    "    display_str_list: list of strings to display in box\n",
    "                      (each to be shown on its own line).\n",
    "    use_normalized_coordinates: If True (default), treat coordinates\n",
    "      ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat\n",
    "      coordinates as absolute.\n",
    "  \"\"\"\n",
    "  draw = ImageDraw.Draw(image)\n",
    "  im_width, im_height = image.size\n",
    "  if use_normalized_coordinates:\n",
    "    (left, right, top, bottom) = (xmin * im_width, xmax * im_width,\n",
    "                                  ymin * im_height, ymax * im_height)\n",
    "  else:\n",
    "    (left, right, top, bottom) = (xmin, xmax, ymin, ymax)\n",
    "  draw.line([(left, top), (left, bottom), (right, bottom),\n",
    "             (right, top), (left, top)], width=thickness, fill=color)\n",
    "  try:\n",
    "    font = ImageFont.truetype(\"arial.ttf\", 24)\n",
    "  except IOError:\n",
    "    font = ImageFont.load_default()\n",
    "\n",
    "  # If the total height of the display strings added to the top of the bounding\n",
    "  # box exceeds the top of the image, stack the strings below the bounding box\n",
    "  # instead of above.\n",
    "  display_str_heights = [font.getsize(ds)[1] for ds in display_str_list]\n",
    "  # Each display_str has a top and bottom margin of 0.05x.\n",
    "  total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights)\n",
    "\n",
    "  if top > total_display_str_height:\n",
    "    text_bottom = top\n",
    "  else:\n",
    "    text_bottom = bottom + total_display_str_height\n",
    "  # Reverse list and print from bottom to top.\n",
    "  for display_str in display_str_list[::-1]:\n",
    "    text_left = min(5, left)\n",
    "    text_width, text_height = font.getsize(display_str)\n",
    "    margin = np.ceil(0.05 * text_height)\n",
    "    draw.rectangle(\n",
    "        [(left, text_bottom - text_height - 2 * margin), (left + text_width,\n",
    "                                                          text_bottom)],\n",
    "        fill=color)\n",
    "    draw.text(\n",
    "        (left + margin, text_bottom - text_height - margin),\n",
    "        display_str,\n",
    "        fill=\"black\",\n",
    "        font=font)\n",
    "    text_bottom -= text_height - 2 * margin\n",
    "\n",
    "def draw_bounding_box_on_image_array(image,\n",
    "                                     ymin,\n",
    "                                     xmin,\n",
    "                                     ymax,\n",
    "                                     xmax,\n",
    "                                     color=\"red\",\n",
    "                                     thickness=4,\n",
    "                                     display_str_list=(),\n",
    "                                     use_normalized_coordinates=True):\n",
    "  \"\"\"Adds a bounding box to an image (numpy array).\n",
    "\n",
    "  Bounding box coordinates can be specified in either absolute (pixel) or\n",
    "  normalized coordinates by setting the use_normalized_coordinates argument.\n",
    "\n",
    "  Args:\n",
    "    image: a numpy array with shape [height, width, 3].\n",
    "    ymin: ymin of bounding box.\n",
    "    xmin: xmin of bounding box.\n",
    "    ymax: ymax of bounding box.\n",
    "    xmax: xmax of bounding box.\n",
    "    color: color to draw bounding box. Default is red.\n",
    "    thickness: line thickness. Default value is 4.\n",
    "    display_str_list: list of strings to display in box\n",
    "                      (each to be shown on its own line).\n",
    "    use_normalized_coordinates: If True (default), treat coordinates\n",
    "      ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat\n",
    "      coordinates as absolute.\n",
    "  \"\"\"\n",
    "  image_pil = Image.fromarray(np.uint8(image)).convert(\"RGB\")\n",
    "  draw_bounding_box_on_image(image_pil, ymin, xmin, ymax, xmax, color,\n",
    "                             thickness, display_str_list,\n",
    "                             use_normalized_coordinates)\n",
    "  np.copyto(image, np.array(image_pil))\n",
    "\n",
    "\n",
    "def draw_mask_on_image_array(image, mask, color=\"red\", alpha=0.4):\n",
    "  \"\"\"Draws mask on an image.\n",
    "\n",
    "  Args:\n",
    "    image: uint8 numpy array with shape (img_height, img_height, 3)\n",
    "    mask: a uint8 numpy array of shape (img_height, img_height) with\n",
    "      values between either 0 or 1.\n",
    "    color: color to draw the keypoints with. Default is red.\n",
    "    alpha: transparency value between 0 and 1. (default: 0.4)\n",
    "\n",
    "  Raises:\n",
    "    ValueError: On incorrect data type for image or masks.\n",
    "  \"\"\"\n",
    "  if image.dtype != np.uint8:\n",
    "    raise ValueError(\"`image` not of type np.uint8\")\n",
    "  if mask.dtype != np.uint8:\n",
    "    raise ValueError(\"`mask` not of type np.uint8\")\n",
    "  if np.any(np.logical_and(mask != 1, mask != 0)):\n",
    "    raise ValueError(\"`mask` elements should be in [0, 1]\")\n",
    "  if image.shape[:2] != mask.shape:\n",
    "    raise ValueError(\"The image has spatial dimensions %s but the mask has \"\n",
    "                     \"dimensions %s\" % (image.shape[:2], mask.shape))\n",
    "  rgb = ImageColor.getrgb(color)\n",
    "  pil_image = Image.fromarray(image)\n",
    "\n",
    "  solid_color = np.expand_dims(\n",
    "      np.ones_like(mask), axis=2) * np.reshape(list(rgb), [1, 1, 3])\n",
    "  pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert(\"RGBA\")\n",
    "  pil_mask = Image.fromarray(np.uint8(255.0*alpha*mask)).convert(\"L\")\n",
    "  pil_image = Image.composite(pil_solid_color, pil_image, pil_mask)\n",
    "  np.copyto(image, np.array(pil_image.convert(\"RGB\")))\n",
    "\n",
    "def visualize_boxes_and_labels_on_image_array(\n",
    "    image,\n",
    "    boxes,\n",
    "    classes,\n",
    "    scores,\n",
    "    category_index,\n",
    "    instance_masks=None,\n",
    "    instance_boundaries=None,\n",
    "    use_normalized_coordinates=False,\n",
    "    max_boxes_to_draw=20,\n",
    "    min_score_thresh=.5,\n",
    "    agnostic_mode=False,\n",
    "    line_thickness=1,\n",
    "    groundtruth_box_visualization_color=\"black\",\n",
    "    skip_scores=False,\n",
    "    skip_labels=False,\n",
    "    mask_alpha=0.4,\n",
    "    plot_color=None,\n",
    "):\n",
    "  \"\"\"Overlay labeled boxes on an image with formatted scores and label names.\n",
    "\n",
    "  This function groups boxes that correspond to the same location\n",
    "  and creates a display string for each detection and overlays these\n",
    "  on the image. Note that this function modifies the image in place, and returns\n",
    "  that same image.\n",
    "\n",
    "  Args:\n",
    "    image: uint8 numpy array with shape (img_height, img_width, 3)\n",
    "    boxes: a numpy array of shape [N, 4]\n",
    "    classes: a numpy array of shape [N]. Note that class indices are 1-based,\n",
    "      and match the keys in the label map.\n",
    "    scores: a numpy array of shape [N] or None.  If scores=None, then\n",
    "      this function assumes that the boxes to be plotted are groundtruth\n",
    "      boxes and plot all boxes as black with no classes or scores.\n",
    "    category_index: a dict containing category dictionaries (each holding\n",
    "      category index `id` and category name `name`) keyed by category indices.\n",
    "    instance_masks: a numpy array of shape [N, image_height, image_width] with\n",
    "      values ranging between 0 and 1, can be None.\n",
    "    instance_boundaries: a numpy array of shape [N, image_height, image_width]\n",
    "      with values ranging between 0 and 1, can be None.\n",
    "    use_normalized_coordinates: whether boxes is to be interpreted as\n",
    "      normalized coordinates or not.\n",
    "    max_boxes_to_draw: maximum number of boxes to visualize.  If None, draw\n",
    "      all boxes.\n",
    "    min_score_thresh: minimum score threshold for a box to be visualized\n",
    "    agnostic_mode: boolean (default: False) controlling whether to evaluate in\n",
    "      class-agnostic mode or not.  This mode will display scores but ignore\n",
    "      classes.\n",
    "    line_thickness: integer (default: 4) controlling line width of the boxes.\n",
    "    groundtruth_box_visualization_color: box color for visualizing groundtruth\n",
    "      boxes\n",
    "    skip_scores: whether to skip score when drawing a single detection\n",
    "    skip_labels: whether to skip label when drawing a single detection\n",
    "\n",
    "  Returns:\n",
    "    uint8 numpy array with shape (img_height, img_width, 3) with overlaid boxes.\n",
    "  \"\"\"\n",
    "  # Create a display string (and color) for every box location, group any boxes\n",
    "  # that correspond to the same location.\n",
    "  box_to_display_str_map = collections.defaultdict(list)\n",
    "  box_to_color_map = collections.defaultdict(str)\n",
    "  box_to_instance_masks_map = {}\n",
    "  box_to_score_map = {}\n",
    "  box_to_instance_boundaries_map = {}\n",
    "  \n",
    "  if not max_boxes_to_draw:\n",
    "    max_boxes_to_draw = boxes.shape[0]\n",
    "  for i in range(min(max_boxes_to_draw, boxes.shape[0])):\n",
    "    if scores is None or scores[i] > min_score_thresh:\n",
    "      box = tuple(boxes[i].tolist())\n",
    "      if instance_masks is not None:\n",
    "        box_to_instance_masks_map[box] = instance_masks[i]\n",
    "      if instance_boundaries is not None:\n",
    "        box_to_instance_boundaries_map[box] = instance_boundaries[i]\n",
    "      if scores is None:\n",
    "        box_to_color_map[box] = groundtruth_box_visualization_color\n",
    "      else:\n",
    "        display_str = \"\"\n",
    "        if not skip_labels:\n",
    "          if not agnostic_mode:\n",
    "            if classes[i] in list(category_index.keys()):\n",
    "              class_name = category_index[classes[i]][\"name\"]\n",
    "            else:\n",
    "              class_name = \"N/A\"\n",
    "            display_str = str(class_name)\n",
    "        if not skip_scores:\n",
    "          if not display_str:\n",
    "            display_str = \"{}%\".format(int(100*scores[i]))\n",
    "          else:\n",
    "            float_score = (\"%.2f\" % scores[i]).lstrip(\"0\")\n",
    "            display_str = \"{}: {}\".format(display_str, float_score)\n",
    "          box_to_score_map[box] = int(100*scores[i])\n",
    "\n",
    "        box_to_display_str_map[box].append(display_str)\n",
    "        if plot_color is not None:\n",
    "          box_to_color_map[box] = plot_color\n",
    "        elif agnostic_mode:\n",
    "          box_to_color_map[box] = \"DarkOrange\"\n",
    "        else:\n",
    "          box_to_color_map[box] = STANDARD_COLORS[\n",
    "              classes[i] % len(STANDARD_COLORS)]\n",
    "\n",
    "  # Handle the case when box_to_score_map is empty.\n",
    "  if box_to_score_map:\n",
    "    box_color_iter = sorted(\n",
    "        box_to_color_map.items(), key=lambda kv: box_to_score_map[kv[0]])\n",
    "  else:\n",
    "    box_color_iter = box_to_color_map.items()\n",
    "\n",
    "  # Draw all boxes onto image.\n",
    "  for box, color in box_color_iter:\n",
    "    ymin, xmin, ymax, xmax = box\n",
    "    if instance_masks is not None:\n",
    "      draw_mask_on_image_array(\n",
    "          image,\n",
    "          box_to_instance_masks_map[box],\n",
    "          color=color,\n",
    "          alpha=mask_alpha\n",
    "      )\n",
    "    if instance_boundaries is not None:\n",
    "      draw_mask_on_image_array(\n",
    "          image,\n",
    "          box_to_instance_boundaries_map[box],\n",
    "          color=\"red\",\n",
    "          alpha=1.0\n",
    "      )\n",
    "    draw_bounding_box_on_image_array(\n",
    "        image,\n",
    "        ymin,\n",
    "        xmin,\n",
    "        ymax,\n",
    "        xmax,\n",
    "        color=color,\n",
    "        thickness=line_thickness,\n",
    "        display_str_list=box_to_display_str_map[box],\n",
    "        use_normalized_coordinates=use_normalized_coordinates)\n",
    "    \n",
    "  return image\n",
    "\n",
    "\n",
    "def paste_instance_masks(masks,\n",
    "                         detected_boxes,\n",
    "                         image_height,\n",
    "                         image_width):\n",
    "  \"\"\"Paste instance masks to generate the image segmentation results.\n",
    "\n",
    "  Args:\n",
    "    masks: a numpy array of shape [N, mask_height, mask_width] representing the\n",
    "      instance masks w.r.t. the `detected_boxes`.\n",
    "    detected_boxes: a numpy array of shape [N, 4] representing the reference\n",
    "      bounding boxes.\n",
    "    image_height: an integer representing the height of the image.\n",
    "    image_width: an integer representing the width of the image.\n",
    "\n",
    "  Returns:\n",
    "    segms: a numpy array of shape [N, image_height, image_width] representing\n",
    "      the instance masks *pasted* on the image canvas.\n",
    "  \"\"\"\n",
    "\n",
    "  def expand_boxes(boxes, scale):\n",
    "    \"\"\"Expands an array of boxes by a given scale.\"\"\"\n",
    "    # Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/utils/boxes.py#L227  # pylint: disable=line-too-long\n",
    "    # The `boxes` in the reference implementation is in [x1, y1, x2, y2] form,\n",
    "    # whereas `boxes` here is in [x1, y1, w, h] form\n",
    "    w_half = boxes[:, 2] * .5\n",
    "    h_half = boxes[:, 3] * .5\n",
    "    x_c = boxes[:, 0] + w_half\n",
    "    y_c = boxes[:, 1] + h_half\n",
    "\n",
    "    w_half *= scale\n",
    "    h_half *= scale\n",
    "\n",
    "    boxes_exp = np.zeros(boxes.shape)\n",
    "    boxes_exp[:, 0] = x_c - w_half\n",
    "    boxes_exp[:, 2] = x_c + w_half\n",
    "    boxes_exp[:, 1] = y_c - h_half\n",
    "    boxes_exp[:, 3] = y_c + h_half\n",
    "\n",
    "    return boxes_exp\n",
    "\n",
    "  # Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/core/test.py#L812  # pylint: disable=line-too-long\n",
    "  # To work around an issue with cv2.resize (it seems to automatically pad\n",
    "  # with repeated border values), we manually zero-pad the masks by 1 pixel\n",
    "  # prior to resizing back to the original image resolution. This prevents\n",
    "  # \"top hat\" artifacts. We therefore need to expand the reference boxes by an\n",
    "  # appropriate factor.\n",
    "  _, mask_height, mask_width = masks.shape\n",
    "  scale = max((mask_width + 2.0) / mask_width,\n",
    "              (mask_height + 2.0) / mask_height)\n",
    "\n",
    "  ref_boxes = expand_boxes(detected_boxes, scale)\n",
    "  ref_boxes = ref_boxes.astype(np.int32)\n",
    "  padded_mask = np.zeros((mask_height + 2, mask_width + 2), dtype=np.float32)\n",
    "  segms = []\n",
    "  for mask_ind, mask in enumerate(masks):\n",
    "    im_mask = np.zeros((image_height, image_width), dtype=np.uint8)\n",
    "    # Process mask inside bounding boxes.\n",
    "    padded_mask[1:-1, 1:-1] = mask[:, :]\n",
    "\n",
    "    ref_box = ref_boxes[mask_ind, :]\n",
    "    w = ref_box[2] - ref_box[0] + 1\n",
    "    h = ref_box[3] - ref_box[1] + 1\n",
    "    w = np.maximum(w, 1)\n",
    "    h = np.maximum(h, 1)\n",
    "\n",
    "    mask = cv2.resize(padded_mask, (w, h))\n",
    "    mask = np.array(mask > 0.5, dtype=np.uint8)\n",
    "\n",
    "    x_0 = min(max(ref_box[0], 0), image_width)\n",
    "    x_1 = min(max(ref_box[2] + 1, 0), image_width)\n",
    "    y_0 = min(max(ref_box[1], 0), image_height)\n",
    "    y_1 = min(max(ref_box[3] + 1, 0), image_height)\n",
    "\n",
    "    im_mask[y_0:y_1, x_0:x_1] = mask[\n",
    "        (y_0 - ref_box[1]):(y_1 - ref_box[1]),\n",
    "        (x_0 - ref_box[0]):(x_1 - ref_box[0])\n",
    "    ]\n",
    "    segms.append(im_mask)\n",
    "\n",
    "  segms = np.array(segms)\n",
    "  assert masks.shape[0] == segms.shape[0]\n",
    "  return segms  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "kKKH_wTak3F_"
   },
   "outputs": [],
   "source": [
    "#@markdown Plot instance masks.\n",
    "def plot_mask(color, alpha, original_image, mask):\n",
    "  rgb = ImageColor.getrgb(color)\n",
    "  pil_image = Image.fromarray(original_image)\n",
    "\n",
    "  solid_color = np.expand_dims(\n",
    "      np.ones_like(mask), axis=2) * np.reshape(list(rgb), [1, 1, 3])\n",
    "  pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert(\"RGBA\")\n",
    "  pil_mask = Image.fromarray(np.uint8(255.0*alpha*mask)).convert(\"L\")\n",
    "  pil_image = Image.composite(pil_solid_color, pil_image, pil_mask)\n",
    "  img_w_mask = np.array(pil_image.convert(\"RGB\"))\n",
    "  return img_w_mask\n",
    "\n",
    "%matplotlib inline\n",
    "def display_image(path_or_array, size=(10, 10)):\n",
    "  if isinstance(path_or_array, str):\n",
    "    image = np.asarray(Image.open(open(image_path, \"rb\")).convert(\"RGB\"))\n",
    "  else:\n",
    "    image = path_or_array\n",
    "  \n",
    "  plt.figure(figsize=size)\n",
    "  plt.imshow(image)\n",
    "  plt.axis(\"off\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. ViLD Forward Pass Cell\n",
    "- **Purpose**: Main model execution logic\n",
    "- **Key Steps**:\n",
    " 1. Preprocess categories and parameters\n",
    " 2. Run object detection\n",
    " 3. Filter and process boxes\n",
    " 4. Compute text embeddings\n",
    " 5. Match detections with categories\n",
    " 6. Return found objects and visualizations\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "X2-SsyEolAr8"
   },
   "outputs": [],
   "source": [
    "#@markdown Define ViLD forward pass.\n",
    "\n",
    "def vild(image_path, category_name_string, params, plot_on=True, prompt_swaps=[]):\n",
    "  #################################################################\n",
    "  # Preprocessing categories and get params\n",
    "  for a, b in prompt_swaps:\n",
    "    category_name_string = category_name_string.replace(a, b)\n",
    "  category_names = [x.strip() for x in category_name_string.split(\";\")]\n",
    "  category_names = [\"background\"] + category_names\n",
    "  categories = [{\"name\": item, \"id\": idx+1,} for idx, item in enumerate(category_names)]\n",
    "  category_indices = {cat[\"id\"]: cat for cat in categories}\n",
    "  \n",
    "  max_boxes_to_draw, nms_threshold, min_rpn_score_thresh, min_box_area, max_box_area = params\n",
    "  fig_size_h = min(max(5, int(len(category_names) / 2.5) ), 10)\n",
    "\n",
    "\n",
    "  #################################################################\n",
    "  # Obtain results and read image\n",
    "  roi_boxes, roi_scores, detection_boxes, scores_unused, box_outputs, detection_masks, visual_features, image_info = session.run(\n",
    "        [\"RoiBoxes:0\", \"RoiScores:0\", \"2ndStageBoxes:0\", \"2ndStageScoresUnused:0\", \"BoxOutputs:0\", \"MaskOutputs:0\", \"VisualFeatOutputs:0\", \"ImageInfo:0\"],\n",
    "        feed_dict={\"Placeholder:0\": [image_path,]})\n",
    "  \n",
    "  roi_boxes = np.squeeze(roi_boxes, axis=0)  # squeeze\n",
    "  # no need to clip the boxes, already done\n",
    "  roi_scores = np.squeeze(roi_scores, axis=0)\n",
    "\n",
    "  detection_boxes = np.squeeze(detection_boxes, axis=(0, 2))\n",
    "  scores_unused = np.squeeze(scores_unused, axis=0)\n",
    "  box_outputs = np.squeeze(box_outputs, axis=0)\n",
    "  detection_masks = np.squeeze(detection_masks, axis=0)\n",
    "  visual_features = np.squeeze(visual_features, axis=0)\n",
    "\n",
    "  image_info = np.squeeze(image_info, axis=0)  # obtain image info\n",
    "  image_scale = np.tile(image_info[2:3, :], (1, 2))\n",
    "  image_height = int(image_info[0, 0])\n",
    "  image_width = int(image_info[0, 1])\n",
    "\n",
    "  rescaled_detection_boxes = detection_boxes / image_scale # rescale\n",
    "\n",
    "  # Read image\n",
    "  image = np.asarray(Image.open(open(image_path, \"rb\")).convert(\"RGB\"))\n",
    "  assert image_height == image.shape[0]\n",
    "  assert image_width == image.shape[1]\n",
    "\n",
    "\n",
    "  #################################################################\n",
    "  # Filter boxes\n",
    "\n",
    "  # Apply non-maximum suppression to detected boxes with nms threshold.\n",
    "  nmsed_indices = nms(\n",
    "      detection_boxes,\n",
    "      roi_scores,\n",
    "      thresh=nms_threshold\n",
    "      )\n",
    "\n",
    "  # Compute RPN box size.\n",
    "  box_sizes = (rescaled_detection_boxes[:, 2] - rescaled_detection_boxes[:, 0]) * (rescaled_detection_boxes[:, 3] - rescaled_detection_boxes[:, 1])\n",
    "\n",
    "  # Filter out invalid rois (nmsed rois)\n",
    "  valid_indices = np.where(\n",
    "      np.logical_and(\n",
    "        np.isin(np.arange(len(roi_scores), dtype=np.int32), nmsed_indices),\n",
    "        np.logical_and(\n",
    "            np.logical_not(np.all(roi_boxes == 0., axis=-1)),\n",
    "            np.logical_and(\n",
    "              roi_scores >= min_rpn_score_thresh,\n",
    "              np.logical_and(\n",
    "                box_sizes > min_box_area,\n",
    "                box_sizes < max_box_area\n",
    "                )\n",
    "              )\n",
    "        )    \n",
    "      )\n",
    "  )[0]\n",
    "\n",
    "  detection_roi_scores = roi_scores[valid_indices][:max_boxes_to_draw, ...]\n",
    "  detection_boxes = detection_boxes[valid_indices][:max_boxes_to_draw, ...]\n",
    "  detection_masks = detection_masks[valid_indices][:max_boxes_to_draw, ...]\n",
    "  detection_visual_feat = visual_features[valid_indices][:max_boxes_to_draw, ...]\n",
    "  rescaled_detection_boxes = rescaled_detection_boxes[valid_indices][:max_boxes_to_draw, ...]\n",
    "\n",
    "\n",
    "  #################################################################\n",
    "  # Compute text embeddings and detection scores, and rank results\n",
    "  text_features = build_text_embedding(categories)\n",
    "  \n",
    "  raw_scores = detection_visual_feat.dot(text_features.T)\n",
    "  if FLAGS.use_softmax:\n",
    "    scores_all = softmax(FLAGS.temperature * raw_scores, axis=-1)\n",
    "  else:\n",
    "    scores_all = raw_scores\n",
    "\n",
    "  indices = np.argsort(-np.max(scores_all, axis=1))  # Results are ranked by scores\n",
    "  indices_fg = np.array([i for i in indices if np.argmax(scores_all[i]) != 0])\n",
    "\n",
    "  \n",
    "  #################################################################\n",
    "  # Print found_objects\n",
    "  found_objects = []\n",
    "  for a, b in prompt_swaps:\n",
    "    category_names = [name.replace(b, a) for name in category_names]  # Extra prompt engineering.\n",
    "  for anno_idx in indices[0:int(rescaled_detection_boxes.shape[0])]:\n",
    "    scores = scores_all[anno_idx]\n",
    "    if np.argmax(scores) == 0:\n",
    "      continue\n",
    "    found_object = category_names[np.argmax(scores)]\n",
    "    if found_object == \"background\":\n",
    "      continue\n",
    "    print(\"Found a\", found_object, \"with score:\", np.max(scores))\n",
    "    found_objects.append(category_names[np.argmax(scores)])\n",
    "  if not plot_on:\n",
    "    return found_objects\n",
    "  \n",
    "\n",
    "  #################################################################\n",
    "  # Plot detected boxes on the input image.\n",
    "  ymin, xmin, ymax, xmax = np.split(rescaled_detection_boxes, 4, axis=-1)\n",
    "  processed_boxes = np.concatenate([xmin, ymin, xmax - xmin, ymax - ymin], axis=-1)\n",
    "  segmentations = paste_instance_masks(detection_masks, processed_boxes, image_height, image_width)\n",
    "\n",
    "  if len(indices_fg) == 0:\n",
    "    display_image(np.array(image), size=overall_fig_size)\n",
    "    print(\"ViLD does not detect anything belong to the given category\")\n",
    "\n",
    "  else:\n",
    "    image_with_detections = visualize_boxes_and_labels_on_image_array(\n",
    "        np.array(image),\n",
    "        rescaled_detection_boxes[indices_fg],\n",
    "        valid_indices[:max_boxes_to_draw][indices_fg],\n",
    "        detection_roi_scores[indices_fg],    \n",
    "        numbered_category_indices,\n",
    "        instance_masks=segmentations[indices_fg],\n",
    "        use_normalized_coordinates=False,\n",
    "        max_boxes_to_draw=max_boxes_to_draw,\n",
    "        min_score_thresh=min_rpn_score_thresh,\n",
    "        skip_scores=False,\n",
    "        skip_labels=True)\n",
    "\n",
    "    # plt.figure(figsize=overall_fig_size)\n",
    "    plt.imshow(image_with_detections)\n",
    "    # plt.axis(\"off\")\n",
    "    plt.title(\"ViLD detected objects and RPN scores.\")\n",
    "    plt.show()\n",
    "\n",
    "  return found_objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8jwAERXZl2Yl"
   },
   "outputs": [],
   "source": [
    "category_names = ['blue block',\n",
    "                  'red block',\n",
    "                  'green block',\n",
    "                  'orange block',\n",
    "                  'yellow block',\n",
    "                  'purple block',\n",
    "                  'pink block',\n",
    "                  'cyan block',\n",
    "                  'brown block',\n",
    "                  'gray block',\n",
    "\n",
    "                  'blue bowl',\n",
    "                  'red bowl',\n",
    "                  'green bowl',\n",
    "                  'orange bowl',\n",
    "                  'yellow bowl',\n",
    "                  'purple bowl',\n",
    "                  'pink bowl',\n",
    "                  'cyan bowl',\n",
    "                  'brown bowl',\n",
    "                  'gray bowl']\n",
    "image_path = 'tmp.jpg'\n",
    "\n",
    "#@markdown ViLD settings.\n",
    "category_name_string = \";\".join(category_names)\n",
    "max_boxes_to_draw = 8 #@param {type:\"integer\"}\n",
    "\n",
    "# Extra prompt engineering: swap A with B for every (A, B) in list.\n",
    "prompt_swaps = [('block', 'cube')]\n",
    "\n",
    "nms_threshold = 0.4 #@param {type:\"slider\", min:0, max:0.9, step:0.05}\n",
    "min_rpn_score_thresh = 0.4  #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
    "min_box_area = 10 #@param {type:\"slider\", min:0, max:10000, step:1.0}\n",
    "max_box_area = 3000  #@param {type:\"slider\", min:0, max:10000, step:1.0}\n",
    "vild_params = max_boxes_to_draw, nms_threshold, min_rpn_score_thresh, min_box_area, max_box_area\n",
    "found_objects = vild(image_path, category_name_string, vild_params, plot_on=True, prompt_swaps=prompt_swaps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yDM9CH_qXGKH"
   },
   "source": [
    "## **Scripted Expert**\n",
    "Scripted pick and place oracle to collect expert demonstrations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oOcKdshBXGKH"
   },
   "outputs": [],
   "source": [
    "class ScriptedPolicy():\n",
    "\n",
    "  def __init__(self, env):\n",
    "    self.env = env\n",
    "\n",
    "  def step(self, text, obs):\n",
    "    print(f'Input: {text}')\n",
    "\n",
    "    # Parse pick and place targets.\n",
    "    pick_text, place_text = text.split('and')\n",
    "    pick_target, place_target = None, None\n",
    "    for name in PICK_TARGETS.keys():\n",
    "      if name in pick_text:\n",
    "        pick_target = name\n",
    "        break\n",
    "    for name in PLACE_TARGETS.keys():\n",
    "      if name in place_text:\n",
    "        place_target = name\n",
    "        break\n",
    "\n",
    "    # Admissable targets only.\n",
    "    assert pick_target is not None\n",
    "    assert place_target is not None\n",
    "\n",
    "    pick_id = self.env.obj_name_to_id[pick_target]\n",
    "    pick_pose = pybullet.getBasePositionAndOrientation(pick_id)\n",
    "    pick_position = np.float32(pick_pose[0])\n",
    "\n",
    "    if place_target in self.env.obj_name_to_id:\n",
    "      place_id = self.env.obj_name_to_id[place_target]\n",
    "      place_pose = pybullet.getBasePositionAndOrientation(place_id)\n",
    "      place_position = np.float32(place_pose[0])\n",
    "    else:\n",
    "      place_position = np.float32(PLACE_TARGETS[place_target])\n",
    "\n",
    "    # Add some noise to pick and place positions.\n",
    "    # pick_position[:2] += np.random.normal(scale=0.01)\n",
    "    place_position[:2] += np.random.normal(scale=0.01)\n",
    "\n",
    "    act = {'pick': pick_position, 'place': place_position}\n",
    "    return act"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GCvl1WmvW3lK"
   },
   "source": [
    "## **Dataset**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lKMSa_ryW3lL"
   },
   "outputs": [],
   "source": [
    "#@markdown Collect demonstrations with a scripted expert, or download a pre-generated dataset.\n",
    "load_pregenerated = True  #@param {type:\"boolean\"}\n",
    "\n",
    "# Load pre-existing dataset.\n",
    "if load_pregenerated:\n",
    "  if not os.path.exists('dataset-9999.pkl'):\n",
    "    # !gdown --id 1TECwTIfawxkRYbzlAey0z1mqXKcyfPc-\n",
    "    !gdown --id 1yCz6C-6eLWb4SFYKdkM-wz5tlMjbG2h8\n",
    "  dataset = pickle.load(open('dataset-9999.pkl', 'rb'))  # ~10K samples.\n",
    "  dataset_size = len(dataset['text'])\n",
    "\n",
    "# Generate new dataset.\n",
    "else:\n",
    "  dataset = {}\n",
    "  dataset_size = 2  # Size of new dataset.\n",
    "  dataset['image'] = np.zeros((dataset_size, 224, 224, 3), dtype=np.uint8)\n",
    "  dataset['pick_yx'] = np.zeros((dataset_size, 2), dtype=np.int32)\n",
    "  dataset['place_yx'] = np.zeros((dataset_size, 2), dtype=np.int32)\n",
    "  dataset['text'] = []\n",
    "  policy = ScriptedPolicy(env)\n",
    "  data_idx = 0\n",
    "  while data_idx < dataset_size:\n",
    "    np.random.seed(data_idx)\n",
    "    num_pick, num_place = 3, 3\n",
    "\n",
    "    # Select random objects for data collection.\n",
    "    pick_items = list(PICK_TARGETS.keys())\n",
    "    pick_items = np.random.choice(pick_items, size=num_pick, replace=False)\n",
    "    place_items = list(PLACE_TARGETS.keys())\n",
    "    for pick_item in pick_items:  # For simplicity: place items != pick items.\n",
    "      place_items.remove(pick_item)\n",
    "    place_items = np.random.choice(place_items, size=num_place, replace=False)\n",
    "    config = {'pick': pick_items, 'place': place_items}\n",
    "\n",
    "    # Initialize environment with selected objects.\n",
    "    obs = env.reset(config)\n",
    "\n",
    "    # Create text prompts.\n",
    "    prompts = []\n",
    "    for i in range(len(pick_items)):\n",
    "      pick_item = pick_items[i]\n",
    "      place_item = place_items[i]\n",
    "      prompts.append(f'Pick the {pick_item} and place it on the {place_item}.')\n",
    "\n",
    "    # Execute 3 pick and place actions.\n",
    "    for prompt in prompts:\n",
    "      act = policy.step(prompt, obs)\n",
    "      dataset['text'].append(prompt)\n",
    "      dataset['image'][data_idx, ...] = obs['image'].copy()\n",
    "      dataset['pick_yx'][data_idx, ...] = xyz_to_pix(act['pick'])\n",
    "      dataset['place_yx'][data_idx, ...] = xyz_to_pix(act['place'])\n",
    "      data_idx += 1\n",
    "      obs, _, _, _ = env.step(act)\n",
    "      debug_clip = ImageSequenceClip(env.cache_video, fps=25)\n",
    "      display(debug_clip.ipython_display(autoplay=1, loop=1))\n",
    "      env.cache_video = []\n",
    "      if data_idx >= dataset_size:\n",
    "        break\n",
    "\n",
    "  pickle.dump(dataset, open(f'dataset-{dataset_size}.pkl', 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "-wVcfLWlYm73"
   },
   "outputs": [],
   "source": [
    "#@markdown Show a demonstration example from the dataset.\n",
    "\n",
    "img = dataset['image'][0]\n",
    "pick_yx = dataset['pick_yx'][0]\n",
    "place_yx = dataset['place_yx'][0]\n",
    "text = dataset['text'][0]\n",
    "plt.title(text)\n",
    "plt.imshow(img)\n",
    "plt.arrow(pick_yx[1], pick_yx[0], place_yx[1]-pick_yx[1], place_yx[0]-pick_yx[0], color='w', head_starts_at_zero=False, head_width=7, length_includes_head=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YoB7QKvOJaWp"
   },
   "source": [
    "## **Model**\n",
    "Define a variant of [CLIPort](https://cliport.github.io/) architecture ([Transporter Nets](https://transporternets.github.io/) + [CLIP](https://arxiv.org/pdf/2103.00020.pdf))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "irhuu2KAJaWt"
   },
   "outputs": [],
   "source": [
    "class ResNetBlock(nn.Module):\n",
    "  \"\"\"ResNet pre-Activation block. https://arxiv.org/pdf/1603.05027.pdf\"\"\"\n",
    "  features: int\n",
    "  stride: int = 1\n",
    "\n",
    "  def setup(self):\n",
    "    self.conv0 = nn.Conv(self.features // 4, (1, 1), (self.stride, self.stride))\n",
    "    self.conv1 = nn.Conv(self.features // 4, (3, 3))\n",
    "    self.conv2 = nn.Conv(self.features, (1, 1))\n",
    "    self.conv3 = nn.Conv(self.features, (1, 1), (self.stride, self.stride))\n",
    "\n",
    "  def __call__(self, x):\n",
    "    y = self.conv0(nn.relu(x))\n",
    "    y = self.conv1(nn.relu(y))\n",
    "    y = self.conv2(nn.relu(y))\n",
    "    if x.shape != y.shape:\n",
    "      x = self.conv3(nn.relu(x))\n",
    "    return x + y\n",
    "\n",
    "\n",
    "class UpSample(nn.Module):\n",
    "  \"\"\"Simple 2D 2x bilinear upsample.\"\"\"\n",
    "\n",
    "  def __call__(self, x):\n",
    "    B, H, W, C = x.shape\n",
    "    new_shape = (B, H * 2, W * 2, C)\n",
    "    return jax.image.resize(x, new_shape, 'bilinear')\n",
    "\n",
    "\n",
    "class ResNet(nn.Module):\n",
    "  \"\"\"Hourglass 53-layer ResNet with 8-stride.\"\"\"\n",
    "  out_dim: int\n",
    "\n",
    "  def setup(self):\n",
    "    self.dense0 = nn.Dense(8)\n",
    "\n",
    "    self.conv0 = nn.Conv(64, (3, 3), (1, 1))\n",
    "    self.block0 = ResNetBlock(64)\n",
    "    self.block1 = ResNetBlock(64)\n",
    "    self.block2 = ResNetBlock(128, stride=2)\n",
    "    self.block3 = ResNetBlock(128)\n",
    "    self.block4 = ResNetBlock(256, stride=2)\n",
    "    self.block5 = ResNetBlock(256)\n",
    "    self.block6 = ResNetBlock(512, stride=2)\n",
    "    self.block7 = ResNetBlock(512)\n",
    "\n",
    "    self.block8 = ResNetBlock(256)\n",
    "    self.block9 = ResNetBlock(256)\n",
    "    self.upsample0 = UpSample()\n",
    "    self.block10 = ResNetBlock(128)\n",
    "    self.block11 = ResNetBlock(128)\n",
    "    self.upsample1 = UpSample()\n",
    "    self.block12 = ResNetBlock(64)\n",
    "    self.block13 = ResNetBlock(64)\n",
    "    self.upsample2 = UpSample()\n",
    "    self.block14 = ResNetBlock(16)\n",
    "    self.block15 = ResNetBlock(16)\n",
    "    self.conv1 = nn.Conv(self.out_dim, (3, 3), (1, 1))\n",
    "\n",
    "  def __call__(self, x, text):\n",
    "\n",
    "    # # Project and concatenate CLIP features (early fusion).\n",
    "    # text = self.dense0(text)\n",
    "    # text = jnp.expand_dims(text, axis=(1, 2))\n",
    "    # text = jnp.broadcast_to(text, x.shape[:3] + (8,))\n",
    "    # x = jnp.concatenate((x, text), axis=-1)\n",
    "\n",
    "    x = self.conv0(x)\n",
    "    x = self.block0(x)\n",
    "    x = self.block1(x)\n",
    "    x = self.block2(x)\n",
    "    x = self.block3(x)\n",
    "    x = self.block4(x)\n",
    "    x = self.block5(x)\n",
    "    x = self.block6(x)\n",
    "    x = self.block7(x)\n",
    "\n",
    "    # Concatenate CLIP features (mid-fusion).\n",
    "    text = jnp.expand_dims(text, axis=(1, 2))\n",
    "    text = jnp.broadcast_to(text, x.shape)\n",
    "    x = jnp.concatenate((x, text), axis=-1)\n",
    "\n",
    "    x = self.block8(x)\n",
    "    x = self.block9(x)\n",
    "    x = self.upsample0(x)\n",
    "    x = self.block10(x)\n",
    "    x = self.block11(x)\n",
    "    x = self.upsample1(x)\n",
    "    x = self.block12(x)\n",
    "    x = self.block13(x)\n",
    "    x = self.upsample2(x)\n",
    "    x = self.block14(x)\n",
    "    x = self.block15(x)\n",
    "    x = self.conv1(x)\n",
    "    return x\n",
    "\n",
    "\n",
    "class TransporterNets(nn.Module):\n",
    "  \"\"\"TransporterNet with 3 ResNets (translation only).\"\"\"\n",
    "\n",
    "  def setup(self):\n",
    "    # Picking affordances.\n",
    "    self.pick_net = ResNet(1)\n",
    "\n",
    "    # Pick-conditioned placing affordances.\n",
    "    self.q_net = ResNet(3)  # Query (crop around pick location).\n",
    "    self.k_net = ResNet(3)  # Key (place features).\n",
    "    self.crop_size = 64\n",
    "    self.crop_conv = nn.Conv(features=1, kernel_size=(self.crop_size, self.crop_size), use_bias=False, dtype=jnp.float32, padding='SAME')\n",
    "\n",
    "  def __call__(self, x, text, p=None, train=True):\n",
    "    B, H, W, C = x.shape\n",
    "    pick_out = self.pick_net(x, text)  # (B, H, W, 1)\n",
    "\n",
    "    # Get key features.\n",
    "    k = self.k_net(x, text)\n",
    "\n",
    "    # Add 0-padding before cropping.\n",
    "    h = self.crop_size // 2\n",
    "    x_crop = jnp.pad(x, [(0, 0), (h, h), (h, h), (0, 0)], 'maximum')\n",
    "\n",
    "    # Get query features and convolve them over key features.\n",
    "    place_out = jnp.zeros((0, H, W, 1), jnp.float32)\n",
    "    for b in range(B):\n",
    "\n",
    "      # Get coordinates at center of crop.\n",
    "      if p is None:\n",
    "        pick_out_b = pick_out[b, ...]  # (H, W, 1)\n",
    "        pick_out_b = pick_out_b.flatten()  # (H * W,)\n",
    "        amax_i = jnp.argmax(pick_out_b)\n",
    "        v, u = jnp.unravel_index(amax_i, (H, W))\n",
    "      else:\n",
    "        v, u = p[b, :]\n",
    "\n",
    "      # Get query crop.\n",
    "      x_crop_b = jax.lax.dynamic_slice(x_crop, (b, v, u, 0), (1, self.crop_size, self.crop_size, x_crop.shape[3]))\n",
    "      # x_crop_b = x_crop[b:b+1, v:(v + self.crop_size), u:(u + self.crop_size), ...]\n",
    "\n",
    "      # Convolve q (query) across k (key).\n",
    "      q = self.q_net(x_crop_b, text[b:b+1, :])  # (1, H, W, 3)\n",
    "      q = jnp.transpose(q, (1, 2, 3, 0))  # (H, W, 3, 1)\n",
    "      place_out_b = self.crop_conv.apply({'params': {'kernel': q}}, k[b:b+1, ...])  # (1, H, W, 1)\n",
    "      scale = 1 / (self.crop_size * self.crop_size)  # For higher softmax temperatures.\n",
    "      place_out_b *= scale\n",
    "      place_out = jnp.concatenate((place_out, place_out_b), axis=0)\n",
    "\n",
    "    return pick_out, place_out\n",
    "\n",
    "\n",
    "def n_params(params):\n",
    "  return jnp.sum(jnp.int32([n_params(v) if isinstance(v, dict) or isinstance(v, flax.core.frozen_dict.FrozenDict) else np.prod(v.shape) for v in params.values()]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yzDKqAwFJaWt"
   },
   "source": [
    "## **Training**\n",
    "Train CLIP-conditioned Transporter Nets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CLIP Text Feature Generation\n",
    "Processes the dataset text through CLIP to generate embeddings used for training and inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aGCsZK5ZJaWu"
   },
   "outputs": [],
   "source": [
    "#@markdown Compute CLIP features for text in the dataset.\n",
    "\n",
    "# Precompute CLIP features for all text in training dataset.\n",
    "text_tokens = clip.tokenize(dataset['text']).to(device)  # Using consistent device\n",
    "text_i = 0\n",
    "data_text_feats = np.zeros((0, 512), dtype=np.float32)\n",
    "while text_i < len(text_tokens):\n",
    "  batch_size = min(len(text_tokens) - text_i, 512)\n",
    "  text_batch = text_tokens[text_i:text_i+batch_size]\n",
    "  with torch.no_grad():\n",
    "    batch_feats = clip_model.encode_text(text_batch).float()\n",
    "  batch_feats /= batch_feats.norm(dim=-1, keepdim=True)\n",
    "  batch_feats = np.float32(batch_feats.cpu())\n",
    "  data_text_feats = np.concatenate((data_text_feats, batch_feats), axis=0)\n",
    "  text_i += batch_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Functions\n",
    "Defines JAX-based training and evaluation functions for the CLIPort model:\n",
    "- InfoNCE loss implementation\n",
    "- Pick/place prediction training\n",
    "- Evaluation step function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lFpqAE86JaWu"
   },
   "outputs": [],
   "source": [
    "#@markdown Define Transporter Nets train and eval functions\n",
    "# Train with InfoNCE loss over pick and place positions.\n",
    "@jax.jit\n",
    "def train_step(optimizer, batch):\n",
    "    def loss_fn(params):\n",
    "        batch_size = batch['img'].shape[0]\n",
    "        pick_logits, place_logits = TransporterNets().apply({'params': params}, batch['img'], batch['text'], batch['pick_yx'])\n",
    "        # InfoNCE pick loss.\n",
    "        pick_logits = pick_logits.reshape(batch_size, -1)\n",
    "        pick_onehot = batch['pick_onehot'].reshape(batch_size, -1)\n",
    "        pick_loss = jnp.mean(optax.softmax_cross_entropy(logits=pick_logits, labels=pick_onehot), axis=0)\n",
    "        # InfoNCE place loss.\n",
    "        place_logits = place_logits.reshape(batch_size, -1)\n",
    "        place_onehot = batch['place_onehot'].reshape(batch_size, -1)\n",
    "        place_loss = jnp.mean(optax.softmax_cross_entropy(logits=place_logits, labels=place_onehot), axis=0)\n",
    "        \n",
    "        loss = pick_loss + place_loss\n",
    "        return loss, (pick_logits, place_logits)\n",
    "    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n",
    "    (loss, logits), grad = grad_fn(optimizer.target)\n",
    "    optimizer = optimizer.apply_gradient(grad)\n",
    "    return optimizer, loss, grad, logits\n",
    "\n",
    "@jax.jit\n",
    "def eval_step(params, batch):\n",
    "    # Convert params to the right type if needed\n",
    "    params = jax.tree_map(lambda x: jnp.asarray(x, dtype=jnp.float32), params)\n",
    "    # Ensure batch inputs are float32\n",
    "    img = jnp.asarray(batch['img'], dtype=jnp.float32)\n",
    "    text = jnp.asarray(batch['text'], dtype=jnp.float32)\n",
    "    # Create a new frozen dict with the right types\n",
    "    model_params = {'params': params}\n",
    "    pick_logits, place_logits = TransporterNets().apply(model_params, img, text)\n",
    "    return pick_logits, place_logits\n",
    "\n",
    "# Coordinate map remains the same\n",
    "coord_x, coord_y = np.meshgrid(np.linspace(-1, 1, 224), np.linspace(-1, 1, 224), sparse=False, indexing='ij')\n",
    "coords = np.concatenate((coord_x[..., None], coord_y[..., None]), axis=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## TensorBoard Setup (Optional)\n",
    "Optional TensorBoard configuration for training visualization and monitoring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4tHOKjGM0_wh"
   },
   "outputs": [],
   "source": [
    "# #@markdown **TensorBoard:** Displays an interactive TensorBoard.\n",
    "# name = datetime.datetime.now().strftime(f'%Y-%m-%d-%H:%M:%S-cliport')\n",
    "# logdir = os.path.join(\"logs\", name)\n",
    "# writer = SummaryWriter(logdir)\n",
    "# %tensorboard --logdir logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checkpoint Management\n",
    "Functions for saving and loading model checkpoints with explicit type handling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "def save_checkpoint(directory, target, step, prefix='checkpoint_', overwrite=False):\n",
    "    \"\"\"Save checkpoint with explicit type handling.\"\"\"\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "    \n",
    "    checkpoint_path = os.path.join(directory, f'{prefix}{step}.npz')\n",
    "    \n",
    "    if os.path.exists(checkpoint_path) and not overwrite:\n",
    "        return checkpoint_path\n",
    "    \n",
    "    try:\n",
    "        # Convert target to numpy arrays with explicit float32 type\n",
    "        arrays = {}\n",
    "        for k, v in target.items():\n",
    "            if isinstance(v, dict):\n",
    "                # Handle nested dictionaries\n",
    "                arrays[k] = {subk: np.array(subv, dtype=np.float32) \n",
    "                           for subk, subv in v.items()}\n",
    "            else:\n",
    "                arrays[k] = np.array(v, dtype=np.float32)\n",
    "        \n",
    "        np.savez(checkpoint_path, **arrays)\n",
    "        print(f\"Saved checkpoint to {checkpoint_path}\")\n",
    "        return checkpoint_path\n",
    "    except Exception as e:\n",
    "        print(f\"Failed to save checkpoint: {e}\")\n",
    "        return None\n",
    "\n",
    "def restore_checkpoint(checkpoint_path, target=None):\n",
    "    \"\"\"Restore checkpoint with explicit type handling.\"\"\"\n",
    "    if not os.path.exists(checkpoint_path):\n",
    "        return target\n",
    "        \n",
    "    try:\n",
    "        with np.load(checkpoint_path) as data:\n",
    "            restored = {}\n",
    "            for k in data.files:\n",
    "                # Ensure all arrays are float32\n",
    "                restored[k] = np.asarray(data[k], dtype=np.float32)\n",
    "        return restored\n",
    "    except Exception as e:\n",
    "        print(f\"Failed to load checkpoint: {e}\")\n",
    "        return target"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Training/Loading\n",
    "Option to either:\n",
    "- Train a new model from scratch\n",
    "- Load pretrained weights for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "UD0GtBiLJaWu"
   },
   "outputs": [],
   "source": [
    "#@markdown Train your own model, or load a pretrained one.\n",
    "load_pretrained = True  #@param {type:\"boolean\"}\n",
    "\n",
    "# Initialize model weights using dummy tensors\n",
    "rng = jax.random.PRNGKey(0)\n",
    "rng, key = jax.random.split(rng)\n",
    "init_img = jnp.ones((4, 224, 224, 5), jnp.float32)\n",
    "init_text = jnp.ones((4, 512), jnp.float32)\n",
    "init_pix = jnp.zeros((4, 2), np.int32)\n",
    "init_params = TransporterNets().init(key, init_img, init_text, init_pix)['params']\n",
    "print(f'Model parameters: {n_params(init_params):,}')\n",
    "optim = flax.optim.Adam(learning_rate=1e-4).create(init_params)\n",
    "\n",
    "if load_pretrained:\n",
    "    ckpt_path = 'logs/ckpt_40000.npz'\n",
    "    if os.path.exists(ckpt_path):\n",
    "        try:\n",
    "            # Load and explicitly convert to float32\n",
    "            restored_params = restore_checkpoint(ckpt_path)\n",
    "            if restored_params is not None:\n",
    "                # Convert to JAX arrays before freezing\n",
    "                jax_params = jax.tree_map(\n",
    "                    lambda x: jnp.array(x, dtype=jnp.float32),\n",
    "                    restored_params\n",
    "                )\n",
    "                # Update the optimizer with restored parameters\n",
    "                optim = optim.replace(target=flax.core.frozen_dict.freeze(jax_params))\n",
    "                print(f\"Successfully loaded weights from {ckpt_path}!\")\n",
    "            else:\n",
    "                print(\"Failed to restore checkpoint\")\n",
    "        except Exception as e:\n",
    "            print(f\"Failed to load weights: {e}\")\n",
    "            print(\"\\nStarting training from random initialization\")\n",
    "    else:\n",
    "        print(f\"\\nCould not find checkpoint at {ckpt_path}. Starting with random initialization.\")\n",
    "else:\n",
    "    print(\"Starting training from scratch...\")\n",
    "    # Training loop.\n",
    "    batch_size = 8\n",
    "    losses = []  # Track losses for plotting later\n",
    "    for train_iter in range(1, 40001):\n",
    "        batch_i = np.random.randint(dataset_size, size=batch_size)\n",
    "        text_feat = data_text_feats[batch_i, ...]\n",
    "        img = dataset['image'][batch_i, ...] / 255\n",
    "        img = np.concatenate((img, np.broadcast_to(coords[None, ...], (batch_size,) + coords.shape)), axis=3)\n",
    "\n",
    "        # Get onehot label maps.\n",
    "        pick_yx = np.zeros((batch_size, 2), dtype=np.int32)\n",
    "        pick_onehot = np.zeros((batch_size, 224, 224), dtype=np.float32)\n",
    "        place_onehot = np.zeros((batch_size, 224, 224), dtype=np.float32)\n",
    "        for i in range(len(batch_i)):\n",
    "            pick_y, pick_x = dataset['pick_yx'][batch_i[i], :]\n",
    "            place_y, place_x = dataset['place_yx'][batch_i[i], :]\n",
    "            pick_onehot[i, pick_y, pick_x] = 1\n",
    "            place_onehot[i, place_y, place_x] = 1\n",
    "            \n",
    "            # Data augmentation (random translation).\n",
    "            roll_y, roll_x = np.random.randint(-112, 112, size=2)\n",
    "            img[i, ...] = np.roll(img[i, ...], roll_y, axis=0)\n",
    "            img[i, ...] = np.roll(img[i, ...], roll_x, axis=1)\n",
    "            pick_onehot[i, ...] = np.roll(pick_onehot[i, ...], roll_y, axis=0)\n",
    "            pick_onehot[i, ...] = np.roll(pick_onehot[i, ...], roll_x, axis=1)\n",
    "            place_onehot[i, ...] = np.roll(place_onehot[i, ...], roll_y, axis=0)\n",
    "            place_onehot[i, ...] = np.roll(place_onehot[i, ...], roll_x, axis=1)\n",
    "            pick_yx[i, 0] = pick_y + roll_y\n",
    "            pick_yx[i, 1] = pick_x + roll_x\n",
    "\n",
    "        # Backpropagate.\n",
    "        batch = {}\n",
    "        batch['img'] = jnp.float32(img)\n",
    "        batch['text'] = jnp.float32(text_feat)\n",
    "        batch['pick_yx'] = jnp.int32(pick_yx)\n",
    "        batch['pick_onehot'] = jnp.float32(pick_onehot)\n",
    "        batch['place_onehot'] = jnp.float32(place_onehot)\n",
    "        rng, batch['rng'] = jax.random.split(rng)\n",
    "        optim, loss, _, _ = train_step(optim, batch)\n",
    "        \n",
    "        # Convert JAX array to numpy for logging\n",
    "        loss_np = np.array(loss)\n",
    "        losses.append(loss_np)\n",
    "\n",
    "        if train_iter % np.power(10, min(4, np.floor(np.log10(train_iter)))) == 0:\n",
    "            avg_loss = np.mean(losses[-100:]) if len(losses) > 100 else np.mean(losses)\n",
    "            print(f'Train Step: {train_iter} Loss: {loss_np:.4f} Avg Loss (100): {avg_loss:.4f}')\n",
    "        \n",
    "        if train_iter % 1000 == 0:\n",
    "            # Save checkpoint\n",
    "            save_checkpoint('.', optim.target, train_iter, prefix='ckpt_', overwrite=True)\n",
    "            \n",
    "            # Plot loss curve every 1000 steps\n",
    "            if len(losses) > 0:\n",
    "                plt.figure(figsize=(10, 4))\n",
    "                plt.plot(losses)\n",
    "                plt.title(f'Training Loss (Step {train_iter})')\n",
    "                plt.xlabel('Step')\n",
    "                plt.ylabel('Loss')\n",
    "                plt.yscale('log')\n",
    "                plt.grid(True)\n",
    "                plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "odcVMlL3JaWu"
   },
   "source": [
    "## **Demo:** CLIPort variant\n",
    "Test a variant of CLIPort: text-conditioned translation-only Transporter Nets.\n",
    "\n",
    "Text must generally be in the form: \"Pick the **{object}** and place it on the **{location}**.\"\n",
    "\n",
    "Admissible **objects:** \"{color} block\" (e.g. \"blue block\")\n",
    "\n",
    "Admissible **colors:** red, orange, yellow, green, blue, purple, pink, cyan, gray, brown\n",
    "\n",
    "Admissible **locations:** \"{color} block\" or \"{color} bowl\" or \"top/bottom/left/right side\" or \"top/bottom left/right corner\" or \"middle\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U0aKxhs6JaWu"
   },
   "outputs": [],
   "source": [
    "# Define and reset environment.\n",
    "config = {'pick':  ['yellow block', 'green block', 'blue block'],\n",
    "          'place': ['yellow bowl', 'green bowl', 'blue bowl']}\n",
    "\n",
    "np.random.seed(42)\n",
    "obs = env.reset(config)\n",
    "img = env.get_camera_image()\n",
    "plt.imshow(img)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x-h0Sg1vJaWu"
   },
   "outputs": [],
   "source": [
    "# Show camera image before pick and place.\n",
    "\n",
    "def run_cliport(obs, text):\n",
    "    \"\"\"Original run_cliport with added error handling\"\"\"\n",
    "    try:\n",
    "        if obs is None:\n",
    "            print(\"Warning: Received None observation\")\n",
    "            return None\n",
    "            \n",
    "        before = env.get_camera_image()\n",
    "        prev_obs = obs['image'].copy()\n",
    "\n",
    "        # Tokenize text and get CLIP features\n",
    "        text_tokens = clip.tokenize([text]).cuda()\n",
    "        with torch.no_grad():\n",
    "            text_feats = clip_model.encode_text(text_tokens).float()\n",
    "        text_feats /= text_feats.norm(dim=-1, keepdim=True)\n",
    "        text_feats = text_feats.cpu().numpy().astype(np.float32)\n",
    "        \n",
    "        # Normalize image and add batch dimension\n",
    "        img = obs['image'][None, ...] / 255\n",
    "        img = np.concatenate((img, coords[None, ...]), axis=3).astype(np.float32)\n",
    "\n",
    "        # Create batch\n",
    "        batch = {\n",
    "            'img': jnp.asarray(img, dtype=jnp.float32),\n",
    "            'text': jnp.asarray(text_feats, dtype=jnp.float32)\n",
    "        }\n",
    "        \n",
    "        # Debug prints\n",
    "        print(\"Final shapes:\", {k: v.shape for k, v in batch.items()})\n",
    "        print(\"Final dtypes:\", {k: v.dtype for k, v in batch.items()})\n",
    "        \n",
    "        pick_map, place_map = eval_step(optim.target, batch)\n",
    "        pick_map, place_map = np.array(pick_map), np.array(place_map)\n",
    "\n",
    "        pick_max = np.argmax(pick_map).squeeze()\n",
    "        pick_yx = (pick_max // 224, pick_max % 224)\n",
    "        pick_yx = np.clip(pick_yx, 20, 204)\n",
    "        pick_xyz = obs['xyzmap'][pick_yx[0], pick_yx[1]]\n",
    "\n",
    "        place_max = np.argmax(place_map).squeeze()\n",
    "        place_yx = (place_max // 224, place_max % 224)\n",
    "        place_yx = np.clip(place_yx, 20, 204)\n",
    "        place_xyz = obs['xyzmap'][place_yx[0], place_yx[1]]\n",
    "\n",
    "        # Step environment\n",
    "        act = {'pick': pick_xyz, 'place': place_xyz}\n",
    "        new_obs, reward, done, info = env.step(act)\n",
    "\n",
    "        # Visualization\n",
    "        plt.title(text)\n",
    "        plt.imshow(prev_obs)\n",
    "        plt.arrow(pick_yx[1], pick_yx[0], place_yx[1]-pick_yx[1], place_yx[0]-pick_yx[0], \n",
    "                 color='w', head_starts_at_zero=False, head_width=7, length_includes_head=True)\n",
    "        plt.show()\n",
    "\n",
    "        plt.subplot(1, 2, 1)\n",
    "        plt.title('Pick Heatmap')\n",
    "        plt.imshow(pick_map.reshape(224, 224))\n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.title('Place Heatmap')\n",
    "        plt.imshow(place_map.reshape(224, 224))\n",
    "        plt.show()\n",
    "\n",
    "        return new_obs\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Error in run_cliport: {e}\")\n",
    "        return obs\n",
    "\n",
    "user_input = 'Pick the yellow block and place it on the blue bowl.'  #@param {type:\"string\"}\n",
    "\n",
    "# Run the demo\n",
    "run_cliport(obs, user_input)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execution Tracking Setup\n",
    "\n",
    "This section implements tracking for:\n",
    "- Plan generation success rates\n",
    "- Execution success rates\n",
    "- Historical record of plans and executions\n",
    "- Plan-to-execution correlation\n",
    "\n",
    "The ExecutionTracker class maintains metrics about the robot's performance in carrying out generated plans."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance Monitoring Setup\n",
    "\n",
    "This section sets up tracking for:\n",
    "- API call latency\n",
    "- Token usage and costs\n",
    "- Success/failure rates\n",
    "- Resource utilization\n",
    "\n",
    "The PerformanceMonitor class handles metrics collection for the language model interactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict # Default dictionary for metrics\n",
    "\n",
    "#@title Performance Monitoring Setup\n",
    "class PerformanceMonitor:\n",
    "    def __init__(self):\n",
    "        self.metrics = defaultdict(list)\n",
    "        self.total_cost = 0\n",
    "        # GPT-3.5-turbo-instruct pricing per 1K tokens\n",
    "        self.price_per_token = {\n",
    "            'prompt': 0.0015 / 1000,  # Input tokens\n",
    "            'completion': 0.002 / 1000 # Output tokens\n",
    "        }\n",
    "    \n",
    "    def log_api_call(self, start_time, response, call_type=\"scoring\"):\n",
    "        duration = time.time() - start_time\n",
    "        prompt_tokens = response['usage']['prompt_tokens']\n",
    "        completion_tokens = response['usage']['completion_tokens']\n",
    "        total_tokens = prompt_tokens + completion_tokens\n",
    "        \n",
    "        # Calculate cost\n",
    "        call_cost = (prompt_tokens * self.price_per_token['prompt'] + \n",
    "                    completion_tokens * self.price_per_token['completion'])\n",
    "        self.total_cost += call_cost\n",
    "        \n",
    "        self.metrics['duration'].append(duration)\n",
    "        self.metrics['prompt_tokens'].append(prompt_tokens)\n",
    "        self.metrics['completion_tokens'].append(completion_tokens)\n",
    "        self.metrics['total_tokens'].append(total_tokens)\n",
    "        self.metrics['cost'].append(call_cost)\n",
    "        self.metrics['call_type'].append(call_type)\n",
    "    \n",
    "    def get_summary(self):\n",
    "        return {\n",
    "            'total_calls': len(self.metrics['duration']),\n",
    "            'total_cost': self.total_cost,\n",
    "            'avg_duration': sum(self.metrics['duration']) / len(self.metrics['duration']),\n",
    "            'total_tokens': sum(self.metrics['total_tokens']),\n",
    "            'avg_tokens_per_call': sum(self.metrics['total_tokens']) / len(self.metrics['total_tokens'])\n",
    "        }\n",
    "    \n",
    "    def print_summary(self):\n",
    "        summary = self.get_summary()\n",
    "        print(f\"Performance Summary:\")\n",
    "        print(f\"Total API Calls: {summary['total_calls']}\")\n",
    "        print(f\"Total Cost: ${summary['total_cost']:.4f}\")\n",
    "        print(f\"Average Duration: {summary['avg_duration']:.2f} seconds\")\n",
    "        print(f\"Total Tokens Used: {summary['total_tokens']}\")\n",
    "        print(f\"Average Tokens per Call: {summary['avg_tokens_per_call']:.1f}\")\n",
    "\n",
    "performance_monitor = PerformanceMonitor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Execution Tracking Setup\n",
    "class ExecutionTracker:\n",
    "    def __init__(self):\n",
    "        self.metrics = {\n",
    "            'total_plans': 0,\n",
    "            'successful_plans': 0,\n",
    "            'total_executions': 0,\n",
    "            'successful_executions': 0,\n",
    "            'plan_history': [],\n",
    "            'execution_history': []\n",
    "        }\n",
    "    \n",
    "    def log_plan(self, instruction, plan, success):\n",
    "        self.metrics['total_plans'] += 1\n",
    "        if success:\n",
    "            self.metrics['successful_plans'] += 1\n",
    "        self.metrics['plan_history'].append({\n",
    "            'instruction': instruction,\n",
    "            'plan': plan,\n",
    "            'success': success\n",
    "        })\n",
    "    \n",
    "    def log_execution(self, instruction, plan, success):\n",
    "        self.metrics['total_executions'] += 1\n",
    "        if success:\n",
    "            self.metrics['successful_executions'] += 1\n",
    "        self.metrics['execution_history'].append({\n",
    "            'instruction': instruction,\n",
    "            'plan': plan,\n",
    "            'success': success\n",
    "        })\n",
    "    \n",
    "    def print_summary(self):\n",
    "        print(f\"Execution Summary:\")\n",
    "        print(f\"Planning Success Rate: {self.metrics['successful_plans']/max(1,self.metrics['total_plans'])*100:.1f}%\")\n",
    "        print(f\"Execution Success Rate: {self.metrics['successful_executions']/max(1,self.metrics['total_executions'])*100:.1f}%\")\n",
    "\n",
    "execution_tracker = ExecutionTracker()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "38Iyw8Zs3Sn5"
   },
   "source": [
    "## Setup SayCan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hPk9A3lQgtcy"
   },
   "outputs": [],
   "source": [
    "#@title LLM Cache\n",
    "overwrite_cache = True\n",
    "if overwrite_cache:\n",
    "  LLM_CACHE = {}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LLM Scoring Functions\n",
    "Implements GPT-3 based scoring for action selection:\n",
    "- Option generation\n",
    "- Action scoring\n",
    "- Planning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Cn53B4YsII3a"
   },
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "# Initialize OpenAI client\n",
    "client = OpenAI(api_key=openai.api_key)\n",
    "\n",
    "def gpt3_call(engine=\"gpt-3.5-turbo-instruct\", prompt=\"\", max_tokens=128, temperature=0):\n",
    "    \"\"\"Updated GPT-3 call using new OpenAI API format\"\"\"\n",
    "    full_query = \"\"\n",
    "    for p in prompt:\n",
    "        full_query += p\n",
    "    id = tuple((engine, full_query, max_tokens, temperature))\n",
    "    \n",
    "    if id in LLM_CACHE.keys():\n",
    "        print('cache hit, returning')\n",
    "        response = LLM_CACHE[id]\n",
    "    else:\n",
    "        start_time = time.time()\n",
    "        response = client.completions.create(\n",
    "            model=engine,\n",
    "            prompt=prompt,\n",
    "            max_tokens=max_tokens,\n",
    "            temperature=temperature\n",
    "        )\n",
    "        # Convert response to dict for compatibility with existing code\n",
    "        response_dict = {\n",
    "            \"choices\": [{\"text\": choice.text} for choice in response.choices],\n",
    "            \"usage\": {\n",
    "                \"prompt_tokens\": response.usage.prompt_tokens,\n",
    "                \"completion_tokens\": response.usage.completion_tokens,\n",
    "                \"total_tokens\": response.usage.total_tokens\n",
    "            }\n",
    "        }\n",
    "        performance_monitor.log_api_call(start_time, response_dict)\n",
    "        LLM_CACHE[id] = response_dict\n",
    "        response = response_dict\n",
    "    return response\n",
    "\n",
    "def gpt3_scoring(query, options, engine=\"gpt-3.5-turbo-instruct\", limit_num_options=None, \n",
    "                option_start=\"\\n\", verbose=False, print_tokens=False):\n",
    "    \"\"\"Score options using GPT-3\"\"\"\n",
    "    if limit_num_options:\n",
    "        options = options[:limit_num_options]\n",
    "    verbose and print(\"Scoring\", len(options), \"options\")\n",
    "    \n",
    "    # Create prompts that ask the model to rate each option\n",
    "    gpt3_prompt_options = []\n",
    "    for option in options:\n",
    "        prompt = f\"{query}{option}\\nHow likely is this to succeed on a scale of 0-10?\"\n",
    "        gpt3_prompt_options.append(prompt)\n",
    "    \n",
    "    # Process in batches of 20\n",
    "    BATCH_SIZE = 20\n",
    "    scores = {}\n",
    "    for i in range(0, len(gpt3_prompt_options), BATCH_SIZE):\n",
    "        batch = gpt3_prompt_options[i:i + BATCH_SIZE]\n",
    "        batch_options = options[i:i + BATCH_SIZE]\n",
    "        \n",
    "        response = gpt3_call(\n",
    "            engine=engine, \n",
    "            prompt=batch, \n",
    "            max_tokens=5,\n",
    "            temperature=0)\n",
    "        \n",
    "        # Parse responses for this batch\n",
    "        for option, choice in zip(batch_options, response[\"choices\"]):\n",
    "            try:\n",
    "                score_text = choice[\"text\"].strip()\n",
    "                score = float(next(s for s in score_text.split() if s.replace('.','',1).isdigit()))\n",
    "                scores[option] = score\n",
    "            except:\n",
    "                scores[option] = 0\n",
    "                \n",
    "        # Optional: Add a small delay between batches\n",
    "        time.sleep(0.1)\n",
    "    \n",
    "    if verbose:\n",
    "        for i, option in enumerate(sorted(scores.items(), key=lambda x : -x[1])):\n",
    "            print(option[1], \"\\t\", option[0])\n",
    "            if i >= 10:\n",
    "                break\n",
    "                \n",
    "    return scores, response\n",
    "\n",
    "def make_options(pick_targets=None, place_targets=None, options_in_api_form=True, termination_string=\"done()\"):\n",
    "  if not pick_targets:\n",
    "    pick_targets = PICK_TARGETS\n",
    "  if not place_targets:\n",
    "    place_targets = PLACE_TARGETS\n",
    "  options = []\n",
    "  for pick in pick_targets:\n",
    "    for place in place_targets:\n",
    "      if options_in_api_form:\n",
    "        option = \"robot.pick_and_place({}, {})\".format(pick, place)\n",
    "      else:\n",
    "        option = \"Pick the {} and place it on the {}.\".format(pick, place)\n",
    "      options.append(option)\n",
    "\n",
    "  options.append(termination_string)\n",
    "  print(\"Considering\", len(options), \"options\")\n",
    "  return options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Shp3CFFDzDie"
   },
   "outputs": [],
   "source": [
    "query = \"To pick the blue block and put it on the red block, I should:\\n\"\n",
    "options = make_options(PICK_TARGETS, PLACE_TARGETS)\n",
    "scores, response = gpt3_scoring(query, options, engine=ENGINE, limit_num_options=5, option_start='\\n', verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Helper Functions\n",
    "Utility functions for:\n",
    "- Scene description\n",
    "- Score normalization\n",
    "- Visualization\n",
    "- Step conversion\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "w16kuljCGd2o"
   },
   "outputs": [],
   "source": [
    "#@title Helper Functions\n",
    "\n",
    "def build_scene_description(found_objects, block_name=\"box\", bowl_name=\"circle\"):\n",
    "  scene_description = f\"objects = {found_objects}\"\n",
    "  scene_description = scene_description.replace(block_name, \"block\")\n",
    "  scene_description = scene_description.replace(bowl_name, \"bowl\")\n",
    "  scene_description = scene_description.replace(\"'\", \"\")\n",
    "  return scene_description\n",
    "\n",
    "def step_to_nlp(step):\n",
    "  step = step.replace(\"robot.pick_and_place(\", \"\")\n",
    "  step = step.replace(\")\", \"\")\n",
    "  pick, place = step.split(\", \")\n",
    "  return \"Pick the \" + pick + \" and place it on the \" + place + \".\"\n",
    "\n",
    "def normalize_scores(scores):\n",
    "  max_score = max(scores.values())  \n",
    "  normed_scores = {key: np.clip(scores[key] / max_score, 0, 1) for key in scores}\n",
    "  return normed_scores\n",
    "\n",
    "def plot_saycan(llm_scores, affordance_scores, combined_scores, task, show_top=None):\n",
    "    \"\"\"More robust plotting function\"\"\"\n",
    "    try:\n",
    "        if show_top:\n",
    "            # Get top options\n",
    "            top_options = nlargest(show_top, combined_scores, key=combined_scores.get)\n",
    "            top_llm_options = nlargest(show_top // 2, llm_scores, key=llm_scores.get)\n",
    "            for llm_option in top_llm_options:\n",
    "                if llm_option not in top_options:\n",
    "                    top_options.append(llm_option)\n",
    "            \n",
    "            # Filter scores to top options\n",
    "            llm_scores = {option: llm_scores[option] for option in top_options}\n",
    "            affordance_scores = {option: affordance_scores[option] for option in top_options}\n",
    "            combined_scores = {option: combined_scores[option] for option in top_options}\n",
    "\n",
    "        sorted_keys = dict(sorted(combined_scores.items()))\n",
    "        keys = list(sorted_keys.keys())\n",
    "        positions = np.arange(len(combined_scores.items()))\n",
    "        width = 0.3\n",
    "\n",
    "        fig = plt.figure(figsize=(12, 6))\n",
    "        ax1 = fig.add_subplot(1,1,1)\n",
    "\n",
    "        # Normalize scores\n",
    "        plot_llm_scores = {k: np.exp(v) for k, v in llm_scores.items()}\n",
    "        plot_llm_scores = normalize_scores(plot_llm_scores)\n",
    "        plot_llm_scores = [plot_llm_scores[k] for k in sorted_keys]\n",
    "        plot_affordance_scores = [affordance_scores[k] for k in sorted_keys]\n",
    "        plot_combined_scores = [combined_scores[k] for k in sorted_keys]\n",
    "        \n",
    "        # Plot bars\n",
    "        ax1.bar(positions, plot_combined_scores, 3 * width, alpha=0.6, color=\"#93CE8E\", label=\"combined\")\n",
    "        ax1.bar(positions - width / 2, plot_affordance_scores, width, color=\"#ea9999ff\", label=\"vfs\")\n",
    "        ax1.bar(positions + width / 2, plot_llm_scores, width, color=\"#a4c2f4ff\", label=\"language\")\n",
    "\n",
    "        plt.xticks(positions, [k.replace(\"robot.pick_and_place\", \"\").replace(\", \", \" to \")\n",
    "                             .replace(\"(\", \"\").replace(\")\", \"\") for k in keys], \n",
    "                  rotation=\"vertical\")\n",
    "        ax1.set_ylim(0.0, 1.0)\n",
    "        ax1.grid(True)\n",
    "        ax1.legend()\n",
    "        \n",
    "        plt.title(task)\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Error in plotting: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Affordance Scoring\n",
    "Implements object-detection based affordance scoring without RL policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TyQ8GVvzyxP5"
   },
   "outputs": [],
   "source": [
    "#@title Affordance Scoring\n",
    "#@markdown Given this environment does not have RL-trained policies or an asscociated value function, we use affordances through an object detector.\n",
    "\n",
    "def affordance_scoring(options, found_objects, verbose=False, block_name=\"box\", bowl_name=\"circle\", termination_string=\"done()\"):\n",
    "  affordance_scores = {}\n",
    "  found_objects = [\n",
    "                   found_object.replace(block_name, \"block\").replace(bowl_name, \"bowl\") \n",
    "                   for found_object in found_objects + list(PLACE_TARGETS.keys())[-5:]]\n",
    "  verbose and print(\"found_objects\", found_objects)\n",
    "  for option in options:\n",
    "    if option == termination_string:\n",
    "      affordance_scores[option] = 0.2\n",
    "      continue\n",
    "    pick, place = option.replace(\"robot.pick_and_place(\", \"\").replace(\")\", \"\").split(\", \")\n",
    "    affordance = 0\n",
    "    found_objects_copy = found_objects.copy()\n",
    "    if pick in found_objects_copy:\n",
    "      found_objects_copy.remove(pick)\n",
    "      if place in found_objects_copy:\n",
    "        affordance = 1\n",
    "    affordance_scores[option] = affordance\n",
    "    verbose and print(affordance, '\\t', option)\n",
    "  return affordance_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Cell\n",
    "Demonstrates basic SayCan functionality with test case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3xZm83muyRWR"
   },
   "outputs": [],
   "source": [
    "#@title Test\n",
    "termination_string = \"done()\"\n",
    "query = \"To pick the blue block and put it on the red block, I should:\\n\"\n",
    "\n",
    "options = make_options(PICK_TARGETS, PLACE_TARGETS, termination_string=termination_string)\n",
    "llm_scores, _ = gpt3_scoring(query, options, verbose=True, engine=ENGINE)\n",
    "\n",
    "affordance_scores = affordance_scoring(options, found_objects, block_name=\"box\", bowl_name=\"circle\", verbose=False, termination_string=termination_string)\n",
    "\n",
    "combined_scores = {option: np.exp(llm_scores[option]) * affordance_scores[option] for option in options}\n",
    "combined_scores = normalize_scores(combined_scores)\n",
    "selected_task = max(combined_scores, key=combined_scores.get)\n",
    "print(\"Selecting: \", selected_task)\n",
    "\n",
    "# Added performance summary\n",
    "print(\"\\nPerformance for test run:\")\n",
    "performance_monitor.print_summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zJVQIUvI77bF"
   },
   "source": [
    "## **SayCan**\n",
    "\n",
    "Here we implement SayCan with LLM scoring and robotic affordances from ViLD (in the absence of a trained value function from an RL policy)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SayCan Demos\n",
    "Final demo cells showing:\n",
    "1. Full SayCan implementation\n",
    "2. Socratic Model variant\n",
    "Each including:\n",
    "- Prompt setup\n",
    "- Task configuration\n",
    "- Scene setup\n",
    "- Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yys24Oj1796Y"
   },
   "outputs": [],
   "source": [
    "#@title Prompt\n",
    "\n",
    "termination_string = \"done()\"\n",
    "\n",
    "gpt3_context = \"\"\"\n",
    "objects = [red block, yellow block, blue block, green bowl]\n",
    "# move all the blocks to the top left corner.\n",
    "robot.pick_and_place(blue block, top left corner)\n",
    "robot.pick_and_place(red block, top left corner)\n",
    "robot.pick_and_place(yellow block, top left corner)\n",
    "done()\n",
    "\n",
    "objects = [red block, yellow block, blue block, green bowl]\n",
    "# put the yellow one the green thing.\n",
    "robot.pick_and_place(yellow block, green bowl)\n",
    "done()\n",
    "\n",
    "objects = [yellow block, blue block, red block]\n",
    "# move the light colored block to the middle.\n",
    "robot.pick_and_place(yellow block, middle)\n",
    "done()\n",
    "\n",
    "objects = [blue block, green bowl, red block, yellow bowl, green block]\n",
    "# stack the blocks.\n",
    "robot.pick_and_place(green block, blue block)\n",
    "robot.pick_and_place(red block, green block)\n",
    "done()\n",
    "\n",
    "objects = [red block, blue block, green bowl, blue bowl, yellow block, green block]\n",
    "# group the blue objects together.\n",
    "robot.pick_and_place(blue block, blue bowl)\n",
    "done()\n",
    "\n",
    "objects = [green bowl, red block, green block, red bowl, yellow bowl, yellow block]\n",
    "# sort all the blocks into their matching color bowls.\n",
    "robot.pick_and_place(green block, green bowl)\n",
    "robot.pick_and_place(red block, red bowl)\n",
    "robot.pick_and_place(yellow block, yellow bowl)\n",
    "done()\n",
    "\"\"\"\n",
    "\n",
    "use_environment_description = False\n",
    "gpt3_context_lines = gpt3_context.split(\"\\n\")\n",
    "gpt3_context_lines_keep = []\n",
    "for gpt3_context_line in gpt3_context_lines:\n",
    "  if \"objects =\" in gpt3_context_line and not use_environment_description:\n",
    "    continue\n",
    "  gpt3_context_lines_keep.append(gpt3_context_line)\n",
    "\n",
    "gpt3_context = \"\\n\".join(gpt3_context_lines_keep)\n",
    "print(gpt3_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eCJckkuwFzDS"
   },
   "outputs": [],
   "source": [
    "#@title Task and Config\n",
    "only_plan = False\n",
    "\n",
    "raw_input = \"put all the blocks in different corners.\" \n",
    "config = {\"pick\":  [\"red block\", \"yellow block\", \"green block\", \"blue block\"],\n",
    "          \"place\": [\"red bowl\"]}\n",
    "\n",
    "# raw_input = \"move the block to the bowl.\"\n",
    "# config = {'pick':  ['red block'],\n",
    "#           'place': ['green bowl']}\n",
    "\n",
    "# raw_input = \"put any blocks on their matched colored bowls.\"\n",
    "# config = {'pick':  ['yellow block', 'green block', 'blue block'],\n",
    "#           'place': ['yellow bowl', 'green bowl', 'blue bowl']}\n",
    "          \n",
    "# raw_input = \"put all the blocks in the green bowl.\"\n",
    "# config = {'pick':  ['yellow block', 'green block', 'red block'],\n",
    "#           'place': ['yellow bowl', 'green bowl']}\n",
    "\n",
    "# raw_input = \"stack all the blocks.\"\n",
    "# config = {'pick':  ['yellow block', 'blue block', 'red block'],\n",
    "#           'place': ['blue bowl', 'red bowl']}\n",
    "\n",
    "# raw_input = \"make the highest block stack.\"\n",
    "# config = {'pick':  ['yellow block', 'blue block', 'red block'],\n",
    "#           'place': ['blue bowl', 'red bowl']}\n",
    "\n",
    "# raw_input = \"stack all the blocks.\"\n",
    "# config = {'pick':  ['green block', 'blue block', 'red block'],\n",
    "#           'place': ['yellow bowl', 'green bowl']}\n",
    "\n",
    "# raw_input = \"put the block in all the corners.\" \n",
    "# config = {'pick':  ['red block'],\n",
    "#           'place': ['red bowl', 'green bowl']}\n",
    "\n",
    "# raw_input = \"clockwise, move the block through all the corners.\"\n",
    "# config = {'pick':  ['red block'],\n",
    "#           'place': ['red bowl', 'green bowl', 'yellow bowl']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DX18LmbIIGfz"
   },
   "outputs": [],
   "source": [
    "#@title Setup Scene\n",
    "image_path = \"./2db.png\"\n",
    "np.random.seed(2)\n",
    "if config is None:\n",
    "  pick_items = list(PICK_TARGETS.keys())\n",
    "  pick_items = np.random.choice(pick_items, size=np.random.randint(1, 5), replace=False)\n",
    "\n",
    "  place_items = list(PLACE_TARGETS.keys())[:-9]\n",
    "  place_items = np.random.choice(place_items, size=np.random.randint(1, 6 - len(pick_items)), replace=False)\n",
    "  config = {\"pick\":  pick_items,\n",
    "            \"place\": place_items}\n",
    "  print(pick_items, place_items)\n",
    "\n",
    "obs = env.reset(config)\n",
    "\n",
    "img_top = env.get_camera_image_top()\n",
    "img_top_rgb = cv2.cvtColor(img_top, cv2.COLOR_BGR2RGB)\n",
    "plt.imshow(img_top)\n",
    "\n",
    "imageio.imsave(image_path, img_top)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m4Yh-PqZ-03n"
   },
   "outputs": [],
   "source": [
    "#@title Runner\n",
    "plot_on = True\n",
    "max_tasks = 5\n",
    "\n",
    "try:\n",
    "    options = make_options(PICK_TARGETS, PLACE_TARGETS, termination_string=termination_string)\n",
    "    found_objects = vild(image_path, category_name_string, vild_params, plot_on=False)\n",
    "    scene_description = build_scene_description(found_objects)\n",
    "    env_description = scene_description\n",
    "\n",
    "    print(scene_description)\n",
    "\n",
    "    gpt3_prompt = gpt3_context\n",
    "    if use_environment_description:\n",
    "        gpt3_prompt += \"\\n\" + env_description\n",
    "    gpt3_prompt += \"\\n# \" + raw_input + \"\\n\"\n",
    "\n",
    "    all_llm_scores = []\n",
    "    all_affordance_scores = []\n",
    "    all_combined_scores = []\n",
    "    affordance_scores = affordance_scoring(options, found_objects, block_name=\"box\", \n",
    "                                         bowl_name=\"circle\", verbose=False)\n",
    "    num_tasks = 0\n",
    "    selected_task = \"\"\n",
    "    steps_text = []\n",
    "    \n",
    "    while not selected_task == termination_string:\n",
    "        num_tasks += 1\n",
    "        if num_tasks > max_tasks:\n",
    "            break\n",
    "\n",
    "        llm_scores, _ = gpt3_scoring(gpt3_prompt, options, verbose=True, \n",
    "                                   engine=ENGINE, print_tokens=False)\n",
    "        combined_scores = {option: np.exp(llm_scores[option]) * affordance_scores[option] \n",
    "                         for option in options}\n",
    "        combined_scores = normalize_scores(combined_scores)\n",
    "        selected_task = max(combined_scores, key=combined_scores.get)\n",
    "        if selected_task and selected_task != termination_string:\n",
    "            steps_text.append(selected_task)\n",
    "            print(f\"{num_tasks} Selecting: {selected_task}\")\n",
    "            gpt3_prompt += selected_task + \"\\n\"\n",
    "\n",
    "            all_llm_scores.append(llm_scores)\n",
    "            all_affordance_scores.append(affordance_scores)\n",
    "            all_combined_scores.append(combined_scores)\n",
    "\n",
    "    if plot_on:\n",
    "        for llm_scores, affordance_scores, combined_scores, step in zip(\n",
    "                all_llm_scores, all_affordance_scores, all_combined_scores, steps_text):\n",
    "            plot_saycan(llm_scores, affordance_scores, combined_scores, step, show_top=10)\n",
    "\n",
    "    print('**** Solution ****')\n",
    "    print(env_description)\n",
    "    print('# ' + raw_input)\n",
    "    for i, step in enumerate(steps_text):\n",
    "        if step == '' or step == termination_string:\n",
    "            break\n",
    "        print('Step ' + str(i) + ': ' + step)\n",
    "\n",
    "    if not only_plan:\n",
    "        print('Initial state:')\n",
    "        plt.imshow(env.get_camera_image())\n",
    "        plt.show()\n",
    "\n",
    "        for i, step in enumerate(steps_text):\n",
    "            if step == '' or step == termination_string:\n",
    "                break\n",
    "            nlp_step = step_to_nlp(step)\n",
    "            print('GPT-3 says next step:', nlp_step)\n",
    "\n",
    "            new_obs = run_cliport(obs, nlp_step)\n",
    "            if new_obs is None:\n",
    "                print(\"Failed to execute step, stopping execution\")\n",
    "                break\n",
    "            obs = new_obs\n",
    "\n",
    "        # Show camera image after task.\n",
    "        print('Final state:')\n",
    "        plt.imshow(env.get_camera_image())\n",
    "        plt.show()\n",
    "\n",
    "        # Add summaries\n",
    "        print(\"\\nOverall Performance Metrics:\")\n",
    "        performance_monitor.print_summary()\n",
    "        execution_tracker.print_summary()\n",
    "\n",
    "except Exception as e:\n",
    "    print(f\"Error in main execution: {e}\")\n",
    "    import traceback\n",
    "    traceback.print_exc()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YWQwjuiDi1nC"
   },
   "source": [
    "## Socratic Model: VILD, GPT3, CLIPort\n",
    "\n",
    "This implements a version of LLM planning shown in [Socratic Models](https://socraticmodels.github.io/), without the grounding, but with a scene description. For this relatively simple environment with clear robotic affordances, the scene description is generally sufficient. This mirrors the implementation attached to the paper [here](https://github.com/google-research/google-research/tree/master/socraticmodels)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "trXP07uypWL8"
   },
   "outputs": [],
   "source": [
    "#@title Prompt\n",
    "\n",
    "gpt3_context = \"\"\"\n",
    "objects = [red block, yellow block, blue block, green bowl]\n",
    "# move all the blocks to the top left corner.\n",
    "robot.pick_and_place(blue block, top left corner)\n",
    "robot.pick_and_place(red block, top left corner)\n",
    "robot.pick_and_place(yellow block, top left corner)\n",
    "done()\n",
    "\n",
    "objects = [red block, yellow block, blue block, green bowl]\n",
    "# put the yellow one the green thing.\n",
    "robot.pick_and_place(yellow block, green bowl)\n",
    "done()\n",
    "\n",
    "objects = [yellow block, blue block, red block]\n",
    "# move the light colored block to the middle.\n",
    "robot.pick_and_place(yellow block, middle)\n",
    "done()\n",
    "\n",
    "objects = [blue block, green bowl, red block, yellow bowl, green block]\n",
    "# stack the blocks.\n",
    "robot.pick_and_place(green block, blue block)\n",
    "robot.pick_and_place(red block, green block)\n",
    "done()\n",
    "\n",
    "objects = [red block, blue block, green bowl, blue bowl, yellow block, green block]\n",
    "# group the blue objects together.\n",
    "robot.pick_and_place(blue block, blue bowl)\n",
    "done()\n",
    "\n",
    "objects = [green bowl, red block, green block, red bowl, yellow bowl, yellow block]\n",
    "# sort all the blocks into their matching color bowls.\n",
    "robot.pick_and_place(green block, green bowl)\n",
    "robot.pick_and_place(red block, red bowl)\n",
    "robot.pick_and_place(yellow block, yellow bowl)\n",
    "done()\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sPx--Ea7IZ69"
   },
   "outputs": [],
   "source": [
    "#@title Queries and Configs\n",
    "\n",
    "only_plan = False\n",
    "\n",
    "raw_input = \"put all the blocks in different corners.\" \n",
    "config = {'pick':  ['red block', 'yellow block', 'green block', 'blue block'],\n",
    "          'place': ['red bowl']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k6TSgRItpPC4"
   },
   "outputs": [],
   "source": [
    "#@title Runner\n",
    "\n",
    "env_description = ''\n",
    "image_path = './2db.png'\n",
    "\n",
    "np.random.seed(2)\n",
    "\n",
    "if config is None:\n",
    "  pick_items = list(PICK_TARGETS.keys())\n",
    "  pick_items = np.random.choice(pick_items, size=np.random.randint(1, 5), replace=False)\n",
    "\n",
    "  place_items = list(PLACE_TARGETS.keys())[:-9]\n",
    "  place_items = np.random.choice(place_items, size=np.random.randint(1, 6 - len(pick_items)), replace=False)\n",
    "  config = {'pick':  pick_items,\n",
    "            'place': place_items}\n",
    "  print(pick_items, place_items)\n",
    "obs = env.reset(config)\n",
    "\n",
    "img_top = env.get_camera_image_top()\n",
    "img_top_rgb = cv2.cvtColor(img_top, cv2.COLOR_BGR2RGB)\n",
    "plt.imshow(img_top_rgb)\n",
    "\n",
    "imageio.imsave(image_path, img_top)\n",
    "\n",
    "found_objects = vild(image_path, category_name_string, vild_params, plot_on=False)\n",
    "scene_description = build_scene_description(found_objects)\n",
    "print(scene_description)\n",
    "\n",
    "env_description = scene_description\n",
    "\n",
    "gpt3_prompt = gpt3_context\n",
    "gpt3_prompt += \"\\n\" + env_description + \"\\n\"\n",
    "gpt3_prompt += \"# \" + raw_input\n",
    "response = gpt3_call(engine=ENGINE, prompt=gpt3_prompt, max_tokens=128, temperature=0)\n",
    "steps_text = [text.strip().strip() for text in response[\"choices\"][0][\"text\"].strip().split(\"#\")[0].split(\"\\n\")][:-1]\n",
    "print('**** Solution ****')\n",
    "print(env_description)\n",
    "print('# ' + raw_input)\n",
    "for i, step in enumerate(steps_text):\n",
    "  if step == '' or step == termination_string:\n",
    "    break\n",
    "  print('Step ' + str(i) + ': ' + step)\n",
    "  nlp_step = step_to_nlp(step)\n",
    "\n",
    "if not only_plan:\n",
    "  print('Initial state:')\n",
    "  plt.imshow(env.get_camera_image())\n",
    "\n",
    "  for i, step in enumerate(steps_text):\n",
    "    if step == '' or step == termination_string:\n",
    "      break\n",
    "    nlp_step = step_to_nlp(step)\n",
    "    print('GPT-3 says next step:', nlp_step)\n",
    "\n",
    "    obs = run_cliport(obs, nlp_step)\n",
    "\n",
    "  # Show camera image after task.\n",
    "  print('Final state:')\n",
    "  plt.imshow(env.get_camera_image())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Josh's SayCan Demo Cell\n",
    "\n",
    "## Overview\n",
    "This demonstration shows the complete SayCan pipeline converting natural language instructions into robot actions through LLM planning and affordance scoring.\n",
    "\n",
    "## Pipeline Steps\n",
    "1. **Vision Analysis**\n",
    "   * Scene detection using ViLD\n",
    "   * Object identification and location mapping\n",
    "\n",
    "2. **Language Planning**\n",
    "   * GPT-3 converts high-level instruction to steps\n",
    "   * Each possible action gets scored for relevance\n",
    "\n",
    "3. **Action Selection**\n",
    "   * Combines language scores with physical affordances \n",
    "   * Selects optimal action at each step\n",
    "\n",
    "4. **Execution**\n",
    "   * Runs selected action through CLIPort\n",
    "   * Shows visual confirmation of results\n",
    "\n",
    "## Example Usage\n",
    "\n",
    "```python\n",
    "# Simple Instructions\n",
    "instruction = \"stack the red block on the blue block\"\n",
    "run_saycan_demo_v2(instruction)\n",
    "\n",
    "# Complex Instructions\n",
    "instruction = \"put all blocks in their matching colored bowls\"\n",
    "run_saycan_demo_v2(instruction)\n",
    "```\n",
    "\n",
    "## Visualization Features\n",
    "* Initial scene inspection\n",
    "* Action scoring visualization\n",
    "* Before/after images per step\n",
    "* Performance metrics display\n",
    "\n",
    "## Try These Instructions\n",
    "* `\"pick up the red block\"`\n",
    "* `\"put the blue block in the green bowl\"`\n",
    "* `\"sort blocks into matching colored bowls\"`\n",
    "* `\"put blocks in different corners\"`\n",
    "\n",
    "## Notes\n",
    "* Each step shows both scoring and execution\n",
    "* Visual confirmation helps track progress\n",
    "* Performance metrics show timing and success rates\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Josh's SayCan Demo - With Scoring\n",
    "def run_saycan_with_scoring(instruction):\n",
    "    \"\"\"SayCan with explicit scoring of each possible action\"\"\"\n",
    "    print(f\"🤖 Running SayCan for: {instruction}\")\n",
    "    \n",
    "    # Reset environment\n",
    "    config = {\n",
    "        'pick': ['red block', 'blue block', 'green block'],\n",
    "        'place': ['red bowl', 'blue bowl', 'green bowl']\n",
    "    }\n",
    "    obs = env.reset(config)\n",
    "    \n",
    "    # Show initial scene\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.imshow(env.get_camera_image())\n",
    "    plt.title(\"Initial Scene\")\n",
    "    plt.show()\n",
    "\n",
    "    # Get scene description\n",
    "    image_path = 'tmp.jpg'\n",
    "    img_top = env.get_camera_image_top()\n",
    "    imageio.imwrite(image_path, img_top)\n",
    "    found_objects = vild(image_path, category_name_string, vild_params, plot_on=False)\n",
    "    scene_description = build_scene_description(found_objects)\n",
    "    print(\"\\n🔍 Scene Description:\", scene_description)\n",
    "\n",
    "    # Generate and execute plan with scoring\n",
    "    options = make_options(PICK_TARGETS, PLACE_TARGETS, termination_string=\"done()\")\n",
    "    gpt3_prompt = gpt3_context + \"\\n\" + scene_description + \"\\n# \" + instruction + \"\\n\"\n",
    "    \n",
    "    num_tasks = 0\n",
    "    selected_task = \"\"\n",
    "    while not selected_task == \"done()\":\n",
    "        num_tasks += 1\n",
    "        if num_tasks > 5:  # Limit number of steps\n",
    "            break\n",
    "            \n",
    "        # Score all possible actions\n",
    "        print(f\"\\n🎯 Step {num_tasks}:\")\n",
    "        llm_scores, _ = gpt3_scoring(gpt3_prompt, options, verbose=False, engine=ENGINE)\n",
    "        affordance_scores = affordance_scoring(options, found_objects)\n",
    "        \n",
    "        # Combine scores and select best action\n",
    "        combined_scores = {option: np.exp(llm_scores[option]) * affordance_scores[option] \n",
    "                         for option in options}\n",
    "        combined_scores = normalize_scores(combined_scores)\n",
    "        selected_task = max(combined_scores, key=combined_scores.get)\n",
    "        \n",
    "        if selected_task and selected_task != \"done()\":\n",
    "            print(f\"Selected action: {selected_task}\")\n",
    "            \n",
    "            # Plot scores for this step\n",
    "            plot_saycan(llm_scores, affordance_scores, combined_scores, \n",
    "                       f\"Step {num_tasks} Scores\", show_top=5)\n",
    "            \n",
    "            # Execute action\n",
    "            nlp_step = step_to_nlp(selected_task)\n",
    "            print(f\"Executing: {nlp_step}\")\n",
    "            obs = run_cliport(obs, nlp_step)\n",
    "            gpt3_prompt += selected_task + \"\\n\"\n",
    "            \n",
    "    # Show final scene\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.imshow(env.get_camera_image())\n",
    "    plt.title(\"Final Scene\")\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Instructions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Try Different Instructions\n",
    "# Simple instruction\n",
    "instruction = \"stack the red block on the blue block\" #@param {type:\"string\"}\n",
    "run_saycan_with_scoring(instruction)\n",
    "\n",
    "# More complex instruction\n",
    "instruction = \"put all blocks in their matching colored bowls\" #@param {type:\"string\"}\n",
    "run_saycan_with_scoring(instruction)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "SayCan_Robot_Pick_Place.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "saycan",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}