add visualize_arrow()

airo_drake/visualization/arrow.py

@@ -0,0 +1,88 @@
+import numpy as np
+from airo_typing import RotationMatrixType, Vector3DType
+from pydrake.geometry import Cylinder, Meshcat, MeshcatCone, Rgba
+from pydrake.math import RigidTransform, RollPitchYaw
+
+
+def orientation_with_Z_axis(Z: Vector3DType) -> RotationMatrixType:
+    Z = np.array(Z)
+    Z = Z / np.linalg.norm(Z)  # normalize just to be sure
+
+    X = np.array([1, 0, 0])
+
+    # handle case where direction is parallel to X
+    if np.all(np.abs(Z - X) < 1e-6):
+        X = np.array([0, 1, 0])
+
+    # make X orthogonal to Z
+    X = X - Z * np.dot(X, Z)
+    X = X / np.linalg.norm(X)
+
+    # make Y orthogonal to Z and X
+    Y = np.cross(Z, X)
+
+    return np.column_stack([X, Y, Z])
+
+
+def visualize_arrow(
+    meshcat: Meshcat,
+    name: str,
+    start: Vector3DType,
+    direction: Vector3DType | None = None,
+    end: Vector3DType | None = None,
+    length: float = 1.0,
+    radius: float = 0.005,
+    color: Rgba | None = None,
+) -> None:
+    """Visualizes an arrow in meshcat.
+
+    Args:
+        meshcat: The MeshCat instance to add the visualization to.
+        name: A name or path for the arrow in the MeshCat hierarchy, can be used to delete it later.
+        start: The starting point of the arrow to visualization.
+        direction: The direction of the arrow to visualize.
+        end: The end point of the arrow to visualize, mutually exclusive with direction.
+        length: Length of the arrow, not used if end is provided.
+        radius: Radius of the cylinder.
+        color: Color of the arrow.
+    """
+    start = np.array(start)
+
+    if direction is None and end is None:
+        raise ValueError("To visualize an arrow either direction or end must be provided.")
+
+    if direction is not None and end is not None:
+        raise ValueError("To visualize an arrow either direction or end must be provided, not both.")
+
+    if direction is None:
+        direction = np.array(end) - start
+        length = float(np.linalg.norm(direction))
+
+    direction = np.array(direction)
+
+    cone_height = 4 * radius  # rule of thumb
+    cone_width = cone_height / 2  # rule of thumb
+    cylinder_height = length - cone_height
+
+    orientation = orientation_with_Z_axis(direction)
+
+    pose = np.identity(4)
+    pose[:3, :3] = orientation
+    pose[:3, 3] = start
+
+    if color is None:
+        color_from_direction = np.abs(direction)
+        color_from_direction = color_from_direction / np.linalg.norm(color_from_direction)
+        color_from_direction = color_from_direction / np.max(color_from_direction)
+        color = Rgba(*color_from_direction, 1)  # type: ignore
+
+    arrow_transform = RigidTransform(pose)
+    meshcat.SetTransform(f"{name}/arrow/", arrow_transform)
+
+    cylinder_transform = RigidTransform(p=[0, 0, cylinder_height / 2])  # type: ignore
+    meshcat.SetObject(f"{name}/arrow/cylinder", Cylinder(radius, cylinder_height), color)
+    meshcat.SetTransform(f"{name}/arrow/cylinder", cylinder_transform)
+
+    cone_transform = RigidTransform(rpy=RollPitchYaw(np.deg2rad([180, 0, 0])), p=[0, 0, cylinder_height + cone_height])  # type: ignore
+    meshcat.SetTransform(f"{name}/arrow/cone/", cone_transform)
+    meshcat.SetObject(f"{name}/arrow/cone", MeshcatCone(cone_height, cone_width, cone_width), color)

airo_drake/visualization/frame.py

