BotShot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip3 install manipulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip3 install plyfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Import Packages and Start Visualizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from copy import deepcopy, copy\n",
    "from plyfile import PlyData\n",
    "from enum import Enum\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AbstractValue,\n",
    "    PointCloud,\n",
    "    PortSwitch,\n",
    "    Concatenate,\n",
    "    ConstantVectorSource,\n",
    "    MeshcatVisualizerParams,\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    GetScopedFrameByName,\n",
    "    VectorLogSink,\n",
    "    DiagramBuilder,\n",
    "    MeshcatVisualizer,\n",
    "    Parser,\n",
    "    Solve,\n",
    "    Rgba,\n",
    "    BaseField,\n",
    "    Concatenate,\n",
    "    Fields,\n",
    "    RigidTransform,\n",
    "    RollPitchYaw,\n",
    "    Simulator,\n",
    "    LeafSystem,\n",
    "    StartMeshcat,\n",
    "    PiecewisePose,\n",
    "    ImageDepth32F, \n",
    "    ImageRgba8U,\n",
    "    InputPortIndex,\n",
    "    RotationMatrix,\n",
    "    PiecewisePolynomial,\n",
    "    PiecewiseQuaternionSlerp,\n",
    ")\n",
    "from pydrake.geometry import StartMeshcat\n",
    "from pydrake.systems.analysis import Simulator\n",
    "from pydrake.multibody.parsing import Parser\n",
    "from pydrake.multibody import inverse_kinematics\n",
    "\n",
    "from manipulation.meshcat_utils import AddMeshcatTriad\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.station import LoadScenario, MakeHardwareStation, MakeMultibodyPlant\n",
    "from manipulation.utils import ConfigureParser\n",
    "from manipulation.meshcat_utils import StopButton\n",
    "from manipulation.systems import AddIiwaDifferentialIK, ExtractPose\n",
    "from manipulation.clutter import GenerateAntipodalGraspCandidate\n",
    "from manipulation.icp import IterativeClosestPoint\n",
    "from manipulation.scenarios import AddMultibodyTriad\n",
    "\n",
    "\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:drake:Meshcat listening for connections at http://localhost:7001\n"
     ]
    }
   ],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Leaf Systems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 IIwa Controller"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ControllerState(Enum):\n",
    "    WAIT_FOR_OBJECTS_TO_SETTLE = 1\n",
    "    PICKING_FROM_TABLE = 2\n",
    "    GO_HOME = 3\n",
    "    PREP_THROW = 4\n",
    "    THROWING = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def piecewise_traj(initial_pose, final_pose, start_time, end_time):\n",
    "    rotation = PiecewiseQuaternionSlerp()\n",
    "    rotation.Append(start_time, initial_pose.rotation())\n",
    "    rotation.Append(end_time, final_pose.rotation())\n",
    "    translation = PiecewisePolynomial.FirstOrderHold(\n",
    "        [start_time, end_time],\n",
    "        np.vstack(\n",
    "            [\n",
    "                [initial_pose.translation()],\n",
    "                [final_pose.translation()],\n",
    "            ]\n",
    "        ).T,\n",
    "    )\n",
    "    return rotation, translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ControlSystem(LeafSystem):\n",
    "    def __init__(self, plant):\n",
    "        LeafSystem.__init__(self)\n",
    "\n",
    "        self._gripper_body_index = plant.GetBodyByName(\"body\").index()\n",
    "        # 1. Input related to body_poses\n",
    "        self.DeclareAbstractInputPort(\n",
    "            \"body_poses\", AbstractValue.Make([RigidTransform()])\n",
    "        )\n",
    "\n",
    "        # 2. Input for selected grasp\n",
    "        self._grasp_index = self.DeclareAbstractInputPort(\n",
    "            \"grasp\", AbstractValue.Make((np.inf, RigidTransform()))\n",
    "        ).get_index()\n",
    "\n",
    "        # 3. Input for pose of the cup\n",
    "        self._cup_pose_index = self.DeclareAbstractInputPort(\n",
    "            \"cup_X_WO\", AbstractValue.Make(RigidTransform())\n",
    "        ).get_index()\n",
    "\n",
    "        # 4. Input for schunk position\n",
    "        self._wsg_state_index = self.DeclareVectorInputPort(\"wsg_state\", 2).get_index()\n",
    "\n",
    "        # 5. Ouput for calculated end-effector pose\n",
    "        self.DeclareAbstractOutputPort(\n",
    "            \"X_WG\",\n",
    "            lambda: AbstractValue.Make(RigidTransform()),\n",
    "            self.CalcGripperPose,\n",
    "        )   \n",
    "\n",
    "        # 6. Output for calculated schunk position\n",
    "        self.DeclareVectorOutputPort(\"wsg_position\", 1, self.CalcWsgPosition)\n",
    "        \n",
    "        # 7. Declare States\n",
    "        self._traj_X_G_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(PiecewisePose())\n",
    "        )\n",
    "        self._traj_wsg_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(PiecewisePolynomial())\n",
    "        )\n",
    "        self._mode_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(ControllerState.WAIT_FOR_OBJECTS_TO_SETTLE)\n",
    "        )\n",
    "        self._times_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make({\"initial\": 0.0})\n",
    "        )\n",
    "\n",
    "        # For GoHome mode\n",
    "        num_positions = 7\n",
    "        self._iiwa_position_index = self.DeclareVectorInputPort(\n",
    "            \"iiwa_position\", num_positions\n",
    "        ).get_index()\n",
    "        self.DeclareAbstractOutputPort(\n",
    "            \"control_mode\",\n",
    "            lambda: AbstractValue.Make(InputPortIndex(0)),\n",
    "            self.CalcControlMode,\n",
    "        )\n",
    "        self.DeclareAbstractOutputPort(\n",
    "            \"reset_diff_ik\",\n",
    "            lambda: AbstractValue.Make(False),\n",
    "            self.CalcDiffIKReset,\n",
    "        )\n",
    "        self._q0_index = self.DeclareDiscreteState(num_positions)  # for q0\n",
    "        self._traj_q_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(PiecewisePolynomial())\n",
    "        )\n",
    "        self.DeclareVectorOutputPort(\n",
    "            \"iiwa_position_command\", num_positions, self.CalcIiwaPosition\n",
    "        )\n",
    "        self.DeclareInitializationDiscreteUpdateEvent(self.Initialize)\n",
    "        self.DeclarePeriodicUnrestrictedUpdateEvent(0.1, 0.0, self.Update)\n",
    "\n",
    "        # # Add this back if we need go home mode or control by velocity\n",
    "        # self.DeclareVectorOutputPort(\n",
    "        #     \"iiwa_position_command\", num_positions, self.CalcIiwaPosition\n",
    "        # )\n",
    "\n",
    "    \n",
    "    def Initialize(self, context, discrete_state):\n",
    "        discrete_state.set_value(\n",
    "            int(self._q0_index),\n",
    "            self.get_input_port(int(self._iiwa_position_index)).Eval(context),\n",
    "        )\n",
    "\n",
    "    def Update(self, context, state):\n",
    "        mode = context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "\n",
    "        current_time = context.get_time()\n",
    "        times = context.get_abstract_state(int(self._times_index)).get_value()\n",
    "\n",
    "        if mode == ControllerState.WAIT_FOR_OBJECTS_TO_SETTLE:\n",
    "            # if context.get_time() - times[\"initial\"] > 0.5:\n",
    "            self.Plan(context, state)\n",
    "            return\n",
    "        elif mode == ControllerState.GO_HOME:\n",
    "            traj_q = context.get_mutable_abstract_state(\n",
    "                int(self._traj_q_index)\n",
    "            ).get_value()\n",
    "            if not traj_q.is_time_in_range(current_time):\n",
    "                self.Plan(context, state)\n",
    "            return\n",
    "        elif mode == ControllerState.PICKING_FROM_TABLE:\n",
    "            traj_X_G = context.get_abstract_state(int(self._traj_X_G_index)).get_value()\n",
    "            if not traj_X_G.is_time_in_range(current_time):\n",
    "                state.get_mutable_abstract_state(int(self._mode_index)).set_value(ControllerState.PREP_THROW)\n",
    "                #self.GoHome(context, state)\n",
    "                self.Plan(context, state)\n",
    "                return\n",
    "            X_G = self.get_input_port(0).Eval(context)[int(self._gripper_body_index)]\n",
    "            if (\n",
    "                np.linalg.norm(\n",
    "                    traj_X_G.GetPose(current_time).translation() - X_G.translation()\n",
    "                )\n",
    "                > 0.2\n",
    "            ):\n",
    "                # If my trajectory tracking has gone this wrong, then I'd better\n",
    "                # stop and replan.\n",
    "                self.GoHome(context, state)\n",
    "                return\n",
    "        elif mode == ControllerState.PREP_THROW: \n",
    "            #\n",
    "            traj_X_G = context.get_abstract_state(int(self._traj_X_G_index)).get_value()\n",
    "            if not traj_X_G.is_time_in_range(current_time):\n",
    "                #print(\"hello from if statement PREP_THROW in update\")\n",
    "                state.get_mutable_abstract_state(int(self._mode_index)).set_value(ControllerState.THROWING)\n",
    "                self.Plan(context, state)\n",
    "                #self.GoHome(context, state)\n",
    "                \n",
    "                return\n",
    "        elif mode == ControllerState.THROWING: \n",
    "            #print(\"aloha from THROWING\")\n",
    "            traj_X_G = context.get_abstract_state(int(self._traj_X_G_index)).get_value()\n",
    "            if not traj_X_G.is_time_in_range(current_time):\n",
    "                state.get_mutable_abstract_state(int(self._mode_index)).set_value(ControllerState.GO_HOME)\n",
    "                self.GoHome(context, state)\n",
    "                #self.Plan(context, state)\n",
    "                return\n",
    "\n",
    "        # Add this back once we pick up the ball\n",
    "        # If we are between pick and place and the gripper is closed, then\n",
    "        # we've missed or dropped the object.  Time to replan.\n",
    "        if \"pick_end\" in times and current_time > times[\"pick_end\"]:\n",
    "            if \"release\" not in times or current_time < times[\"release\"]:\n",
    "                wsg_state = self.get_input_port(self._wsg_state_index).Eval(context)\n",
    "                if wsg_state[0] < 0.01:\n",
    "                    print(\"Ball dropped. Go back to ball pick up.\")\n",
    "                    state.get_mutable_abstract_state(int(self._mode_index)).set_value(\n",
    "                        ControllerState.WAIT_FOR_OBJECTS_TO_SETTLE\n",
    "                    )\n",
    "                    times = {\"initial\": current_time}\n",
    "                    state.get_mutable_abstract_state(int(self._times_index)).set_value(\n",
    "                        times\n",
    "                    )\n",
    "                    X_G = self.get_input_port(0).Eval(context)[\n",
    "                        int(self._gripper_body_index)\n",
    "                    ]\n",
    "                    state.get_mutable_abstract_state(int(self._traj_X_G_index)).set_value(\n",
    "                        PiecewisePose.MakeLinear([current_time, np.inf], [X_G, X_G])\n",
    "                    )\n",
    "                    return\n",
    "\n",
    "    def GoHome(self, context, state):\n",
    "        print(\"Replanning due to large tracking error.\")\n",
    "        state.get_mutable_abstract_state(int(self._mode_index)).set_value(\n",
    "            ControllerState.GO_HOME\n",
    "        )\n",
    "        q = self.get_input_port(self._iiwa_position_index).Eval(context)\n",
    "        q0 = copy(context.get_discrete_state(self._q0_index).get_value())\n",
    "        q0[0] = q[0]  # Safer to not reset the first joint.\n",
    "\n",
    "        current_time = context.get_time()\n",
    "        q_traj = PiecewisePolynomial.FirstOrderHold(\n",
    "            [current_time, current_time + 5.0], np.vstack((q, q0)).T\n",
    "        )\n",
    "        state.get_mutable_abstract_state(int(self._traj_q_index)).set_value(q_traj)\n",
    "\n",
    "    def Plan(self, context, state):\n",
    "        mode = copy(state.get_mutable_abstract_state(int(self._mode_index)).get_value())\n",
    "\n",
    "        initial_pose = self.get_input_port(0).Eval(context)[\n",
    "                int(self._gripper_body_index)\n",
    "            ]\n",
    "\n",
    "        # Pose of the cup, use for calculating toss position!\n",
    "        cup_pose = self.get_input_port(self._cup_pose_index).Eval(context)\n",
    "        times = context.get_abstract_state(int(self._times_index)).get_value()\n",
    "\n",
    "        # print(\"cup pose position:\", cup_pose.translation())\n",
    "        # print(\"cup pose orientation (rpy):\", RollPitchYaw(cup_pose.rotation()).vector())\n",
    "\n",
    "        if mode == ControllerState.WAIT_FOR_OBJECTS_TO_SETTLE:\n",
    "            cost, pick_up_pose = self.get_input_port(self._grasp_index).Eval(context)\n",
    "            assert not np.isinf(cost), \"Could not find a valid grasp\"\n",
    "            state.get_mutable_abstract_state(int(self._mode_index)).set_value(ControllerState.PICKING_FROM_TABLE)\n",
    "            mode = ControllerState.PICKING_FROM_TABLE  # Update the mode variable.\n",
    "            \n",
    "        if mode == ControllerState.PICKING_FROM_TABLE:\n",
    "            cost, pick_up_pose = self.get_input_port(self._grasp_index).Eval(context)\n",
    "            # Prepare prepick pose\n",
    "            pre_pick_pose = pick_up_pose @ RigidTransform([0, -0.1, 0.0])\n",
    "\n",
    "            # print(\"Pre-pick pose position:\", pre_pick_pose.translation())\n",
    "            # print(\"Pre-pick pose orientation (rpy):\", RollPitchYaw(pre_pick_pose.rotation()).vector())\n",
    "\n",
    "            start_time = context.get_time()\n",
    "            times['start_robot'] = start_time\n",
    "            times['prepick'] = start_time + 5\n",
    "            times['pick_start'] = start_time + 7\n",
    "            times['pick_end'] = start_time + 9\n",
    "            times['pick_up_end'] = start_time + 11\n",
    "\n",
    "            # Interpolate pose for entry\n",
    "            entry_rot, entry_tran = piecewise_traj(initial_pose, pre_pick_pose, times['start_robot'], times['prepick'])\n",
    "            # Interpolate pose for pickup\n",
    "            pre_pick_rot, pre_pick_tran = piecewise_traj(pre_pick_pose, pick_up_pose, times['prepick'], times['pick_start'])\n",
    "            # Interpolate pose for ready position\n",
    "            pick_rot, pick_tran = piecewise_traj(pick_up_pose, pick_up_pose, times['pick_start'], times['pick_end'])\n",
    "            # Interpolate pose for ready position\n",
    "            ready_rot, ready_tran = piecewise_traj(pick_up_pose, pre_pick_pose, times['pick_end'], times['pick_up_end'])\n",
    "            \n",
    "            t_lst = np.linspace(times['start_robot'], times['pick_up_end'], 30)\n",
    "            pose_lst = []\n",
    "            for t in t_lst:\n",
    "                if t <= times['prepick']: # Prepick\n",
    "                    rot_traj = entry_rot\n",
    "                    tran_traj = entry_tran\n",
    "                elif t > times['prepick'] and t <= times['pick_start']:\n",
    "                    rot_traj = pre_pick_rot\n",
    "                    tran_traj = pre_pick_tran\n",
    "                elif t > times['pick_start'] and t <= times['pick_end']:\n",
    "                    rot_traj = pick_rot\n",
    "                    tran_traj = pick_tran\n",
    "                elif t > times['pick_end'] and t <= times['pick_up_end']:\n",
    "                    rot_traj = ready_rot\n",
    "                    tran_traj = ready_tran\n",
    "                else:\n",
    "                    assert False, \"Unexpected time\"\n",
    "                pose_lst.append(RigidTransform(\n",
    "                    RotationMatrix(rot_traj.orientation(t)),\n",
    "                    tran_traj.value(t),\n",
    "                ))\n",
    "                #AddMeshcatTriad(meshcat, path=str(t), X_PT=pose_lst[-1], opacity=0.2)\n",
    "            \n",
    "            # Visualize our end-effector nominal trajectory.\n",
    "            # AddMeshcatTriad(meshcat, path=str(t_lst[-1]), X_PT=pick_up_pose, opacity=0.2)\n",
    "\n",
    "            # q_knots = np.array(create_q_knots(pose_lst))\n",
    "            # q_traj = PiecewisePolynomial.CubicShapePreserving(t_lst, q_knots[:, 0:7].T)\n",
    "            q_traj = PiecewisePose.MakeLinear(t_lst, pose_lst)\n",
    "\n",
    "            opened = np.array([0.107])\n",
    "            closed = np.array([0.0])\n",
    "\n",
    "            traj_wsg_command = PiecewisePolynomial.FirstOrderHold(\n",
    "                [times['start_robot'], times['pick_start']],\n",
    "                np.hstack([[opened], [opened]]),\n",
    "            )\n",
    "            traj_wsg_command.AppendFirstOrderSegment(times['pick_end'], closed)\n",
    "            traj_wsg_command.AppendFirstOrderSegment(times['pick_up_end'] + 1, closed)\n",
    "\n",
    "            state.get_mutable_abstract_state(int(self._times_index)).set_value(times)\n",
    "            state.get_mutable_abstract_state(int(self._traj_X_G_index)).set_value(q_traj)\n",
    "            state.get_mutable_abstract_state(int(self._traj_wsg_index)).set_value(traj_wsg_command)\n",
    "\n",
    "        if mode == ControllerState.PREP_THROW: \n",
    "                # Get current gripper pose in world frame\n",
    "            X_current = self.get_input_port(0).Eval(context)[int(self._gripper_body_index)]\n",
    "\n",
    "            # Desired end-effector orientation (R) and position (p)\n",
    "            R_target = RotationMatrix(RollPitchYaw([-0.7853981, 0.04184365, 0.0])) #-0.6040839\n",
    "            target_pose = RigidTransform(R_target, [0.0, 0.80, 0.45]) #0.63893282\n",
    "\n",
    "            #AddMeshcatTriad(meshcat, path=\"temp\", X_PT=target_pose, opacity=0.8)\n",
    "\n",
    "            current_time = context.get_time()\n",
    "            times['prep_throw_start'] = current_time\n",
    "            times['pre_throw_end'] = times['prep_throw_start'] + 6.0\n",
    "            \n",
    "            # Create a pose trajectory from current pose to the desired target pose\n",
    "            traj_X_G = PiecewisePose.MakeLinear(\n",
    "                [current_time, times['pre_throw_end']],\n",
    "                [X_current, target_pose]\n",
    "            )\n",
    "\n",
    "            state.get_mutable_abstract_state(int(self._traj_X_G_index)).set_value(traj_X_G)\n",
    "\n",
    "            # For the gripper, you can keep it closed if you're holding the ball\n",
    "            closed = np.array([0.0])\n",
    "            traj_wsg = PiecewisePolynomial.FirstOrderHold([current_time, times['pre_throw_end'] + 1.0],\n",
    "                                                        np.hstack([[closed], [closed]]))\n",
    "            state.get_mutable_abstract_state(int(self._traj_wsg_index)).set_value(traj_wsg)\n",
    "            state.get_mutable_abstract_state(int(self._times_index)).set_value(times)           \n",
    "\n",
    "        if mode == ControllerState.THROWING: \n",
    "\n",
    "            X_current = self.get_input_port(0).Eval(context)[int(self._gripper_body_index)]\n",
    "\n",
    "            # Desired release pose and orientation\n",
    "            # Orientation can remain similar to current or slightly angled\n",
    "            R_target = RotationMatrix(RollPitchYaw([-2.356194, 0.04184365, 0.0])) \n",
    "            X_release = RigidTransform(R_target, [0.0, 0.5, 0.57]) #z = 0.55\n",
    "\n",
    "            # We'll give ourselves some time to move to the release pose.\n",
    "            current_time = context.get_time()\n",
    "            release_duration = 0.2 # time to move hand into release position\n",
    "            times['throw_start'] = current_time\n",
    "            times['release'] = times['throw_start'] + release_duration\n",
    "\n",
    "            # Create a pose trajectory from current pose to the release pose\n",
    "            # traj_X_G = PiecewisePose.MakeLinear(\n",
    "            #     [current_time, release_time],\n",
    "            #     [X_current, X_release]\n",
    "            # )\n",
    "\n",
    "            # Here, let's just approximate by adding another keyframe slightly after release_time.\n",
    "            # If delta_t is small, position change ~ velocity * delta_t.\n",
    "            delta_t = 0.05\n",
    "            desired_v = np.array([0.0, -0.1, .25])  # m/s\n",
    "            next_pose = RigidTransform(\n",
    "                R_target,\n",
    "                X_release.translation() + desired_v * delta_t\n",
    "            )\n",
    "\n",
    "            times_t = [times['throw_start'], times['release'], times['release'] + delta_t]\n",
    "            poses = [X_current, X_release, next_pose]\n",
    "            traj_X_G = PiecewisePose.MakeLinear(times_t, poses)\n",
    "\n",
    "\n",
    "            closed = np.array([0.0])\n",
    "            opened = np.array([0.107])\n",
    "            traj_wsg = PiecewisePolynomial.FirstOrderHold(\n",
    "                [times['throw_start'], times['release']], np.hstack([[closed], [closed]])\n",
    "            )\n",
    "            # Open right after release\n",
    "            traj_wsg.AppendFirstOrderSegment(times['release'] + 0.01, opened)\n",
    "\n",
    "            # Update state\n",
    "            state.get_mutable_abstract_state(int(self._traj_X_G_index)).set_value(traj_X_G)\n",
    "            state.get_mutable_abstract_state(int(self._traj_wsg_index)).set_value(traj_wsg)\n",
    "            state.get_mutable_abstract_state(int(self._times_index)).set_value(times)\n",
    "\n",
    "            sample_times = np.linspace(times['throw_start'], times['release'] + delta_t, 10)\n",
    "            for i, t_i in enumerate(sample_times):\n",
    "                X_t = traj_X_G.GetPose(t_i)\n",
    "                AddMeshcatTriad(meshcat, path=f\"throwing_traj/sample_{i}\", X_PT=X_t, length=0.05, radius=0.002, opacity=0.4)\n",
    "\n",
    "\n",
    "    def CalcControlMode(self, context, output):\n",
    "        mode = context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "\n",
    "        if mode == ControllerState.GO_HOME :\n",
    "            output.set_value(InputPortIndex(2))  # iiwa_position_command\n",
    "        else:\n",
    "            output.set_value(InputPortIndex(1))  # Diff IK\n",
    "\n",
    "    def CalcDiffIKReset(self, context, output):\n",
    "        mode = context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "\n",
    "        if mode == ControllerState.GO_HOME:\n",
    "            output.set_value(True)\n",
    "        else:\n",
    "            output.set_value(False)\n",
    "\n",
    "    def CalcGripperPose(self, context, output):\n",
    "        context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "\n",
    "        traj_X_G = context.get_abstract_state(int(self._traj_X_G_index)).get_value()\n",
    "        if traj_X_G.get_number_of_segments() > 0 and traj_X_G.is_time_in_range(\n",
    "            context.get_time()\n",
    "        ):\n",
    "            # Evaluate the trajectory at the current time, and write it to the\n",
    "            # output port.\n",
    "            output.set_value(\n",
    "                context.get_abstract_state(int(self._traj_X_G_index))\n",
    "                .get_value()\n",
    "                .GetPose(context.get_time())\n",
    "            )\n",
    "            return\n",
    "\n",
    "        # Command the current position (note: this is not particularly good if the velocity is non-zero)\n",
    "        output.set_value(\n",
    "            self.get_input_port(0).Eval(context)[int(self._gripper_body_index)]\n",
    "        )\n",
    "        \n",
    "    def CalcWsgPosition(self, context, output):\n",
    "        mode = context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "        opened = np.array([0.107])\n",
    "        closed = np.array([0.0])\n",
    "\n",
    "        if mode == ControllerState.GO_HOME:\n",
    "            # Command the open position\n",
    "            output.SetFromVector([opened])\n",
    "            return\n",
    "\n",
    "        traj_wsg = context.get_abstract_state(int(self._traj_wsg_index)).get_value()\n",
    "        if traj_wsg.get_number_of_segments() > 0 and traj_wsg.is_time_in_range(\n",
    "            context.get_time()\n",
    "        ):\n",
    "            # Evaluate the trajectory at the current time, and write it to the\n",
    "            # output port.\n",
    "            output.SetFromVector(traj_wsg.value(context.get_time()))\n",
    "            return\n",
    "\n",
    "        # Command the open position\n",
    "        output.SetFromVector([opened])\n",
    "\n",
    "    def CalcIiwaPosition(self, context, output):\n",
    "        traj_q = context.get_mutable_abstract_state(int(self._traj_q_index)).get_value()\n",
    "        t = context.get_time()\n",
    "\n",
    "        if traj_q.get_number_of_segments() == 0 or not traj_q.is_time_in_range(t):\n",
    "            #q0 = context.get_discrete_state(int(self._q0_index)).get_value()\n",
    "            #output.SetFromVector(q0)\n",
    "            pass\n",
    "        else : \n",
    "            output.SetFromVector(traj_q.value(context.get_time()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Grasp Selector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Another diagram for the objects the robot \"knows about\": gripper, cameras and table \n",
    "# Think of this as the model in the robot's head.\n",
    "def make_internal_model():\n",
    "    builder = DiagramBuilder()\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    parser.package_map().AddPackageXml(filename=\"./package.xml\")\n",
    "    parser.AddModelsFromUrl(\"package://bot_shot/models/pingpong_planning.dmd.yaml\")\n",
    "    plant.Finalize()\n",
    "    return builder.Build()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_grasp_candidate(X_G, prefix=\"gripper\", draw_frames=True):\n",
    "    builder = DiagramBuilder()\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    parser.AddModelsFromUrl(\"package://manipulation/schunk_wsg_50_welded_fingers.sdf\")\n",
    "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"body\"), X_G)\n",
    "    AddMeshcatTriad(meshcat, path=str(5), X_PT=X_G, opacity=0.7)\n",
    "    plant.Finalize()\n",
    "\n",
    "    # frames_to_draw = {\"gripper\": {\"body\"}} if draw_frames else {}\n",
    "    params = MeshcatVisualizerParams()\n",
    "    params.prefix = prefix\n",
    "    params.delete_prefix_on_initialization_event = False\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat, params)\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    diagram.ForcedPublish(context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Takes ball point clouds (in world coordinates) as input, and outputs and estimated pose for the mustard bottle.\n",
    "class GraspSelector(LeafSystem):\n",
    "    def __init__(self):\n",
    "        LeafSystem.__init__(self)\n",
    "        self.DeclareAbstractInputPort(\"ball_pcd\", AbstractValue.Make(PointCloud(0)))\n",
    "        port = self.DeclareAbstractOutputPort(\n",
    "            \"grasp_selection\",\n",
    "            lambda: AbstractValue.Make((np.inf, RigidTransform())),\n",
    "            self.SelectGrasp,\n",
    "        )\n",
    "        port.disable_caching_by_default()\n",
    "\n",
    "        self._internal_model = make_internal_model()\n",
    "        self._internal_model_context = self._internal_model.CreateDefaultContext()\n",
    "        \n",
    "        self._rng = np.random.default_rng(27)\n",
    "\n",
    "    def SelectGrasp(self, context, output):\n",
    "        point_cloud = self.get_input_port().Eval(context)\n",
    "\n",
    "        min_cost = np.inf\n",
    "        best_X_G = None\n",
    "        for i in range(100):\n",
    "            cost, X_G = GenerateAntipodalGraspCandidate(\n",
    "                self._internal_model,\n",
    "                self._internal_model_context, \n",
    "                point_cloud, \n",
    "                self._rng)\n",
    "            if np.isfinite(cost) and cost < min_cost:\n",
    "                min_cost = cost\n",
    "                best_X_G = X_G\n",
    "        \n",
    "        if best_X_G is None:\n",
    "            assert False, \"Cannot find a valid antipodal grasp\"\n",
    "        else:\n",
    "            #draw_grasp_candidate(best_X_G, draw_frames=False)\n",
    "            # Update gripper pose since in the scenario file, the gripper has been rotated\n",
    "            # print(best_X_G.rotation().ToRollPitchYaw())\n",
    "            R_G = RotationMatrix.MakeXRotation(\n",
    "                -np.pi / 2.0\n",
    "            ) @ RotationMatrix.MakeZRotation(np.pi / 2.0)\n",
    "            #best_X_G = best_X_G @ RigidTransform(R_G, [0, 0, 0])\n",
    "            output.set_value((min_cost, best_X_G))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Pose Estimation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PoseEstimator(LeafSystem):\n",
    "    # Use point cloud from camera to estimate pose of the cup in the world frame\n",
    "    def __init__(self, prefix):\n",
    "        LeafSystem.__init__(self)\n",
    "        self.DeclareAbstractInputPort(f\"{prefix}_pcd\", AbstractValue.Make(PointCloud(0)))\n",
    "        self.DeclareAbstractOutputPort(\"X_WO\", \n",
    "                                       lambda: AbstractValue.Make(RigidTransform()),\n",
    "                                       self.get_pose_estimate)\n",
    "        \n",
    "    def get_pose_estimate(self, context, output):\n",
    "        # Scene pcd\n",
    "        scene_pcd = self.get_input_port().Eval(context)\n",
    "\n",
    "        # Read the point cloud from a .ply file - Model pcd\n",
    "        plydata = PlyData.read(\"./models/cup/cup.ply\")\n",
    "\n",
    "        # Access point coordinates\n",
    "        x = plydata['vertex']['x']\n",
    "        y = plydata['vertex']['y']\n",
    "        z = plydata['vertex']['z']\n",
    "        model_pcd = np.stack([x, y, z], axis=-1).T\n",
    "        \n",
    "        cloud = PointCloud(model_pcd.shape[1], Fields(BaseField.kXYZs | BaseField.kRGBs))\n",
    "        cloud.mutable_xyzs()[:] = model_pcd\n",
    "        cloud = cloud.VoxelizedDownSample(voxel_size=0.005)\n",
    "        \n",
    "        X_WOhat, chat = IterativeClosestPoint(\n",
    "            cloud.xyzs(),\n",
    "            scene_pcd.xyzs(),\n",
    "            meshcat=meshcat,\n",
    "            meshcat_scene_path=\"icp_scene\",\n",
    "        )\n",
    "        #print(X_WOhat)\n",
    "        #X_WOhat = RigidTransform() #TODO remove and uncomment above\n",
    "        output.set_value(X_WOhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Camera System"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from urllib.request import urlretrieve\n",
    "import torch\n",
    "import torch.utils.data\n",
    "import torchvision  \n",
    "import torchvision.transforms.functional as Tf\n",
    "\n",
    "from torchvision.models.detection import MaskRCNN_ResNet50_FPN_Weights\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\n",
    "\n",
    "import ssl\n",
    "ssl._create_default_https_context = ssl._create_unverified_context\n",
    "CUP_IDX = 6\n",
    "BALL_IDX = 7\n",
    "\n",
    "\n",
    "model_file = \"clutter_maskrcnn_model.pt\"\n",
    "# if not os.path.exists(model_file):\n",
    "#     urlretrieve(\n",
    "#         \"https://groups.csail.mit.edu/locomotion/clutter_maskrcnn_model.pt\",\n",
    "#         model_file,\n",
    "#     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MaskRCNN(\n",
       "  (transform): GeneralizedRCNNTransform(\n",
       "      Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
       "      Resize(min_size=(800,), max_size=1333, mode='bilinear')\n",
       "  )\n",
       "  (backbone): BackboneWithFPN(\n",
       "    (body): IntermediateLayerGetter(\n",
       "      (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "      (bn1): FrozenBatchNorm2d(64, eps=0.0)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "      (layer1): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (2): Bottleneck(\n",
       "          (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "      (layer2): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (2): Bottleneck(\n",
       "          (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (3): Bottleneck(\n",
       "          (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "      (layer3): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (2): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (3): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (4): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (5): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "      (layer4): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(2048, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): FrozenBatchNorm2d(2048, eps=0.0)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(2048, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (2): Bottleneck(\n",
       "          (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(2048, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (fpn): FeaturePyramidNetwork(\n",
       "      (inner_blocks): ModuleList(\n",
       "        (0): Conv2dNormActivation(\n",
       "          (0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (1): Conv2dNormActivation(\n",
       "          (0): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (2): Conv2dNormActivation(\n",
       "          (0): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (3): Conv2dNormActivation(\n",
       "          (0): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (layer_blocks): ModuleList(\n",
       "        (0-3): 4 x Conv2dNormActivation(\n",
       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (extra_blocks): LastLevelMaxPool()\n",
       "    )\n",
       "  )\n",
       "  (rpn): RegionProposalNetwork(\n",
       "    (anchor_generator): AnchorGenerator()\n",
       "    (head): RPNHead(\n",
       "      (conv): Sequential(\n",
       "        (0): Conv2dNormActivation(\n",
       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "      (cls_logits): Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (bbox_pred): Conv2d(256, 12, kernel_size=(1, 1), stride=(1, 1))\n",
       "    )\n",
       "  )\n",
       "  (roi_heads): RoIHeads(\n",
       "    (box_roi_pool): MultiScaleRoIAlign(featmap_names=['0', '1', '2', '3'], output_size=(7, 7), sampling_ratio=2)\n",
       "    (box_head): TwoMLPHead(\n",
       "      (fc6): Linear(in_features=12544, out_features=1024, bias=True)\n",
       "      (fc7): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "    )\n",
       "    (box_predictor): FastRCNNPredictor(\n",
       "      (cls_score): Linear(in_features=1024, out_features=9, bias=True)\n",
       "      (bbox_pred): Linear(in_features=1024, out_features=36, bias=True)\n",
       "    )\n",
       "    (mask_roi_pool): MultiScaleRoIAlign(featmap_names=['0', '1', '2', '3'], output_size=(14, 14), sampling_ratio=2)\n",
       "    (mask_head): MaskRCNNHeads(\n",
       "      (0): Conv2dNormActivation(\n",
       "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): ReLU(inplace=True)\n",
       "      )\n",
       "      (1): Conv2dNormActivation(\n",
       "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): ReLU(inplace=True)\n",
       "      )\n",
       "      (2): Conv2dNormActivation(\n",
       "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): ReLU(inplace=True)\n",
       "      )\n",
       "      (3): Conv2dNormActivation(\n",
       "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (1): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (mask_predictor): MaskRCNNPredictor(\n",
       "      (conv5_mask): ConvTranspose2d(256, 256, kernel_size=(2, 2), stride=(2, 2))\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (mask_fcn_logits): Conv2d(256, 9, kernel_size=(1, 1), stride=(1, 1))\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_instance_segmentation_model(num_classes):\n",
    "    # load an instance segmentation model pre-trained on COCO\n",
    "    model = torchvision.models.detection.maskrcnn_resnet50_fpn(\n",
    "        weights=MaskRCNN_ResNet50_FPN_Weights.DEFAULT, progress=False\n",
    "    )\n",
    "\n",
    "    # get the number of input features for the classifier\n",
    "    in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
    "    # replace the pre-trained head with a new one\n",
    "    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n",
    "\n",
    "    # now get the number of input features for the mask classifier\n",
    "    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels\n",
    "    hidden_layer = 256\n",
    "    # and replace the mask predictor with a new one\n",
    "    model.roi_heads.mask_predictor = MaskRCNNPredictor(\n",
    "        in_features_mask, hidden_layer, num_classes\n",
    "    )\n",
    "\n",
    "    return model\n",
    "\n",
    "num_classes = 9\n",
    "model = get_instance_segmentation_model(num_classes)\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "model.load_state_dict(torch.load(\"clutter_maskrcnn_model.pt\", map_location=device))\n",
    "model.eval()\n",
    "\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def AddExtractPose(scenario, station, builder, poses_output_port=None):\n",
    "    \"\"\"\n",
    "    Adds one ExtractPose system to the `builder` for each camera in `scenario`, and connects it to the respective camera station output ports.\n",
    "\n",
    "    Args:\n",
    "        scenario: A Scenario structure, populated using the `LoadScenario` method.\n",
    "        station: A HardwareStation system (e.g. from MakeHardwareStation) that has already been added to `builder`.\n",
    "        builder: The DiagramBuilder containing `station` into which the new systems will be added.\n",
    "        poses_output_port: (optional) HardwareStation will have a body_poses output port iff it was created with `hardware=False`. Alternatively, one could create a MultibodyPositionsToGeometryPoses system to consume the position measurements; this optional input can be used to support that workflow.\n",
    "\n",
    "    Returns:\n",
    "        A mapping from camera name to the ExtractPose system.\n",
    "    \"\"\"\n",
    "    to_extract_pose = dict()\n",
    "    for _, config in scenario.cameras.items():\n",
    "        if not config.depth:\n",
    "            return\n",
    "\n",
    "        plant = station.GetSubsystemByName(\"plant\")\n",
    "        # frame names in local variables:\n",
    "        # P for parent frame, B for base frame, C for camera frame.\n",
    "\n",
    "        # Extract the camera extrinsics from the config struct.\n",
    "        P = (\n",
    "            GetScopedFrameByName(plant, config.X_PB.base_frame)\n",
    "            if config.X_PB.base_frame\n",
    "            else plant.world_frame()\n",
    "        )\n",
    "        X_PC = config.X_PB.GetDeterministicValue()\n",
    "\n",
    "        # convert mbp frame to geometry frame\n",
    "        body = P.body()\n",
    "        plant.GetBodyFrameIdIfExists(body.index())\n",
    "        # assert body_frame_id.has_value()\n",
    "\n",
    "        X_BP = P.GetFixedPoseInBodyFrame()\n",
    "        X_BC = X_BP @ X_PC\n",
    "\n",
    "        if poses_output_port is None:\n",
    "            # Note: this is a cheat port; it will only work in single process\n",
    "            # mode.\n",
    "            poses_output_port = station.GetOutputPort(\"body_poses\")\n",
    "\n",
    "        camera_pose = builder.AddSystem(ExtractPose(int(body.index()), X_BC))\n",
    "        to_extract_pose[config.name] = camera_pose\n",
    "        camera_pose.set_name(f\"{config.name}.pose\")\n",
    "        builder.Connect(\n",
    "            poses_output_port,\n",
    "            camera_pose.get_input_port(),\n",
    "        )\n",
    "\n",
    "    return to_extract_pose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from manipulation.mustard_depth_camera_example import MustardPointCloud\n",
    "class CameraSystem(LeafSystem):\n",
    "    def __init__(self, start_idx, num_of_cameras, output_prefix, station):\n",
    "        LeafSystem.__init__(self)\n",
    "        rgb_img = AbstractValue.Make(ImageRgba8U())\n",
    "        depth_img = AbstractValue.Make(ImageDepth32F())\n",
    "        for i in range(start_idx, start_idx + num_of_cameras):\n",
    "            self.DeclareAbstractInputPort(f\"camera{i}_depth_image\", depth_img)\n",
    "            self.DeclareAbstractInputPort(f\"camera{i}_rgb_image\", rgb_img)\n",
    "            self.DeclareAbstractInputPort(f\"camera{i}_pose\", AbstractValue.Make(RigidTransform()))\n",
    "\n",
    "        pcd_port = self.DeclareAbstractOutputPort(\n",
    "            f\"{output_prefix}_pcd\",\n",
    "            lambda: AbstractValue.Make(PointCloud(0)),\n",
    "            self.get_merged_masked_pcd,\n",
    "        )\n",
    "\n",
    "        self.station = station\n",
    "        self.output_prefix = output_prefix\n",
    "        self.num_of_cameras = num_of_cameras\n",
    "        self.start_idx = start_idx\n",
    "        self.depth_imgs = []\n",
    "        self.rgb_imgs = []\n",
    "        self.X_WCs = []\n",
    "        self.cam_infos = []\n",
    "\n",
    "    def project_depth_to_pC(self, idx, depth_pixel):\n",
    "        \"\"\"\n",
    "        project depth pixels to points in camera frame\n",
    "        using pinhole camera model\n",
    "        Input:\n",
    "            depth_pixels: numpy array of (nx3) or (3,)\n",
    "        Output:\n",
    "            pC: 3D point in camera frame, numpy array of (nx3)\n",
    "        \"\"\"\n",
    "        # switch u,v due to python convention\n",
    "        cam_info = self.cam_infos[idx]\n",
    "        v = depth_pixel[:, 0]\n",
    "        u = depth_pixel[:, 1]\n",
    "        Z = depth_pixel[:, 2]\n",
    "        cx = cam_info.center_x()\n",
    "        cy = cam_info.center_y()\n",
    "        fx = cam_info.focal_x()\n",
    "        fy = cam_info.focal_y()\n",
    "        X = (u - cx) * Z / fx\n",
    "        Y = (v - cy) * Z / fy\n",
    "        pC = np.c_[X, Y, Z]\n",
    "        return pC\n",
    "\n",
    "    def read_images(self, cam_context):\n",
    "        # Read images\n",
    "        for idx in range(self.start_idx, self.start_idx + self.num_of_cameras):\n",
    "            depth_im_read = (\n",
    "                self.GetInputPort(\"camera{}_depth_image\".format(idx))\n",
    "                .Eval(cam_context)\n",
    "                .data.squeeze()\n",
    "            )\n",
    "            depth_im = deepcopy(depth_im_read)\n",
    "            depth_im[depth_im == np.inf] = 10.0\n",
    "            self.depth_imgs.append(depth_im)\n",
    "            self.rgb_imgs.append(\n",
    "                self.GetInputPort(\"camera{}_rgb_image\".format(idx)).Eval(cam_context).data\n",
    "            )\n",
    "\n",
    "            # # Visualize\n",
    "            # point_cloud = diagram.GetOutputPort(\"camera{}_point_cloud\".format(idx)).Eval(\n",
    "            #     diagram_context\n",
    "            # )\n",
    "            # meshcat.SetObject(f\"Camera {idx} point cloud\", point_cloud)\n",
    "\n",
    "            # Get other info about the camera\n",
    "            cam = self.station.GetSubsystemByName(\"rgbd_sensor_camera\" + str(idx))\n",
    "            self.X_WCs.append(self.GetInputPort(\"camera{}_pose\".format(idx)).Eval(cam_context))\n",
    "            self.cam_infos.append(cam.default_depth_render_camera().core().intrinsics())\n",
    "\n",
    "    def predict(self):\n",
    "        with torch.no_grad():\n",
    "            predictions = []\n",
    "            for i in range(self.num_of_cameras):\n",
    "                predictions.append(\n",
    "                    model([Tf.to_tensor(self.rgb_imgs[i][:, :, :3]).to(device)])\n",
    "                )\n",
    "        for i in range(self.num_of_cameras):\n",
    "            for k in predictions[i][0].keys():\n",
    "                if k == \"masks\":\n",
    "                    predictions[i][0][k] = (\n",
    "                        predictions[i][0][k].mul(255).byte().cpu().numpy()\n",
    "                    )\n",
    "                else:\n",
    "                    predictions[i][0][k] = predictions[i][0][k].cpu().numpy()\n",
    "        return predictions\n",
    "    \n",
    "    def get_merged_masked_pcd(self, context, output):\n",
    "        \"\"\"\n",
    "        predictions: The output of the trained network (one for each camera)\n",
    "        rgb_ims: RGBA images from each camera\n",
    "        depth_ims: Depth images from each camera\n",
    "        project_depth_to_pC_funcs: Functions that perform the pinhole camera operations to convert pixels\n",
    "            into points. See the analogous function in problem 5.2 to see how to use it.\n",
    "        X_WCs: Poses of the cameras in the world frame\n",
    "        \"\"\"\n",
    "        pcd = []\n",
    "        self.read_images(context)\n",
    "        predictions = self.predict()\n",
    "        i = 0\n",
    "        if self.output_prefix == \"cup\":\n",
    "            pred_idx = CUP_IDX\n",
    "        elif self.output_prefix == \"ball\":\n",
    "            pred_idx = BALL_IDX\n",
    "        else:\n",
    "            assert False\n",
    "        for prediction, rgb_im, depth_im, X_WC in zip(\n",
    "            predictions, self.rgb_imgs, self.depth_imgs, self.X_WCs\n",
    "        ):\n",
    "            mask_idx = np.argmax(prediction[0][\"labels\"] == pred_idx)\n",
    "            prediction_mask = prediction[0][\"masks\"][mask_idx, 0] > 250 # mask_threshold\n",
    "            # if self.output_prefix == \"cup\" and i == 1:\n",
    "            #     plt.imshow(prediction_mask)\n",
    "\n",
    "            non_zero_idx = np.array(list(zip(*np.nonzero(prediction_mask))))\n",
    "            num_of_masked_points = np.sum(prediction_mask)\n",
    "\n",
    "            depth = depth_im[prediction_mask].reshape(-1, 1)\n",
    "            depth_pixels = np.concatenate((non_zero_idx, depth), axis=1)\n",
    "            \n",
    "            spatial_points = self.project_depth_to_pC(i, depth_pixels).T\n",
    "            spatial_points = X_WC @ spatial_points\n",
    "            rgb_points = rgb_im[prediction_mask, :-1].T\n",
    "\n",
    "            # You get an unhelpful RunTime error if your arrays are the wrong\n",
    "            # shape, so we'll check beforehand that they're the correct shapes.\n",
    "            assert (\n",
    "                len(spatial_points.shape) == 2\n",
    "            ), \"Spatial points is the wrong size -- should be 3 x N\"\n",
    "            assert (\n",
    "                spatial_points.shape[0] == 3\n",
    "            ), \"Spatial points is the wrong size -- should be 3 x N\"\n",
    "            assert (\n",
    "                len(rgb_points.shape) == 2\n",
    "            ), \"RGB points is the wrong size -- should be 3 x N\"\n",
    "            assert (\n",
    "                rgb_points.shape[0] == 3\n",
    "            ), \"RGB points is the wrong size -- should be 3 x N\"\n",
    "            assert rgb_points.shape[1] == spatial_points.shape[1]\n",
    "\n",
    "            N = spatial_points.shape[1]\n",
    "            pcd.append(PointCloud(N, Fields(BaseField.kXYZs | BaseField.kRGBs)))\n",
    "            pcd[-1].mutable_xyzs()[:] = spatial_points\n",
    "            pcd[-1].mutable_rgbs()[:] = rgb_points\n",
    "            # Estimate normals\n",
    "            pcd[-1].EstimateNormals(radius=0.1, num_closest=30)\n",
    "            # Flip normals toward camera\n",
    "            pcd[-1].FlipNormalsTowardPoint(X_WC.translation())\n",
    "            i += 1\n",
    "\n",
    "        # Merge point clouds.\n",
    "        merged_pcd = Concatenate(pcd)\n",
    "        # Voxelize down-sample.  (Note that the normals still look reasonable)\n",
    "        downsampled_pcd = merged_pcd.VoxelizedDownSample(voxel_size=0.005)\n",
    "        # print(downsampled_pcd)\n",
    "        # Visualize\n",
    "        \n",
    "        output.set_value(downsampled_pcd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Build RoboToss Station"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RoboTossStationSim:\n",
    "    def __init__(self):\n",
    "        builder = DiagramBuilder()\n",
    "        scenario = LoadScenario(filename=\"models/pingpong.scenario.yaml\")\n",
    "        self.station = builder.AddSystem(\n",
    "            MakeHardwareStation(scenario, meshcat=meshcat,\n",
    "                                package_xmls=[\"./package.xml\"])\n",
    "        )\n",
    "        self.plant = self.station.GetSubsystemByName(\"plant\")\n",
    "        self.scene_graph = self.station.GetSubsystemByName(\"scene_graph\")\n",
    "        AddMultibodyTriad(self.plant.GetFrameByName(\"body\"), self.scene_graph)\n",
    "\n",
    "        self.query_output_port = self.scene_graph.GetOutputPort(\"query\")\n",
    "\n",
    "        # Add system to extract camera poses\n",
    "        self.extract_pose = AddExtractPose(scenario=scenario, station=self.station, builder=builder)\n",
    "\n",
    "        # Add point cloud processing system\n",
    "        self.num_of_cameras = 3\n",
    "        self.camera_system_ball = builder.AddSystem(CameraSystem(0, self.num_of_cameras, \"ball\", self.station))\n",
    "        self.camera_system_ball.set_name(\"ball_camera_system\")\n",
    "        for i in range(self.num_of_cameras):\n",
    "            builder.Connect(\n",
    "                self.station.GetOutputPort(f\"camera{i}.depth_image\"),\n",
    "                self.camera_system_ball.GetInputPort(f\"camera{i}_depth_image\"),\n",
    "            )\n",
    "            builder.Connect(\n",
    "                self.station.GetOutputPort(f\"camera{i}.rgb_image\"),\n",
    "                self.camera_system_ball.GetInputPort(f\"camera{i}_rgb_image\"),\n",
    "            )\n",
    "            extract_pose = self.extract_pose[f\"camera{i}\"]\n",
    "            builder.Connect(\n",
    "                extract_pose.GetOutputPort(\"pose\"),\n",
    "                self.camera_system_ball.GetInputPort(f\"camera{i}_pose\"),\n",
    "            )\n",
    "        builder.ExportOutput(\n",
    "            self.camera_system_ball.GetOutputPort(\"ball_pcd\"),\n",
    "            \"ball_pcd\",\n",
    "        )\n",
    "\n",
    "        # Connect to grasp selector\n",
    "        self.grasph_selector = builder.AddSystem(GraspSelector())\n",
    "        builder.Connect(\n",
    "            self.camera_system_ball.GetOutputPort(\"ball_pcd\"), self.grasph_selector.get_input_port()\n",
    "        )       \n",
    "\n",
    "        # Add point cloud processing system for cup\n",
    "        self.camera_system_cup = builder.AddSystem(CameraSystem(self.num_of_cameras, self.num_of_cameras, \"cup\", self.station))\n",
    "        self.camera_system_cup.set_name(\"cup_camera_system\")\n",
    "        for j in range(self.num_of_cameras):\n",
    "            i = j + self.num_of_cameras\n",
    "            builder.Connect(\n",
    "                self.station.GetOutputPort(f\"camera{i}.depth_image\"),\n",
    "                self.camera_system_cup.GetInputPort(f\"camera{i}_depth_image\"),\n",
    "            )\n",
    "            builder.Connect(\n",
    "                self.station.GetOutputPort(f\"camera{i}.rgb_image\"),\n",
    "                self.camera_system_cup.GetInputPort(f\"camera{i}_rgb_image\"),\n",
    "            )\n",
    "            extract_pose = self.extract_pose[f\"camera{i}\"]\n",
    "            builder.Connect(\n",
    "                extract_pose.GetOutputPort(\"pose\"),\n",
    "                self.camera_system_cup.GetInputPort(f\"camera{i}_pose\"),\n",
    "            )\n",
    "        builder.ExportOutput(\n",
    "            self.camera_system_cup.GetOutputPort(\"cup_pcd\"),\n",
    "            \"cup_pcd\",\n",
    "        )\n",
    "\n",
    "        # Add controller to iiwa joint inputs\n",
    "        self.iiwa_controller = builder.AddSystem(ControlSystem(self.plant))\n",
    "        builder.Connect(\n",
    "            self.station.GetOutputPort(\"body_poses\"), self.iiwa_controller.GetInputPort(\"body_poses\")\n",
    "        )\n",
    "        builder.Connect(\n",
    "            self.grasph_selector.get_output_port(), self.iiwa_controller.GetInputPort(\"grasp\"),\n",
    "        )\n",
    "        builder.Connect(\n",
    "            self.station.GetOutputPort(\"wsg.state_measured\"),\n",
    "            self.iiwa_controller.GetInputPort(\"wsg_state\"),\n",
    "        )\n",
    "        builder.Connect(\n",
    "            self.station.GetOutputPort(\"iiwa.position_measured\"),\n",
    "            self.iiwa_controller.GetInputPort(\"iiwa_position\"),\n",
    "        )\n",
    "\n",
    "        # Connect to planner*********************\n",
    "         # Connect to pose estimator\n",
    "        self.icp = builder.AddSystem(PoseEstimator(\"cup\"))\n",
    "        builder.Connect(\n",
    "            self.camera_system_cup.GetOutputPort(\"cup_pcd\"), self.icp.get_input_port()\n",
    "        )\n",
    "        builder.Connect(self.icp.GetOutputPort(\"X_WO\"), self.iiwa_controller.GetInputPort(\"cup_X_WO\"))\n",
    "        builder.ExportOutput(\n",
    "            self.icp.GetOutputPort(\"X_WO\"),\n",
    "            \"X_WO\",\n",
    "        )\n",
    "        \n",
    "        # Add controller to wsg inputs\n",
    "        builder.Connect(\n",
    "            self.iiwa_controller.GetOutputPort(\"wsg_position\"),\n",
    "            self.station.GetInputPort(\"wsg.position\"),\n",
    "        )\n",
    "\n",
    "        # Set up differential inverse kinematics\n",
    "        robot = self.station.GetSubsystemByName(\"iiwa_controller_plant_pointer_system\").get() #Modify???\n",
    "        self.diff_ik = AddIiwaDifferentialIK(builder, robot, frame=robot.GetFrameByName(\"body\"))\n",
    "        ############################\n",
    "        # Modify the parameters to remove effective velocity constraints\n",
    "        params = self.diff_ik.get_mutable_parameters()\n",
    "\n",
    "        # Set extremely large velocity limits \n",
    "        iiwa14_velocity_limits = np.array([14000, 14000, 17000, 13000, 23000, 23000, 23000])\n",
    "        params.set_joint_velocity_limits((-iiwa14_velocity_limits, iiwa14_velocity_limits))\n",
    "\n",
    "        # Set high end-effector speed limits\n",
    "        params.set_end_effector_angular_speed_limit(20000)\n",
    "        params.set_end_effector_translational_velocity_limits(\n",
    "            [-20000, -20000, -20000],\n",
    "            [20000, 20000, 20000]\n",
    "        )\n",
    "\n",
    "        #########################    \n",
    "\n",
    "        builder.Connect(self.iiwa_controller.GetOutputPort(\"X_WG\"), self.diff_ik.get_input_port(0))\n",
    "        builder.Connect(\n",
    "            self.station.GetOutputPort(\"iiwa.state_estimated\"),\n",
    "            self.diff_ik.GetInputPort(\"robot_state\"),\n",
    "        )\n",
    "        # builder.Connect(\n",
    "        #     self.diff_ik.get_output_port(),\n",
    "        #     self.station.GetInputPort(\"iiwa.position\"),\n",
    "        # )\n",
    "\n",
    "        # Add this back if we need go home mode or control by velocity\n",
    "        builder.Connect(\n",
    "            self.iiwa_controller.GetOutputPort(\"reset_diff_ik\"),\n",
    "            self.diff_ik.GetInputPort(\"use_robot_state\"),\n",
    "        )      \n",
    "        \n",
    "        # The DiffIK and the direct position-control modes go through a PortSwitch\n",
    "        switch = builder.AddSystem(PortSwitch(7))\n",
    "        builder.Connect(self.diff_ik.get_output_port(), switch.DeclareInputPort(\"diff_ik\"))\n",
    "        builder.Connect(\n",
    "            self.iiwa_controller.GetOutputPort(\"iiwa_position_command\"),\n",
    "            switch.DeclareInputPort(\"position\"),\n",
    "        )\n",
    "        builder.Connect(switch.get_output_port(), self.station.GetInputPort(\"iiwa.position\"))\n",
    "        builder.Connect(\n",
    "            self.iiwa_controller.GetOutputPort(\"control_mode\"),\n",
    "            switch.get_port_selector_input_port(),\n",
    "        )\n",
    "\n",
    "        # # Add controller to wsg inputs\n",
    "        # wsg_position = builder.AddSystem(ConstantVectorSource([0.06]))\n",
    "        # builder.Connect(\n",
    "        #     wsg_position.get_output_port(), self.station.GetInputPort(\"wsg.position\")\n",
    "        # )\n",
    "\n",
    "        # Add logger\n",
    "        logger = builder.AddSystem(VectorLogSink(7))\n",
    "        builder.Connect(switch.get_output_port(), logger.get_input_port(0))\n",
    "        \n",
    "        # Add stop button\n",
    "        builder.AddSystem(StopButton(meshcat))\n",
    "        self.diagram = builder.Build()\n",
    "\n",
    "        # Store contexts\n",
    "        self.context_diagram = self.diagram.CreateDefaultContext()\n",
    "        self.context_station = self.diagram.GetSubsystemContext(\n",
    "            self.station, self.context_diagram\n",
    "        )\n",
    "        self.context_scene_graph = self.station.GetSubsystemContext(\n",
    "            self.scene_graph, self.context_station\n",
    "        )\n",
    "        self.context_plant = self.station.GetMutableSubsystemContext(\n",
    "            self.plant, self.context_station\n",
    "        )\n",
    "         \n",
    "        # Set initial position of the ball\n",
    "        ball_instance = self.plant.GetModelInstanceByName(\"ball\")\n",
    "        ball_body = self.plant.GetBodyByName(\"ball\", ball_instance)\n",
    "        self.plant.SetFreeBodyPose(self.context_plant, ball_body, RigidTransform(RollPitchYaw(0, 0, 0), [0, 0.55, 0.2825]))\n",
    "\n",
    "        soup_instance = self.plant.GetModelInstanceByName(\"soup_can\")\n",
    "        soup_body = self.plant.GetBodyByName(\"base_link_soup\", soup_instance)\n",
    "        self.plant.SetFreeBodyPose(self.context_plant, soup_body, RigidTransform(RollPitchYaw(0, 0, 15), [0.2, 0.35, 0.3])) # modify planning dmd too\n",
    "\n",
    "        cup_instance = self.plant.GetModelInstanceByName(\"cup\")\n",
    "        cup_body = self.plant.GetBodyByName(\"cup\", cup_instance)\n",
    "        self.plant.SetFreeBodyPose(self.context_plant, cup_body, RigidTransform(RollPitchYaw(0, 0, 0), [0, -0.5, 0.25]))\n",
    "\n",
    "        cracker_instance = self.plant.GetModelInstanceByName(\"cracker_box\")\n",
    "        cracker_body = self.plant.GetBodyByName(\"base_link_cracker\", cracker_instance)\n",
    "        self.plant.SetFreeBodyPose(self.context_plant, cracker_body, RigidTransform(RollPitchYaw(0, 0, 45), [0, -0.2, 0.3])) # modify planning dmd too\n",
    "\n",
    "        # import pydot\n",
    "        # from IPython.display import SVG, display\n",
    "\n",
    "        # display(SVG(pydot.graph_from_dot_data(\n",
    "        #     self.diagram.GetGraphvizString(max_depth=2))[0].create_svg()))\n",
    "        \n",
    "        simulator = Simulator(self.diagram, self.context_diagram)\n",
    "        self.logger = logger.FindLog(self.context_diagram)\n",
    "        simulator.set_target_realtime_rate(1.0)\n",
    "        meshcat.StartRecording()\n",
    "        simulator.AdvanceTo(24 if running_as_notebook else 0.1)\n",
    "        meshcat.PublishRecording()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:drake:Drake materials have been assigned to a glTF file. glTF defines its own materials, so post hoc materials will be ignored and should be removed from the model specification. glTF file: '/private/var/tmp/.cache_aprilhu/drake/package_map/bacc01fba8f324b8dce1d15ae98083fa7e61463d0d278df46f198c2e3ae14abd-d542c1cfb057a27398ab0b818609a2edb14ec9c621b58dc2d78af6daa96f8056/wsg_50_description/meshes/finger_with_tip.gltf'\n",
      "WARNING:drake:Drake materials have been assigned to a glTF file. glTF defines its own materials, so post hoc materials will be ignored and should be removed from the model specification. glTF file: '/private/var/tmp/.cache_aprilhu/drake/package_map/bacc01fba8f324b8dce1d15ae98083fa7e61463d0d278df46f198c2e3ae14abd-d542c1cfb057a27398ab0b818609a2edb14ec9c621b58dc2d78af6daa96f8056/wsg_50_description/meshes/wsg_body.gltf'\n",
      "WARNING:drake:Drake materials have been assigned to a glTF file. glTF defines its own materials, so post hoc materials will be ignored and should be removed from the model specification. glTF file: '/private/var/tmp/.cache_aprilhu/drake/package_map/bacc01fba8f324b8dce1d15ae98083fa7e61463d0d278df46f198c2e3ae14abd-d542c1cfb057a27398ab0b818609a2edb14ec9c621b58dc2d78af6daa96f8056/wsg_50_description/meshes/finger_with_tip.gltf'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Press Escape to stop the simulation\n",
      "fixing improper rotation\n",
      "fixing improper rotation\n",
      "Replanning due to large tracking error.\n",
      "fixing improper rotation\n",
      "fixing improper rotation\n",
      "fixing improper rotation\n",
      "fixing improper rotation\n",
      "fixing improper rotation\n",
      "fixing improper rotation\n",
      "fixing improper rotation\n",
      "fixing improper rotation\n",
      "fixing improper rotation\n",
      "[-3.34689636e-01  5.00000000e-01 -2.77555756e-17  1.30000000e+00\n",
      "  1.38777878e-17 -1.50000000e+00 -1.11022302e-16]\n"
     ]
    }
   ],
   "source": [
    "meshcat.Delete()\n",
    "env = RoboTossStationSim()\n",
    "\n",
    "np.set_printoptions(threshold=np.inf)\n",
    "traj = env.logger.data()\n",
    "print(traj[:, -1])\n",
    "\n",
    "downsampled_pcd = env.diagram.GetOutputPort(\"cup_pcd\").Eval(env.context_diagram)\n",
    "meshcat.SetObject(\"masked_cloud\", downsampled_pcd, point_size=0.002)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Backups"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}