notebooks/05_tactile_sensor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tactile Sensor Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import airo_models\n",
    "import numpy as np\n",
    "from airo_models.primitives.sphere import sphere_geometry_dict\n",
    "\n",
    "def sensor_urdf_dict():\n",
    "\n",
    "    base_link = {\"@name\": \"base_link\"}\n",
    "\n",
    "    links = [base_link]\n",
    "    joints = []\n",
    "\n",
    "    # add a grid of spheres as links to the sensor\n",
    "    sphere_size = 0.01 # diameter\n",
    "    sphere_radius = sphere_size / 2\n",
    "\n",
    "    # sensor is a 2x8x4 grid of spheres\n",
    "    x_centers = np.linspace(-0.005, 0.005, 2) \n",
    "    y_centers = np.linspace(-0.035, 0.035, 8)\n",
    "    z_centers = np.linspace(-0.015, 0.015, 4)\n",
    "\n",
    "    sphere_geometry = sphere_geometry_dict(radius=sphere_radius)\n",
    "\n",
    "    for i, x in enumerate(x_centers):\n",
    "        for j, y in enumerate(y_centers):\n",
    "            for k, z in enumerate(z_centers):\n",
    "\n",
    "                tixel_link = {\n",
    "                    \"@name\": f\"tixel_{i}_{j}_{k}\",\n",
    "                    \"visual\": {\n",
    "                        \"geometry\": sphere_geometry,\n",
    "                    },\n",
    "                    \"collision\": {\n",
    "                        \"geometry\": sphere_geometry,\n",
    "                    }\n",
    "                }\n",
    "\n",
    "                joint = {\n",
    "                    \"@name\": f\"joint_{i}_{j}_{k}\",\n",
    "                    \"@type\": \"fixed\",\n",
    "                    \"parent\": { \"@link\": \"base_link\" },\n",
    "                    \"child\": { \"@link\": f\"tixel_{i}_{j}_{k}\" },\n",
    "                    \"origin\": {\n",
    "                        \"@xyz\": f\"{x} {y} {z}\"\n",
    "                    }\n",
    "                }\n",
    "\n",
    "                links.append(tixel_link)\n",
    "                joints.append(joint)\n",
    "                \n",
    "    urdf_dict = {\n",
    "        \"robot\": {\n",
    "            \"@name\": \"sensor\",\n",
    "            \"link\": links,\n",
    "            \"joint\": joints\n",
    "        }\n",
    "    }\n",
    "\n",
    "    return urdf_dict\n",
    "\n",
    "sensor_urdf_dict_ = sensor_urdf_dict()\n",
    "sensor_urdf_path = airo_models.urdf.write_urdf_to_tempfile(sensor_urdf_dict_, prefix=\"tactile_sensor_\")\n",
    "sensor_urdf_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.planning import RobotDiagramBuilder\n",
    "from airo_drake import add_manipulator, add_floor, add_meshcat, finish_build\n",
    "from airo_drake import SingleArmScene\n",
    "from pydrake.math import RigidTransform\n",
    "\n",
    "\n",
    "robot_diagram_builder = RobotDiagramBuilder() \n",
    "\n",
    "meshcat = add_meshcat(robot_diagram_builder)\n",
    "arm_index, gripper_index = add_manipulator(robot_diagram_builder, \"ur5e\", \"robotiq_2f_85\", static_gripper=True)\n",
    "add_floor(robot_diagram_builder)\n",
    "\n",
    "\n",
    "# adding sensor area bbox\n",
    "plant = robot_diagram_builder.plant()\n",
    "parser = robot_diagram_builder.parser()\n",
    "\n",
    "sensor_transform = RigidTransform(p=np.array([0, 0, 0.13]))\n",
    "sensor_index = parser.AddModels(sensor_urdf_path)[0]\n",
    "sensor_frame = plant.GetFrameByName(\"base_link\", sensor_index)\n",
    "\n",
    "# welding the sensor to the gripper\n",
    "gripper_frame = plant.GetFrameByName(\"base_link\", gripper_index)\n",
    "plant.WeldFrames(gripper_frame, sensor_frame, sensor_transform)\n",
    "\n",
    "\n",
    "# adding a thin (x) obstable\n",
    "\n",
    "obstacle_urdf_path = airo_models.box_urdf_path((0.01, 0.5, 0.25), \"obstacle\")\n",
    "obstacle_transform = RigidTransform(p=np.array([0.3, 0, 0.15]))\n",
    "obstacle_index = parser.AddModels(obstacle_urdf_path)[0]\n",
    "obstacle_frame = plant.GetFrameByName(\"base_link\", obstacle_index)\n",
    "plant.WeldFrames(plant.world_frame(), obstacle_frame, obstacle_transform)\n",
    "\n",
    "\n",
    "\n",
    "robot_diagram, context = finish_build(robot_diagram_builder, meshcat)\n",
    "del robot_diagram_builder # no longer needed\n",
    "\n",
    "scene = SingleArmScene(robot_diagram, arm_index, gripper_index, meshcat)\n",
    "scene"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import visualize_frame\n",
    "\n",
    "plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "gripper_body = plant.GetBodyByName(\"base_link\", gripper_index)\n",
    "gripper_pose = plant.EvalBodyPoseInWorld(plant_context, gripper_body)\n",
    "\n",
    "visualize_frame(scene.meshcat, \"tcp_frame\", gripper_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plant = scene.robot_diagram.plant()\n",
    "joints_1 = np.deg2rad([0, -64, -120, -90, 90, 90])\n",
    "\n",
    "plant_context = plant.GetMyContextFromRoot(context)\n",
    "plant.SetPositions(plant_context, scene.arm_index, joints_1)\n",
    "scene.robot_diagram.ForcedPublish(context) # updates the meshcat visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.planning import SceneGraphCollisionChecker\n",
    "\n",
    "collision_checker = SceneGraphCollisionChecker(\n",
    "    model=scene.robot_diagram,\n",
    "    robot_model_instances=[scene.arm_index, scene.gripper_index, sensor_index],\n",
    "    edge_step_size=0.125,  # Arbitrary value: we don't use the CheckEdgeCollisionFree\n",
    "    env_collision_padding=0.0,\n",
    "    self_collision_padding=0.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import list_collisions_between_bodies\n",
    "\n",
    "def print_robot_collisions(q):\n",
    "    colliding_bodies = list_collisions_between_bodies(collision_checker, q)\n",
    "\n",
    "    for colliding_body_1, colliding_body_2, is_self_collision in colliding_bodies:\n",
    "            bodyA_name = plant.get_body(colliding_body_1).scoped_name()\n",
    "            bodyB_name = plant.get_body(colliding_body_2).scoped_name()\n",
    "            print(f\"Body {bodyA_name} is colliding with body {bodyB_name} (is self collision: {is_self_collision})\")\n",
    "\n",
    "print_robot_collisions(joints_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_sensor(q): \n",
    "    colliding_bodies = list_collisions_between_bodies(collision_checker, q)\n",
    "    readout = np.zeros((2, 8, 4))\n",
    "\n",
    "    for colliding_body_1, colliding_body_2, is_self_collision in colliding_bodies:\n",
    "            bodyA_name = plant.get_body(colliding_body_1).scoped_name()\n",
    "            # bodyB_name = plant.get_body(colliding_body_2).scoped_name()\n",
    "\n",
    "            if \"tixel_\" not in str(bodyA_name): # TODO check if this assumption is OK\n",
    "                    continue\n",
    "            \n",
    "            i, j, k = [int(x) for x in str(bodyA_name).split(\"_\")[1:]]\n",
    "            readout[i, j, k] = 1\n",
    "\n",
    "    return readout\n",
    "            \n",
    "            \n",
    "readout = read_sensor(joints_1)\n",
    "readout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Create an empty color array with the same shape as readout\n",
    "colors = np.empty((8,8,8), dtype=object)\n",
    "\n",
    "for i in range(2):\n",
    "    for j in range(8):\n",
    "        for k in range(4):\n",
    "            if readout[i, j, k]:\n",
    "                colors[i, j, k] = 'red'\n",
    "            else:\n",
    "                colors[i, j, k] = 'green'\n",
    "\n",
    "# and plot everything\n",
    "ax = plt.figure().add_subplot(projection='3d')\n",
    "ax.voxels(colors, facecolors=colors, edgecolor='k', alpha=0.9)\n",
    "\n",
    "ax.set_xlabel('X')\n",
    "ax.set_ylabel('Y')\n",
    "ax.set_zlabel('Z')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "readout = read_sensor(joints_1)\n",
    "plt.imshow(readout.sum(axis=0))\n",
    "# label axes\n",
    "plt.xlabel(\"X\")\n",
    "plt.ylabel(\"Y\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cloth-competition-dev",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}