tactile sensor example

notebooks/05_tactile_sensor.ipynb

@@ -0,0 +1,299 @@
+{
+ "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
+}