@@ -6,8 +6,8 @@
 
 def visualize_frame(
     meshcat: Meshcat,
+    name: str,
     transform: RigidTransform | HomogeneousMatrixType,
-    path: str,
     length: float = 0.1,
     radius: float = 0.002,
     opacity: float = 1.0,
@@ -16,14 +16,14 @@ def visualize_frame(
 
     Args:
         transform: The 6D pose / transform of the frame to visualize.
+        name: A name or path for the frame in the MeshCat hierarchy, can be used to delete it later.
         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)
+    meshcat.SetTransform(name, transform)
     axis_labels = ["x-axis", "y-axis", "z-axis"]
     axis_colors = [
         Rgba(1, 0, 0, opacity),
@@ -37,5 +37,5 @@ def visualize_frame(
     ]
 
     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)
+        meshcat.SetTransform(f"{name}/{label}", transform)
+        meshcat.SetObject(f"{name}/{label}", Cylinder(radius, length), color)

notebooks/02_visualization.ipynb

@@ -118,7 +118,7 @@
    "source": [
     "from airo_drake import visualize_frame\n",
     "\n",
-    "visualize_frame(scene.meshcat, tcp_pose, \"tcp_frame\")"
+    "visualize_frame(scene.meshcat, \"tcp_frame\", tcp_pose)"
    ]
   },
   {
@@ -127,7 +127,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "visualize_frame(scene.meshcat, RigidTransform(), \"world_frame\", length=1.0, radius=0.01, opacity=0.2)"
+    "visualize_frame(scene.meshcat, \"world_frame\", RigidTransform(), length=1.0, radius=0.01, opacity=0.2)"
    ]
   },
   {

notebooks/03_collision_checking.ipynb

@@ -275,7 +275,9 @@
    "source": [
     "## Checking TCP paths 📏\n",
     "\n",
-    "For the cloth competition, I've had the problem that my `pregrasp_pose` pose is collision-free, but movely linearly to the `grasp_pose` make the robot elbow collide with the table.\n",
+    "### Use case 1: Grasping a hanging shirt 👕\n",
+    "\n",
+    "For the cloth competition, I've had the problem that my `pregrasp_pose` pose is collision-free, but moving linearly to the `grasp_pose` make the robot elbow collide with the table.\n",
     "\n",
     "Below is a recreation of this scenario, and a demonstration of how checking the TCP path can help us avoid this problem."
    ]
@@ -306,8 +308,8 @@
     "pregrasp_pose = grasp_pose.copy()\n",
     "pregrasp_pose[0:3, 3] -= 0.4 * gripper_forward_direction\n",
     "\n",
-    "visualize_frame(scene.meshcat, pregrasp_pose, \"pregrasp_pose\")\n",
-    "visualize_frame(scene.meshcat, grasp_pose, \"grasp_pose\")"
+    "visualize_frame(scene.meshcat, \"pregrasp_pose\", pregrasp_pose)\n",
+    "visualize_frame(scene.meshcat, \"grasp_pose\", grasp_pose)"
    ]
   },
   {
@@ -387,7 +389,7 @@
     "tcp_path = interpolate_pose_path_positions(pregrasp_pose, grasp_pose, n=20)\n",
     "\n",
     "for i, tcp_pose in enumerate(tcp_path):\n",
-    "    visualize_frame(scene.meshcat, tcp_pose, f\"tcp_path/pose_{i}\", length=0.05, opacity=0.1)\n",
+    "    visualize_frame(scene.meshcat, f\"tcp_path/pose_{i}\", tcp_pose, length=0.05, opacity=0.1)\n",
     "\n",
     "\n",
     "path_joints_solutions = []\n",
@@ -396,8 +398,10 @@
     "    joint_solutions = ur5e.inverse_kinematics_with_tcp(tcp_pose, tcp_transform)\n",
     "    path_joints_solutions.append(np.array(joint_solutions).squeeze())\n",
     "\n",
-    "    \n",
-    "path_joints_solutions = np.array(path_joints_solutions) # note this is not always possible, because the number of solutions can change\n",
+    "\n",
+    "path_joints_solutions = np.array(\n",
+    "    path_joints_solutions\n",
+    ")  # note this is not always possible, because the number of solutions can change\n",
     "path_joints_solutions.shape"
    ]
   },
@@ -511,6 +515,101 @@
     "animate_joint_configurations(scene.meshcat, scene.robot_diagram, scene.arm_index, paths_without_collisions[0])"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from airo_drake import time_parametrize_toppra\n",
+    "\n",
+    "\n",
+    "trajectories = [time_parametrize_toppra(scene.robot_diagram.plant(), path) for path in paths_without_collisions]\n",
+    "durations = [trajectory.end_time() for trajectory in trajectories]\n",
+    "\n",
+    "print(durations)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fastest_trajectory = trajectories[np.argmin(durations)]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from airo_drake import animate_joint_trajectory\n",
+    "\n",
+    "\n",
+    "animate_joint_trajectory(scene.meshcat, scene.robot_diagram, scene.arm_index, fastest_trajectory, context)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Use case two: Closing a Nepresso machine's lever ☕"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lever_axis_point = np.array([0.3, 0.0, 0.5])\n",
+    "lever_axis_direction = np.array([0, -1, 0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from airo_drake.visualization.arrow import visualize_arrow\n",
+    "\n",
+    "visualize_arrow(meshcat, \"arrow\", lever_axis_point, lever_axis_direction)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "start = (1, 0, 0)\n",
+    "direction = (0, 1, 0)\n",
+    "visualize_arrow(meshcat, \"arrow2\", start, direction)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "start = (1, 0, 0)\n",
+    "direction = (0, 0, -1)\n",
+    "visualize_arrow(meshcat, \"arrow3\", start, direction)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "visualize_arrow(meshcat, \"arrow4\", (0, 1, 0), end=(1, 1, 1))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,