start visualization notebook

README.md

@@ -35,4 +35,4 @@ See the [`airo-mono`](https://github.com/airo-ugent/airo-mono) developer guide.
 A very similar process and tools are used for this package.
 
 ### Releasing 🏷️
-See [`airo-models`](https://github.com/airo-ugent/airo-models/tree/main),releasing `airo-drake` works the same way.
+See [`airo-models`](https://github.com/airo-ugent/airo-models/tree/main), releasing `airo-drake` works the same way.

airo_drake/__init__.py

@@ -8,6 +8,7 @@
 from airo_drake.building.manipulator import add_manipulator
 from airo_drake.building.meshcat import add_meshcat
 from airo_drake.scene import DualArmScene, SingleArmScene
+from airo_drake.visualization.frame import visualize_frame
 
 __all__ = [
     "add_floor",
@@ -16,4 +17,5 @@
     "finish_build",
     "SingleArmScene",
     "DualArmScene",
+    "visualize_frame",
 ]

airo_drake/building/manipulator.py

@@ -10,7 +10,8 @@
 # X_URBASE_ROSBASE is the 180 rotation between ROS URDF base and the UR control box base
 X_URBASE_ROSBASE = RigidTransform(rpy=RollPitchYaw([0, 0, np.pi]), p=[0, 0, 0])  # type: ignore
 
-X_URTOOL0_ROBOTIQ = RigidTransform(rpy=RollPitchYaw([0, 0, np.pi / 2]), p=[0, 0, 0])  # type: ignore
+# 180 degree rotation and 1 cm (estimate) offset for the coupling / flange
+X_URTOOL0_ROBOTIQ = RigidTransform(rpy=RollPitchYaw([0, 0, np.pi / 2]), p=[0, 0, 0.01])  # type: ignore
 
 
 def add_manipulator(

airo_drake/visualization/frame.py

@@ -0,0 +1,41 @@
+import numpy as np
+from airo_typing import HomogeneousMatrixType
+from pydrake.geometry import Cylinder, Meshcat, Rgba
+from pydrake.math import RigidTransform, RotationMatrix
+
+
+def visualize_frame(
+    meshcat: Meshcat,
+    transform: RigidTransform | HomogeneousMatrixType,
+    path: str,
+    length: float = 0.1,
+    radius: float = 0.002,
+    opacity: float = 1.0,
+) -> None:
+    """Visualizes a static frame in meshcat as an RGB (for xyz) colored triad.
+
+    Args:
+        transform: The 6D pose / transform of the frame to visualize.
+        meshcat: The MeshCat instance to add the visualization to.
+        path: A name or path for the frame in the MeshCat hierarchy, can be used to delete it later.
+        length: Length of each axis in the triad.
+        radius: Radius of each axis cylinder.
+        opacity: Opacity for the axis colors.
+    """
+
+    meshcat.SetTransform(path, transform)
+    axis_labels = ["x-axis", "y-axis", "z-axis"]
+    axis_colors = [
+        Rgba(1, 0, 0, opacity),
+        Rgba(0, 1, 0, opacity),
+        Rgba(0, 0, 1, opacity),
+    ]
+    axis_transforms = [
+        RigidTransform(RotationMatrix.MakeYRotation(np.pi / 2), [length / 2.0, 0, 0]),  # type: ignore
+        RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2), [0, length / 2.0, 0]),  # type: ignore
+        RigidTransform([0, 0, length / 2.0]),  # type: ignore
+    ]
+
+    for transform, color, label in zip(axis_transforms, axis_colors, axis_labels):
+        meshcat.SetTransform(f"{path}/{label}", transform)
+        meshcat.SetObject(f"{path}/{label}", Cylinder(radius, length), color)

notebooks/02_visualization.ipynb

@@ -0,0 +1,191 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Visualization tutorial 🎨\n",
+    "\n",
+    "This notebook shows the visualization tools of `airo-drake`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Building the Scene 🏗️"
+   ]
+  },
+  {
+   "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",
+    "\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\")\n",
+    "add_floor(robot_diagram_builder)\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": [
+    "joints = np.deg2rad([0, -45, -90, -90, 90, 0])\n",
+    "\n",
+    "plant = scene.robot_diagram.plant()\n",
+    "plant_context = plant.GetMyContextFromRoot(context)\n",
+    "\n",
+    "plant.SetPositions(plant_context, scene.arm_index, joints)\n",
+    "scene.robot_diagram.ForcedPublish(context) # updates the meshcat visualization"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plant.GetPositions(plant_context, scene.gripper_index)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "🧙‍♂️ Closing the gripper is a bit more tricky because the mimic joints are not supported in Drake: "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "v = 0.81 # magic joint value for to close the gripper\n",
+    "plant.SetPositions(plant_context, scene.gripper_index, np.array([v, -v, v] * 2))\n",
+    "scene.robot_diagram.ForcedPublish(context) # updates the meshcat visualization"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Visualizing frames 🖼️\n",
+    "\n",
+    "⏩ Let's start by calculating the forward kinematics of the robot and visualizing the TCP frame:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ur_analytic_ik import ur5e\n",
+    "\n",
+    "tcp_transform = np.identity(4)\n",
+    "tcp_transform[2, 3] = 0.175 # 175 mm in z\n",
+    "\n",
+    "tcp_pose = ur5e.forward_kinematics(*joints) @ tcp_transform\n",
+    "\n",
+    "with np.printoptions(precision=3, suppress=True):\n",
+    "    print(tcp_pose)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from airo_drake import visualize_frame\n",
+    "\n",
+    "visualize_frame(scene.meshcat, tcp_pose, \"tcp_frame\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pydrake.math import RigidTransform\n",
+    "\n",
+    "visualize_frame(scene.meshcat, RigidTransform(), \"world_frame\", length=1.0, opacity=0.2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "scene.meshcat.Delete(\"tcp_frame\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Visualizing IK solutions 🦾"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Visualizing paths 🛤️"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Visualizing trajectories 🌠"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "cloth-competition",
+   "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.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

pyproject.toml

@@ -21,6 +21,7 @@ classifiers = [
 dependencies = [
     "airo-models",
     "drake",
+    "ur-analytic-ik"
 ]
 
 [project.urls]